Assembler // Continue patching memory management issues.

Author: ShikiSuen

Sources/Megrez/2_Assembler.swift

@@ -128,15 +128,16 @@ extension Megrez.Compositor {
     let sortedStates = visited.sorted { $0.distance > $1.distance }
     let statesToRemove = Array(sortedStates.dropFirst(keepRecentCount))
 
-    // 先從 Set 中移除，再清理參據（避免 hash 不一致）
+    // 先從 Set 中移除（避免 hash 不一致）
     for state in statesToRemove {
       visited.remove(state)
-      state.node = nil
-      state.prev = nil
     }
+
+    // 採用鏈條逆向清理策略清理要移除的狀態
+    Self.cleanSearchStateChains(allStates: Set(statesToRemove))
   }
 
-  /// 即時清理策略：直接清理各個資料結構，避免額外的 Set 集合
+  /// 即時清理策略：採用鏈條逆向清理避免循環引用
   /// - Parameters:
   ///   - visited: 已訪問的狀態集合
   ///   - openSet: 優先序列中剩餘的狀態
@@ -146,23 +147,33 @@ extension Megrez.Compositor {
     openSet: inout HybridPriorityQueue<PrioritizedState>,
     leadingState: SearchState
   ) {
-    // 策略1: 直接清理 visited set 中的所有狀態
+    // 收集所有需要清理的 SearchState 物件
+    var allStates = Set<SearchState>()
+
+    // 1. 新增 visited set 中的所有狀態
     for state in visited {
-      state.node = nil
-      state.prev = nil
+      allStates.insert(state)
     }
 
-    // 策略2: 直接清理 openSet 中剩餘的所有狀態
+    // 2. 新增 openSet 中剩餘的所有狀態
     while !openSet.isEmpty {
       if let prioritizedState = openSet.dequeue() {
-        prioritizedState.state.node = nil
-        prioritizedState.state.prev = nil
+        allStates.insert(prioritizedState.state)
       }
     }
 
-    // 策略3: 清理 leadingState
-    leadingState.node = nil
-    leadingState.prev = nil
+    // 3. 新增 leadingState
+    allStates.insert(leadingState)
+
+    // 4. 採用鏈條逆向清理策略
+    Self.cleanSearchStateChains(allStates: allStates)
+  }
+
+  /// 採用 SearchState 自身的清理方法進行批次清理
+  /// - Parameter allStates: 所有需要清理的狀態集合
+  private static func cleanSearchStateChains(allStates: Set<SearchState>) {
+    // 使用 SearchState 的靜態批次清理方法
+    SearchState.batchClean(Array(allStates))
   }
 }
 
@@ -196,9 +207,9 @@ extension Megrez.Compositor {
 
     // MARK: Internal
 
-    var node: Megrez.Node? // 當前節點（可變，用於清理）
+    weak var node: Megrez.Node? // 當前節點（可變，用於清理）
     let position: Int // 在輸入串中的位置
-    var prev: SearchState? // 前一個狀態
+    weak var prev: SearchState? // 前一個狀態
     var distance: Double // 累計分數
 
     // MARK: - Hashable 協定實作
@@ -207,13 +218,62 @@ extension Megrez.Compositor {
       lhs.originalNodeRef === rhs.originalNodeRef && lhs.position == rhs.position
     }
 
-    /// 清理單一 SearchState 的參據
-    /// 注意：由於新的清理策略是直接清理各個集合，這個方法現在主要用於向下相容
-    func cleanState() {
+    /// 靜態方法：批次清理多個 SearchState
+    static func batchClean(_ states: [SearchState]) {
+      var visited = Set<ObjectIdentifier>()
+
+      for state in states {
+        let stateId = ObjectIdentifier(state)
+        guard !visited.contains(stateId) else { continue }
+
+        // 清理此狀態及其鏈條
+        state.cleanChainRecursively()
+
+        // 標記整個鏈條為已處理
+        var current: SearchState? = state
+        while let currentState = current {
+          visited.insert(ObjectIdentifier(currentState))
+          current = currentState.prev
+        }
+      }
+    }
+
+    /// 清理自身的所有引用，立即釋放記憶體
+    func cleanSelf() {
       node = nil
       prev = nil
     }
 
+    /// 清理整個 SearchState 鏈條，從當前節點開始向後遞歸清理
+    /// 採用深度優先遍歷策略，確保每個節點都被清理
+    func cleanChainRecursively() {
+      var visited = Set<ObjectIdentifier>()
+      var stack = [Self]()
+      stack.append(self)
+
+      while !stack.isEmpty {
+        let current = stack.removeLast()
+        let currentId = ObjectIdentifier(current)
+
+        // 避免重複清理和無限循環
+        guard !visited.contains(currentId) else { continue }
+        visited.insert(currentId)
+
+        // 在清理前，將 prev 加入堆疊（如果存在）
+        if let prevState = current.prev {
+          stack.append(prevState)
+        }
+
+        // 清理當前節點
+        current.cleanSelf()
+      }
+    }
+
+    /// 向下相容的清理方法
+    func cleanState() {
+      cleanSelf()
+    }
+
     func hash(into hasher: inout Hasher) {
       hasher.combine(stableHash)
     }