miden_ace_codegen/dag/
ir.rs

1use core::sync::atomic::{AtomicUsize, Ordering};
2
3use miden_crypto::{
4    field::TwoAdicField,
5    stark::dft::{NaiveDft, TwoAdicSubgroupDft},
6};
7
8use crate::layout::InputKey;
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
11pub(crate) struct DagId(usize);
12
13impl DagId {
14    pub(crate) fn fresh() -> Self {
15        static NEXT_DAG_ID: AtomicUsize = AtomicUsize::new(0);
16
17        Self(NEXT_DAG_ID.fetch_add(1, Ordering::Relaxed))
18    }
19}
20
21/// Identifier for a node in the DAG.
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
23pub struct NodeId {
24    pub(super) dag_id: DagId,
25    pub(super) index: usize,
26}
27
28impl NodeId {
29    /// Return the underlying node index.
30    pub const fn index(self) -> usize {
31        self.index
32    }
33
34    pub(super) const fn in_dag(index: usize, dag_id: DagId) -> Self {
35        Self { dag_id, index }
36    }
37}
38
39/// Node kinds in the DAG.
40///
41/// These nodes mirror the verifier expression tree after lowering:
42/// inputs are read via `InputKey`, constants are lifted into the DAG, and
43/// arithmetic nodes capture the evaluation order.
44#[derive(Debug, Clone, PartialEq, Eq, Hash)]
45pub enum NodeKind<EF> {
46    /// Layout-addressable input (public, OOD, aux, etc.).
47    Input(InputKey),
48    /// Constant extension-field value.
49    Constant(EF),
50    /// Addition node.
51    Add(NodeId, NodeId),
52    /// Subtraction node.
53    Sub(NodeId, NodeId),
54    /// Multiplication node.
55    Mul(NodeId, NodeId),
56    /// Negation node (modeled as 0 - x when emitting ops).
57    Neg(NodeId),
58}
59
60/// Precomputed periodic column data for DAG construction.
61#[derive(Debug, Clone)]
62pub struct PeriodicColumnData<EF> {
63    /// Per-column coefficient vectors (highest-degree first).
64    coeffs: Vec<Vec<EF>>,
65}
66
67impl<EF> PeriodicColumnData<EF> {
68    /// Convert periodic columns (evaluations) into coefficient form for DAG building.
69    ///
70    /// Applies an inverse DFT so the DAG can evaluate them at `z_k` inside the circuit.
71    pub fn from_periodic_columns<F>(periodic_columns: Vec<Vec<F>>) -> Self
72    where
73        F: TwoAdicField,
74        EF: From<F>,
75    {
76        if periodic_columns.is_empty() {
77            return Self { coeffs: Vec::new() };
78        }
79
80        let dft = NaiveDft;
81        let mut coeffs = Vec::with_capacity(periodic_columns.len());
82        for col in periodic_columns {
83            assert!(!col.is_empty(), "periodic column must not be empty");
84            assert!(col.len().is_power_of_two(), "periodic column length must be a power of two");
85            let values = dft.idft(col);
86            let coeff_row = values.into_iter().map(EF::from).collect();
87            coeffs.push(coeff_row);
88        }
89
90        Self { coeffs }
91    }
92
93    /// Number of periodic columns.
94    pub fn num_columns(&self) -> usize {
95        self.coeffs.len()
96    }
97
98    /// Maximum periodic column length (used to align powers).
99    pub fn max_period(&self) -> usize {
100        self.coeffs.iter().map(|c| c.len()).max().unwrap_or(0)
101    }
102
103    /// Iterate over the per-column coefficient vectors.
104    pub fn columns(&self) -> &[Vec<EF>] {
105        &self.coeffs
106    }
107}
108
109/// A built DAG with a designated root.
110#[derive(Debug)]
111pub struct AceDag<EF> {
112    dag_id: DagId,
113    /// Topologically ordered nodes.
114    pub nodes: Vec<NodeKind<EF>>,
115    /// Root node of the verifier equation.
116    pub root: NodeId,
117}
118
119/// Exported DAG data that preserves the source DAG id across imports.
120#[derive(Debug, Clone)]
121pub struct DagSnapshot<EF> {
122    nodes: Vec<NodeKind<EF>>,
123    root: NodeId,
124    source_dag_id: DagId,
125}
126
127impl<EF> AceDag<EF> {
128    pub(crate) fn from_parts(dag_id: DagId, nodes: Vec<NodeKind<EF>>, root: NodeId) -> Self {
129        Self { dag_id, nodes, root }
130    }
131
132    pub(crate) fn nodes(&self) -> &[NodeKind<EF>] {
133        &self.nodes
134    }
135
136    pub(crate) fn into_nodes(self) -> Vec<NodeKind<EF>> {
137        self.nodes
138    }
139
140    pub(crate) fn dag_id(&self) -> DagId {
141        self.dag_id
142    }
143
144    pub fn root(&self) -> NodeId {
145        self.root
146    }
147
148    /// Consume the DAG and return an exported snapshot that can be re-imported later.
149    pub fn into_snapshot(self) -> DagSnapshot<EF> {
150        DagSnapshot {
151            nodes: self.nodes,
152            root: self.root,
153            source_dag_id: self.dag_id,
154        }
155    }
156}
157
158impl<EF> DagSnapshot<EF> {
159    /// Root node of the verifier equation.
160    pub fn root(&self) -> NodeId {
161        self.root
162    }
163
164    pub(super) fn into_parts(self) -> (DagId, Vec<NodeKind<EF>>, NodeId) {
165        (self.source_dag_id, self.nodes, self.root)
166    }
167}
miden_ace_codegen/dag/ir.rs

miden_ace_codegen/dag/
ir.rs
