arnaucube
arnaucube/shockwave-plusPublic Self-hosted repository
shockwave-plus/tensor_pcs/src/tree.rs
tensor_pcs/src/tree.rsRaw
 1use crate::{utils::hash_all, FieldExt};
 2use serde::{Deserialize, Serialize};
 3
 4#[derive(Clone)]
 5pub struct CommittedMerkleTree<F> {
 6    pub column_roots: Vec<[u8; 32]>,
 7    pub leaves: Vec<F>,
 8    pub num_cols: usize,
 9    pub root: [u8; 32],
10}
11
12impl<F: FieldExt> CommittedMerkleTree<F> {
13    pub fn from_leaves(leaves: Vec<F>, num_cols: usize) -> Self {
14        let n = leaves.len();
15        debug_assert!(n.is_power_of_two());
16        let num_rows = n / num_cols;
17        assert!(num_rows & 1 == 0); // Number of rows must be even
18
19        let leaf_bytes = leaves
20            .iter()
21            .map(|x| x.to_repr())
22            .collect::<Vec<[u8; 32]>>();
23
24        let mut column_roots = Vec::with_capacity(num_cols);
25        for col in 0..num_cols {
26            let column_leaves = leaf_bytes[col * num_rows..(col + 1) * num_rows].to_vec();
27            let column_root = hash_all(&column_leaves);
28            column_roots.push(column_root);
29        }
30
31        let root = hash_all(&column_roots);
32
33        Self {
34            column_roots,
35            leaves,
36            root,
37            num_cols,
38        }
39    }
40
41    pub fn root(&self) -> [u8; 32] {
42        self.root
43    }
44
45    pub fn leaves(&self) -> Vec<F> {
46        self.leaves.clone()
47    }
48}
49
50#[derive(Serialize, Deserialize, Clone, Debug)]
51pub struct BaseOpening {
52    pub hashes: Vec<[u8; 32]>,
53}
54
55impl BaseOpening {
56    pub fn verify(&self, root: [u8; 32]) -> bool {
57        let r = hash_all(&self.hashes);
58        root == r
59    }
60}