concordium_std/test_infrastructure/
trie.rs

1use super::{TestStateEntry, TestStateEntryData};
2use crate::{
3    cell::{Cell, RefCell},
4    collections::{btree_map, BTreeMap, HashMap as Map, VecDeque},
5    rc::Rc,
6    Box, StateEntryId, StateError, Vec,
7};
8use core::convert::TryInto;
9
10const BRANCHING_FACTOR: usize = 16;
11
12pub(crate) type Index = usize;
13
14#[derive(Debug)]
15pub(crate) struct StateTrie {
16    nodes:           Node,
17    next_entry_id:   Cell<StateEntryId>,
18    entry_map:       RefCell<Map<StateEntryId, Vec<Index>>>,
19    iterator_counts: RefCell<BTreeMap<Vec<Index>, u32>>,
20}
21
22impl Default for StateTrie {
23    /// Creates an empty state trie.
24    fn default() -> Self { Self::new() }
25}
26
27fn to_indexes(key: &[u8]) -> Vec<Index> {
28    let mut indexes = Vec::new();
29    for byte in key {
30        indexes.push(((byte & 0b_11_11_00_00) >> 4) as Index);
31        indexes.push((byte & 0b_00_00_11_11) as Index);
32    }
33    indexes
34}
35
36/// The inverse of `to_indexes`.
37fn from_indexes(indexes: &[Index]) -> Vec<u8> {
38    let mut key = Vec::new();
39    for chunk in indexes.chunks(2) {
40        let n = (chunk[0] << 4 | chunk[1]) as u8;
41        key.push(n);
42    }
43    key
44}
45
46impl StateTrie {
47    pub(crate) fn new() -> Self {
48        Self {
49            nodes:           Node::new(),
50            next_entry_id:   Cell::new(0),
51            entry_map:       RefCell::new(Map::default()),
52            iterator_counts: Default::default(),
53        }
54    }
55
56    /// Construct a `TestStateEntry` and use interior mutation to add increment
57    /// next_entry_id and add the entry to the entry_map.
58    fn construct_state_entry_test(
59        &self,
60        indexes: Vec<Index>,
61        data: Rc<RefCell<TestStateEntryData>>,
62        key: Vec<u8>,
63    ) -> TestStateEntry {
64        // Get the current next_entry_id
65        let state_entry_id = self.next_entry_id.get();
66
67        // Add the id and indexes to the map and increment the next_entry_id
68        self.entry_map.borrow_mut().insert(state_entry_id, indexes);
69        self.next_entry_id.set(state_entry_id + 1);
70
71        TestStateEntry::open(data, key, state_entry_id)
72    }
73
74    pub(crate) fn delete_prefix(&mut self, prefix: &[u8]) -> Result<bool, StateError> {
75        let indexes = to_indexes(prefix);
76        if self.is_locked(&indexes) {
77            return Err(StateError::SubtreeLocked);
78        }
79
80        // Unwrapping is safe, because the subtree isn't locked.
81        // TODO: Getting an iterator is not OK here due to the bound.
82        let iterator = match self.iterator(prefix) {
83            Err(StateError::SubtreeWithPrefixNotFound) => {
84                return Ok(false);
85            }
86            Err(e) => crate::fail!("Unexpected error in delete_prefix: {:?}", e),
87            Ok(v) => v,
88        };
89
90        // Invalidate all the data in the deleted entries such that reading and writing
91        // them will fail.
92        // Invalidating is necessary because the data is kept alive in any entries given
93        // out due to the Rc. This uses the queue iter because Iterator isn't
94        // implemented for &Iter and we need to delete the iterator afterwards.
95        for entry in iterator.queue.iter() {
96            *entry.cursor.data.borrow_mut() = TestStateEntryData::EntryDeleted;
97        }
98        self.delete_iterator(iterator);
99
100        // Delete the nodes in the tree.
101        self.nodes.delete_prefix(&indexes, false)?;
102        Ok(true)
103    }
104
105    /// Returns true if the subtree corresponding to the given key is
106    /// already locked by an existing iterator, false otherwise.
107    fn is_locked(&self, prefix: &[usize]) -> bool {
108        self.iterator_counts.borrow().keys().any(|p| {
109            let shortest = crate::cmp::min(p.len(), prefix.len());
110            p[..shortest] == prefix[..shortest]
111        })
112    }
113
114    pub(crate) fn create_entry(&mut self, key: &[u8]) -> Result<TestStateEntry, StateError> {
115        let indexes = to_indexes(key);
116        if self.is_locked(&indexes) {
117            return Err(StateError::SubtreeLocked);
118        }
119        let data = self.nodes.create(&indexes);
120        let entry = self.construct_state_entry_test(indexes, data, key.to_vec());
121        Ok(entry)
122    }
123
124    pub(crate) fn lookup(&self, key: &[u8]) -> Option<TestStateEntry> {
125        let indexes = to_indexes(key);
126        self.nodes
127            .lookup(&indexes)
128            .map(|data| self.construct_state_entry_test(indexes, data, key.to_vec()))
129    }
130
131    pub(crate) fn delete_entry(&mut self, entry: TestStateEntry) -> Result<(), StateError> {
132        let indexes = to_indexes(&entry.key);
133        if self.is_locked(&indexes) {
134            return Err(StateError::SubtreeLocked);
135        }
136        match self.entry_map.borrow_mut().remove(&entry.state_entry_id) {
137            Some(indexes) => self.nodes.delete_data(&indexes),
138            None => Err(StateError::EntryNotFound), /* Entry did not exist. Only happens
139                                                     * when entry was deleted using
140                                                     * delete_prefix. */
141        }
142    }
143
144    pub(crate) fn iterator(&self, prefix: &[u8]) -> Result<TestStateIter, StateError> {
145        let index_prefix = to_indexes(prefix);
146
147        // Try to find the root_node for the prefix.
148        let node =
149            self.nodes.lookup_node(&index_prefix).ok_or(StateError::SubtreeWithPrefixNotFound)?;
150
151        // Keep track of the number of iterators given out.
152        match self.iterator_counts.borrow_mut().entry(index_prefix.clone()) {
153            btree_map::Entry::Vacant(vac) => {
154                let _ = vac.insert(1);
155            }
156            btree_map::Entry::Occupied(ref mut occ) => {
157                if *occ.get() == u32::MAX {
158                    return Err(StateError::IteratorLimitForPrefixExceeded);
159                }
160                *occ.get_mut() += 1;
161            }
162        }
163
164        let iter = TestStateIter::new(self, index_prefix, node);
165        Ok(iter)
166    }
167
168    pub(crate) fn delete_iterator(&mut self, iterator: TestStateIter) {
169        match self.iterator_counts.borrow_mut().entry(iterator.prefix) {
170            btree_map::Entry::Vacant(_) => crate::fail!(), // Internal error: Should never happen.
171            btree_map::Entry::Occupied(mut occ) => {
172                if *occ.get() == 1 {
173                    // Deleting last iterator for the prefix.
174                    occ.remove();
175                } else {
176                    *occ.get_mut() -= 1;
177                }
178            }
179        }
180    }
181
182    /// Makes a deep clone of the trie. Used for rollbacks.
183    pub(crate) fn clone_deep(&self) -> Self {
184        Self {
185            nodes:           self.nodes.clone_deep(),
186            next_entry_id:   self.next_entry_id.clone(),
187            entry_map:       self.entry_map.clone(),
188            iterator_counts: self.iterator_counts.clone(),
189        }
190    }
191}
192
193#[derive(Debug)]
194pub struct TestStateIter {
195    // Only used when deleting the iterator.
196    prefix: Vec<Index>,
197    queue:  VecDeque<TestStateEntry>,
198}
199
200impl TestStateIter {
201    fn new(trie: &StateTrie, mut root_index: Vec<Index>, root_of_iter: &Node) -> Self {
202        let mut queue = VecDeque::new();
203        let prefix = root_index.clone();
204
205        fn build_queue(
206            trie: &StateTrie,
207            queue: &mut VecDeque<TestStateEntry>,
208            indexes: &mut Vec<Index>,
209            node: &Node,
210        ) {
211            for idx in 0..BRANCHING_FACTOR {
212                if let Some(child) = &node.children[idx] {
213                    // Push current index.
214                    indexes.push(idx);
215
216                    if let Some(data) = &child.data {
217                        let state_entry = trie.construct_state_entry_test(
218                            indexes.clone(),
219                            Rc::clone(data),
220                            from_indexes(indexes),
221                        );
222                        queue.push_back(state_entry);
223                    }
224                    build_queue(trie, queue, indexes, child);
225
226                    // Pop current index again.
227                    indexes.pop();
228                }
229            }
230        }
231
232        build_queue(trie, &mut queue, &mut root_index, root_of_iter);
233        Self {
234            prefix,
235            queue,
236        }
237    }
238}
239
240impl Iterator for TestStateIter {
241    type Item = TestStateEntry;
242
243    fn next(&mut self) -> Option<Self::Item> { self.queue.pop_front() }
244}
245
246#[derive(Debug)]
247struct Node {
248    data:     Option<Rc<RefCell<TestStateEntryData>>>,
249    children: [Option<Box<Node>>; BRANCHING_FACTOR],
250}
251
252impl Node {
253    fn new() -> Self {
254        Self {
255            data:     None,
256            children: Default::default(),
257        }
258    }
259
260    /// Tries to find the data in a node with the given index.
261    /// Returns `None` if the node doesn't exist or if it doesn't have any data.
262    /// Note: If `Some` is returned, it will _always_ be a
263    /// `TestStateEntryData::EntryExists(..)`.
264    fn lookup(&self, indexes: &[Index]) -> Option<Rc<RefCell<TestStateEntryData>>> {
265        self.lookup_node(indexes).and_then(|node| node.data.as_ref().map(Rc::clone))
266    }
267
268    /// Tries to find the node with the given index.
269    /// Returns `None` if the node doesn't exist.
270    fn lookup_node(&self, indexes: &[Index]) -> Option<&Self> {
271        match indexes.first() {
272            Some(idx) => {
273                self.children[*idx].as_ref().and_then(|node| node.lookup_node(&indexes[1..]))
274            }
275            None => Some(self),
276        }
277    }
278
279    /// Create a new entry.
280    /// It will always return `TestStateEntryData::EntryExists(..)`.
281    fn create(&mut self, indexes: &[Index]) -> Rc<RefCell<TestStateEntryData>> {
282        match indexes.first() {
283            Some(idx) => {
284                self.children[*idx].get_or_insert(Box::new(Self::new())).create(&indexes[1..])
285            }
286            None => {
287                let new_data = Rc::new(RefCell::new(TestStateEntryData::new()));
288                let new_data_clone = Rc::clone(&new_data);
289                self.data = Some(new_data);
290                new_data_clone
291            }
292        }
293    }
294
295    fn delete_data(&mut self, indexes: &[Index]) -> Result<(), StateError> {
296        self.delete_prefix(indexes, true)
297    }
298
299    /// Delete nodes with the given prefix. If `exact`, then it only deletes the
300    /// data in the node with that specific key (prefix).
301    fn delete_prefix(&mut self, prefix: &[Index], exact: bool) -> Result<(), StateError> {
302        match prefix.first() {
303            Some(idx) => match &mut self.children[*idx] {
304                Some(child) => {
305                    let something_was_deleted = child.delete_prefix(&prefix[1..], exact);
306                    if child.is_empty() {
307                        self.children[*idx] = None;
308                    }
309                    something_was_deleted
310                }
311                None => {
312                    if exact {
313                        Err(StateError::EntryNotFound)
314                    } else {
315                        Err(StateError::SubtreeWithPrefixNotFound)
316                    }
317                }
318            },
319            None => {
320                // If `exact` and we found a non-leaf node, then do nothing and return false.
321                if exact && self.data.is_none() {
322                    // Make no changes and return false.
323                    return Ok(());
324                }
325
326                // If not `exact` delete the children, as we are deleting the whole prefix.
327                if !exact {
328                    self.children.iter_mut().for_each(|child| {
329                        *child = None;
330                    });
331                }
332
333                self.data = None;
334                Ok(())
335            }
336        }
337    }
338
339    /// Check whether a node is empty.
340    /// A node is considered empty when it has no data and no children.
341    fn is_empty(&self) -> bool { self.data.is_none() && self.children.iter().all(|x| x.is_none()) }
342
343    /// Make a deep clone of the node. Used for rollbacks.
344    fn clone_deep(&self) -> Self {
345        Self {
346            data:     self.data.as_ref().map(|d| Rc::new(RefCell::new(d.borrow().clone()))),
347            children: self
348                .children
349                .iter()
350                .map(|child| child.as_ref().map(|child| Box::new(child.clone_deep())))
351                .collect::<Vec<_>>()
352                .try_into()
353                .unwrap(), // This is safe because we know it has the right size and type.
354        }
355    }
356}
357
358#[cfg(test)]
359mod tests {
360    use crate::test_infrastructure::{trie::StateTrie, TestStateEntry};
361    use concordium_contracts_common::{to_bytes, Deserial, Read, Seek, SeekFrom, Write};
362
363    /// Create an entry and unwrap the result.
364    fn create_entry(trie: &mut StateTrie, key: &[u8]) -> TestStateEntry {
365        trie.create_entry(key).expect("Failed to create entry")
366    }
367
368    /// Delete an entry and unwrap the result.
369    fn delete_entry(trie: &mut StateTrie, entry: TestStateEntry) {
370        trie.delete_entry(entry).expect("Failed to delete entry")
371    }
372
373    #[test]
374    fn insert_get_test() {
375        let expected_value = "hello";
376        let key = [0, 1, 2];
377        let mut trie = StateTrie::new();
378
379        create_entry(&mut trie, &key)
380            .write_all(&to_bytes(&expected_value))
381            .expect("Writing to state failed.");
382
383        let mut entry = trie.lookup(&key).expect("Entry not found");
384        let actual_value = String::deserial(&mut entry).unwrap();
385        assert_eq!(&expected_value, &actual_value);
386    }
387
388    #[test]
389    fn delete_entry_test() {
390        let key1 = [0];
391        let key2 = [0, 0]; // A leaf, which is the child of the key1 node.
392        let mut trie = StateTrie::new();
393        create_entry(&mut trie, &key1);
394        create_entry(&mut trie, &key2);
395
396        // Both entries exist in the tree.
397        let entry1 = trie.lookup(&key1).expect("entry1 not found");
398        assert!(trie.lookup(&key2).is_some());
399
400        delete_entry(&mut trie, entry1); // Delete the data in the parent node.
401        assert!(trie.lookup(&key1).is_none());
402        assert!(trie.lookup(&key2).is_some()); // The child should still exist.
403    }
404
405    #[test]
406    fn delete_prefix_test() {
407        let key1 = [0];
408        let key2 = [0, 0];
409        let key3 = [0, 0, 0];
410
411        let mut trie = StateTrie::new();
412        create_entry(&mut trie, &key1);
413        create_entry(&mut trie, &key2);
414        create_entry(&mut trie, &key3);
415
416        assert!(trie.delete_prefix(&key2).is_ok());
417
418        assert!(trie.lookup(&key1).is_some());
419        assert!(trie.lookup(&key2).is_none());
420        assert!(trie.lookup(&key3).is_none());
421    }
422
423    #[test]
424    fn double_create_overwrites_data() {
425        let key = [];
426        let mut trie = StateTrie::new();
427        create_entry(&mut trie, &key)
428            .write_all(&to_bytes(&"hello"))
429            .expect("Writing to state failed");
430
431        // Creating again overwrites the old data.
432        let res = String::deserial(&mut create_entry(&mut trie, &key));
433
434        assert!(res.is_err())
435    }
436
437    #[test]
438    fn iterator_test() {
439        let mut trie = StateTrie::new();
440
441        create_entry(&mut trie, b"a").write_u8(42).unwrap();
442        create_entry(&mut trie, b"ab").write_u8(43).unwrap();
443        let mut entry_abd = create_entry(&mut trie, b"abd");
444        let mut entry_abdf = create_entry(&mut trie, b"abdf");
445        let mut entry_abdg = create_entry(&mut trie, b"abdg");
446        let mut entry_abe = create_entry(&mut trie, b"abe");
447        create_entry(&mut trie, b"ac").write_u8(44).unwrap();
448
449        entry_abd.write_u8(0).unwrap();
450        entry_abdf.write_u8(1).unwrap();
451        entry_abdg.write_u8(2).unwrap();
452        entry_abe.write_u8(3).unwrap();
453
454        // Get an iterator of the trie.
455        let mut iter = trie.iterator(b"ab").unwrap();
456        assert_eq!(u8::deserial(&mut iter.next().unwrap()), Ok(0));
457        assert_eq!(u8::deserial(&mut iter.next().unwrap()), Ok(1));
458        assert_eq!(u8::deserial(&mut iter.next().unwrap()), Ok(2));
459        assert_eq!(u8::deserial(&mut iter.next().unwrap()), Ok(3));
460        assert!(iter.next().is_none());
461
462        // Delete some entries.
463        trie.delete_iterator(iter);
464        assert!(trie.delete_entry(entry_abd).is_ok());
465        delete_entry(&mut trie, entry_abdf);
466        delete_entry(&mut trie, entry_abe);
467
468        // Only "abdg" is left.
469        let mut new_trie = trie.iterator(b"ab").unwrap();
470        assert_eq!(u8::deserial(&mut new_trie.next().unwrap()), Ok(2));
471        assert!(new_trie.next().is_none());
472    }
473
474    #[test]
475    fn index_conversion() {
476        let expected_key1 = [1, 2, 3, 4, 5, 6, 7];
477        let expected_key2 = [92, 255, 23, 5];
478        let index1 = super::to_indexes(&expected_key1);
479        let index2 = super::to_indexes(&expected_key2);
480        let actual_key1 = super::from_indexes(&index1);
481        let actual_key2 = super::from_indexes(&index2);
482        assert_eq!(expected_key1, &actual_key1[..]);
483        assert_eq!(expected_key2, &actual_key2[..]);
484    }
485
486    #[test]
487    fn write_to_deleted_entry_should_fail() {
488        let mut trie = StateTrie::new();
489        let mut entry = create_entry(&mut trie, b"ab");
490        assert!(entry.write_u8(1).is_ok());
491        trie.delete_prefix(&[]).unwrap();
492        assert!(entry.write_u8(1).is_err());
493    }
494
495    #[test]
496    fn seek_on_deleted_entry_should_fail() {
497        let mut trie = StateTrie::new();
498        let mut entry = create_entry(&mut trie, b"ab");
499        assert!(entry.write_u8(1).is_ok());
500        trie.delete_prefix(&[]).unwrap();
501        assert!(entry.seek(SeekFrom::Start(0)).is_err());
502    }
503
504    #[test]
505    fn read_from_deleted_entry_should_fail() {
506        let mut trie = StateTrie::new();
507        let mut entry = create_entry(&mut trie, b"ab");
508        assert!(entry.write_u8(1).is_ok());
509        trie.delete_prefix(&[]).unwrap();
510        // Manually reset offset, since Seek will also fail.
511        entry.cursor.offset = 0;
512        assert!(entry.read_u8().is_err());
513    }
514
515    #[test]
516    fn read_from_deleted_aliased_entry_should_fail() {
517        let mut trie = StateTrie::new();
518        let mut entry = create_entry(&mut trie, b"ab");
519        let mut alias_entry = create_entry(&mut trie, b"ab");
520        assert!(entry.write_u8(1).is_ok());
521        trie.delete_prefix(&[]).unwrap();
522        assert!(alias_entry.read_u8().is_err());
523    }
524
525    #[test]
526    fn test_deep_clone() {
527        let mut trie = StateTrie::new();
528
529        // Create two entries
530        let mut e1 = create_entry(&mut trie, b"ab");
531        let mut e2 = create_entry(&mut trie, b"ac");
532
533        // Insert data in both
534        e1.write_u32(10001).unwrap();
535        e2.write_u64(10002).unwrap();
536
537        // Get iterator
538        let _iterator_1 = trie.iterator(b"a");
539
540        // Save stats
541        let next_entry_id = trie.next_entry_id.get();
542
543        // Clone deeply
544        let trie_clone = trie.clone_deep();
545
546        // ---------------
547        // Modify original
548        e1.write_u32(20001).unwrap(); // Also moves cursor
549        e2.write_u32(20002).unwrap(); // Also moves cursor
550        let _e3 = create_entry(&mut trie, b"qq");
551        let _iterator_2 = trie.iterator(&[]);
552        // ---------------
553
554        // Ensure that clone matches stats
555        assert_eq!(trie_clone.entry_map.borrow().len(), 4); // Only has [e1, e2, e1, e2] (the last two from the iterator)
556        assert_eq!(trie_clone.iterator_counts.borrow().len(), 1); // Only has [iterator_1]
557        assert_eq!(trie_clone.next_entry_id.get(), next_entry_id);
558
559        let mut e1_clone = trie_clone.lookup(b"ab").expect("Entry 'ab' is missing.");
560        let mut e2_clone = trie_clone.lookup(b"ac").expect("Entry 'ac' is missing.");
561
562        assert_eq!(Ok(10001), e1_clone.read_u32());
563        assert_eq!(Ok(10002), e2_clone.read_u64());
564
565        assert!(trie_clone.lookup(b"qq").is_none());
566    }
567}
concordium_std/test_infrastructure/trie.rs

concordium_std/test_infrastructure/
trie.rs
