Swift での Copy on Write の実装方法の解説

Swift の Array や Dictionary は Copy on Write になっていてとても使いやすいです。この記事では、そのような Copy on Write の実装方法を解説します。

CoW の動作

Copy on Write はよく CoW と略されるので以後そのように呼びます。さて、 CoW はその名の通り、書き込みが生じる時にコピーをするものです。まずはその動作について説明します。

以下のコードを見てください。

var a = [1, 2, 3]
var b = a

変数 a に [1, 2, 3] を代入した後、 b にコピーしています。変数 a の配列は 3 つの値が入ったデータ領域を持っていますが、 a が b にコピーされたときには、そのデータ領域はコピーされません。その代わりに、 b は a の持っているデータ領域を共有します。だから、 a の配列の要素がどれだけ多くても、このコピーは高速で消費メモリもとても小さいです。

ただし、このままデータ領域が共有されているだけであれば a に変更を加えたら、 b の値も変化してしまいます。

a.append(4) // a に値 4 を追加する。
print(a) // "1, 2, 3, 4" が出力される。
print(b) // "1, 2, 3, 4" が出力されてしまうでしょうか？

実際に Swift を使ってみれば、 b の値はちゃんと [1, 2, 3] のままであることが確認できます。この動作は append を呼び出した時にデータ領域のコピーが行われることで実現されています。つまり append を呼び出した時点で a と b のデータ領域の共有が解除され、それぞれが固有のデータ領域を持った状態に変化します。 append は a に対する変更操作ですから、変更操作 (=Write) のタイミングでデータ領域のコピー (Copy) が行われています。これが、 Copy on Write の動作です。

さてこの時点で、 a は [1, 2, 3, 4] が入ったデータ領域、 b は [1, 2, 3] が入ったデータ領域を保持しています。ここで、さらに a に値を追加する事を考えてみます。

a.append(5)

もし、 append メソッドを呼び出すたびにデータ領域をコピーしていたら、この append の呼び出し時にも、 [1, 2, 3, 4] の入ったデータ領域をコピーしてから、その後で 5 を追加する事になってしまいますが、そのコピーは不要です。現在のデータ領域に 5 を追加するだけで良いからです。実際、この場合はコピーは行われません。

つまり、変更操作のタイミングならいつでもデータ領域のコピーを行うわけではなく、そこには更に条件があります。すなわち、変更操作が生じたとき、データ領域が他のオブジェクトと共有されているならば、データ領域のコピーを行う、というのがより正確な CoW の動作になります。

CoW の設計

この章では、 CoW をもった配列を自作するためにその設計を示します。まず、配列それ自体の型を用意します。これは値型にします。名前は CoWArray とします。インターフェースを以下に示します。

public struct CoWArray<T> {
    public init()
    public var count: Int { get }
    public subscript(index: Int) -> T { get set }
    public mutating func append(_ element: T)
    public mutating func remove(at index: Int)
}

機能は要素数の取得とサブスクリプトでのアクセス、要素の追加と削除を用意します。サブスクリプトアクセスのセッターと、要素の追加、削除は mutating メソッドになります。

次に、データ領域を表す型を用意します。これは他の配列と共有する必要があるので参照型にします。名前は CowArrayStorage とします。インターフェースを以下に示します。

internal class CoWArrayStorage<T> {
    public init()
    public init(copy: CoWArrayStorage<T>)
    public var count: Int { get }
    public subscript(index: Int) -> T { get set }
    public func append(_ element: T)
    public func remove(at index: Int)
}

定義を CoWArray と見比べると、 init(copy:) が追加されている以外は同じになっています。これがポイントです。

CoWArray は内部でプロパティとして CoWArrayStorage を保持します。そして、全てのメソッドはそのストレージの同名メソッドを呼び出すだけにします。ただし、変更が発生するメソッド、つまり mutating なメソッドの場合は、ストレージのメソッドを呼び出す前に、ストレージが他のオブジェクトから共有されているかチェックして、もし共有されている場合はコピーをする処理をはさみます。これが CoW を実現する肝になります。

