Port PastConfigs

src/lib.rs

@@ -20,7 +20,7 @@ use blocks::{measure_blocks, unroll_tape};
 use parse::{
     parse as parse_fn, read_slot, read_state, show_instr, show_slot, show_state, tcompile,
 };
-use prover::PastConfig;
+use prover::PastConfigs;
 use rules::{InfiniteRule, RuleLimit, UnknownRule};
 
 #[pymodule]
@@ -39,7 +39,7 @@ fn rust_stuff(py: Python, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(show_instr, m)?)?;
 
     // prover
-    m.add_class::<PastConfig>()?;
+    m.add_class::<PastConfigs>()?;
 
     // rules
     m.add("UnknownRule", py.get_type::<UnknownRule>())?;

src/prover.rs

@@ -1,15 +1,16 @@
+use std::collections::HashMap;
+
 use pyo3::prelude::*;
 
+use crate::instrs::State;
+
 type Cycle = u32;
 
-#[pyclass]
-pub struct PastConfig {
+struct PastConfig {
     cycles: Vec<Cycle>,
 }
 
-#[pymethods]
 impl PastConfig {
-    #[new]
     const fn new() -> Self {
         Self { cycles: Vec::new() }
     }
@@ -55,3 +56,27 @@ impl PastConfig {
         None
     }
 }
+
+#[pyclass]
+pub struct PastConfigs {
+    configs: HashMap<State, PastConfig>,
+}
+
+#[pymethods]
+impl PastConfigs {
+    #[new]
+    fn new() -> Self {
+        Self {
+            configs: HashMap::new(),
+        }
+    }
+
+    fn next_deltas(&mut self, state: State, cycle: Cycle) -> Option<(Cycle, Cycle, Cycle)> {
+        let config = self.configs.entry(state).or_insert_with(PastConfig::new);
+        config.next_deltas(cycle)
+    }
+
+    fn delete_configs(&mut self, state: State) {
+        self.configs.remove(&state);
+    }
+}

tm/prover.py

@@ -3,7 +3,7 @@
 from typing import TYPE_CHECKING
 
 from tm.rules import make_rule
-from tm.rust_stuff import PastConfig
+from tm.rust_stuff import PastConfigs
 
 if TYPE_CHECKING:
     from tm.rules import Rule
@@ -17,28 +17,6 @@ class ConfigLimit(Exception):
     pass
 
 
-class PastConfigs:
-    _configs: dict[State, PastConfig]
-
-    def __init__(self) -> None:
-        self._configs = {}
-
-    def next_deltas(
-            self,
-            state: State,
-            cycle: int,
-    ) -> tuple[int, int, int] | None:
-        try:
-            config = self._configs[state]
-        except KeyError:
-            self._configs[state] = (config := PastConfig())
-
-        return config.next_deltas(cycle)
-
-    def delete_configs(self, state: State) -> None:
-        del self._configs[state]
-
-
 class Prover:
     prog: GetInstr
 

tm/rust_stuff.pyi

@@ -1,9 +1,15 @@
 ## prover ##############################
 
-class PastConfig:
-    cycles: list[int]
+class PastConfigs:
+    def __init__(self) -> None: ...
 
-    def next_deltas(self, cycle: int) -> tuple[int, int, int] | None: ...
+    def next_deltas(
+            self,
+            state: State,
+            cycle: int,
+    ) -> tuple[int, int, int] | None: ...
+
+    def delete_configs(self, state: State) -> None: ...
 
 ## parse ###############################