From bcc4855791bdb11f50c7d2da359bb20ada96f700 Mon Sep 17 00:00:00 2001 From: Muukii Date: Wed, 4 Feb 2026 02:17:59 +0900 Subject: [PATCH] Add generic ClosureBox to fix thunk stack growing in continuous tracking - Create generic ClosureBox in withTracking.swift to eliminate code duplication - Fix thunk stack growing for didChange closure in withContinuousStateGraphTracking by wrapping closures once at entry point and passing them directly in recursive calls - Remove duplicate private ClosureBox definitions from withGraphTrackingGroup.swift and withGraphTrackingMap.swift - Add callAsFunction support for cleaner invocation syntax The previous implementation wrapped/unwrapped closures via UnsafeSendable on each recursive iteration, causing thunk stack growth. Now closures are wrapped in ClosureBox once and passed through all iterations. Co-Authored-By: Claude Opus 4.5 --- .../Observation/withGraphTrackingGroup.swift | 14 +-- .../Observation/withGraphTrackingMap.swift | 20 ++-- .../StateGraph/Observation/withTracking.swift | 95 +++++++++++++------ 3 files changed, 76 insertions(+), 53 deletions(-) diff --git a/Sources/StateGraph/Observation/withGraphTrackingGroup.swift b/Sources/StateGraph/Observation/withGraphTrackingGroup.swift index 3dac4ae..79ff8b7 100644 --- a/Sources/StateGraph/Observation/withGraphTrackingGroup.swift +++ b/Sources/StateGraph/Observation/withGraphTrackingGroup.swift @@ -95,8 +95,8 @@ public func withGraphTrackingGroup( return } - let _handlerBox = OSAllocatedUnfairLock( - uncheckedState: .init(handler: handler) + let _handlerBox = OSAllocatedUnfairLock?>( + uncheckedState: ClosureBox(handler) ) // Create a cancellable for this scope that manages nested tracking @@ -113,7 +113,7 @@ public func withGraphTrackingGroup( // Nested groups/maps will register with this parent via addChild() ThreadLocal.currentCancellable.withValue(scopeCancellable) { _handlerBox.withLock { - $0?.handler() + $0?() } } }, @@ -132,11 +132,3 @@ public func withGraphTrackingGroup( } } - -private struct ClosureBox { - let handler: () -> Void - - init(handler: @escaping () -> Void) { - self.handler = handler - } -} diff --git a/Sources/StateGraph/Observation/withGraphTrackingMap.swift b/Sources/StateGraph/Observation/withGraphTrackingMap.swift index e6ed8d0..259338d 100644 --- a/Sources/StateGraph/Observation/withGraphTrackingMap.swift +++ b/Sources/StateGraph/Observation/withGraphTrackingMap.swift @@ -137,8 +137,8 @@ public func withGraphTrackingMap( var filter = filter - let _handlerBox = OSAllocatedUnfairLock( - uncheckedState: .init(handler: { + let _handlerBox = OSAllocatedUnfairLock?>( + uncheckedState: ClosureBox({ let result = applier() let filtered = filter.send(value: result) if let filtered { @@ -160,7 +160,7 @@ public func withGraphTrackingMap( // Set this scope's cancellable as the current parent for nested tracking // Nested groups/maps will register with this parent via addChild() ThreadLocal.currentCancellable.withValue(scopeCancellable) { - _handlerBox.withLock { $0?.handler() } + _handlerBox.withLock { $0?() } } }, didChange: { @@ -305,8 +305,8 @@ public func withGraphTrackingMap( var filter = filter - let _handlerBox = OSAllocatedUnfairLock( - uncheckedState: .init(handler: { + let _handlerBox = OSAllocatedUnfairLock?>( + uncheckedState: ClosureBox({ guard let dependency = weakDependency else { return } @@ -335,7 +335,7 @@ public func withGraphTrackingMap( // Set this scope's cancellable as the current parent for nested tracking // Nested groups/maps will register with this parent via addChild() ThreadLocal.currentCancellable.withValue(scopeCancellable) { - _handlerBox.withLock { $0?.handler() } + _handlerBox.withLock { $0?() } } }, didChange: { @@ -353,11 +353,3 @@ public func withGraphTrackingMap( subscriptions!.append(AnyCancellable(scopeCancellable)) } } - -private struct ClosureBox { - let handler: () -> Void - - init(handler: @escaping () -> Void) { - self.handler = handler - } -} diff --git a/Sources/StateGraph/Observation/withTracking.swift b/Sources/StateGraph/Observation/withTracking.swift index 9cec6b7..95d18a3 100644 --- a/Sources/StateGraph/Observation/withTracking.swift +++ b/Sources/StateGraph/Observation/withTracking.swift @@ -32,6 +32,49 @@ public enum StateGraphTrackingContinuation: Sendable { case next } +// MARK: - Internal Types + +struct UnsafeSendable: ~Copyable, @unchecked Sendable { + + let _value: V + + init(_ value: consuming V) { + _value = value + } + +} + +/// A box that wraps a closure to prevent thunk stack growing during recursive calls. +/// By wrapping closures in this struct and passing the struct instead of the raw closure, +/// we avoid the overhead of repeatedly wrapping/unwrapping closure types. +struct ClosureBox { + let closure: () -> R + + init(_ closure: @escaping () -> R) { + self.closure = closure + } + + func callAsFunction() -> R { + closure() + } +} + +func perform( + _ closure: () -> Return, + isolation: isolated (any Actor)? = #isolation +) -> Return { + closure() +} + +func perform( + _ box: ClosureBox, + isolation: isolated (any Actor)? = #isolation +) -> Return { + box() +} + +// MARK: - Continuous Tracking + /// Tracks access to the properties of StoredNode or Computed. /// Continuously tracks until `didChange` returns `.stop`. /// It does not provides update of the properties granurarly. some frequency of updates may be aggregated into single event. @@ -40,21 +83,37 @@ func withContinuousStateGraphTracking( didChange: @escaping () -> StateGraphTrackingContinuation, isolation: isolated (any Actor)? = #isolation ) { + // Wrap closures in ClosureBox to prevent thunk stack growing during recursive calls. + // The boxes are created once here and passed through all recursive iterations. + let applyBox = ClosureBox(apply) + let didChangeBox = ClosureBox(didChange) + _withContinuousStateGraphTracking( + apply: applyBox, + didChange: didChangeBox, + isolation: isolation + ) +} - let applyBox = UnsafeSendable(apply) - let didChangeBox = UnsafeSendable(didChange) - - withStateGraphTracking(apply: apply) { - let continuation = perform(didChangeBox._value, isolation: isolation) +/// Private implementation that receives pre-wrapped closures to prevent thunk stack growing. +/// By passing ClosureBox<...> directly in recursive calls, we avoid +/// the cost of re-wrapping/unwrapping closure types on each iteration. +private func _withContinuousStateGraphTracking( + apply: ClosureBox, + didChange: ClosureBox, + isolation: isolated (any Actor)? = #isolation +) { + withStateGraphTracking(apply: apply.closure) { + let continuation = perform(didChange, isolation: isolation) switch continuation { case .stop: break case .next: // continue tracking on next event loop. // It uses isolation and task dispatching to ensure apply closure is called on the same actor. - withContinuousStateGraphTracking( - apply: applyBox._value, - didChange: didChangeBox._value, + // Pass the already-wrapped closures directly to avoid thunk stack growing. + _withContinuousStateGraphTracking( + apply: apply, + didChange: didChange, isolation: isolation ) } @@ -239,23 +298,3 @@ public final class TrackingRegistration: Sendable, Hashable { } } - -struct UnsafeSendable: ~Copyable, @unchecked Sendable { - - let _value: V - - init(_ value: consuming V) { - _value = value - } - -} - -func perform( - _ closure: () -> Return, - isolation: isolated (any Actor)? = #isolation -) - -> Return -{ - closure() -} -