Extend ConstrainedGenerationSession with rollback, jump-forward, and direct bitmask fill

- rollback(_:) — unwind grammar state by N tokens (budget: 64)
- findJumpForwardString() — peek at deterministic continuations
- fillBitmask(into:) → BitmaskResult — write bitmask directly into
  a caller-provided pointer (e.g. GPU-visible MTLBuffer), avoiding
  array allocation per token
- BitmaskResult enum: .terminated, .unconstrained, .constrained
- C bridge additions: XGrammarRollback, XGrammarFindJumpForwardString,
  XGrammarFillNextTokenBitmask

Part 1 of 4 for GPU-based constrained sampling (#114).

Author: stikves

swift/Sources/CoreAILanguageModels/GuidedGeneration/ConstrainedGenerationSession.swift

@@ -93,7 +93,7 @@ public struct ConstrainedGenerationSession: ~Copyable {
         self.tokenizerInfo = tokenizerInfo
         self.compiler = GrammarCompiler(tokenizerInfo: tokenizerInfo)
         self.compiledGrammar = try compiler.compileJSONSchema(jsonSchema)
-        self.matcher = GrammarMatcher(compiledGrammar: compiledGrammar)
+        self.matcher = GrammarMatcher(compiledGrammar: compiledGrammar, maxRollbackTokens: 64)
         self.vocabularySize = tokenizerInfo.vocabularySize
         self.bitmaskSize = (vocabularySize + 31) / 32
         self.bitmaskBuffer = Array(repeating: 0, count: bitmaskSize)
@@ -193,6 +193,49 @@ public struct ConstrainedGenerationSession: ~Copyable {
         matcher.reset()
         allTokensBlocked = false
     }
+
+    /// Rollback the grammar state by N tokens. Returns false if rollback failed
+    /// (e.g., exceeds maxRollbackTokens budget).
+    @discardableResult
+    public mutating func rollback(_ numTokens: Int = 1) -> Bool {
+        matcher.rollback(numTokens)
+    }
+
+    /// Find the longest deterministic string from the current grammar state.
+    /// Does not change the matcher state. Returns nil if no jump-forward is possible.
+    public func findJumpForwardString() -> String? {
+        matcher.findJumpForwardString()
+    }
+
+    /// Result of filling a bitmask for the next token.
+    public enum BitmaskResult: Equatable {
+        /// Grammar is terminated or all tokens are blocked — generation should stop.
+        case terminated
+        /// All tokens are allowed — no mask needed, generate unconstrained.
+        case unconstrained
+        /// Bitmask was written — apply it to constrain sampling.
+        case constrained
+    }
+
+    /// Fill the bitmask directly into a caller-provided buffer (e.g., a GPU-visible MTLBuffer).
+    ///
+    /// The caller must ensure the pointer has room for at least `(vocabularySize + 31) / 32`
+    /// Int32 words.
+    public mutating func fillBitmask(into pointer: UnsafeMutablePointer<Int32>) -> BitmaskResult {
+        if isTerminated { return .terminated }
+
+        let needsApplication = matcher.fillNextTokenBitmask(pointer)
+        if !needsApplication {
+            // xgrammar signals all tokens are allowed — no mask needed
+            return .unconstrained
+        }
+        // Check for all-zeros (no tokens allowed — grammar done)
+        for i in 0..<bitmaskSize {
+            if pointer[i] != 0 { return .constrained }
+        }
+        allTokensBlocked = true
+        return .terminated
+    }
 }
 
 // MARK: - Float16 Masking

swift/Sources/CoreAILanguageModels/GuidedGeneration/XGrammarWrapper.swift

@@ -136,9 +136,29 @@ public final class GrammarMatcher {
         return xgrammar_matcher_is_terminated(handle)
     }
 
+    public var isCompleted: Bool {
+        return xgrammar_matcher_is_completed(handle)
+    }
+
     public func reset() {
         xgrammar_matcher_reset(handle)
     }
+
+    @discardableResult
+    public func rollback(_ numTokens: Int = 1) -> Bool {
+        xgrammar_matcher_rollback(handle, Int32(numTokens))
+    }
+
+    /// Returns the longest deterministic string from the current grammar state,
+    /// or nil if no jump-forward is possible. Does not change matcher state.
+    public func findJumpForwardString() -> String? {
+        guard let cStr = xgrammar_matcher_find_jump_forward_string(handle) else {
+            return nil
+        }
+        let result = String(cString: cStr)
+        free(UnsafeMutablePointer(mutating: cStr))
+        return result.isEmpty ? nil : result
+    }
 }
 
 // MARK: - Errors

