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        let num_rows = n / num_cols;
16        assert!(num_rows & 1 == 0); // Number of rows must be even
17
18        let leaf_bytes = leaves
19            .iter()
20            .map(|x| x.to_repr())
21            .collect::<Vec<[u8; 32]>>();
22
23        let mut column_roots = Vec::with_capacity(num_cols);
24        for col in 0..num_cols {
25            let column_leaves = leaf_bytes[col * num_rows..(col + 1) * num_rows].to_vec();
26            let column_root = hash_all(&column_leaves);
27            column_roots.push(column_root);
28        }
29
30        let root = hash_all(&column_roots);
31
32        Self {
33            column_roots,
34            leaves,
35            root,
36            num_cols,
37        }
38    }
39
40    pub fn root(&self) -> [u8; 32] {
41        self.root
42    }
43
44    pub fn leaves(&self) -> Vec<F> {
45        self.leaves.clone()
46    }
47}
48
49#[derive(Serialize, Deserialize, Clone, Debug)]
50pub struct BaseOpening {
51    pub hashes: Vec<[u8; 32]>,
52}
53
54impl BaseOpening {
55    pub fn verify(&self, root: [u8; 32]) -> bool {
56        let r = hash_all(&self.hashes);
57        root == r
58    }
59}