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