Swift Digest
SE-0245 | Swift Evolution

Add an Array Initializer with Access to Uninitialized Storage

Proposal
SE-0245
Authors
Nate Cook
Review Manager
Ted Kremenek
Status
Implemented (Swift 5.1)

01 何が問題だったのか

固定長のバッファに対してアルゴリズムを実装したい場合、標準ライブラリの Array では「指定したサイズの未初期化の領域」を直接扱う手段がありませんでした。Array を作成する手段は、リテラル、repeating:count: のように全要素を初期化するもの、または既存の Sequence からコピーするものに限られ、どれも要素をあらかじめ初期化する必要があります。

たとえば、配列に対して stable partition(順序を保ったまま要素を2グループに振り分ける操作)を O(n) で実装したい場合、次のようなアルゴリズムを使いたくなります。

  1. 元配列と同じサイズの新しい配列を用意する。
  2. 元配列を走査し、条件に合う要素は新配列の先頭側から、合わない要素は末尾側から書き込む。
  3. 最後に末尾側に集めたスライスを反転させる。

ところが、このためには「サイズだけ決まっていて中身は未初期化の Array」が必要です。Swift では UnsafeMutableBufferPointer で未初期化のメモリを確保することはできますが、そのバッファを Array の内部ストレージとして直接受け渡す方法が無く、結局は要素をコピーすることになります。結果として、不要な初期化やコピーが避けられず、アルゴリズムを本来の効率で書けませんでした。

同様の制約は、C の API と組み合わせるときにも現れます。バッファに要素を書き込み、書き込んだ個数を返してくる類の API(たとえば vDSP_vsadd のような Accelerate 系の関数)では、呼び出し側が結果を受け取るバッファを用意する必要があります。このバッファを Array として扱いたい場合、

  • あらかじめ 0 などで初期化した Array を渡す(不要な初期化のコストがかかる)、
  • UnsafeMutableBufferPointer に書き込んでから Array にコピーする(コピーのコストがかかる)、

のどちらかを選ぶしかなく、せっかくの高速な API の性能上の利点を打ち消してしまいます。

要するに、「指定したサイズの未初期化領域を持つ Array を作り、その領域を直接初期化して、そのまま Array として扱う」ことが Swift では表現できませんでした。

02 どのように解決されるのか

ArrayContiguousArray に、未初期化のストレージへ直接アクセスできる新しいイニシャライザが追加されます。確保した容量を受け取るクロージャの中で要素を初期化し、実際に初期化した個数を書き戻すことで、不要な初期化やコピーを経ずに Array を作れます。

var myArray = Array<Int>(unsafeUninitializedCapacity: 10) { buffer, initializedCount in
    for x in 1..<5 {
        buffer[x] = x
    }
    buffer[0] = 10
    initializedCount = 5
}
// myArray == [10, 1, 2, 3, 4]

シグネチャは次のとおりです。

public init(
    unsafeUninitializedCapacity: Int,
    initializingWith initializer: (
        _ buffer: inout UnsafeMutableBufferPointer<Element>,
        _ initializedCount: inout Int
    ) throws -> Void
) rethrows

クロージャには次の2つが渡されます。

  • buffer: 指定した容量ぶんの未初期化メモリを指す UnsafeMutableBufferPointercountunsafeUninitializedCapacity と一致します(実際の Array の容量はそれ以上になる可能性がありますが、クロージャから見えるバッファはちょうど指定した個数です)。
  • initializedCount: クロージャ終了時までに「実際に初期化した要素の個数」を書き込むための inout パラメータ。

呼び出し側が守るべき事後条件

クロージャを抜けるとき、バッファは次の状態になっている必要があります。

  • buffer[0..<initializedCount] はすべて初期化済み。
  • buffer[initializedCount...] はすべて未初期化(deinitialize 済み)。

この約束はイニシャライザの名前にある通り unsafe で、メモリ管理はユーザーの責任です。後続の Array としての扱い(要素数、destructor の呼び出しなど)はこの事後条件を前提に動きます。

エラーをスローする場合

クロージャが途中で throw する場合も、上記の事後条件を満たす必要があります。つまりスローする前に、次のどちらかをしなければなりません。

  1. すでに初期化してしまった要素をすべて deinitialize してから throw する(initializedCount はそのまま)。
  2. initializedCount を現時点で実際に初期化済みの個数に更新してから throw する。

スロー時点の initializedCount が正しいものとして扱われ、Array 側はその範囲だけを生きている要素として解放します。

使用例: stable partition

このイニシャライザを使うと、冒頭の stable partition を余計なコピーなしで書けます。lowhigh の2つのポインタを端から進めながら要素を直接書き込み、最後に末尾側のスライスを反転するだけです。エラー時は、それまでに初期化した両端の要素を deinitialize してから rethrow します。

func stablyPartitioned(by belongsInFirstPartition: (Element) throws -> Bool) rethrows -> [Element] {
    return try Array<Element>(unsafeUninitializedCapacity: count) {
        buffer, initializedCount in
        var low = buffer.baseAddress!
        var high = low + buffer.count
        do {
            for element in self {
                if try belongsInFirstPartition(element) {
                    low.initialize(to: element)
                    low += 1
                } else {
                    high -= 1
                    high.initialize(to: element)
                }
            }

            let highIndex = high - buffer.baseAddress!
            buffer[highIndex...].reverse()
            initializedCount = buffer.count
        } catch {
            let lowCount = low - buffer.baseAddress!
            let highCount = (buffer.baseAddress! + buffer.count) - high
            buffer.baseAddress!.deinitialize(count: lowCount)
            high.deinitialize(count: highCount)
            throw error
        }
    }
}

使用例: C API への橋渡し

バッファに結果を書き込んでくる C 関数のラッパーも、Array を 0 で埋めるコストなしに書けます。

extension Array where Element == Float {
    func dspAdd(scalar: Float) -> [Float] {
        let n = self.count
        return self.withUnsafeBufferPointer { buf in
            var scalar = scalar
            return Array<Float>(unsafeUninitializedCapacity: n) { rbuf, count in
                vDSP_vsadd(buf.baseAddress!, 1, &scalar, rbuf.baseAddress!, 1, UInt(n))
                count = n
            }
        }
    }
}

命名について

引数ラベルが unsafeUninitializedCapacity と長めなのは意図的です。このイニシャライザは unsafe(メモリ管理の責任がユーザー側にある)で、かつ uninitialized なストレージへのアクセスを提供する、という他の withUnsafe... 系 API とは異なる性質を持ちます。トレーリングクロージャで呼ばれたときでもこれらの性質が使用箇所から読み取れるよう、最初のラベルに明示的に含めるという選択です。