Add an Array Initializer with Access to Uninitialized Storage
01 何が問題だったのか
固定長のバッファに対してアルゴリズムを実装したい場合、標準ライブラリの Array では「指定したサイズの未初期化の領域」を直接扱う手段がありませんでした。Array を作成する手段は、リテラル、repeating:count: のように全要素を初期化するもの、または既存の Sequence からコピーするものに限られ、どれも要素をあらかじめ初期化する必要があります。
たとえば、配列に対して stable partition(順序を保ったまま要素を2グループに振り分ける操作)を O(n) で実装したい場合、次のようなアルゴリズムを使いたくなります。
- 元配列と同じサイズの新しい配列を用意する。
- 元配列を走査し、条件に合う要素は新配列の先頭側から、合わない要素は末尾側から書き込む。
- 最後に末尾側に集めたスライスを反転させる。
ところが、このためには「サイズだけ決まっていて中身は未初期化の Array」が必要です。Swift では UnsafeMutableBufferPointer で未初期化のメモリを確保することはできますが、そのバッファを Array の内部ストレージとして直接受け渡す方法が無く、結局は要素をコピーすることになります。結果として、不要な初期化やコピーが避けられず、アルゴリズムを本来の効率で書けませんでした。
同様の制約は、C の API と組み合わせるときにも現れます。バッファに要素を書き込み、書き込んだ個数を返してくる類の API(たとえば vDSP_vsadd のような Accelerate 系の関数)では、呼び出し側が結果を受け取るバッファを用意する必要があります。このバッファを Array として扱いたい場合、
- あらかじめ 0 などで初期化した
Arrayを渡す(不要な初期化のコストがかかる)、 UnsafeMutableBufferPointerに書き込んでからArrayにコピーする(コピーのコストがかかる)、
のどちらかを選ぶしかなく、せっかくの高速な API の性能上の利点を打ち消してしまいます。
要するに、「指定したサイズの未初期化領域を持つ Array を作り、その領域を直接初期化して、そのまま Array として扱う」ことが Swift では表現できませんでした。
02 どのように解決されるのか
Array と ContiguousArray に、未初期化のストレージへ直接アクセスできる新しいイニシャライザが追加されます。確保した容量を受け取るクロージャの中で要素を初期化し、実際に初期化した個数を書き戻すことで、不要な初期化やコピーを経ずに 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: 指定した容量ぶんの未初期化メモリを指すUnsafeMutableBufferPointer。countはunsafeUninitializedCapacityと一致します(実際のArrayの容量はそれ以上になる可能性がありますが、クロージャから見えるバッファはちょうど指定した個数です)。initializedCount: クロージャ終了時までに「実際に初期化した要素の個数」を書き込むためのinoutパラメータ。
呼び出し側が守るべき事後条件
クロージャを抜けるとき、バッファは次の状態になっている必要があります。
buffer[0..<initializedCount]はすべて初期化済み。buffer[initializedCount...]はすべて未初期化(deinitialize 済み)。
この約束はイニシャライザの名前にある通り unsafe で、メモリ管理はユーザーの責任です。後続の Array としての扱い(要素数、destructor の呼び出しなど)はこの事後条件を前提に動きます。
エラーをスローする場合
クロージャが途中で throw する場合も、上記の事後条件を満たす必要があります。つまりスローする前に、次のどちらかをしなければなりません。
- すでに初期化してしまった要素をすべて deinitialize してから
throwする(initializedCountはそのまま)。 initializedCountを現時点で実際に初期化済みの個数に更新してからthrowする。
スロー時点の initializedCount が正しいものとして扱われ、Array 側はその範囲だけを生きている要素として解放します。
使用例: stable partition
このイニシャライザを使うと、冒頭の stable partition を余計なコピーなしで書けます。low と high の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 とは異なる性質を持ちます。トレーリングクロージャで呼ばれたときでもこれらの性質が使用箇所から読み取れるよう、最初のラベルに明示的に含めるという選択です。