1use std::marker::PhantomData;
2
3use alloc::vec::Vec;
4use anyhow::Error;
5use prost::Message;
6use warg_crypto::{
7 hash::{Hash, SupportedDigest},
8 VisitBytes,
9};
10use warg_protobuf::internal as protobuf;
11
12use super::{hash_branch, hash_empty, hash_leaf, node::Node, Checkpoint, LogBuilder};
13
14#[derive(Clone, Debug)]
16pub struct StackLog<D, V>
17where
18 D: SupportedDigest,
19 V: VisitBytes,
20{
21 stack: Vec<(Node, Hash<D>)>,
22 length: usize,
23 _value: PhantomData<V>,
25}
26
27impl<D, V> StackLog<D, V>
28where
29 D: SupportedDigest,
30 V: VisitBytes,
31{
32 pub fn length(&self) -> usize {
34 self.length
35 }
36
37 pub fn is_empty(&self) -> bool {
39 self.length == 0
40 }
41
42 pub fn to_protobuf(self) -> Vec<u8> {
44 let proto: protobuf::StackLog = self.into();
45 proto.encode_to_vec()
46 }
47
48 pub fn from_protobuf(bytes: &[u8]) -> Result<Self, Error> {
50 let proto = protobuf::StackLog::decode(bytes)?;
51 let value = proto.try_into()?;
52 Ok(value)
53 }
54}
55
56impl<D, V> From<StackLog<D, V>> for protobuf::StackLog
57where
58 D: SupportedDigest,
59 V: VisitBytes,
60{
61 fn from(value: StackLog<D, V>) -> Self {
62 let stack = value
63 .stack
64 .into_iter()
65 .map(|(node, hash)| protobuf::HashEntry {
66 index: node.0 as u32,
67 hash: hash.bytes().to_vec(),
68 })
69 .collect();
70 protobuf::StackLog {
71 length: value.length as u32,
72 stack,
73 }
74 }
75}
76
77impl<D, V> TryFrom<protobuf::StackLog> for StackLog<D, V>
78where
79 D: SupportedDigest,
80 V: VisitBytes,
81{
82 type Error = Error;
83
84 fn try_from(value: protobuf::StackLog) -> Result<Self, Self::Error> {
85 let length = value.length as usize;
86 let mut stack = Vec::with_capacity(length);
87 for entry in value.stack {
88 stack.push((Node(entry.index as usize), entry.hash.try_into()?))
89 }
90 let stack_log = StackLog {
91 length,
92 stack,
93 _value: PhantomData,
94 };
95 Ok(stack_log)
96 }
97}
98
99impl<D, V> LogBuilder<D, V> for StackLog<D, V>
100where
101 D: SupportedDigest,
102 V: VisitBytes,
103{
104 fn checkpoint(&self) -> Checkpoint<D> {
105 let root = self
106 .stack
107 .iter()
108 .rev()
109 .map(|(_n, hash)| hash.clone())
110 .reduce(|new, old| hash_branch::<D>(old, new))
111 .unwrap_or_else(hash_empty::<D>);
112
113 Checkpoint {
114 root,
115 length: self.length,
116 }
117 }
118
119 fn push(&mut self, entry: &V) -> Node {
120 let node = Node(self.length * 2);
121
122 let leaf_digest = hash_leaf::<D>(entry);
123
124 self.length += 1;
125 self.stack.push((node, leaf_digest));
126 self.reduce();
127
128 node
129 }
130}
131
132impl<D, V> StackLog<D, V>
133where
134 D: SupportedDigest,
135 V: VisitBytes,
136{
137 fn reduce(&mut self) {
138 while self.stack.len() >= 2 {
139 let (top_node, top_hash) = &self.stack[self.stack.len() - 1];
140 let (second_node, second_hash) = &self.stack[self.stack.len() - 2];
141
142 if top_node.height() == second_node.height() {
143 let new_node = top_node.parent();
144 let new_hash = hash_branch::<D>(second_hash.clone(), top_hash.clone());
145 self.stack.pop();
146 self.stack.pop();
147 self.stack.push((new_node, new_hash));
148 } else {
149 return;
150 }
151 }
152 }
153}
154
155impl<D, V> Default for StackLog<D, V>
156where
157 D: SupportedDigest,
158 V: VisitBytes,
159{
160 fn default() -> Self {
161 Self {
162 stack: Default::default(),
163 length: Default::default(),
164 _value: Default::default(),
165 }
166 }
167}
168
169#[cfg(test)]
170mod tests {
171 use warg_crypto::hash::Sha256;
172
173 use super::super::VecLog;
174 use super::*;
175
176 #[test]
177 fn test_matches_vec_log() {
178 let mut vec_log: VecLog<Sha256, &str> = VecLog::default();
179 let mut stack_log: StackLog<Sha256, &str> = StackLog::default();
180
181 let data: [&str; 25] = [
182 "93", "67", "30", "37", "23", "75", "57", "89", "76", "42", "9", "14", "40", "59",
183 "26", "66", "77", "38", "47", "34", "8", "81", "101", "102", "103",
184 ];
185
186 for leaf in data {
187 vec_log.push(&leaf);
188 stack_log.push(&leaf);
189
190 assert_eq!(vec_log.checkpoint(), stack_log.checkpoint());
191 }
192 }
193}