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,65 @@ 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)
+  }
+
+  /// 採用鏈條逆向清理策略，從最遠端的 prev 開始逐一退行清理
+  /// - Parameter allStates: 所有需要清理的狀態集合
+  private static func cleanSearchStateChains(allStates: Set<SearchState>) {
+    var processedStates = Set<SearchState>()
+
+    // 對每個狀態進行鏈條逆向清理
+    for state in allStates {
+      guard !processedStates.contains(state) else { continue }
+
+      // 找到鏈條的末端（最遠端的 prev）
+      var chainStates = [SearchState]()
+      var current: SearchState? = state
+
+      // 追溯整個鏈條，收集所有狀態
+      while let currentState = current {
+        // 避免無限循環（檢測循環引用）
+        if chainStates.contains(where: { $0 === currentState }) {
+          break
+        }
+        chainStates.append(currentState)
+        current = currentState.prev
+      }
+
+      // 從鏈條末端開始逐一清理（逆向清理）
+      for i in stride(from: chainStates.count - 1, through: 0, by: -1) {
+        let stateToClean = chainStates[i]
+
+        // 首先清理 node 引用
+        stateToClean.node = nil
+
+        // 然後清理 prev 引用（從末端開始，確保沒有懸掛指針）
+        stateToClean.prev = nil
+
+        // 標記為已處理
+        processedStates.insert(stateToClean)
+      }
+    }
   }
 }
 
@@ -196,9 +239,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,11 +250,33 @@ extension Megrez.Compositor {
       lhs.originalNodeRef === rhs.originalNodeRef && lhs.position == rhs.position
     }
 
-    /// 清理單一 SearchState 的參據
+    /// 清理單一 SearchState 的參據，採用鏈條逆向清理策略
     /// 注意：由於新的清理策略是直接清理各個集合，這個方法現在主要用於向下相容
     func cleanState() {
-      node = nil
-      prev = nil
+      // 對於單一狀態，我們使用簡化的鏈條清理
+      var chainStates = [Self]()
+      var current: SearchState? = self
+
+      // 追溯整個鏈條，收集所有狀態
+      while let currentState = current {
+        // 避免無限循環（檢測循環引用）
+        if chainStates.contains(where: { $0 === currentState }) {
+          break
+        }
+        chainStates.append(currentState)
+        current = currentState.prev
+      }
+
+      // 從鏈條末端開始逐一清理（逆向清理）
+      for i in stride(from: chainStates.count - 1, through: 0, by: -1) {
+        let stateToClean = chainStates[i]
+
+        // 首先清理 node 引用
+        stateToClean.node = nil
+
+        // 然後清理 prev 引用（從末端開始，確保沒有懸掛指針）
+        stateToClean.prev = nil
+      }
     }
 
     func hash(into hasher: inout Hasher) {