1use super::{
  2    BTreeMap, EmptySubtreeRoots, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, StarkField,
  3    Vec, Word, EMPTY_WORD,
  4};
  5
  6#[cfg(test)]
  7mod tests;
  8
  9// TIERED SPARSE MERKLE TREE
 10// ================================================================================================
 11
 12#[derive(Debug, Clone, PartialEq, Eq)]
 13pub struct TieredSmt {
 14    root: RpoDigest,
 15    nodes: BTreeMap<NodeIndex, RpoDigest>,
 16    upper_leaves: BTreeMap<NodeIndex, RpoDigest>,
 17    bottom_leaves: BTreeMap<u64, Vec<RpoDigest>>,
 18    values: BTreeMap<RpoDigest, Word>,
 19}
 20
 21impl TieredSmt {
 22    // CONSTANTS
 23    // --------------------------------------------------------------------------------------------
 24
 25    const MAX_DEPTH: u8 = 64;
 26
 27    // CONSTRUCTORS
 28    // --------------------------------------------------------------------------------------------
 29
 30    pub fn new() -> Self {
 31        Self {
 32            root: EmptySubtreeRoots::empty_hashes(Self::MAX_DEPTH)[0],
 33            nodes: BTreeMap::new(),
 34            upper_leaves: BTreeMap::new(),
 35            bottom_leaves: BTreeMap::new(),
 36            values: BTreeMap::new(),
 37        }
 38    }
 39
 40    // PUBLIC ACCESSORS
 41    // --------------------------------------------------------------------------------------------
 42
 43    pub const fn root(&self) -> RpoDigest {
 44        self.root
 45    }
 46
 47    pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
 48        if index.is_root() {
 49            return Err(MerkleError::DepthTooSmall(index.depth()));
 50        } else if index.depth() > Self::MAX_DEPTH {
 51            return Err(MerkleError::DepthTooBig(index.depth() as u64));
 52        } else if !self.is_node_available(index) {
 53            todo!()
 54        }
 55
 56        Ok(self.get_branch_node(&index))
 57    }
 58
 59    pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
 60        if index.is_root() {
 61            return Err(MerkleError::DepthTooSmall(index.depth()));
 62        } else if index.depth() > Self::MAX_DEPTH {
 63            return Err(MerkleError::DepthTooBig(index.depth() as u64));
 64        } else if !self.is_node_available(index) {
 65            todo!()
 66        }
 67
 68        let mut path = Vec::with_capacity(index.depth() as usize);
 69        for _ in 0..index.depth() {
 70            let node = self.get_branch_node(&index.sibling());
 71            path.push(node.into());
 72            index.move_up();
 73        }
 74        Ok(path.into())
 75    }
 76
 77    pub fn get_value(&self, key: RpoDigest) -> Result<Word, MerkleError> {
 78        match self.values.get(&key) {
 79            Some(value) => Ok(*value),
 80            None => Ok(EMPTY_WORD),
 81        }
 82    }
 83
 84    // STATE MUTATORS
 85    // --------------------------------------------------------------------------------------------
 86
 87    pub fn insert(&mut self, key: RpoDigest, value: Word) -> Result<Word, MerkleError> {
 88        let (mut index, leaf_key) = self.get_insert_location(&key);
 89
 90        if let Some(other_key) = leaf_key {
 91            if other_key != key {
 92                let common_prefix_len = get_common_prefix_length(&key, &other_key);
 93                let depth = common_prefix_len + 16;
 94
 95                let other_index = key_to_index(&other_key, depth);
 96                self.move_leaf_node(other_key, index, other_index);
 97
 98                index = key_to_index(&key, depth);
 99            }
100        }
101
102        let old_value = self.values.insert(key, value).unwrap_or(EMPTY_WORD);
103        if value != old_value {
104            self.upper_leaves.insert(index, key);
105            let new_node = build_leaf_node(key, value, index.depth().into());
106            self.root = self.update_path(index, new_node);
107        }
108
109        Ok(old_value)
110    }
111
112    // HELPER METHODS
113    // --------------------------------------------------------------------------------------------
114
115    fn is_node_available(&self, index: NodeIndex) -> bool {
116        match index.depth() {
117            32 => true,
118            48 => true,
119            _ => true,
120        }
121    }
122
123    fn get_branch_node(&self, index: &NodeIndex) -> RpoDigest {
124        match self.nodes.get(index) {
125            Some(node) => *node,
126            None => EmptySubtreeRoots::empty_hashes(Self::MAX_DEPTH)[index.depth() as usize],
127        }
128    }
129
130    fn get_insert_location(&self, key: &RpoDigest) -> (NodeIndex, Option<RpoDigest>) {
131        let mse = Word::from(key)[3].as_int();
132        for depth in (16..64).step_by(16) {
133            let index = NodeIndex::new(depth, mse >> (Self::MAX_DEPTH - depth)).unwrap();
134            if let Some(leaf_key) = self.upper_leaves.get(&index) {
135                return (index, Some(*leaf_key));
136            } else if self.nodes.contains_key(&index) {
137                continue;
138            } else {
139                return (index, None);
140            }
141        }
142
143        // TODO: handle bottom tier
144        unimplemented!()
145    }
146
147    fn move_leaf_node(&mut self, key: RpoDigest, old_index: NodeIndex, new_index: NodeIndex) {
148        self.upper_leaves.remove(&old_index).unwrap();
149        self.upper_leaves.insert(new_index, key);
150        let value = *self.values.get(&key).unwrap();
151        let new_node = build_leaf_node(key, value, new_index.depth().into());
152        self.update_path(new_index, new_node);
153    }
154
155    fn update_path(&mut self, mut index: NodeIndex, mut node: RpoDigest) -> RpoDigest {
156        for _ in 0..index.depth() {
157            self.nodes.insert(index, node);
158            let sibling = self.get_branch_node(&index.sibling());
159            node = Rpo256::merge(&index.build_node(node, sibling));
160            index.move_up();
161        }
162        node
163    }
164}
165
166impl Default for TieredSmt {
167    fn default() -> Self {
168        Self::new()
169    }
170}
171
172// HELPER FUNCTIONS
173// ================================================================================================
174
175fn get_remaining_path(key: RpoDigest, depth: u32) -> RpoDigest {
176    let mut key = Word::from(key);
177    let remaining = (key[3].as_int() << depth) >> depth;
178    key[3] = remaining.into();
179    key.into()
180}
181
182fn build_leaf_node(key: RpoDigest, value: Word, depth: u32) -> RpoDigest {
183    let remaining_path = get_remaining_path(key, depth);
184    Rpo256::merge_in_domain(&[remaining_path, value.into()], depth.into())
185}
186
187fn get_common_prefix_length(key1: &RpoDigest, key2: &RpoDigest) -> u8 {
188    let e1 = Word::from(key1)[3].as_int();
189    let e2 = Word::from(key2)[3].as_int();
190
191    if e1 == e2 {
192        64
193    } else if e1 >> 16 == e2 >> 16 {
194        48
195    } else if e1 >> 32 == e2 >> 32 {
196        32
197    } else if e1 >> 48 == e2 >> 48 {
198        16
199    } else {
200        0
201    }
202}
203
204fn key_to_index(key: &RpoDigest, depth: u8) -> NodeIndex {
205    let mse = Word::from(key)[3].as_int();
206    let value = match depth {
207        16 | 32 | 48 => mse >> (depth as u32),
208        _ => unreachable!("invalid depth: {depth}"),
209    };
210
211    // TODO: use unchecked version?
212    NodeIndex::new(depth, value).unwrap()
213}