sos_core/commit/
tree.rs

1use crate::{Error, Result};
2use rs_merkle::{algorithms::Sha256, Hasher, MerkleTree};
3
4use super::{CommitHash, CommitProof, CommitState, Comparison, TreeHash};
5
6/// Encapsulates a Merkle tree and provides functions
7/// for generating and comparing proofs.
8#[derive(Default)]
9pub struct CommitTree {
10    tree: MerkleTree<Sha256>,
11    maybe_last_commit: Option<TreeHash>,
12    last_commit: Option<TreeHash>,
13}
14
15impl CommitTree {
16    /// Create a new empty commit tree.
17    pub fn new() -> Self {
18        Self {
19            tree: MerkleTree::<Sha256>::new(),
20            maybe_last_commit: None,
21            last_commit: None,
22        }
23    }
24
25    /// Compute the Sha256 hash of some data.
26    pub fn hash(data: &[u8]) -> TreeHash {
27        Sha256::hash(data)
28    }
29
30    /// Number of leaves in the tree.
31    pub fn len(&self) -> usize {
32        self.tree.leaves_len()
33    }
34
35    /// Determine if this commit tree is empty.
36    pub fn is_empty(&self) -> bool {
37        self.len() == 0
38    }
39
40    /// Insert a commit hash into the tree,
41    pub fn insert(&mut self, hash: TreeHash) -> &mut Self {
42        self.maybe_last_commit = Some(hash);
43        self.tree.insert(hash);
44        self
45    }
46
47    /// Append a collections of commit hashes to the tree.
48    pub fn append(&mut self, hashes: &mut Vec<TreeHash>) -> &mut Self {
49        self.maybe_last_commit = hashes.last().cloned();
50        self.tree.append(hashes);
51        self
52    }
53
54    /// Commit changes to the tree to compute the root.
55    pub fn commit(&mut self) {
56        self.tree.commit();
57        if self.maybe_last_commit.is_some() {
58            self.last_commit = self.maybe_last_commit.take();
59        }
60    }
61
62    /// Revert changes to the tree.
63    pub fn rollback(&mut self) {
64        self.tree.rollback();
65        self.maybe_last_commit = None;
66        if let Some(leaves) = self.leaves() {
67            self.last_commit = leaves.last().cloned();
68        } else {
69            self.last_commit = None;
70        }
71    }
72
73    /// Leaves of the tree.
74    pub fn leaves(&self) -> Option<Vec<TreeHash>> {
75        self.tree.leaves()
76    }
77
78    /// Root hash and a proof of the last leaf node.
79    pub fn head(&self) -> Result<CommitProof> {
80        if self.is_empty() {
81            return Err(Error::NoRootCommit);
82        }
83        self.proof(&[self.tree.leaves_len() - 1])
84    }
85
86    /// Proof for the given indices.
87    pub fn proof(&self, leaf_indices: &[usize]) -> Result<CommitProof> {
88        let root = self.root().ok_or(Error::NoRootCommit)?;
89        let proof = self.tree.proof(leaf_indices);
90        Ok(CommitProof {
91            root,
92            proof,
93            length: self.len(),
94            indices: leaf_indices.to_vec(),
95        })
96    }
97
98    /// Compare this tree against another root hash and merkle proof.
99    pub fn compare(&self, proof: &CommitProof) -> Result<Comparison> {
100        let CommitProof {
101            root: other_root,
102            proof,
103            length,
104            indices: indices_to_prove,
105        } = proof;
106        let root = self.root().ok_or(Error::NoRootCommit)?;
107        if &root == other_root {
108            Ok(Comparison::Equal)
109        } else {
110            let leaves = self.tree.leaves().unwrap_or_default();
111            let leaves_to_prove = indices_to_prove
112                .into_iter()
113                .filter_map(|i| leaves.get(*i).cloned())
114                .collect::<Vec<_>>();
115            if leaves_to_prove.len() == indices_to_prove.len() {
116                if proof.verify(
117                    other_root.into(),
118                    indices_to_prove.as_slice(),
119                    leaves_to_prove.as_slice(),
120                    *length,
121                ) {
122                    Ok(Comparison::Contains(indices_to_prove.to_vec()))
123                } else {
124                    Ok(Comparison::Unknown)
125                }
126            } else {
127                Ok(Comparison::Unknown)
128            }
129        }
130    }
131
132    /// Compute the first commit state.
133    ///
134    /// Must have at least one commit.
135    pub fn first_commit(&self) -> Result<CommitState> {
136        let leaves = self.tree.leaves().ok_or(Error::NoRootCommit)?;
137
138        // Compute a proof for the first commit
139        let leaf = *leaves.first().ok_or(Error::NoRootCommit)?;
140        let leaves = vec![leaf];
141        let tree = MerkleTree::<Sha256>::from_leaves(&leaves);
142        let proof = tree.proof(&[0]);
143        let root = tree.root().map(CommitHash).ok_or(Error::NoRootCommit)?;
144        let first_proof = CommitProof {
145            root,
146            proof,
147            length: leaves.len(),
148            indices: vec![0],
149        };
150
151        let first_commit = CommitHash(leaf);
152        Ok(CommitState(first_commit, first_proof))
153    }
154
155    /// Last commit hash in the underlying merkle tree.
156    pub fn last_commit(&self) -> Option<CommitHash> {
157        self.last_commit.map(CommitHash)
158    }
159
160    /// Commit state of this tree.
161    ///
162    /// The tree must already have some commits.
163    pub fn commit_state(&self) -> Result<CommitState> {
164        let last_commit = self.last_commit().ok_or(Error::NoLastCommit)?;
165        Ok(CommitState(last_commit, self.head()?))
166    }
167
168    /// Root hash of the underlying merkle tree.
169    pub fn root(&self) -> Option<CommitHash> {
170        self.tree.root().map(CommitHash)
171    }
172
173    /// Root hash of the underlying merkle tree as hexadecimal.
174    pub fn root_hex(&self) -> Option<String> {
175        self.tree.root().map(hex::encode)
176    }
177
178    /// Given a proof from another tree determine if this
179    /// tree contains the other proof.
180    pub fn contains(
181        &self,
182        other_proof: &CommitProof,
183    ) -> Result<Option<CommitProof>> {
184        Ok(match self.compare(other_proof)? {
185            Comparison::Contains(indices) => Some(self.proof(&indices)?),
186            _ => None,
187        })
188    }
189}
190
191/*
192#[cfg(debug_assertions)]
193impl From<&crate::vault::Vault> for CommitTree {
194    fn from(value: &crate::vault::Vault) -> Self {
195        let mut commit_tree = CommitTree::new();
196        for (_, commit) in value.commits() {
197            commit_tree.tree.insert(commit.into());
198        }
199        commit_tree.tree.commit();
200        commit_tree
201    }
202}
203*/
sos_core/commit/tree.rs

sos_core/commit/
tree.rs