swift/Sources/lib/CXGrammar/include/xgrammar_c_bridge.h

@@ -101,6 +101,16 @@ bool xgrammar_matcher_is_terminated(const XGrammarMatcher* matcher);
 // Reset matcher
 void xgrammar_matcher_reset(XGrammarMatcher* matcher);
 
+// Rollback N tokens (restores grammar state to N tokens ago)
+bool xgrammar_matcher_rollback(XGrammarMatcher* matcher, int num_tokens);
+
+// Find the longest deterministic string from the current grammar state.
+// Returns a malloc'd C string (caller must free), or NULL if no jump-forward is possible.
+const char* xgrammar_matcher_find_jump_forward_string(XGrammarMatcher* matcher);
+
+// Check if the grammar's root rule has been fully matched (without requiring stop token)
+bool xgrammar_matcher_is_completed(const XGrammarMatcher* matcher);
+
 // Free grammar matcher
 void xgrammar_matcher_free(XGrammarMatcher* matcher);
 

swift/Sources/lib/CXGrammar/xgrammar_c_bridge.cpp

@@ -222,6 +222,39 @@ void xgrammar_matcher_reset(XGrammarMatcher* matcher) {
     } catch (...) {}
 }
 
+bool xgrammar_matcher_rollback(XGrammarMatcher* matcher, int num_tokens) {
+    try {
+        if (matcher) {
+            matcher->cpp_obj.Rollback(num_tokens);
+            return true;
+        }
+    } catch (...) {}
+    return false;
+}
+
+const char* xgrammar_matcher_find_jump_forward_string(XGrammarMatcher* matcher) {
+    try {
+        if (!matcher) return nullptr;
+        std::string result = matcher->cpp_obj.FindJumpForwardString();
+        if (result.empty()) return nullptr;
+        char* buf = static_cast<char*>(malloc(result.size() + 1));
+        if (!buf) return nullptr;
+        memcpy(buf, result.c_str(), result.size() + 1);
+        return buf;
+    } catch (...) {
+        return nullptr;
+    }
+}
+
+bool xgrammar_matcher_is_completed(const XGrammarMatcher* matcher) {
+    try {
+        if (!matcher) return false;
+        return matcher->cpp_obj.IsCompleted();
+    } catch (...) {
+        return false;
+    }
+}
+
 void xgrammar_matcher_free(XGrammarMatcher* matcher) {
     delete matcher;
 }

swift/Tests/GuidedGenerationTests/ConstrainedGenerationSessionTests.swift

@@ -84,6 +84,43 @@ struct ConstrainedGenerationSessionTests {
         #expect(allowedTokenIDs(session: &session).isEmpty)
     }
 
+    @Test func fillBitmaskIntoPointerMatchesNextTokenBitmask() throws {
+        var session = try createTestSession()
+
+        // Get bitmask via the array-returning method
+        let arrayBitmask = session.nextTokenBitmask()
+        #expect(arrayBitmask != nil)
+
+        // Reset and get bitmask via the pointer method
+        session.reset()
+        let bitmaskSize = (session.vocabularySize + 31) / 32
+        var pointerBitmask = [Int32](repeating: 0, count: bitmaskSize)
+        let result = pointerBitmask.withUnsafeMutableBufferPointer { buf in
+            session.fillBitmask(into: buf.baseAddress!)
+        }
+        #expect(result == .constrained)
+
+        // They should be identical
+        #expect(arrayBitmask!.count == pointerBitmask.count)
+        for i in 0..<arrayBitmask!.count {
+            #expect(
+                arrayBitmask![i] == pointerBitmask[i],
+                "Mismatch at word \(i): array=\(arrayBitmask![i]) pointer=\(pointerBitmask[i])")
+        }
+    }
+
+    @Test func fillBitmaskReturnsTerminatedWhenDone() throws {
+        var session = try createTestSession()
+        driveToCompletion(session: &session)
+
+        let bitmaskSize = (session.vocabularySize + 31) / 32
+        var buffer = [Int32](repeating: 0, count: bitmaskSize)
+        let result = buffer.withUnsafeMutableBufferPointer { buf in
+            session.fillBitmask(into: buf.baseAddress!)
+        }
+        #expect(result == .terminated)
+    }
+
     // MARK: - Token Acceptance Tests
 
     @Test func acceptToken() throws {