metrale_circuit/
state_ops.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: The state programs of a circuit (M5, LIFECYCLE-DESIGN.md section 3.5): the small
4//! programs that move recurrent state between its places around a speculative verify, and reset
5//! a slot, as nodes over the circuit's declared states. Legacy runs them as host code that sizes
6//! each copy itself (`verify_a_ssm.rs`, `async_chkpt.rs`, `ssm_pool_slots.rs`), with the
7//! GatedDeltaNet conv formula, which is 0 bytes for a Mamba2 layer; here every node is sized from
8//! its state's declaration, so a Mamba2 layer's copies are its own.
9//!
10//! Owner: metrale-circuit.
11//! Invariants:
12//! - The programs are generic over recurrent state kinds: every recurrent state of the target
13//!   (`Section::Main`) gets one node in each program, in state order.
14//! - A node's bytes are one unit of its state ([`StatePlan`] at one slot), so a program's bytes
15//!   are what one sequence's copies move.
16
17use std::collections::BTreeMap;
18
19use crate::ir::{Circuit, Section};
20use crate::state::{Holding, StateDtype, StateError, StateInputs, StateKind, StatePlan};
21
22/// 2026-09-30: Where one sequence's unit of a recurrent state lives.
23#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
24pub enum StatePlace {
25    /// 2026-09-30: The slot's live state.
26    Live,
27    /// 2026-09-30: The pre-verify checkpoint.
28    Checkpoint,
29    /// 2026-09-30: The verify intermediate after the last accepted row.
30    AcceptedStep,
31}
32
33/// 2026-09-30: What a state-program node does.
34#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
35pub enum StateOp {
36    /// 2026-09-30: Copy one unit from `from` to `to`.
37    Copy { from: StatePlace, to: StatePlace },
38    /// 2026-09-30: Zero the live unit.
39    Zero,
40}
41
42/// 2026-09-30: The state programs.
43#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
44pub enum StateProgramId {
45    /// 2026-09-30: Before a verify: live to checkpoint (`checkpoint_ssm_states_dispatch`).
46    VerifyCheckpoint,
47    /// 2026-09-30: A verify that accepted nothing: checkpoint to live (`rollback_ssm_states`).
48    VerifyRollback,
49    /// 2026-09-30: A verify that accepted a prefix: the accepted row's intermediate to live
50    /// (`commit_accepted_prefix_dispatch`).
51    CommitAccepted,
52    /// 2026-09-30: A slot handed to a new sequence: zero its live state (`zero_slot`).
53    SlotZero,
54}
55
56impl StateProgramId {
57    /// 2026-09-30: Every program, in this order.
58    pub const ALL: [StateProgramId; 4] = [
59        Self::VerifyCheckpoint,
60        Self::VerifyRollback,
61        Self::CommitAccepted,
62        Self::SlotZero,
63    ];
64
65    /// 2026-09-30: The op each of its nodes runs.
66    pub fn op(self) -> StateOp {
67        match self {
68            Self::VerifyCheckpoint => StateOp::Copy {
69                from: StatePlace::Live,
70                to: StatePlace::Checkpoint,
71            },
72            Self::VerifyRollback => StateOp::Copy {
73                from: StatePlace::Checkpoint,
74                to: StatePlace::Live,
75            },
76            Self::CommitAccepted => StateOp::Copy {
77                from: StatePlace::AcceptedStep,
78                to: StatePlace::Live,
79            },
80            Self::SlotZero => StateOp::Zero,
81        }
82    }
83
84    /// 2026-09-30: The spelling used in node ids.
85    pub fn name(self) -> &'static str {
86        match self {
87            Self::VerifyCheckpoint => "verify_checkpoint",
88            Self::VerifyRollback => "verify_rollback",
89            Self::CommitAccepted => "commit_accepted",
90            Self::SlotZero => "slot_zero",
91        }
92    }
93}
94
95/// 2026-09-30: One node of a state program.
96#[derive(Debug, Clone, PartialEq, Eq)]
97pub struct StateOpNode {
98    /// 2026-09-30: `<program>.<state id>` (`verify_rollback.l4.mamba.conv`).
99    pub id: String,
100    /// 2026-09-30: Index into [`Circuit::states`].
101    pub state: usize,
102    /// 2026-09-30: What it does.
103    pub op: StateOp,
104}
105
106/// 2026-09-30: One state program.
107#[derive(Debug, Clone, PartialEq, Eq)]
108pub struct StateProgram {
109    pub id: StateProgramId,
110    pub nodes: Vec<StateOpNode>,
111}
112
113impl StateProgram {
114    /// 2026-09-30: Bytes one sequence's run of the program moves (or zeroes): one unit of each
115    /// node's state, its keyed format from `formats`.
116    pub fn bytes(
117        &self,
118        circuit: &Circuit,
119        formats: &BTreeMap<String, StateDtype>,
120    ) -> Result<u64, StateError> {
121        let decls: Vec<_> = self
122            .nodes
123            .iter()
124            .map(|n| circuit.states[n.state].clone())
125            .collect();
126        let plan = StatePlan::new(
127            &decls,
128            &StateInputs {
129                formats: formats.clone(),
130                slots: 1,
131                verify: None,
132                kv: None,
133                draft_kv: None,
134            },
135        )?;
136        Ok(plan.bytes_where(|t| t.holding == Holding::Live))
137    }
138}
139
140/// 2026-09-30: Every state program of `circuit`, one node per recurrent target state.
141pub fn state_programs(circuit: &Circuit) -> Vec<StateProgram> {
142    StateProgramId::ALL
143        .into_iter()
144        .map(|id| StateProgram {
145            id,
146            nodes: circuit
147                .states
148                .iter()
149                .enumerate()
150                .filter(|(_, s)| s.kind == StateKind::Recurrent && s.section == Section::Main)
151                .map(|(i, s)| StateOpNode {
152                    id: format!("{}.{}", id.name(), s.id),
153                    state: i,
154                    op: id.op(),
155                })
156                .collect(),
157        })
158        .collect()
159}
160
161#[cfg(test)]
162#[path = "state_ops_tests.rs"]
163mod state_ops_tests;
metrale_circuit/state_ops.rs

metrale_circuit/
state_ops.rs