ストレージが共有されているかどうかは、 Swift 標準ライブラリの isKnownUniquelyReferenced 関数を使います。この関数は、引数で与えた参照型の値が、ユニークであれば、つまり共有されていなければ true を返します。ユニークでなければ、つまり共有されていれば false を返します。内部的には、参照型のオブジェクトが内部に持っている参照カウンタの値を調べて、それが 1 か 2以上かどうかで判定しています。

設計は以上です。

CoW の実装

この章では実装を行います。まず、配列の機能本体を提供する CoWArrayStorage の実装を示し、次に配列自体を提供する CoWArray の実装を示します。前者の詳細は CoW の実現とは特に関係が無いので読み飛ばしても良いです。後者は CoW の重要なポイントになります。

この手の配列の実装では、 replaceSubrange を実装すれば、 append, insert, remove, subscript set はいずれも replaceSubrange の呼び出しで済ませられますが、今回はシンプルにするために append と remove の直接の実装だけを行います。

データ領域の実装

データ構造

CoWArrayStorage はメモリ領域を示すポインタ memory と、そのメモリ領域の大きさを示す capacity 、配列に入っている要素の個数を表す count をプロパティとして持ちます。 capacity と count の 2 つを用意することで、配列の要素数よりも余分にメモリ領域を持っておく事ができます。そうすると、 append する時にメモリ領域の拡張を省略できる場合があります。ここまでの実装を示します。

internal class CoWArrayStorage<T> {
    public init() {
        self.memory = nil
        self.capacity = 0
        self.count = 0
    }

    private var memory: UnsafeMutablePointer<T>?
    private var capacity: Int
    private var count: Int
}

なお、 capacity が 0 のときは、 memory は nil にするという規約にします。 deinit ではデータ領域の解放が必要ですが、それは後で書きます。

subscript

subscript と count については素直に実装します。

    public private(set) var count: Int
    
    public subscript(index: Int) -> T {
        get {
            return memory![index]
        }
        set {
            memory![index] = newValue
        }
    }

subscript の範囲チェックは省略しました。メモリ安全にしたければチェックしてください。

reserveCapacity

メモリ領域のリサイズを行う関数、 reserveCapacity を実装します。

この関数は、はじめに新しい領域を確保し、そこに古い領域のデータをコピーして、さいごに古い領域を解放します。この時に注意しなければならないのが、ポインタの 3 つの状態です。ポインタには、 未確保、未初期化、初期化済の 3種類の状態があります。今回のポインタは連続領域なので、まず領域を指しているかどうかで未確保か確保済みの 2 つがあり、確保済みの場合において、領域の 1 要素ごとに、未初期化か初期化済があります。また、新しいメモリ領域は、大きくなる場合と小さくなる場合があり、小さくなる場合には 0 になる場合もあります。このポインタの状態については、過去に 解説記事 を書いたので参考にしてください。

まず、大きくなる場合は、古いの領域の要素を全て新しい領域に移す事ができます。これは、新しい領域に対しては 未初期化領域を初期化する操作なので、 initialize 操作になります。古い領域に対しては、初期化済み領域を未初期化に戻す操作なので deinitialize 操作になります。2 つの領域があって、片方を deinitialize しながら別の領域に書き込む操作は move 操作になります。なので moveInitialize を使うと一発で記述できます。

小さくなる場合には、古い領域の要素のうち、新しい領域に入る分と、入り切らない分が出てきます。入りきる分については同様に moveInitialize で移動させて、入り切らない分については破棄します。破棄するためには deinitialize が使えます。この、前半を moveInitialize 、 後半を deinitialize するロジックは、その境界となる要素番号を使って書いてやれば、大きくなる場合とひとまとめに書くことができます。

0 になる場合については、新しい領域のポインタは nil にしておく以外は小さくなる場合と同じです。

最後に 古い領域を deallocate で解放して、 memory プロパティと capacity プロパティを更新して完了です。

なお、 count プロパティについてはこの関数では関与しない規約とします。

