[Performance] LOUDSのinit処理を高速化 (#30)

* simplify calculation * cleanup * add code fro debug * fix bug * cleanup * cleanup * add binary search * cleanup binary search * cleanup binary search * cleanup init * cleanup init * debugging: flatten indices * flatten indices
2025-12-03 02:58:27 +00:00 · 2024-01-07 15:48:19 +09:00
parent 1bbb1ec03e
commit f2c5c3c58d
1 changed files with 73 additions and 35 deletions
--- a/Sources/KanaKanjiConverterModule/LOUDS/LOUDS.swift
+++ b/Sources/KanaKanjiConverterModule/LOUDS/LOUDS.swift
@@ -15,7 +15,11 @@ struct LOUDS: Sendable {
    private static let uExp = 6

    private let bits: [Unit]
-    private let char2nodeIndices: [[Int]]
+    /// indexを並べてflattenしたArray。
+    ///  - seealso: flatChar2nodeIndicesIndex
+    private let flatChar2nodeIndices: [Int]
+    /// 256個の値を入れるArray。`flatChar2nodeIndices[flatChar2nodeIndicesIndex[char - 1] ..< flatChar2nodeIndicesIndex[char]]`が`nodeIndices`になる
+    private let flatChar2nodeIndicesIndex: [Int]
    /// 0の数（1の数ではない）
    ///
    /// LOUDSのサイズが4GBまでは`UInt32`で十分
@@ -23,12 +27,42 @@ struct LOUDS: Sendable {

    @inlinable init(bytes: [UInt64], nodeIndex2ID: [UInt8]) {
        self.bits = bytes
-        self.char2nodeIndices = nodeIndex2ID.enumerated().reduce(into: .init(repeating: [], count: 1 << 8)) { list, data in
-            list[Int(data.element)].append(data.offset)
+        // flatChar2nodeIndicesIndexを構築する
+        // これは、どのcharがどれだけの長さのnodeIndicesを持つかを知るために行う
+        var flatChar2nodeIndicesIndex = [Int](repeating: 0, count: 256)
+        flatChar2nodeIndicesIndex.withUnsafeMutableBufferPointer { buffer in
+            for value in nodeIndex2ID {
+                buffer[Int(value)] += 1
+            }
+            // 累積和にする
+            for i in 1 ..< 256 {
+                buffer[i] = buffer[i - 1] + buffer[i]
+            }
        }
-        self.rankLarge = bytes.reduce(into: [0]) {
-            $0.append($0.last! &+ UInt32(Self.unit &- $1.nonzeroBitCount))
+        // flatChar2nodeIndicesを構築する
+        // すでに開始位置はflatChar2nodeIndicesIndexで分かるので、もう一度countsを構築しながら適切な場所にindexを入れていく
+        var counts = [Int](repeating: 0, count: 256)
+        self.flatChar2nodeIndices = counts.withUnsafeMutableBufferPointer { countsBuffer in
+            var flatChar2nodeIndices = [Int](repeating: 0, count: nodeIndex2ID.count)
+            for (i, value) in zip(nodeIndex2ID.indices, nodeIndex2ID) {
+                if value == .zero {
+                    flatChar2nodeIndices[countsBuffer[Int(value)]] = i
+                } else {
+                    flatChar2nodeIndices[flatChar2nodeIndicesIndex[Int(value) - 1] + countsBuffer[Int(value)]] = i
+                }
+                countsBuffer[Int(value)] += 1
+            }
+            return flatChar2nodeIndices
        }
+        self.flatChar2nodeIndicesIndex = consume flatChar2nodeIndicesIndex
+
+        var rankLarge: [UInt32] = .init(repeating: 0, count: bytes.count + 1)
+        rankLarge.withUnsafeMutableBufferPointer { buffer in
+            for (i, byte) in zip(bytes.indices, bytes) {
+                buffer[i + 1] = buffer[i] &+ UInt32(Self.unit &- byte.nonzeroBitCount)
+            }
+        }
+        self.rankLarge = consume rankLarge
    }

    /// parentNodeIndex個の0を探索し、その次から1個増えるまでのIndexを返す。
@@ -43,34 +77,33 @@ struct LOUDS: Sendable {
        // rankLargeは左側の0の数を示すので、difを取っている
        // まず最低限の絞り込みを行う。leftを探索する。
        // 探しているのは、startIndexが含まれるbitsのindex `i`
-        var left = (parentNodeIndex >> Self.uExp) &- 1
-        while true {
-            let dif = parentNodeIndex &- Int(self.rankLarge[Int(left) &+ 1])
-            if dif >= Self.unit {
-                left &+= dif >> Self.uExp
+        var left = parentNodeIndex >> Self.uExp
+        var right = self.rankLarge.endIndex - 1
+        var i = self.rankLarge.endIndex
+        while left <= right {
+            let mid = (left + right) / 2
+            if self.rankLarge[mid] >= parentNodeIndex {
+                i = mid
+                right = mid - 1
            } else {
-                break
+                left = mid + 1
            }
        }
-        var i: Int?
-        for index in left &+ 1 ..< self.bits.endIndex where self.rankLarge[index &+ 1] >= parentNodeIndex {
-            i = index
-            break
-        }
-        guard let i else {
+        guard i != self.rankLarge.endIndex else {
            return 0 ..< 0
        }
-        // 探索パート②
-        // 目標はparentNodeIndex番目の0の位置である`k`の発見
-        let byte = self.bits[i]
-        var k = 0
-        for _ in  0 ..< parentNodeIndex - Int(self.rankLarge[i]) {
-            k = (~(byte << k)).leadingZeroBitCount &+ k &+ 1
-        }
-        let start = (i << Self.uExp) &+ k &- parentNodeIndex &+ 1
-        // ちょうどparentNodeIndex個の0がi番目にあるかどうか
-        if self.rankLarge[i &+ 1] == parentNodeIndex {
-            return self.bits.withUnsafeBufferPointer {(buffer: UnsafeBufferPointer<Unit>) -> Range<Int> in
+        i -= 1
+        return self.bits.withUnsafeBufferPointer {(buffer: UnsafeBufferPointer<Unit>) -> Range<Int> in
+            // 探索パート②
+            // 目標はparentNodeIndex番目の0の位置である`k`の発見
+            let byte = buffer[i]
+            var k = 0
+            for _ in  0 ..< parentNodeIndex - Int(self.rankLarge[i]) {
+                k = (~(byte << k)).leadingZeroBitCount &+ k &+ 1
+            }
+            let start = (i << Self.uExp) &+ k &- parentNodeIndex &+ 1
+            // ちょうどparentNodeIndex個の0がi番目にあるかどうか
+            if self.rankLarge[i &+ 1] == parentNodeIndex {
                var j = i &+ 1
                while buffer[j] == Unit.max {
                    j &+= 1
@@ -81,13 +114,13 @@ struct LOUDS: Sendable {
                let byte2 = buffer[j]
                let a = (~byte2).leadingZeroBitCount % Self.unit
                return start ..< (j << Self.uExp) &+ a &- parentNodeIndex &+ 1
+            } else {
+                // difが0以上の場合、k番目以降の初めての0を発見したい
+                // 例えばk=1の場合
+                // Ex. 1011_1101 => 0111_1010 => 1000_0101 => 1 => 2
+                let a = ((~(byte << k)).leadingZeroBitCount &+ k) % Self.unit
+                return start ..< (i << Self.uExp) &+ a &- parentNodeIndex &+ 1
            }
-        } else {
-            // difが0以上の場合、k番目以降の初めての0を発見したい
-            // 例えばk=1の場合
-            // Ex. 1011_1101 => 0111_1010 => 1000_0101 => 1 => 2
-            let a = ((~(byte << k)).leadingZeroBitCount &+ k) % Self.unit
-            return start ..< (i << Self.uExp) &+ a &- parentNodeIndex &+ 1
        }
    }

@@ -96,7 +129,12 @@ struct LOUDS: Sendable {
    @inlinable func searchCharNodeIndex(from parentNodeIndex: Int, char: UInt8) -> Int? {
        // char2nodeIndicesには単調増加性があるので二分探索が成立する
        let childNodeIndices = self.childNodeIndices(from: parentNodeIndex)
-        let nodeIndices = self.char2nodeIndices[Int(char)]
+        let nodeIndices: ArraySlice<Int> = if char == .zero {
+            self.flatChar2nodeIndices[0 ..< self.flatChar2nodeIndicesIndex[Int(char)]]
+        } else {
+            self.flatChar2nodeIndices[self.flatChar2nodeIndicesIndex[Int(char - 1)] ..< self.flatChar2nodeIndicesIndex[Int(char)]]
+        }
+
        var left = nodeIndices.startIndex
        var right = nodeIndices.endIndex
        while left < right {