openvm_continuations/circuit/root/def_paths/
trace.rs

1use std::borrow::BorrowMut;
2
3use openvm_circuit::{
4    arch::{instructions::DEFERRAL_AS, POSEIDON2_WIDTH},
5    system::memory::dimensions::MemoryDimensions,
6};
7use openvm_cpu_backend::CpuBackend;
8use openvm_stark_backend::{prover::AirProvingContext, StarkProtocolConfig};
9use openvm_stark_sdk::config::baby_bear_poseidon2::{
10    poseidon2_compress_with_capacity, DIGEST_SIZE, F,
11};
12use p3_field::PrimeCharacteristicRing;
13use p3_matrix::dense::RowMajorMatrix;
14
15use crate::{
16    circuit::root::def_paths::air::DeferralAccMerklePathsCols, utils::digests_to_poseidon2_input,
17};
18
19fn bit_at(value: usize, bit: usize) -> bool {
20    (value >> bit) & 1 == 1
21}
22
23fn build_path_nodes(
24    acc_hash: [F; DIGEST_SIZE],
25    proof: &[[F; DIGEST_SIZE]],
26    is_right_child_bits: &[bool],
27    skip_depth: usize,
28) -> Vec<[F; DIGEST_SIZE]> {
29    let num_layers = proof.len() + 1;
30    let mut nodes = vec![[F::ZERO; DIGEST_SIZE]; num_layers];
31
32    if skip_depth >= num_layers {
33        return nodes;
34    }
35
36    nodes[skip_depth] = acc_hash;
37    for row_idx in skip_depth..proof.len() {
38        let sibling = proof[row_idx];
39        let is_right_child = is_right_child_bits[row_idx];
40        let left = if is_right_child {
41            sibling
42        } else {
43            nodes[row_idx]
44        };
45        let right = if is_right_child {
46            nodes[row_idx]
47        } else {
48            sibling
49        };
50        nodes[row_idx + 1] = poseidon2_compress_with_capacity(left, right).0;
51    }
52
53    nodes
54}
55
56pub fn generate_proving_input<SC: StarkProtocolConfig<F = F>>(
57    initial_acc_hash: [F; DIGEST_SIZE],
58    final_acc_hash: [F; DIGEST_SIZE],
59    initial_merkle_proof: &[[F; DIGEST_SIZE]],
60    final_merkle_proof: &[[F; DIGEST_SIZE]],
61    memory_dimensions: MemoryDimensions,
62    depth: usize,
63    is_unset: bool,
64) -> (AirProvingContext<CpuBackend<SC>>, Vec<[F; POSEIDON2_WIDTH]>) {
65    assert_eq!(
66        initial_merkle_proof.len(),
67        final_merkle_proof.len(),
68        "initial/final Merkle proofs must have the same depth"
69    );
70
71    let proof_len = initial_merkle_proof.len();
72    let num_layers = proof_len + 1;
73    let height = num_layers.next_power_of_two();
74    let width = DeferralAccMerklePathsCols::<u8>::width();
75
76    // Matches AccMerklePathsAir::new().
77    let expected_branch_bits = usize::try_from(memory_dimensions.label_to_index((DEFERRAL_AS, 0)))
78        .expect("label index must fit in usize");
79    let address_height = memory_dimensions.address_height;
80
81    let skip_depth = if is_unset {
82        0
83    } else {
84        assert!(
85            depth > 0 && depth <= address_height,
86            "depth must be in 1..=address_height when is_unset is false"
87        );
88        depth
89    };
90
91    // is_within_deferral_as tracks the DEFERRAL_AS prefix outside the Merkle-rooted part
92    let untouched_cut = address_height.min(proof_len.saturating_sub(1));
93
94    let mut is_right_child_bits = vec![false; num_layers];
95    let mut row_branch_bits = vec![0usize; num_layers];
96    let mut branch_acc = 0usize;
97    for row_idx in 0..num_layers {
98        let is_right_child = if row_idx < skip_depth {
99            false
100        } else if is_unset && row_idx == 0 {
101            true
102        } else {
103            bit_at(expected_branch_bits, row_idx)
104        };
105        is_right_child_bits[row_idx] = is_right_child;
106        if is_right_child {
107            branch_acc += 1usize << row_idx;
108        }
109        row_branch_bits[row_idx] = branch_acc;
110    }
111
112    let initial_start_hash = if is_unset {
113        poseidon2_compress_with_capacity(initial_acc_hash, [F::ZERO; DIGEST_SIZE]).0
114    } else {
115        initial_acc_hash
116    };
117    let final_start_hash = if is_unset {
118        poseidon2_compress_with_capacity(final_acc_hash, [F::ZERO; DIGEST_SIZE]).0
119    } else {
120        final_acc_hash
121    };
122
123    let initial_nodes = build_path_nodes(
124        initial_start_hash,
125        initial_merkle_proof,
126        &is_right_child_bits,
127        skip_depth,
128    );
129    let final_nodes = build_path_nodes(
130        final_start_hash,
131        final_merkle_proof,
132        &is_right_child_bits,
133        skip_depth,
134    );
135    let final_siblings = final_merkle_proof.to_vec();
136
137    let mut trace = vec![F::ZERO; height * width];
138    let mut poseidon2_inputs = Vec::with_capacity(proof_len * 2);
139
140    for row_idx in 0..num_layers {
141        let row = &mut trace[row_idx * width..(row_idx + 1) * width];
142        let cols: &mut DeferralAccMerklePathsCols<F> = row.borrow_mut();
143
144        cols.is_valid = if row_idx == 0 { F::TWO } else { F::ONE };
145        cols.depth = F::from_usize(row_idx);
146        cols.is_skip = F::from_bool(row_idx < skip_depth);
147        cols.is_within_deferral_as = F::from_bool(row_idx < untouched_cut);
148        cols.is_unset = F::from_bool(is_unset);
149
150        cols.is_right_child = F::from_bool(is_right_child_bits[row_idx]);
151        cols.row_idx_exp_2 = F::from_usize(1usize << row_idx);
152        cols.merkle_path_branch_bits = F::from_usize(row_branch_bits[row_idx]);
153
154        cols.initial_node_commit = initial_nodes[row_idx];
155        cols.final_node_commit = final_nodes[row_idx];
156        if row_idx < proof_len {
157            cols.initial_sibling = initial_merkle_proof[row_idx];
158            cols.final_sibling = final_siblings[row_idx];
159        }
160    }
161
162    for row_idx in 0..proof_len {
163        if row_idx < skip_depth {
164            continue;
165        }
166        let is_right_child = is_right_child_bits[row_idx];
167
168        let init_left = if is_right_child {
169            initial_merkle_proof[row_idx]
170        } else {
171            initial_nodes[row_idx]
172        };
173        let init_right = if is_right_child {
174            initial_nodes[row_idx]
175        } else {
176            initial_merkle_proof[row_idx]
177        };
178        poseidon2_inputs.push(digests_to_poseidon2_input(init_left, init_right));
179
180        let final_left = if is_right_child {
181            final_siblings[row_idx]
182        } else {
183            final_nodes[row_idx]
184        };
185        let final_right = if is_right_child {
186            final_nodes[row_idx]
187        } else {
188            final_siblings[row_idx]
189        };
190        poseidon2_inputs.push(digests_to_poseidon2_input(final_left, final_right));
191    }
192
193    (
194        AirProvingContext::simple_no_pis(RowMajorMatrix::new(trace, width)),
195        poseidon2_inputs,
196    )
197}
openvm_continuations/circuit/root/def_paths/trace.rs

openvm_continuations/circuit/root/def_paths/
trace.rs
