ic_storage/adapter/
tree.rs

1//! Owned AVL ordered storage. Small mutations stage affected nodes and allocator
2//! metadata, reserve capacity once, then publish the protected root last.
3//! Power-of-two reusable blocks bound churn without whole-map heap snapshots.
4use super::{
5    Memory, Storable,
6    error::{CollectionError as Error, Result, ensure},
7    record,
8};
9use serde::{Deserialize, Serialize};
10use std::{
11    borrow::Cow,
12    cell::{Cell, RefCell},
13    cmp::Ordering,
14    collections::BTreeMap,
15    marker::PhantomData,
16    ops::{Bound, RangeBounds},
17};
18
19pub const MAGIC: [u8; 8] = *b"ICMAP001";
20const NODE: [u8; 8] = *b"ICNODE01";
21const FREE: [u8; 8] = *b"ICFREE01";
22const BINS: usize = 32;
23const MAX_DEPTH: usize = 96;
24
25#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
26struct Pointer {
27    offset: u64,
28    order: u8,
29}
30impl Pointer {
31    fn capacity(self) -> u64 {
32        1_u64 << self.order
33    }
34    fn check(self, cursor: u64) -> Result<()> {
35        ensure(
36            (8..32).contains(&self.order),
37            Error::Corrupt("tree block order"),
38        )?;
39        ensure(
40            self.offset >= record::HEADER_BYTES
41                && self
42                    .offset
43                    .checked_add(self.capacity())
44                    .is_some_and(|end| end <= cursor),
45            Error::Corrupt("tree block bounds"),
46        )?;
47        Ok(())
48    }
49}
50
51#[derive(Clone, Serialize, Deserialize)]
52struct Root {
53    top: Option<Pointer>,
54    length: u64,
55    cursor: u64,
56    free: Vec<Option<Pointer>>,
57    key_bound: Option<(u32, bool)>,
58    value_bound: Option<(u32, bool)>,
59}
60
61#[derive(Clone, Serialize, Deserialize)]
62struct Node {
63    left: Option<Pointer>,
64    right: Option<Pointer>,
65    height: u8,
66    key: Vec<u8>,
67    value: Vec<u8>,
68}
69
70pub struct Tree<K, V, M: Memory> {
71    pub memory: M,
72    root: RefCell<Root>,
73    poisoned: Cell<bool>,
74    types: PhantomData<(K, V)>,
75}
76
77impl<K: Storable + Ord + Clone, V: Storable, M: Memory> Tree<K, V, M> {
78    pub fn reset(memory: M) -> Result<Self> {
79        let root = Root {
80            top: None,
81            length: 0,
82            cursor: record::HEADER_BYTES,
83            free: vec![None; BINS],
84            key_bound: record::bound::<K>(),
85            value_bound: record::bound::<V>(),
86        };
87        let bytes = record::encode(MAGIC, &root)?;
88        ensure(
89            bytes.len() as u64 <= record::HEADER_BYTES,
90            Error::Corrupt("map invariant"),
91        )?;
92        record::reserve(&memory, record::HEADER_BYTES)?;
93        memory.write(0, &bytes);
94        Ok(Self {
95            memory,
96            root: RefCell::new(root),
97            poisoned: Cell::new(false),
98            types: PhantomData,
99        })
100    }
101
102    pub fn load(memory: M) -> Result<Self> {
103        let root: Root = record::read(&memory, 0, record::HEADER_BYTES, MAGIC)?;
104        ensure(
105            root.key_bound == record::bound::<K>(),
106            Error::Corrupt("map key bound changed"),
107        )?;
108        ensure(
109            root.value_bound == record::bound::<V>(),
110            Error::Corrupt("map value bound changed"),
111        )?;
112        ensure(
113            root.free.len() == BINS
114                && root.cursor >= record::HEADER_BYTES
115                && root.cursor <= memory.size() * crate::WASM_PAGE_SIZE_BYTES,
116            Error::Corrupt("map allocator bounds"),
117        )?;
118        ensure(
119            root.top.is_none() == (root.length == 0),
120            Error::Corrupt("map root/count disagree"),
121        )?;
122        if let Some(pointer) = root.top {
123            pointer.check(root.cursor)?;
124        }
125        for (order, pointer) in root.free.iter().enumerate() {
126            if let Some(pointer) = pointer {
127                pointer.check(root.cursor)?;
128                ensure(
129                    usize::from(pointer.order) == order,
130                    Error::Corrupt("free block order"),
131                )?;
132            }
133        }
134        Ok(Self {
135            memory,
136            root: RefCell::new(root),
137            poisoned: Cell::new(false),
138            types: PhantomData,
139        })
140    }
141
142    fn healthy(&self) -> Result<()> {
143        ensure(!self.poisoned.get(), Error::Poisoned)
144    }
145    pub fn len(&self) -> u64 {
146        self.healthy().expect("map owner is poisoned");
147        self.root.borrow().length
148    }
149    fn node(&self, pointer: Pointer) -> Result<Node> {
150        pointer.check(self.root.borrow().cursor)?;
151        let node: Node = record::read(&self.memory, pointer.offset, pointer.capacity(), NODE)?;
152        validate_node::<K, V>(&node, self.root.borrow().cursor)?;
153        Ok(node)
154    }
155
156    pub fn get(&self, key: &K) -> Result<Option<V>> {
157        self.healthy()?;
158        let mut current = self.root.borrow().top;
159        for _ in 0..MAX_DEPTH {
160            let Some(pointer) = current else {
161                return Ok(None);
162            };
163            let node = self.node(pointer)?;
164            match key.cmp(&decode::<K>(node.key)?) {
165                Ordering::Less => current = node.left,
166                Ordering::Greater => current = node.right,
167                Ordering::Equal => return decode(node.value).map(Some),
168            }
169        }
170        Err(Error::Corrupt("map path exceeds height bound"))
171    }
172
173    pub fn insert(&mut self, key: K, value: V) -> Result<Option<V>> {
174        self.healthy()?;
175        let key_bytes = key.to_bytes_checked()?.into_owned();
176        let value_bytes = value.into_bytes_checked()?;
177        let mut transaction = Transaction::<K, V, M>::new(self);
178        let (root, previous) =
179            transaction.insert(transaction.root.top, &key, &key_bytes, &value_bytes, 0)?;
180        transaction.root.top = Some(root);
181        drop(key);
182        let previous = previous.map(decode).transpose()?;
183        transaction.commit()?;
184        Ok(previous)
185    }
186
187    pub fn remove(&mut self, key: &K) -> Result<Option<V>> {
188        self.healthy()?;
189        let mut transaction = Transaction::<K, V, M>::new(self);
190        let (root, previous) = transaction.remove(transaction.root.top, key, 0)?;
191        transaction.root.top = root;
192        let previous = previous.map(decode).transpose()?;
193        if previous.is_some() {
194            transaction.commit()?;
195        }
196        Ok(previous)
197    }
198
199    pub fn clear(&mut self) -> Result<()> {
200        self.healthy()?;
201        let mut transaction = Transaction::<K, V, M>::new(self);
202        transaction.root.top = None;
203        transaction.root.length = 0;
204        transaction.root.cursor = record::HEADER_BYTES;
205        transaction.root.free.fill(None);
206        transaction.commit()
207    }
208
209    fn seek_path(&self, bound: &Bound<K>, forward: bool) -> Result<Vec<Pointer>> {
210        self.healthy()?;
211        let mut current = self.root.borrow().top;
212        let mut path = Vec::new();
213        for _ in 0..MAX_DEPTH {
214            let Some(pointer) = current else {
215                return Ok(path);
216            };
217            let node = self.node(pointer)?;
218            let key = decode::<K>(node.key)?;
219            let accepted = match bound {
220                Bound::Unbounded => true,
221                Bound::Included(limit) => {
222                    if forward {
223                        key >= *limit
224                    } else {
225                        key <= *limit
226                    }
227                }
228                Bound::Excluded(limit) => {
229                    if forward {
230                        key > *limit
231                    } else {
232                        key < *limit
233                    }
234                }
235            };
236            if accepted {
237                current = if forward { node.left } else { node.right };
238                path.push(pointer);
239            } else {
240                current = if forward { node.right } else { node.left };
241            }
242        }
243        Err(Error::Corrupt("map path exceeds height bound"))
244    }
245
246    fn advance_path(
247        &self,
248        path: &mut Vec<Pointer>,
249        forward: bool,
250    ) -> Result<Option<MapEntry<K, V>>> {
251        self.healthy()?;
252        let Some(pointer) = path.pop() else {
253            return Ok(None);
254        };
255        let node = self.node(pointer)?;
256        let mut child = if forward { node.right } else { node.left };
257        let mut depth = 0;
258        while let Some(pointer) = child {
259            ensure(
260                depth < MAX_DEPTH && path.len() < MAX_DEPTH,
261                Error::Corrupt("map cursor exceeds height bound"),
262            )?;
263            depth += 1;
264            path.push(pointer);
265            let child_node = self.node(pointer)?;
266            child = if forward {
267                child_node.left
268            } else {
269                child_node.right
270            };
271        }
272        Ok(Some(MapEntry {
273            key: decode(node.key)?,
274            bytes: node.value,
275            value: PhantomData,
276        }))
277    }
278
279    pub fn range(&self, bounds: impl RangeBounds<K>) -> TreeIter<'_, K, V, M> {
280        let lower = bounds.start_bound().cloned();
281        let upper = bounds.end_bound().cloned();
282        let invalid = match (&lower, &upper) {
283            (Bound::Included(a) | Bound::Excluded(a), Bound::Included(b) | Bound::Excluded(b)) => {
284                a > b
285                    || (a == b
286                        && matches!((&lower, &upper), (Bound::Excluded(_), Bound::Excluded(_))))
287            }
288            _ => false,
289        };
290        TreeIter {
291            tree: self,
292            lower,
293            upper,
294            front: None,
295            back: None,
296            exhausted: false,
297            invalid,
298        }
299    }
300}
301
302fn decode<T: Storable>(bytes: Vec<u8>) -> Result<T> {
303    T::check_bounds(&bytes)?;
304    Ok(T::from_bytes(Cow::Owned(bytes))?)
305}
306fn validate_node<K: Storable, V: Storable>(node: &Node, cursor: u64) -> Result<()> {
307    ensure(
308        (1..=96).contains(&node.height),
309        Error::Corrupt("map node height"),
310    )?;
311    K::check_bounds(&node.key)?;
312    V::check_bounds(&node.value)?;
313    for pointer in [node.left, node.right].into_iter().flatten() {
314        pointer.check(cursor)?;
315    }
316    Ok(())
317}
318
319struct Transaction<'a, K, V, M: Memory> {
320    tree: &'a Tree<K, V, M>,
321    root: Root,
322    writes: BTreeMap<u64, Vec<u8>>,
323}
324
325impl<'a, K: Storable + Ord + Clone, V: Storable, M: Memory> Transaction<'a, K, V, M> {
326    fn new(tree: &'a Tree<K, V, M>) -> Self {
327        Self {
328            tree,
329            root: tree.root.borrow().clone(),
330            writes: BTreeMap::new(),
331        }
332    }
333    fn read<T: serde::de::DeserializeOwned>(&self, pointer: Pointer, magic: [u8; 8]) -> Result<T> {
334        pointer.check(self.root.cursor)?;
335        if let Some(bytes) = self.writes.get(&pointer.offset) {
336            record::decode(magic, bytes)
337        } else {
338            record::read(&self.tree.memory, pointer.offset, pointer.capacity(), magic)
339        }
340    }
341    fn node(&self, pointer: Pointer) -> Result<Node> {
342        let node: Node = self.read(pointer, NODE)?;
343        validate_node::<K, V>(&node, self.root.cursor)?;
344        Ok(node)
345    }
346    fn height(&self, pointer: Option<Pointer>) -> Result<u8> {
347        pointer.map_or(Ok(0), |pointer| Ok(self.node(pointer)?.height))
348    }
349    fn free(&mut self, pointer: Pointer) -> Result<()> {
350        let index = usize::from(pointer.order);
351        self.writes.insert(
352            pointer.offset,
353            record::encode(FREE, &self.root.free[index])?,
354        );
355        self.root.free[index] = Some(pointer);
356        Ok(())
357    }
358    fn allocate(&mut self, length: usize) -> Result<Pointer> {
359        let capacity = length
360            .max(256)
361            .checked_next_power_of_two()
362            .ok_or(Error::Budget)?;
363        let order = usize::try_from(capacity.trailing_zeros()).unwrap();
364        ensure(order < BINS, Error::Corrupt("map record too large"))?;
365        for bin in order..BINS {
366            if let Some(pointer) = self.root.free[bin] {
367                let next: Option<Pointer> = self.read(pointer, FREE)?;
368                if let Some(next) = next {
369                    next.check(self.root.cursor)?;
370                    ensure(
371                        next.order == pointer.order,
372                        Error::Corrupt("free block order"),
373                    )?;
374                    ensure(next != pointer, Error::Corrupt("map free-list cycle"))?;
375                }
376                self.root.free[bin] = next;
377                return Ok(pointer);
378            }
379        }
380        let pointer = Pointer {
381            offset: self.root.cursor,
382            order: u8::try_from(order).unwrap(),
383        };
384        self.root.cursor = self
385            .root
386            .cursor
387            .checked_add(pointer.capacity())
388            .ok_or(Error::Budget)?;
389        Ok(pointer)
390    }
391    fn save(&mut self, previous: Option<Pointer>, node: &Node) -> Result<Pointer> {
392        let bytes = record::encode(NODE, node)?;
393        let pointer = if let Some(pointer) =
394            previous.filter(|pointer| pointer.capacity() >= bytes.len() as u64)
395        {
396            pointer
397        } else {
398            if let Some(pointer) = previous {
399                self.free(pointer)?;
400            }
401            self.allocate(bytes.len())?
402        };
403        self.writes.insert(pointer.offset, bytes);
404        Ok(pointer)
405    }
406    fn rotate_left(&mut self, pointer: Pointer, mut node: Node) -> Result<Pointer> {
407        let right_pointer = node.right.ok_or(Error::Corrupt("map rotation child"))?;
408        let mut right = self.node(right_pointer)?;
409        node.right = right.left;
410        node.height = 1 + self.height(node.left)?.max(self.height(node.right)?);
411        right.left = Some(self.save(Some(pointer), &node)?);
412        right.height = 1 + self.height(right.left)?.max(self.height(right.right)?);
413        self.save(Some(right_pointer), &right)
414    }
415    fn rotate_right(&mut self, pointer: Pointer, mut node: Node) -> Result<Pointer> {
416        let left_pointer = node.left.ok_or(Error::Corrupt("map rotation child"))?;
417        let mut left = self.node(left_pointer)?;
418        node.left = left.right;
419        node.height = 1 + self.height(node.left)?.max(self.height(node.right)?);
420        left.right = Some(self.save(Some(pointer), &node)?);
421        left.height = 1 + self.height(left.left)?.max(self.height(left.right)?);
422        self.save(Some(left_pointer), &left)
423    }
424    fn balance(&mut self, pointer: Pointer, mut node: Node) -> Result<Pointer> {
425        let left = self.height(node.left)?;
426        let right = self.height(node.right)?;
427        node.height = 1 + left.max(right);
428        if i16::from(left) - i16::from(right) > 1 {
429            let child = node.left.unwrap();
430            let value = self.node(child)?;
431            if self.height(value.right)? > self.height(value.left)? {
432                node.left = Some(self.rotate_left(child, value)?);
433            }
434            self.rotate_right(pointer, node)
435        } else if i16::from(right) - i16::from(left) > 1 {
436            let child = node.right.unwrap();
437            let value = self.node(child)?;
438            if self.height(value.left)? > self.height(value.right)? {
439                node.right = Some(self.rotate_right(child, value)?);
440            }
441            self.rotate_left(pointer, node)
442        } else {
443            self.save(Some(pointer), &node)
444        }
445    }
446    fn insert(
447        &mut self,
448        pointer: Option<Pointer>,
449        key: &K,
450        key_bytes: &[u8],
451        value: &[u8],
452        depth: usize,
453    ) -> Result<(Pointer, Option<Vec<u8>>)> {
454        ensure(
455            depth < MAX_DEPTH,
456            Error::Corrupt("map path exceeds height bound"),
457        )?;
458        let Some(pointer) = pointer else {
459            let node = Node {
460                left: None,
461                right: None,
462                height: 1,
463                key: key_bytes.to_vec(),
464                value: value.to_vec(),
465            };
466            self.root.length = self.root.length.checked_add(1).ok_or(Error::Budget)?;
467            return Ok((self.save(None, &node)?, None));
468        };
469        let mut node = self.node(pointer)?;
470        let previous = match key.cmp(&decode::<K>(node.key.clone())?) {
471            Ordering::Less => {
472                let (child, old) = self.insert(node.left, key, key_bytes, value, depth + 1)?;
473                if old.is_some() && node.left == Some(child) {
474                    return Ok((pointer, old));
475                }
476                node.left = Some(child);
477                old
478            }
479            Ordering::Greater => {
480                let (child, old) = self.insert(node.right, key, key_bytes, value, depth + 1)?;
481                if old.is_some() && node.right == Some(child) {
482                    return Ok((pointer, old));
483                }
484                node.right = Some(child);
485                old
486            }
487            Ordering::Equal => Some(std::mem::replace(&mut node.value, value.to_vec())),
488        };
489        // Replacement can relocate a record, but never changes subtree height.
490        let pointer = if previous.is_some() {
491            self.save(Some(pointer), &node)?
492        } else {
493            self.balance(pointer, node)?
494        };
495        Ok((pointer, previous))
496    }
497    fn remove(
498        &mut self,
499        pointer: Option<Pointer>,
500        key: &K,
501        depth: usize,
502    ) -> Result<(Option<Pointer>, Option<Vec<u8>>)> {
503        ensure(
504            depth < MAX_DEPTH,
505            Error::Corrupt("map path exceeds height bound"),
506        )?;
507        let Some(pointer) = pointer else {
508            return Ok((None, None));
509        };
510        let mut node = self.node(pointer)?;
511        let previous = match key.cmp(&decode::<K>(node.key.clone())?) {
512            Ordering::Less => {
513                let (child, old) = self.remove(node.left, key, depth + 1)?;
514                node.left = child;
515                old
516            }
517            Ordering::Greater => {
518                let (child, old) = self.remove(node.right, key, depth + 1)?;
519                node.right = child;
520                old
521            }
522            Ordering::Equal => {
523                if node.left.is_none() || node.right.is_none() {
524                    self.free(pointer)?;
525                    self.root.length = self
526                        .root
527                        .length
528                        .checked_sub(1)
529                        .ok_or(Error::Corrupt("map length"))?;
530                    return Ok((node.left.or(node.right), Some(node.value)));
531                }
532                let (right, successor) = self.extract_min(node.right.unwrap(), depth + 1)?;
533                node.right = right;
534                node.key = successor.key;
535                Some(std::mem::replace(&mut node.value, successor.value))
536            }
537        };
538        if previous.is_none() {
539            Ok((Some(pointer), None))
540        } else {
541            Ok((Some(self.balance(pointer, node)?), previous))
542        }
543    }
544    fn extract_min(&mut self, pointer: Pointer, depth: usize) -> Result<(Option<Pointer>, Node)> {
545        ensure(
546            depth < MAX_DEPTH,
547            Error::Corrupt("map path exceeds height bound"),
548        )?;
549        let mut node = self.node(pointer)?;
550        let Some(left) = node.left else {
551            self.free(pointer)?;
552            self.root.length = self
553                .root
554                .length
555                .checked_sub(1)
556                .ok_or(Error::Corrupt("map length"))?;
557            return Ok((node.right, node));
558        };
559        let (child, minimum) = self.extract_min(left, depth + 1)?;
560        node.left = child;
561        Ok((Some(self.balance(pointer, node)?), minimum))
562    }
563    fn commit(self) -> Result<()> {
564        let header = record::encode(MAGIC, &self.root)?;
565        ensure(
566            header.len() as u64 <= record::HEADER_BYTES,
567            Error::Corrupt("map invariant"),
568        )?;
569        record::reserve(&self.tree.memory, self.root.cursor)?;
570        self.tree.poisoned.set(true);
571        for (offset, bytes) in self.writes {
572            self.tree.memory.write(offset, &bytes);
573        }
574        self.tree.memory.write(0, &header);
575        *self.tree.root.borrow_mut() = self.root;
576        self.tree.poisoned.set(false);
577        Ok(())
578    }
579}
580
581///
582/// `MapEntry`
583///
584/// An owned key and lazily decoded value from an ordered scan.
585///
586pub struct MapEntry<K, V> {
587    key: K,
588    bytes: Vec<u8>,
589    value: PhantomData<V>,
590}
591impl<K, V: Storable> MapEntry<K, V> {
592    /// Borrow the decoded key.
593    pub const fn key(&self) -> &K {
594        &self.key
595    }
596    /// Decode the value using its declared codec.
597    ///
598    /// # Errors
599    /// Returns the application codec or declared-bound error for this value.
600    pub fn value(&self) -> Result<V> {
601        decode(self.bytes.clone())
602    }
603    /// Consume the entry and decode its value.
604    ///
605    /// # Errors
606    /// Returns the application codec or declared-bound error for this value.
607    pub fn into_pair(self) -> Result<(K, V)> {
608        Ok((self.key, decode(self.bytes)?))
609    }
610    /// Consume the entry without decoding its value.
611    pub fn into_key(self) -> K {
612        self.key
613    }
614    /// Consume the entry and decode only its value.
615    ///
616    /// # Errors
617    /// Returns the application codec or declared-bound error for this value.
618    pub fn into_value(self) -> Result<V> {
619        decode(self.bytes)
620    }
621}
622
623///
624/// `TreeIter`
625///
626/// Lazy bidirectional traversal with bounded pointer stacks. Ancestor payloads
627/// are not retained; values are decoded only when requested by the caller.
628///
629pub struct TreeIter<'a, K, V, M: Memory> {
630    tree: &'a Tree<K, V, M>,
631    lower: Bound<K>,
632    upper: Bound<K>,
633    front: Option<Vec<Pointer>>,
634    back: Option<Vec<Pointer>>,
635    exhausted: bool,
636    invalid: bool,
637}
638impl<K: Storable + Ord + Clone, V: Storable, M: Memory> TreeIter<'_, K, V, M> {
639    fn step(&mut self, forward: bool) -> Result<Option<MapEntry<K, V>>> {
640        ensure(!self.invalid, Error::Bounds)?;
641        let (path, bound) = if forward {
642            (&mut self.front, &self.lower)
643        } else {
644            (&mut self.back, &self.upper)
645        };
646        if path.is_none() {
647            *path = Some(self.tree.seek_path(bound, forward)?);
648        }
649        let entry = self.tree.advance_path(path.as_mut().unwrap(), forward)?;
650        if let Some(entry) = &entry {
651            let after = match &self.lower {
652                Bound::Unbounded => true,
653                Bound::Included(k) => entry.key >= *k,
654                Bound::Excluded(k) => entry.key > *k,
655            };
656            let before = match &self.upper {
657                Bound::Unbounded => true,
658                Bound::Included(k) => entry.key <= *k,
659                Bound::Excluded(k) => entry.key < *k,
660            };
661            ensure(
662                if forward { after } else { before },
663                Error::Corrupt("map cursor key order"),
664            )?;
665            if !(after && before) {
666                return Ok(None);
667            }
668            if forward {
669                self.lower = Bound::Excluded(entry.key.clone());
670            } else {
671                self.upper = Bound::Excluded(entry.key.clone());
672            }
673        }
674        Ok(entry)
675    }
676    fn advance(&mut self, forward: bool) -> Option<Result<MapEntry<K, V>>> {
677        if self.exhausted {
678            return None;
679        }
680        match self.step(forward) {
681            Ok(Some(entry)) => Some(Ok(entry)),
682            Ok(None) => {
683                self.exhausted = true;
684                None
685            }
686            Err(error) => {
687                self.exhausted = true;
688                Some(Err(error))
689            }
690        }
691    }
692}
693impl<K: Storable + Ord + Clone, V: Storable, M: Memory> Iterator for TreeIter<'_, K, V, M> {
694    type Item = Result<MapEntry<K, V>>;
695    fn next(&mut self) -> Option<Self::Item> {
696        self.advance(true)
697    }
698}
699impl<K: Storable + Ord + Clone, V: Storable, M: Memory> DoubleEndedIterator
700    for TreeIter<'_, K, V, M>
701{
702    fn next_back(&mut self) -> Option<Self::Item> {
703        self.advance(false)
704    }
705}
706impl<K: Storable + Ord + Clone, V: Storable, M: Memory> std::iter::FusedIterator
707    for TreeIter<'_, K, V, M>
708{
709}
ic_storage/adapter/tree.rs

ic_storage/adapter/
tree.rs
