Add accepts, rollback methods on Guide and tests (#212)

Adds two methods on the Guide object:
- `accepts_tokens`: given a sequence of tokens, returns a boolean
indicating if that sequence will end in a valid DFA state.
- `rollback_state`: rollback the state of the Guide `n` tokens ( states
). errors if `n` is greater than the number of states in
self.state_cache

This is needed for speculative decoding compatibility in vLLM. See:
vllm-project/vllm#15975 (comment)

Author: unaidedelf8777

src/python_bindings/mod.rs

@@ -1,5 +1,6 @@
 //! Provides tools and interfaces to integrate the crate's functionality with Python.
 
+use std::collections::VecDeque;
 use std::sync::Arc;
 
 use bincode::{config, Decode, Encode};
@@ -27,16 +28,19 @@ macro_rules! type_name {
 pub struct PyGuide {
     state: StateId,
     index: PyIndex,
+    state_cache: VecDeque<StateId>,
 }
 
 #[pymethods]
 impl PyGuide {
     /// Creates a Guide object based on Index.
     #[new]
-    fn __new__(index: PyIndex) -> Self {
+    #[pyo3(signature = (index, max_rollback=32))]
+    fn __new__(index: PyIndex, max_rollback: usize) -> Self {
         PyGuide {
             state: index.get_initial_state(),
             index,
+            state_cache: VecDeque::with_capacity(max_rollback),
         }
     }
 
@@ -57,6 +61,11 @@ impl PyGuide {
             )))
     }
 
+    /// Get the number of rollback steps available.
+    fn get_allowed_rollback(&self) -> usize {
+        self.state_cache.len()
+    }
+
     /// Guide moves to the next state provided by the token id and returns a list of allowed tokens, unless return_tokens is False.
     #[pyo3(signature = (token_id, return_tokens=None))]
     fn advance(
@@ -66,6 +75,11 @@ impl PyGuide {
     ) -> PyResult<Option<Vec<TokenId>>> {
         match self.index.get_next_state(self.state, token_id) {
             Some(new_state) => {
+                // Free up space in state_cache if needed.
+                if self.state_cache.len() == self.state_cache.capacity() {
+                    self.state_cache.pop_front();
+                }
+                self.state_cache.push_back(self.state);
                 self.state = new_state;
                 if return_tokens.unwrap_or(true) {
                     self.get_tokens().map(Some)
@@ -80,6 +94,41 @@ impl PyGuide {
         }
     }
 
+    /// Rollback the Guide state `n` tokens (states).
+    /// Fails if `n` is greater than stored prior states.
+    fn rollback_state(&mut self, n: usize) -> PyResult<()> {
+        if n == 0 {
+            return Ok(());
+        }
+        if n > self.get_allowed_rollback() {
+            return Err(PyValueError::new_err(format!(
+                "Cannot roll back {n} step(s): only {available} states stored (max_rollback = {cap}). \
+                 You must advance through at least {n} state(s) before rolling back {n} step(s).",
+                 cap = self.state_cache.capacity(),
+                 available = self.get_allowed_rollback(),
+            )));
+        }
+        let mut new_state: u32 = self.state;
+        for _ in 0..n {
+            // unwrap is safe because length is checked above
+            new_state = self.state_cache.pop_back().unwrap();
+        }
+        self.state = new_state;
+        Ok(())
+    }
+
+    // Returns a boolean indicating if the sequence leads to a valid state in the DFA
+    fn accepts_tokens(&self, sequence: Vec<u32>) -> bool {
+        let mut state = self.state;
+        for t in sequence {
+            match self.index.get_next_state(state, t) {
+                Some(s) => state = s,
+                None => return false,
+            }
+        }
+        true
+    }
+
     /// Checks if the automaton is in a final state.
     fn is_finished(&self) -> bool {
         self.index.is_final_state(self.state)

tests/test_guide.py

@@ -181,3 +181,38 @@ def test_write_mask_into_interface(index):
         guide.write_mask_into(0, mask.numel(), mask.element_size())
     with pytest.raises(ValueError, match="Invalid data pointer alignment"):
         guide.write_mask_into(5, mask.numel(), mask.element_size())
+
+
+def test_rollback(index):
+    guide = Guide(index, max_rollback=3)
+
+    first_state = guide.get_state()
+    guide.advance(1)
+
+    # Roll back one token to initial state
+    guide.rollback_state(1)
+    assert not guide.is_finished()
+    # we should now be in the initial state
+    assert guide.get_state() == first_state
+
+
+def test_rollback_interface(index):
+    guide = Guide(index, max_rollback=3)
+
+    # Rolling back more than recorded history must raise
+    with pytest.raises(ValueError, match="Cannot roll back"):
+        guide.rollback_state(5)
+
+
+@pytest.mark.parametrize(
+    "seq, expected",
+    [
+        ([1], True),  # single allowed token: accept
+        ([2], True),  # different allowed token: accept
+        ([1, 1], False),  # too long for r"[1-9]": reject
+        ([2, 3], False),  # extra token: reject
+    ],
+)
+def test_accepts_tokens_correctness(index, seq, expected):
+    guide = Guide(index)
+    assert guide.accepts_tokens(seq) is expected