以上の方針で実装したコードが下記になります。

    private func reserveCapacity(_ newCapacity: Int) {
        let newMemory: UnsafeMutablePointer<T>?
        if newCapacity > 0 {
            newMemory = UnsafeMutablePointer<T>.allocate(capacity: newCapacity)
        } else {
            newMemory = nil
        }
        
        let moveCount = min(count, newCapacity)
        let deinitCount = capacity - moveCount

        if let nmem = newMemory {
            if moveCount > 0 {
                nmem.moveInitialize(from: memory!, count: moveCount)
            }
        }
        if let mem = memory {
            if deinitCount > 0 {
                (mem + moveCount).deinitialize(count: capacity - moveCount)
            }
            mem.deallocate(capacity: capacity)
        }
        
        memory = newMemory
        capacity = newCapacity
    }

deinit での破棄と解放

メモリ確保処理を実装したので、ここで忘れずに deinit でのメモリ解放も記述しておきます。データ領域は先頭から要素数分は初期化済みなので、これの破棄も行います。

    deinit {
        memory?.deinitialize(count: count)
        reserveCapacity(0)
    }

append

append では、最初にメモリ領域が足りているかチェックします。足りていない場合は、先程の reserveCapacity を使ってメモリ確保をします。ここで領域の拡大は要素 1 つ分にせず、一気に領域を 2 倍に広げるようにすると効率的です。連続で append が呼び出される場合に、メモリの再確保の頻度が減らせるからです。しかし、領域の大きさが 0 の場合には 2 倍にしても 0 のままなので、特別に初期値を設定しておきます。メモリ領域の準備ができたら、末尾に要素を追加します。ここもやはり、未初期化領域を初期化する操作なので、 initialize 操作をします。最後に要素数を 1 増やします。以下にコードを示します。

    public func append(_ element: T) {
        if count == capacity {
            if capacity == 0 {
                reserveCapacity(4)
            } else {
                reserveCapacity(capacity * 2)
            }
        }
        
        (memory! + count).initialize(to: element)
        count += 1
    }

remove

remove では、指定されたインデックスの要素を削除した後、後続の要素を 1 つずつ前にずらします。まず、削除する操作は deinitialize です。これによって削除された場所は未初期化になります。そして、 1 つ後ろの要素を、削除された要素のあった場所に移動させるわけですが、これは reserveCapacity のときと同様、移動元が初期化済で、移動先が未初期化なので、 moveInitialize で書けます。最後に要素数を 1 減らします。

moveInitialize は2つのメモリ領域が共有領域を持っていても、要素ごとに move していくのでうまく動きます。

    public func remove(at index: Int) {
        (memory! + index).deinitialize()
        let moveNum = count - index - 1
        if moveNum > 0 {
            (memory! + index).moveInitialize(from: (memory! + index + 1), count: moveNum)
        }
        count -= 1
    }

コピー init

最後に、コピー用の init を作ります。コピーする要素数でメモリ確保して、確保した未初期化領域をまとめて initialize します。

    public init(copy: CoWArrayStorage<T>) {
        self.count = copy.count
        self.capacity = count
        self.memory = UnsafeMutablePointer<T>.allocate(capacity: count)
        
        memory?.initialize(from: copy.memory!, count: count)
    }

まとめ

以上で CoWArrayStorage が実装できました。

CoW 型の実装

いよいよ CoW 型の実装をします。まず、データ構造ですが、ストレージオブジェクトを保持するだけです。

public struct CoWArray<T> {
    public init() {
        self.storage = .init()
    }

    private var storage: CoWArrayStorage<T>
}

次に、 CoW の肝となる、共有されていればストレージをコピーするメソッドを作ります。共有されているかどうかは isKnownUniquelyReference でチェックして、コピーは先程作ったコピー用の init を使います。

    private mutating func copyStorageIfShared() {
        if isKnownUniquelyReferenced(&storage) {
            return
        }
        
        storage = .init(copy: storage)
    }

あとは、全てのメソッドを素通ししつつ、 mutating なものについては事前に copyStorageIfShared を呼び出すだけです。

    public var count: Int {
        return storage.count
    }

    public subscript(index: Int) -> T {
        get {
            return storage[index]
        }
        set {
            copyStorageIfShared()
            storage[index] = newValue
        }
    }

    public mutating func append(_ element: T) {
        copyStorageIfShared()
        storage.append(element)
    }

    public mutating func remove(at index: Int) {
        copyStorageIfShared()
        storage.remove(at: index)
    }

これで CoW な配列が完成しました。

ソース

完成したソースはこちらにアップしてあります。