Assembler // Continue patching memory management issues.

Author: ShikiSuen

Sources/Megrez/2_Assembler.swift

@@ -21,13 +21,17 @@ extension Megrez.Compositor {
     var visited = Set<SearchState>()
     var bestScore = [Int: Double]() // 追蹤每個位置的最佳分數
 
+    var stateCleaningTasks: [() -> ()] = []
+    defer { stateCleaningTasks.forEach { $0() } }
+
     // 初期化起始狀態。
     let leadingNode = Megrez.Node(keyArray: ["$LEADING"])
     let start = SearchState(
       node: leadingNode,
       position: 0,
       prev: nil,
-      distance: 0
+      distance: 0,
+      cleaningTaskRegister: &stateCleaningTasks
     )
     openSet.enqueue(PrioritizedState(state: start))
     bestScore[0] = 0
@@ -67,7 +71,8 @@ extension Megrez.Compositor {
           node: nextNode,
           position: nextPos,
           prev: current,
-          distance: newScore
+          distance: newScore,
+          cleaningTaskRegister: &stateCleaningTasks
         )
 
         bestScore[nextPos] = newScore
@@ -76,18 +81,17 @@ extension Megrez.Compositor {
 
       // 即時記憶體最佳化：當 visited 集合過大時進行部分清理
       if visited.count > 1_000 { // 可調整的閾值
-        Self.partialCleanVisitedStates(visited: &visited, keepRecentCount: 500)
+        // 簡單清理：保留最近的高分狀態，清理較舊的低分狀態
+        let sortedStates = visited.sorted { $0.distance > $1.distance }
+        let statesToRemove = Array(sortedStates.dropFirst(500))
+        for state in statesToRemove {
+          visited.remove(state)
+        }
       }
     }
 
     // 從最佳終止狀態重建路徑。
     guard let finalState = bestFinalState else {
-      // 即使沒有找到最佳狀態，也需要清理所有建立的 SearchState 物件
-      Self.batchCleanAllSearchStates(
-        visited: visited,
-        openSet: &openSet,
-        leadingState: start
-      )
       return []
     }
 
@@ -103,67 +107,8 @@ extension Megrez.Compositor {
       // 備註：此處不需要手動 ASAN，因為沒有參據循環（Retain Cycle）。
     }
     assembledNodes = pathNodes.map(\.copy)
-
-    // 手動 ASAN：批次清理所有 SearchState 物件以防止記憶體洩漏
-    // 包括 visited set 中的所有狀態、openSet 中剩餘的狀態，以及 leadingState
-    Self.batchCleanAllSearchStates(
-      visited: visited,
-      openSet: &openSet,
-      leadingState: start
-    )
     return assembledNodes
   }
-
-  /// 部分清理已訪問狀態集合以控制記憶體使用
-  /// - Parameters:
-  ///   - visited: 已訪問的狀態集合
-  ///   - keepRecentCount: 要保留的最近狀態數量
-  private static func partialCleanVisitedStates(
-    visited: inout Set<SearchState>,
-    keepRecentCount: Int
-  ) {
-    guard visited.count > keepRecentCount else { return }
-
-    // 按距離排序，保留分數較高的狀態
-    let sortedStates = visited.sorted { $0.distance > $1.distance }
-    let statesToRemove = Array(sortedStates.dropFirst(keepRecentCount))
-
-    // 先從 Set 中移除，再清理參據（避免 hash 不一致）
-    for state in statesToRemove {
-      visited.remove(state)
-      state.node = nil
-      state.prev = nil
-    }
-  }
-
-  /// 即時清理策略：直接清理各個資料結構，避免額外的 Set 集合
-  /// - Parameters:
-  ///   - visited: 已訪問的狀態集合
-  ///   - openSet: 優先序列中剩餘的狀態
-  ///   - leadingState: 初始狀態
-  private static func batchCleanAllSearchStates(
-    visited: Set<SearchState>,
-    openSet: inout HybridPriorityQueue<PrioritizedState>,
-    leadingState: SearchState
-  ) {
-    // 策略1: 直接清理 visited set 中的所有狀態
-    for state in visited {
-      state.node = nil
-      state.prev = nil
-    }
-
-    // 策略2: 直接清理 openSet 中剩餘的所有狀態
-    while !openSet.isEmpty {
-      if let prioritizedState = openSet.dequeue() {
-        prioritizedState.state.node = nil
-        prioritizedState.state.prev = nil
-      }
-    }
-
-    // 策略3: 清理 leadingState
-    leadingState.node = nil
-    leadingState.prev = nil
-  }
 }
 
 // MARK: - 搜尋狀態相關定義
@@ -179,11 +124,13 @@ extension Megrez.Compositor {
     ///   - position: 在輸入串中的位置。
     ///   - prev: 前一個狀態。
     ///   - distance: 到達此狀態的累計分數。
+    ///   - cleaningTaskRegister: 登記自毀任務池。
     init(
       node: Megrez.Node?,
       position: Int,
       prev: SearchState?,
-      distance: Double = Double(Int.min)
+      distance: Double = Double(Int.min),
+      cleaningTaskRegister: inout [() -> ()]
     ) {
       self.node = node
       self.position = position
@@ -192,13 +139,16 @@ extension Megrez.Compositor {
       // 使用不可變的標識符來確保 hash 一致性
       self.originalNodeRef = node
       self.stableHash = Self.computeStableHash(node: node, position: position)
+      cleaningTaskRegister.append { [weak self] in
+        self?.cleanChainRecursively()
+      }
     }
 
     // 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 +157,30 @@ extension Megrez.Compositor {
       lhs.originalNodeRef === rhs.originalNodeRef && lhs.position == rhs.position
     }
 
-    /// 清理單一 SearchState 的參據
-    /// 注意：由於新的清理策略是直接清理各個集合，這個方法現在主要用於向下相容
-    func cleanState() {
-      node = nil
-      prev = nil
+    /// 清理整個 SearchState 鏈條，從當前節點開始向後遞歸清理
+    /// 採用深度優先遍歷策略，確保每個節點都被清理
+    func cleanChainRecursively() {
+      var visited = Set<ObjectIdentifier>()
+      var stack: [SearchState] = []
+      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.node = nil
+        current.prev = nil
+      }
     }
 
     func hash(into hasher: inout Hasher) {