openvm_continuations/circuit/subair/merkle_tree/
trace.rs

1use openvm_stark_sdk::config::baby_bear_poseidon2::{
2    poseidon2_compress_with_capacity, DIGEST_SIZE, F,
3};
4use p3_field::PrimeCharacteristicRing;
5
6use super::MerkleTreeCols;
7use crate::circuit::deferral::{utils::def_tagged_compress, DEF_INTERNAL_TAG, DEF_LEAF_TAG};
8
9/// Build Merkle-tree rows from leaf `(left_child, right_child)` inputs using
10/// Poseidon2 compression over BabyBear.
11///
12/// The number of leaf rows must be a power of two and non-zero.
13/// The returned vector has length `2 * num_leaves`, with the final row being the
14/// terminal row (`send_type = receive_type = 0`).
15pub fn generate_cols_from_leaf_children(
16    leaf_children: Vec<([F; DIGEST_SIZE], [F; DIGEST_SIZE])>,
17    tagged: bool,
18) -> Vec<MerkleTreeCols<F>> {
19    let num_leaves = leaf_children.len();
20    debug_assert!(num_leaves > 0);
21    debug_assert!(num_leaves.is_power_of_two());
22
23    let mut rows = Vec::with_capacity(2 * num_leaves);
24    let mut next_layer = Vec::with_capacity(num_leaves);
25    let leaf_tag = tagged.then_some(DEF_LEAF_TAG);
26    let internal_tag = tagged.then_some(DEF_INTERNAL_TAG);
27
28    for (row_idx, (left, right)) in leaf_children.into_iter().enumerate() {
29        let parent = compress_with_optional_tag(left, right, leaf_tag);
30        rows.push(MerkleTreeCols {
31            row_idx: F::from_usize(row_idx),
32            send_type: if num_leaves == 1 { F::TWO } else { F::ONE },
33            receive_type: F::ONE,
34            parent,
35            is_right_child: F::from_bool(row_idx & 1 == 1),
36            left_child: left,
37            right_child: right,
38        });
39        next_layer.push(parent);
40    }
41
42    let mut row_idx = num_leaves;
43    while next_layer.len() > 1 {
44        let parent_layer_len = next_layer.len() >> 1;
45        for parent_idx in 0..parent_layer_len {
46            let left = next_layer[2 * parent_idx];
47            let right = next_layer[2 * parent_idx + 1];
48            let parent = compress_with_optional_tag(left, right, internal_tag);
49
50            rows.push(MerkleTreeCols {
51                row_idx: F::from_usize(row_idx),
52                send_type: if parent_layer_len > 1 { F::ONE } else { F::TWO },
53                receive_type: F::TWO,
54                parent,
55                is_right_child: F::from_bool(row_idx & 1 == 1),
56                left_child: left,
57                right_child: right,
58            });
59
60            row_idx += 1;
61            next_layer[parent_idx] = parent;
62        }
63        next_layer.truncate(parent_layer_len);
64    }
65
66    rows.push(MerkleTreeCols {
67        row_idx: F::from_usize(row_idx),
68        send_type: F::ZERO,
69        receive_type: F::ZERO,
70        parent: [F::ZERO; DIGEST_SIZE],
71        is_right_child: F::from_bool(row_idx & 1 == 1),
72        left_child: [F::ZERO; DIGEST_SIZE],
73        right_child: [F::ZERO; DIGEST_SIZE],
74    });
75
76    rows
77}
78
79fn compress_with_optional_tag(
80    left: [F; DIGEST_SIZE],
81    right: [F; DIGEST_SIZE],
82    tag: Option<[u8; DIGEST_SIZE]>,
83) -> [F; DIGEST_SIZE] {
84    match tag {
85        Some(tag) => def_tagged_compress(tag, left, right).1,
86        None => poseidon2_compress_with_capacity(left, right).0,
87    }
88}
openvm_continuations/circuit/subair/merkle_tree/trace.rs

openvm_continuations/circuit/subair/merkle_tree/
trace.rs
