Skip to main content

proof_stream/viz/
circuit.rs

1//! Helpers for building `ProofEvent::CircuitCompiled` from stwo-ml circuit types.
2//!
3//! This module is called from stwo-ml (feature-gated) — it may import stwo-ml
4//! types directly. The output is a plain `ProofEvent` with no stwo/nightly deps.
5
6use crate::events::{CircuitNodeMeta, LayerKind, ProofEvent};
7
8/// A minimal description of one circuit layer, collected by stwo-ml before
9/// calling `circuit_compiled_event`. This avoids importing stwo-ml types
10/// directly into proof-stream.
11#[derive(Debug, Clone)]
12pub struct CircuitLayerDesc {
13    pub layer_idx: usize,
14    pub node_id: usize,
15    pub kind: LayerKind,
16    pub input_shape: (usize, usize),
17    pub output_shape: (usize, usize),
18    /// Estimated number of trace cells (e.g. m*k*n for MatMul).
19    pub trace_cost: usize,
20    /// Predecessor layer indices.
21    pub input_layers: Vec<usize>,
22}
23
24/// Build `ProofEvent::CircuitCompiled` from a flat description list.
25pub fn circuit_compiled_event(
26    layers: &[CircuitLayerDesc],
27    total_layers: usize,
28    input_shape: (usize, usize),
29    output_shape: (usize, usize),
30) -> ProofEvent {
31    let nodes: Vec<CircuitNodeMeta> = layers
32        .iter()
33        .map(|l| CircuitNodeMeta {
34            layer_idx: l.layer_idx,
35            node_id: l.node_id,
36            kind: l.kind,
37            input_shape: l.input_shape,
38            output_shape: l.output_shape,
39            trace_cost: l.trace_cost,
40            input_layers: l.input_layers.clone(),
41        })
42        .collect();
43
44    ProofEvent::CircuitCompiled {
45        total_layers,
46        input_shape,
47        output_shape,
48        nodes,
49        has_simd: false,
50        simd_num_blocks: 0,
51    }
52}
53
54/// Estimate trace cost for a MatMul layer: m × k × n rounded up to next
55/// power of two (proportional to sumcheck size).
56pub fn matmul_trace_cost(m: usize, k: usize, n: usize) -> usize {
57    let raw = m * k * n;
58    raw.next_power_of_two()
59}