Add bind(into:) for zero-copy state binding with loop-based insert (#156)

State handlers are now classes (AnyObject) that own their NDArrays at
refcount 1. bind(into:) inserts all states into MutableViews in a loop
using stdlib _overrideLifetime to express disjoint element access.
No COW, no switch on state count, no write-back needed.

FixedMTLBufferState also gains bind(into:) for AsyncMutableViews,
eliminating the 3x duplicated switch blocks in the pipelined engine.

Engine call sites: runWithStates (sequential) and encodeWithStates
(pipelined) are now thin wrappers around bind + function.run/encode.

Author: stikves

Package.swift

@@ -57,6 +57,7 @@ let package = Package(
             swiftSettings: [
                 .define("CXGRAMMAR_IMPORT"),
                 .enableUpcomingFeature("MemberImportVisibility"),
+                .enableExperimentalFeature("Lifetimes"),
             ],
             linkerSettings: [
                 .linkedLibrary("c++")
@@ -225,6 +226,9 @@ let package = Package(
             resources: [
                 .copy("Resources/MinimalTokenizer")
             ],
+            swiftSettings: [
+                .enableExperimentalFeature("Lifetimes"),
+            ],
             linkerSettings: [
                 .linkedLibrary("c++")
             ]

swift/Sources/CoreAILanguageModels/Handlers/StateHandler+Encode.swift

@@ -0,0 +1,38 @@
+// Copyright 2026 Apple Inc.
+//
+// Use of this source code is governed by a BSD-3-clause license that can
+// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
+
+import CoreAI
+import Metal
+
+/// Encode an inference step with KV cache states, optional additional MTLBuffer
+/// states, and logits output.
+func encodeWithStates(
+    function: InferenceFunction,
+    inputs: [String: InferenceFunction.AsyncValue],
+    keyState: inout InferenceFunction.AsyncMutableValue,
+    keyCacheName: String,
+    valState: inout InferenceFunction.AsyncMutableValue,
+    valueCacheName: String,
+    additionalStates: FixedMTLBufferState?,
+    logitsBuffer: MTLBuffer,
+    logitsName: String,
+    logitsShape: [Int],
+    logitsStrides: [Int],
+    computeStream: ComputeStream
+) throws {
+    var asyncStates = InferenceFunction.AsyncMutableViews()
+    asyncStates.insert(&keyState, for: keyCacheName)
+    asyncStates.insert(&valState, for: valueCacheName)
+    additionalStates?.bind(into: &asyncStates)
+
+    var logitsOutput = unsafe InferenceFunction.AsyncMutableValue(
+        unsafeBuffer: logitsBuffer, byteOffset: 0,
+        scalarType: .float16, shape: logitsShape, strides: logitsStrides)
+    var asyncOutputs = InferenceFunction.AsyncMutableViews()
+    asyncOutputs.insert(&logitsOutput, for: logitsName)
+    let _ = try function.encode(
+        inputs: inputs, states: consume asyncStates,
+        outputViews: consume asyncOutputs, to: computeStream)
+}

swift/Sources/CoreAILanguageModels/Handlers/StateHandler+MTLBuffer.swift

@@ -7,9 +7,9 @@ import CoreAI
 import CoreAIShared
 import Metal
 
-/// Fixed-size MTLBuffer state for non-truncatable persistent states.
+/// Fixed-size MTLBuffer state for non-truncatable persistent states (pipelined engine).
 /// Allocated once at init, zero-initialized, never grows.
-public struct FixedMTLBufferState {
+public final class FixedMTLBufferState {
     public let stateNames: [String]
     public var stateCount: Int { bindings.count }
 
@@ -29,7 +29,8 @@ public struct FixedMTLBufferState {
             let resolved = desc.resolvingDynamicDimensions(desc.shape)
             let strides = resolved.preferredStrides
             let byteCount = resolved.minimumByteCount
-            guard let buffer = device.makeBuffer(length: max(byteCount, 64), options: .storageModeShared) else {
+            guard let buffer = device.makeBuffer(length: max(byteCount, 64), options: .storageModeShared)
+            else {
                 throw InferenceRuntimeError.bufferAllocationFailed("\(name) (\(byteCount) bytes)")
             }
             memset(buffer.contents(), 0, buffer.length)
@@ -39,17 +40,21 @@ public struct FixedMTLBufferState {
         self.stateNames = states.map(\.name)
     }
 
-    /// Access a state binding by index for AsyncMutableValue construction.
-    /// The engine builds AsyncMutableValue views from these in the same scope as encode().
-    /// Note: MTLBuffer is a reference type — the returned buffer is shared, not copied.
-    public subscript(stateIndex index: Int) -> (
-        name: String, buffer: MTLBuffer, scalarType: NDArray.ScalarType, shape: [Int], strides: [Int]
-    ) {
-        get { bindings[index] }
+    /// Insert all managed states into async mutable views. MTLBuffer is a reference
+    /// type (no COW). Uses _overrideLifetime for disjoint element access in the loop.
+    @_lifetime(views: borrow self)
+    public func bind(into views: inout InferenceFunction.AsyncMutableViews) {
+        for binding in bindings {
+            var value = unsafe InferenceFunction.AsyncMutableValue(
+                unsafeBuffer: binding.buffer, byteOffset: 0,
+                scalarType: binding.scalarType, shape: binding.shape, strides: binding.strides)
+            views.insert(&value, for: binding.name)
+            views = unsafe _overrideLifetime(consume views, borrowing: self)
+        }
     }
 
     /// Zero all state buffers. Caller must ensure no in-flight GPU work references these.
-    public mutating func reset() {
+    public func reset() {
         for (_, buffer, _, _, _) in bindings {
             memset(buffer.contents(), 0, buffer.length)
         }

swift/Sources/CoreAILanguageModels/Handlers/StateHandler+NDArray.swift

@@ -11,56 +11,63 @@ import Darwin
 
 /// Fixed-size state for non-truncatable persistent states.
 /// Allocated at full size on init, zero-initialized. No capacity management needed.
-public struct FixedNDArrayState: SyncStateHandler {
+public final class FixedNDArrayState: SyncStateHandler {
     public let stateNames: [String]
     public let supportsTruncation: Bool = false
     public let currentCapacity: Int = .max
     public var stateCount: Int { arrays.count }
 
-    private var arrays: [(name: String, array: NDArray)]
+    private var arrays: [String: NDArray]
 
     public init(states: [(name: String, descriptor: NDArrayDescriptor)]) {
-        var arrays: [(String, NDArray)] = []
+        var arrays: [String: NDArray] = [:]
         for (name, desc) in states {
             var array = NDArray(descriptor: desc)
             zeroFillNDArray(&array)
-            arrays.append((name, array))
+            arrays[name] = array
         }
         self.arrays = arrays
         self.stateNames = states.map(\.name)
     }
 
-    public mutating func ensureCapacity(forContextLength contextLength: Int) throws -> Bool {
+    public func ensureCapacity(forContextLength contextLength: Int) throws -> Bool {
         false
     }
 
     public subscript(stateIndex index: Int) -> (name: String, array: NDArray) {
-        get { arrays[index] }
-        set { arrays[index] = newValue }
+        get { (stateNames[index], arrays[stateNames[index]]!) }
+        set { arrays[stateNames[index]] = newValue.array }
     }
 
-    public mutating func reset() {
-        for i in arrays.indices {
-            zeroFillNDArray(&arrays[i].array)
+    @_lifetime(views: borrow self)
+    public func bind(into views: inout InferenceFunction.MutableViews) {
+        for name in stateNames {
+            let view = _overrideLifetime(arrays[name]!.mutableRawView(), borrowing: Void())
+            views.insert(view, for: name)
         }
     }
 
-    public mutating func truncate(to tokenCount: Int) {
+    public func reset() {
+        for name in stateNames {
+            zeroFillNDArray(&arrays[name]!)
+        }
+    }
+
+    public func truncate(to tokenCount: Int) {
         preconditionFailure("truncate(to:) called on non-truncatable FixedNDArrayState")
     }
 }
 
 // MARK: - Growing NDArray State
 
-/// Dynamically-growing KV cache state. Starts small and doubles capacity
-/// when more context is needed.
-public struct GrowingNDArrayState: SyncStateHandler {
+/// Dynamically-growing KV cache state. Starts small and doubles capacity.
+public final class GrowingNDArrayState: SyncStateHandler {
     public let stateNames: [String]
     public let supportsTruncation: Bool = true
     public private(set) var currentCapacity: Int
     public var stateCount: Int { arrays.count }
 
-    private var arrays: [(name: String, array: NDArray)]
+    private var arrays: [String: NDArray]
     private let descriptors: [NDArrayDescriptor]
     private let maxCapacity: Int
     private let sequenceDimIndex: Int
@@ -80,16 +87,16 @@ public struct GrowingNDArrayState: SyncStateHandler {
         let capacity = min(initialCapacity, maxCapacity)
         self.currentCapacity = capacity
 
-        var arrays: [(String, NDArray)] = []
+        var arrays: [String: NDArray] = [:]
         for (name, desc) in states {
             let resolved = desc.resolvingDynamicDimensions(
                 desc.shape.map { $0 < 0 ? capacity : $0 })
-            arrays.append((name, NDArray(descriptor: resolved)))
+            arrays[name] = NDArray(descriptor: resolved)
         }
         self.arrays = arrays
     }
 
-    public mutating func ensureCapacity(forContextLength contextLength: Int) throws -> Bool {
+    public func ensureCapacity(forContextLength contextLength: Int) throws -> Bool {
         guard contextLength > currentCapacity else { return false }
         guard contextLength <= maxCapacity else {
             throw InferenceRuntimeError.invalidState(
@@ -101,38 +108,41 @@ public struct GrowingNDArrayState: SyncStateHandler {
             newCapacity = min(newCapacity * 2, maxCapacity)
         }
 
-        for i in arrays.indices {
+        for (i, name) in stateNames.enumerated() {
             let desc = descriptors[i]
             let newShape = desc.shape.map { $0 < 0 ? newCapacity : $0 }
             let resolvedDesc = desc.resolvingDynamicDimensions(newShape)
             var newArray = NDArray(descriptor: resolvedDesc)
-            // Force backing allocation before copy
             _ = newArray.mutableRawView()
-
-            copyCache(from: arrays[i].array, to: &newArray, sequenceDim: sequenceDimIndex)
-            arrays[i].array = newArray
+            copyCache(from: arrays[name]!, to: &newArray, sequenceDim: sequenceDimIndex)
+            arrays[name] = newArray
         }
 
         currentCapacity = newCapacity
         return true
     }
 
     public subscript(stateIndex index: Int) -> (name: String, array: NDArray) {
-        get { arrays[index] }
-        set { arrays[index] = newValue }
+        get { (stateNames[index], arrays[stateNames[index]]!) }
+        set { arrays[stateNames[index]] = newValue.array }
     }
 
-    public mutating func reset() {
-        for i in arrays.indices {
-            zeroFillNDArray(&arrays[i].array)
+    @_lifetime(views: borrow self)
+    public func bind(into views: inout InferenceFunction.MutableViews) {
+        for name in stateNames {
+            let view = _overrideLifetime(arrays[name]!.mutableRawView(), borrowing: Void())
+            views.insert(view, for: name)
         }
     }
 
-    public mutating func truncate(to tokenCount: Int) {
-        // KV cache truncation is a no-op on the backing storage.
-        // The causal mask hides positions beyond processedTokenCount.
+    public func reset() {
+        for name in stateNames {
+            zeroFillNDArray(&arrays[name]!)
+        }
     }
 
+    public func truncate(to tokenCount: Int) {}
+
     // MARK: - Private
 
     private func copyCache(from source: NDArray, to destination: inout NDArray, sequenceDim: Int) {
@@ -175,7 +185,6 @@ public struct GrowingNDArrayState: SyncStateHandler {
 
 // MARK: - Shared Utilities
 
-/// Zero-initialize an NDArray, dispatching on scalar type.
 func zeroFillNDArray(_ array: inout NDArray) {
     let count = array.shape.reduce(1, *)
     switch array.scalarType {

swift/Sources/CoreAILanguageModels/Handlers/StateHandler+Run.swift

@@ -0,0 +1,29 @@
+// Copyright 2026 Apple Inc.
+//
+// Use of this source code is governed by a BSD-3-clause license that can
+// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
+
+import CoreAI
+
+/// Run an inference step with combined primary + secondary states and output.
+/// Zero-copy: bind(into:) uses reference-backed storage and _overrideLifetime.
+func runWithStates(
+    function: InferenceFunction,
+    inputs: [String: NDArray],
+    primary: any SyncStateHandler,
+    secondary: FixedNDArrayState?,
+    outputArray: inout NDArray,
+    outputName: String
+) async throws {
+    var states = InferenceFunction.MutableViews()
+    primary.bind(into: &states)
+    secondary?.bind(into: &states)
+
+    var outputViews = InferenceFunction.MutableViews()
+    outputViews.insert(&outputArray, for: outputName)
+
+    _ = try await function.run(
+        inputs: inputs,
+        states: _unsafeEscapeMutableViews(consume states),
+        outputViews: consume outputViews)
+}

swift/Sources/CoreAILanguageModels/Handlers/StateHandler.swift

@@ -8,16 +8,13 @@ import CoreAI
 /// Persistent model state that the engine carries across inference steps.
 ///
 /// Each handler manages one or more named state tensors (e.g., key_cache + value_cache,
-/// or a single convolution_state). The engine owns an array of handlers and delegates
-/// all state lifecycle to them — allocation, growth, and reset.
+/// or a single convolution_state). The engine owns handlers and delegates all state
+/// lifecycle to them — allocation, growth, and reset.
 ///
-/// ## MutableViews Lifetime Constraint
-///
-/// `InferenceFunction.MutableViews.insert` creates a lifetime dependency on each `inout`
-/// variable. This means state binding CANNOT be abstracted into a method that takes
-/// `inout MutableViews` — the inserts must happen in the same scope as `function.run()`.
-/// Handlers therefore expose their arrays directly via subscript for the engine to bind.
-public protocol SyncStateHandler {
+/// Handlers are classes (AnyObject) so they own their NDArrays at refcount 1 —
+/// `bind(into:)` calls `mutableRawView()` without triggering COW. The loop uses
+/// `_overrideLifetime` to express disjoint element access to the compiler.
+public protocol SyncStateHandler: AnyObject {
     /// Names of the states managed by this handler.
     var stateNames: [String] { get }
 
@@ -35,18 +32,32 @@ public protocol SyncStateHandler {
 
     /// Ensure the state can accommodate `contextLength` tokens.
     /// Returns true if reallocation occurred.
-    mutating func ensureCapacity(forContextLength contextLength: Int) throws -> Bool
+    func ensureCapacity(forContextLength contextLength: Int) throws -> Bool
 
-    /// Access a state array by index for binding into MutableViews.
-    /// The engine calls this to get name + array pairs for insert.
+    /// Access a state array by name+index (value copy — use bind(into:) on hot paths).
     subscript(stateIndex index: Int) -> (name: String, array: NDArray) { get set }
 
+    /// Insert all managed states into `views`. Zero-copy: uses reference-backed
+    /// storage internally, so mutableRawView() never triggers COW.
+    @_lifetime(views: borrow self)
+    func bind(into views: inout InferenceFunction.MutableViews)
+
     /// Full reset — zero all backing storage, rewind to position 0.
-    mutating func reset()
+    func reset()
 
     /// Truncate to a given token position.
-    /// Only valid when `supportsTruncation == true`. For KV cache, this is a no-op
-    /// on the backing storage (causal mask handles visibility); the engine just
-    /// updates its processedTokenCount.
-    mutating func truncate(to tokenCount: Int)
+    func truncate(to tokenCount: Int)
+}
+
+// MARK: - Lifetime Helpers
+
+/// Detach lifetime dependencies from MutableViews so it can cross scope
+/// boundaries (closures, await). Caller must ensure inserted arrays remain valid.
+@inline(__always)
+@_unsafeNonescapableResult
+@_lifetime(immortal)
+func _unsafeEscapeMutableViews(
+    _ views: consuming InferenceFunction.MutableViews
+) -> InferenceFunction.MutableViews {
+    views
 }