ic_storage/adapter/
tree.rs

1//! Owned degree-six B-tree. The multi-entry search, split/borrow/merge design and
2//! selective payload access follow the established B-tree techniques used by
3//! ic-stable-structures 0.7.2 (see docs/storage/upstreams.md). This is a new checked
4//! format, not an implementation of its durable layout or compatibility reader.
5//! Routing records contain small inline values; larger values have separate
6//! checksummed records. Mutations stage bytes and reserve before publication.
7use super::{
8    Memory, Storable,
9    error::{CollectionError as Error, Result, ensure},
10    record,
11};
12use serde::{Deserialize, Serialize};
13use std::{
14    borrow::Cow,
15    cell::{Cell, RefCell},
16    collections::BTreeMap,
17    marker::PhantomData,
18    ops::{Bound, RangeBounds},
19    rc::Rc,
20};
21mod read;
22use read::{NodeCache, ReadNode};
23
24pub const MAGIC: [u8; 8] = *b"ICMAP002";
25const NODE: [u8; 8] = *b"ICBTND01";
26const VALUE: [u8; 8] = *b"ICBTVAL1";
27const FREE: [u8; 8] = *b"ICFREE01";
28const BINS: usize = 32;
29const MAX_DEPTH: usize = 32;
30const DEGREE: usize = 6;
31const MAX_ENTRIES: usize = DEGREE * 2 - 1;
32const MIN_ENTRIES: usize = DEGREE - 1;
33const INLINE: usize = 32;
34
35#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
36struct Pointer {
37    offset: u64,
38    order: u8,
39}
40impl Pointer {
41    fn capacity(self) -> u64 {
42        1_u64 << self.order
43    }
44    fn check(self, cursor: u64) -> Result<()> {
45        ensure(
46            (8..32).contains(&self.order),
47            Error::Corrupt("tree block order"),
48        )?;
49        ensure(
50            self.offset >= record::HEADER_BYTES
51                && self
52                    .offset
53                    .checked_add(self.capacity())
54                    .is_some_and(|end| end <= cursor),
55            Error::Corrupt("tree block bounds"),
56        )
57    }
58    fn put(self, bytes: &mut Vec<u8>) {
59        bytes.extend_from_slice(&self.offset.to_le_bytes());
60        bytes.push(self.order);
61    }
62}
63#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
64struct Root {
65    top: Option<Pointer>,
66    length: u64,
67    cursor: u64,
68    free: Vec<Option<Pointer>>,
69    key_bound: Option<(u32, bool)>,
70    value_bound: Option<(u32, bool)>,
71}
72#[derive(Clone)]
73enum Value {
74    Inline(Vec<u8>),
75    External(Pointer),
76}
77#[derive(Clone)]
78struct Entry {
79    key: Vec<u8>,
80    value: Value,
81}
82#[derive(Clone, Default)]
83struct Node {
84    entries: Vec<Entry>,
85    children: Vec<Pointer>,
86}
87impl Node {
88    fn encode(&self) -> Result<Vec<u8>> {
89        ensure(
90            self.entries.len() <= MAX_ENTRIES,
91            Error::Corrupt("B-tree node capacity"),
92        )?;
93        ensure(
94            self.children.is_empty() || self.children.len() == self.entries.len() + 1,
95            Error::Corrupt("B-tree child count"),
96        )?;
97        let capacity =
98            self.entries
99                .iter()
100                .try_fold(2 + self.children.len() * 9, |size, entry| {
101                    let value_size = match &entry.value {
102                        Value::Inline(v) => v.len(),
103                        Value::External(_) => 9,
104                    };
105                    size.checked_add(5)
106                        .and_then(|n| n.checked_add(entry.key.len()))
107                        .and_then(|n| n.checked_add(value_size))
108                        .ok_or(Error::Budget)
109                })?;
110        ensure(capacity <= record::MAX_RECORD, Error::Budget)?;
111        let mut bytes = Vec::with_capacity(capacity);
112        bytes.extend_from_slice(&[
113            u8::try_from(self.entries.len()).unwrap(),
114            u8::from(!self.children.is_empty()),
115        ]);
116        for p in &self.children {
117            p.put(&mut bytes);
118        }
119        for entry in &self.entries {
120            bytes.extend_from_slice(
121                &u32::try_from(entry.key.len())
122                    .map_err(|_| Error::Budget)?
123                    .to_le_bytes(),
124            );
125            bytes.extend_from_slice(&entry.key);
126            match &entry.value {
127                Value::Inline(value) => {
128                    bytes.push(u8::try_from(value.len()).unwrap());
129                    bytes.extend_from_slice(value);
130                }
131                Value::External(p) => {
132                    bytes.push(255);
133                    p.put(&mut bytes);
134                }
135            }
136        }
137        record::encode_payload(NODE, &bytes)
138    }
139}
140struct Reader<'a>(&'a [u8]);
141impl<'a> Reader<'a> {
142    fn take(&mut self, n: usize) -> Result<&'a [u8]> {
143        ensure(n <= self.0.len(), Error::Corrupt("truncated B-tree node"))?;
144        let (a, b) = self.0.split_at(n);
145        self.0 = b;
146        Ok(a)
147    }
148    fn byte(&mut self) -> Result<u8> {
149        Ok(self.take(1)?[0])
150    }
151    fn pointer(&mut self, cursor: u64) -> Result<Pointer> {
152        let offset = u64::from_le_bytes(self.take(8)?.try_into().unwrap());
153        let p = Pointer {
154            offset,
155            order: self.byte()?,
156        };
157        p.check(cursor)?;
158        Ok(p)
159    }
160}
161
162pub struct Tree<K, V, M: Memory> {
163    pub memory: M,
164    root: RefCell<Root>,
165    poisoned: Cell<bool>,
166    cache: RefCell<NodeCache>,
167    types: PhantomData<(K, V)>,
168}
169impl<K: Storable + Ord + Clone, V: Storable, M: Memory> Tree<K, V, M> {
170    pub fn reset(memory: M) -> Result<Self> {
171        let root = Root {
172            top: None,
173            length: 0,
174            cursor: record::HEADER_BYTES,
175            free: vec![None; BINS],
176            key_bound: record::bound::<K>(),
177            value_bound: record::bound::<V>(),
178        };
179        let bytes = record::encode(MAGIC, &root)?;
180        ensure(bytes.len() as u64 <= record::HEADER_BYTES, Error::Budget)?;
181        record::reserve(&memory, record::HEADER_BYTES)?;
182        memory.write(0, &bytes);
183        Ok(Self {
184            memory,
185            root: RefCell::new(root),
186            poisoned: Cell::new(false),
187            cache: RefCell::default(),
188            types: PhantomData,
189        })
190    }
191    pub fn load(memory: M) -> Result<Self> {
192        let root: Root = record::read(&memory, 0, record::HEADER_BYTES, MAGIC)?;
193        ensure(
194            root.key_bound == record::bound::<K>() && root.value_bound == record::bound::<V>(),
195            Error::Corrupt("map codec bound changed"),
196        )?;
197        ensure(
198            root.free.len() == BINS
199                && root.cursor >= record::HEADER_BYTES
200                && root.cursor
201                    <= memory
202                        .size()
203                        .checked_mul(crate::WASM_PAGE_SIZE_BYTES)
204                        .ok_or(Error::Budget)?,
205            Error::Corrupt("map allocator bounds"),
206        )?;
207        ensure(
208            root.top.is_none() == (root.length == 0),
209            Error::Corrupt("map root/count disagree"),
210        )?;
211        if let Some(p) = root.top {
212            p.check(root.cursor)?;
213        }
214        for (order, p) in root.free.iter().enumerate() {
215            if let Some(p) = p {
216                p.check(root.cursor)?;
217                ensure(
218                    usize::from(p.order) == order,
219                    Error::Corrupt("free block order"),
220                )?;
221            }
222        }
223        Ok(Self {
224            memory,
225            root: RefCell::new(root),
226            poisoned: Cell::new(false),
227            cache: RefCell::default(),
228            types: PhantomData,
229        })
230    }
231    fn healthy(&self) -> Result<()> {
232        ensure(!self.poisoned.get(), Error::Poisoned)?;
233        // A different allocation can poison the shared governed backing. This
234        // also covers iterator entries already held in a routing frame.
235        let capacity = self
236            .memory
237            .size()
238            .checked_mul(crate::WASM_PAGE_SIZE_BYTES)
239            .ok_or(Error::Budget)?;
240        ensure(
241            self.root.borrow().cursor <= capacity,
242            Error::Corrupt("map allocator bounds"),
243        )
244    }
245    pub fn len(&self) -> u64 {
246        self.healthy().expect("map owner is poisoned");
247        self.root.borrow().length
248    }
249    #[cfg(test)]
250    fn node(&self, p: Pointer) -> Result<Node> {
251        Ok(self.read_node(p, MAX_DEPTH)?.owned())
252    }
253    pub fn set_cache_budget(&mut self, bytes: usize) {
254        self.cache.get_mut().configure(bytes);
255    }
256    fn read_node(&self, p: Pointer, depth: usize) -> Result<Rc<ReadNode>> {
257        let cursor = self.root.borrow().cursor;
258        p.check(cursor)?;
259        if let Some(node) = self.cache.borrow_mut().get(p, depth) {
260            return Ok(node);
261        }
262        let node = Rc::new(ReadNode::decode::<K, V>(
263            record::read_bytes(&self.memory, p.offset, p.capacity(), NODE)?,
264            cursor,
265        )?);
266        self.cache.borrow_mut().put(p, Rc::clone(&node), depth);
267        Ok(node)
268    }
269    fn value(&self, v: Value) -> Result<Vec<u8>> {
270        let bytes = match v {
271            Value::Inline(bytes) => bytes,
272            Value::External(p) => {
273                p.check(self.root.borrow().cursor)?;
274                let bytes = record::read_bytes(&self.memory, p.offset, p.capacity(), VALUE)?;
275                record::decode_payload(VALUE, &bytes)?.to_vec()
276            }
277        };
278        V::check_bounds(&bytes)?;
279        Ok(bytes)
280    }
281    fn find(&self, key: &K) -> Result<Option<Value>> {
282        self.healthy()?;
283        let mut p = self.root.borrow().top;
284        for depth in 0..MAX_DEPTH {
285            let Some(current) = p else { return Ok(None) };
286            let node = self.read_node(current, depth)?;
287            match node.search(key)? {
288                Ok(i) => return Ok(Some(node.value(i))),
289                Err(i) => p = node.children.get(i).copied(),
290            }
291        }
292        Err(Error::Corrupt("map path exceeds height bound"))
293    }
294    pub fn get(&self, key: &K) -> Result<Option<V>> {
295        self.find(key)?.map(|v| decode(self.value(v)?)).transpose()
296    }
297    pub fn contains_key(&self, key: &K) -> Result<bool> {
298        Ok(self.find(key)?.is_some())
299    }
300    pub fn insert(&mut self, key: K, value: V) -> Result<Option<V>> {
301        self.healthy()?;
302        let key_bytes = key.to_bytes_checked()?.into_owned();
303        let bytes = value.into_bytes_checked()?;
304        let mut tx = Transaction::new(self);
305        let (top, split, old) = tx.insert(tx.root.top, &key, key_bytes, bytes, 0)?;
306        tx.root.top = Some(if let Some((entry, right)) = split {
307            tx.save(
308                None,
309                &Node {
310                    entries: vec![entry],
311                    children: vec![top, right],
312                },
313            )?
314        } else {
315            top
316        });
317        drop(key);
318        let old = old.map(decode).transpose()?;
319        tx.commit()?;
320        Ok(old)
321    }
322    pub fn remove(&mut self, key: &K) -> Result<Option<V>> {
323        self.healthy()?;
324        let mut tx = Transaction::new(self);
325        let Some(top) = tx.root.top else {
326            return Ok(None);
327        };
328        let (p, old, entries) = tx.remove(top, key, 0)?;
329        let Some(old) = old else { return Ok(None) };
330        let value = decode(old)?;
331        tx.root.top = if entries == 0 {
332            let root = tx.read_node(p, 0)?;
333            tx.free(p)?;
334            root.children.first().copied()
335        } else {
336            Some(p)
337        };
338        tx.root.length -= 1;
339        tx.commit()?;
340        Ok(Some(value))
341    }
342    pub fn clear(&mut self) -> Result<()> {
343        self.healthy()?;
344        let mut tx = Transaction::new(self);
345        tx.root.top = None;
346        tx.root.length = 0;
347        tx.root.cursor = record::HEADER_BYTES;
348        tx.root.free.fill(None);
349        tx.commit()
350    }
351    fn seek_path(&self, bound: &Bound<K>, forward: bool) -> Result<Vec<Frame>> {
352        self.healthy()?;
353        let mut path = Vec::new();
354        let mut pointer = self.root.borrow().top;
355        while let Some(p) = pointer {
356            ensure(
357                path.len() < MAX_DEPTH,
358                Error::Corrupt("map cursor exceeds height bound"),
359            )?;
360            let node = self.read_node(p, path.len())?;
361            let index = match bound {
362                Bound::Unbounded => {
363                    if forward {
364                        0
365                    } else {
366                        node.len()
367                    }
368                }
369                Bound::Included(k) => match node.search(k)? {
370                    Ok(i) => {
371                        if forward {
372                            i
373                        } else {
374                            i + 1
375                        }
376                    }
377                    Err(i) => i,
378                },
379                Bound::Excluded(k) => match node.search(k)? {
380                    Ok(i) => {
381                        if forward {
382                            i + 1
383                        } else {
384                            i
385                        }
386                    }
387                    Err(i) => i,
388                },
389            };
390            pointer = node.children.get(index).copied();
391            path.push(Frame {
392                node,
393                index,
394                descend: false,
395            });
396        }
397        Ok(path)
398    }
399    fn advance_path(&self, path: &mut Vec<Frame>, forward: bool) -> Result<Option<Entry>> {
400        self.healthy()?;
401        while let Some(frame) = path.last_mut() {
402            if frame.descend {
403                frame.descend = false;
404                let mut child = frame.node.children.get(frame.index).copied();
405                while let Some(p) = child {
406                    ensure(
407                        path.len() < MAX_DEPTH,
408                        Error::Corrupt("map cursor exceeds height bound"),
409                    )?;
410                    let node = self.read_node(p, path.len())?;
411                    let index = if forward { 0 } else { node.len() };
412                    child = node.children.get(index).copied();
413                    path.push(Frame {
414                        node,
415                        index,
416                        descend: false,
417                    });
418                }
419                continue;
420            }
421            let index = if forward {
422                frame.index
423            } else {
424                if frame.index == 0 {
425                    path.pop();
426                    continue;
427                }
428                frame.index - 1
429            };
430            if index >= frame.node.len() {
431                path.pop();
432                continue;
433            }
434            let entry = frame.node.entry(index);
435            frame.index = if forward { index + 1 } else { index };
436            // A caller may stop at this entry. Read the following subtree only
437            // when another step in this direction is actually requested.
438            frame.descend = !frame.node.children.is_empty();
439            return Ok(Some(entry));
440        }
441        Ok(None)
442    }
443    pub fn range(&self, bounds: impl RangeBounds<K>) -> TreeIter<'_, K, V, M> {
444        let lower = bounds.start_bound().cloned();
445        let upper = bounds.end_bound().cloned();
446        let invalid = match (&lower, &upper) {
447            (Bound::Included(a) | Bound::Excluded(a), Bound::Included(b) | Bound::Excluded(b)) => {
448                a > b
449                    || (a == b
450                        && matches!((&lower, &upper), (Bound::Excluded(_), Bound::Excluded(_))))
451            }
452            _ => false,
453        };
454        TreeIter {
455            tree: self,
456            lower,
457            upper,
458            front: None,
459            back: None,
460            exhausted: false,
461            invalid,
462            keys_only: false,
463        }
464    }
465    pub fn keys_range(
466        &self,
467        bounds: impl RangeBounds<K>,
468    ) -> impl DoubleEndedIterator<Item = Result<K>> + '_ {
469        let mut iter = self.range(bounds);
470        iter.keys_only = true;
471        iter.map(|e| e.map(MapEntry::into_key))
472    }
473}
474struct Frame {
475    node: Rc<ReadNode>,
476    index: usize,
477    descend: bool,
478}
479fn decode<T: Storable>(bytes: Vec<u8>) -> Result<T> {
480    T::check_bounds(&bytes)?;
481    Ok(T::from_bytes(Cow::Owned(bytes))?)
482}
483
484type Split = Option<(Entry, Pointer)>;
485struct Transaction<'a, K, V, M: Memory> {
486    tree: &'a Tree<K, V, M>,
487    root: Root,
488    writes: BTreeMap<u64, Vec<u8>>,
489}
490impl<'a, K: Storable + Ord + Clone, V: Storable, M: Memory> Transaction<'a, K, V, M> {
491    fn new(tree: &'a Tree<K, V, M>) -> Self {
492        Self {
493            tree,
494            root: tree.root.borrow().clone(),
495            writes: BTreeMap::new(),
496        }
497    }
498    fn bytes(&self, p: Pointer, magic: [u8; 8]) -> Result<Vec<u8>> {
499        p.check(self.root.cursor)?;
500        if let Some(bytes) = self.writes.get(&p.offset) {
501            Ok(bytes.clone())
502        } else {
503            record::read_bytes(&self.tree.memory, p.offset, p.capacity(), magic)
504        }
505    }
506    fn node(&self, p: Pointer, depth: usize) -> Result<Node> {
507        Ok(self.read_node(p, depth)?.owned())
508    }
509    fn read_node(&self, p: Pointer, depth: usize) -> Result<Rc<ReadNode>> {
510        if let Some(bytes) = self.writes.get(&p.offset) {
511            Ok(Rc::new(ReadNode::decode::<K, V>(
512                bytes.clone(),
513                self.root.cursor,
514            )?))
515        } else {
516            self.tree.read_node(p, depth)
517        }
518    }
519    fn value(&self, v: &Value) -> Result<Vec<u8>> {
520        let bytes = match v {
521            Value::Inline(b) => b.clone(),
522            Value::External(p) => record::decode_payload(VALUE, &self.bytes(*p, VALUE)?)?.to_vec(),
523        };
524        V::check_bounds(&bytes)?;
525        Ok(bytes)
526    }
527    fn free(&mut self, p: Pointer) -> Result<()> {
528        let index = usize::from(p.order);
529        self.writes
530            .insert(p.offset, record::encode(FREE, &self.root.free[index])?);
531        self.root.free[index] = Some(p);
532        Ok(())
533    }
534    fn allocate(&mut self, len: usize) -> Result<Pointer> {
535        let size = len
536            .max(256)
537            .checked_next_power_of_two()
538            .ok_or(Error::Budget)?;
539        let order = size.trailing_zeros() as usize;
540        ensure(order < BINS, Error::Budget)?;
541        for bin in order..BINS {
542            if let Some(mut p) = self.root.free[bin] {
543                let next: Option<Pointer> = record::decode(FREE, &self.bytes(p, FREE)?)?;
544                if let Some(next) = next {
545                    next.check(self.root.cursor)?;
546                    ensure(
547                        next.order == p.order && next != p,
548                        Error::Corrupt("map free-list link"),
549                    )?;
550                }
551                self.root.free[bin] = next;
552                // Return only the requested capacity. Each unused half remains
553                // an ordinary free-list block, staged with the same mutation.
554                while usize::from(p.order) > order {
555                    p.order -= 1;
556                    let sibling = Pointer {
557                        offset: p.offset + p.capacity(),
558                        order: p.order,
559                    };
560                    self.free(sibling)?;
561                }
562                return Ok(p);
563            }
564        }
565        let p = Pointer {
566            offset: self.root.cursor,
567            order: u8::try_from(order).unwrap(),
568        };
569        self.root.cursor = self
570            .root
571            .cursor
572            .checked_add(p.capacity())
573            .ok_or(Error::Budget)?;
574        Ok(p)
575    }
576    fn save_bytes(&mut self, previous: Option<Pointer>, bytes: Vec<u8>) -> Result<Pointer> {
577        let p = if let Some(p) = previous.filter(|p| p.capacity() >= bytes.len() as u64) {
578            let needed = bytes.len().max(256).next_power_of_two() as u64;
579            // Reclaim only into already reusable smaller capacity. Splitting the
580            // live block itself loses its large size class and causes unbounded
581            // growth under repeated large/small replacements without coalescing.
582            let smaller_available = (needed.trailing_zeros() as usize..usize::from(p.order))
583                .any(|order| self.root.free[order].is_some());
584            if p.capacity() / needed >= 4 && smaller_available {
585                let smaller = self.allocate(bytes.len())?;
586                self.free(p)?;
587                smaller
588            } else {
589                p
590            }
591        } else {
592            if let Some(p) = previous {
593                self.free(p)?;
594            }
595            self.allocate(bytes.len())?
596        };
597        self.writes.insert(p.offset, bytes);
598        Ok(p)
599    }
600    fn save(&mut self, p: Option<Pointer>, node: &Node) -> Result<Pointer> {
601        self.save_bytes(p, node.encode()?)
602    }
603    fn save_value(&mut self, previous: Option<&Value>, bytes: Vec<u8>) -> Result<Value> {
604        let old = match previous {
605            Some(Value::External(p)) => Some(*p),
606            _ => None,
607        };
608        if bytes.len() <= INLINE {
609            if let Some(p) = old {
610                self.free(p)?;
611            }
612            Ok(Value::Inline(bytes))
613        } else {
614            Ok(Value::External(
615                self.save_bytes(old, record::encode_payload(VALUE, &bytes)?)?,
616            ))
617        }
618    }
619    fn insert(
620        &mut self,
621        p: Option<Pointer>,
622        key: &K,
623        key_bytes: Vec<u8>,
624        bytes: Vec<u8>,
625        depth: usize,
626    ) -> Result<(Pointer, Split, Option<Vec<u8>>)> {
627        ensure(
628            depth < MAX_DEPTH,
629            Error::Corrupt("map path exceeds height bound"),
630        )?;
631        let read = p.map(|p| self.read_node(p, depth)).transpose()?;
632        let search = read.as_ref().map_or(Ok(Err(0)), |node| node.search(key))?;
633        let (mut node, old) = match search {
634            Ok(i) => {
635                let read = read.as_ref().unwrap();
636                let previous = read.value(i);
637                let old = self.value(&previous)?;
638                let value = self.save_value(Some(&previous), bytes)?;
639                let same_pointer =
640                    matches!((&previous,&value),(Value::External(a),Value::External(b)) if a==b);
641                if same_pointer {
642                    return Ok((p.unwrap(), None, Some(old)));
643                }
644                let mut node = read.owned();
645                node.entries[i].value = value;
646                (node, Some(old))
647            }
648            Err(i) => {
649                if let Some(child) = read.as_ref().and_then(|node| node.children.get(i)).copied() {
650                    let (next, split, old) =
651                        self.insert(Some(child), key, key_bytes, bytes, depth + 1)?;
652                    if split.is_none() && next == child {
653                        return Ok((p.unwrap(), None, old));
654                    }
655                    let mut node = read.as_ref().unwrap().owned();
656                    node.children[i] = next;
657                    if let Some((entry, right)) = split {
658                        node.entries.insert(i, entry);
659                        node.children.insert(i + 1, right);
660                    }
661                    (node, old)
662                } else {
663                    let value = self.save_value(None, bytes)?;
664                    let mut node = read
665                        .as_ref()
666                        .map_or_else(Node::default, |node| node.owned());
667                    node.entries.insert(
668                        i,
669                        Entry {
670                            key: key_bytes,
671                            value,
672                        },
673                    );
674                    self.root.length = self.root.length.checked_add(1).ok_or(Error::Budget)?;
675                    (node, None)
676                }
677            }
678        };
679        let split = if node.entries.len() > MAX_ENTRIES {
680            let right_entries = node.entries.split_off(DEGREE + 1);
681            let median = node.entries.pop().unwrap();
682            let children = if node.children.is_empty() {
683                vec![]
684            } else {
685                node.children.split_off(DEGREE + 1)
686            };
687            let right = self.save(
688                None,
689                &Node {
690                    entries: right_entries,
691                    children,
692                },
693            )?;
694            Some((median, right))
695        } else {
696            None
697        };
698        Ok((self.save(p, &node)?, split, old))
699    }
700    fn remove(
701        &mut self,
702        p: Pointer,
703        key: &K,
704        depth: usize,
705    ) -> Result<(Pointer, Option<Vec<u8>>, usize)> {
706        ensure(
707            depth < MAX_DEPTH,
708            Error::Corrupt("map path exceeds height bound"),
709        )?;
710        let read = self.read_node(p, depth)?;
711        let (node, old) = match read.search(key)? {
712            Ok(i) => {
713                let mut node = read.owned();
714                let old = self.value(&node.entries[i].value)?;
715                if let Value::External(p) = node.entries[i].value {
716                    self.free(p)?;
717                }
718                if node.children.is_empty() {
719                    node.entries.remove(i);
720                } else {
721                    let (child, entry, entries) =
722                        self.extract_min(node.children[i + 1], depth + 1)?;
723                    node.entries[i] = entry;
724                    node.children[i + 1] = child;
725                    if entries < MIN_ENTRIES {
726                        self.rebalance(&mut node, i + 1, depth)?;
727                    }
728                }
729                (node, Some(old))
730            }
731            Err(i) => {
732                let Some(child) = read.children.get(i).copied() else {
733                    return Ok((p, None, read.len()));
734                };
735                let (next, old, entries) = self.remove(child, key, depth + 1)?;
736                if old.is_none() {
737                    return Ok((p, None, read.len()));
738                }
739                if next == child && entries >= MIN_ENTRIES {
740                    return Ok((p, old, read.len()));
741                }
742                let mut node = read.owned();
743                node.children[i] = next;
744                if entries < MIN_ENTRIES {
745                    self.rebalance(&mut node, i, depth)?;
746                }
747                (node, old)
748            }
749        };
750        Ok((self.save(Some(p), &node)?, old, node.entries.len()))
751    }
752    fn extract_min(&mut self, p: Pointer, depth: usize) -> Result<(Pointer, Entry, usize)> {
753        ensure(
754            depth < MAX_DEPTH,
755            Error::Corrupt("map path exceeds height bound"),
756        )?;
757        let read = self.read_node(p, depth)?;
758        ensure(read.len() != 0, Error::Corrupt("empty successor node"))?;
759        let (node, entry) = if read.children.is_empty() {
760            let mut node = read.owned();
761            let entry = node.entries.remove(0);
762            (node, entry)
763        } else {
764            let (child, entry, entries) = self.extract_min(read.children[0], depth + 1)?;
765            if read.children[0] == child && entries >= MIN_ENTRIES {
766                return Ok((p, entry, read.len()));
767            }
768            let mut node = read.owned();
769            node.children[0] = child;
770            if entries < MIN_ENTRIES {
771                self.rebalance(&mut node, 0, depth)?;
772            }
773            (node, entry)
774        };
775        Ok((self.save(Some(p), &node)?, entry, node.entries.len()))
776    }
777    fn rebalance(&mut self, parent: &mut Node, i: usize, parent_depth: usize) -> Result<()> {
778        let depth = parent_depth + 1;
779        let p = parent.children[i];
780        let mut child = self.node(p, depth)?;
781        if child.entries.len() >= MIN_ENTRIES {
782            return Ok(());
783        }
784        if i > 0 {
785            let left_p = parent.children[i - 1];
786            let mut left = self.node(left_p, depth)?;
787            ensure(
788                left.children.is_empty() == child.children.is_empty(),
789                Error::Corrupt("B-tree sibling levels"),
790            )?;
791            if left.entries.len() > MIN_ENTRIES {
792                child.entries.insert(
793                    0,
794                    std::mem::replace(&mut parent.entries[i - 1], left.entries.pop().unwrap()),
795                );
796                if let Some(p) = left.children.pop() {
797                    child.children.insert(0, p);
798                }
799                parent.children[i - 1] = self.save(Some(left_p), &left)?;
800                parent.children[i] = self.save(Some(p), &child)?;
801                return Ok(());
802            }
803        }
804        if i + 1 < parent.children.len() {
805            let right_p = parent.children[i + 1];
806            let mut right = self.node(right_p, depth)?;
807            ensure(
808                right.children.is_empty() == child.children.is_empty(),
809                Error::Corrupt("B-tree sibling levels"),
810            )?;
811            if right.entries.len() > MIN_ENTRIES {
812                child.entries.push(std::mem::replace(
813                    &mut parent.entries[i],
814                    right.entries.remove(0),
815                ));
816                if !right.children.is_empty() {
817                    child.children.push(right.children.remove(0));
818                }
819                parent.children[i + 1] = self.save(Some(right_p), &right)?;
820                parent.children[i] = self.save(Some(p), &child)?;
821                return Ok(());
822            }
823        }
824        let left_index = if i > 0 { i - 1 } else { i };
825        ensure(
826            left_index < parent.entries.len(),
827            Error::Corrupt("B-tree missing separator"),
828        )?;
829        let left_p = parent.children[left_index];
830        let right_p = parent.children[left_index + 1];
831        let mut left = self.node(left_p, depth)?;
832        let right = self.node(right_p, depth)?;
833        ensure(
834            left.children.is_empty() == right.children.is_empty(),
835            Error::Corrupt("B-tree sibling levels"),
836        )?;
837        left.entries.push(parent.entries.remove(left_index));
838        left.entries.extend(right.entries);
839        left.children.extend(right.children);
840        self.free(right_p)?;
841        parent.children[left_index] = self.save(Some(left_p), &left)?;
842        parent.children.remove(left_index + 1);
843        Ok(())
844    }
845    fn commit(self) -> Result<()> {
846        let changed = self.root != *self.tree.root.borrow();
847        let header = if changed {
848            Some(record::encode(MAGIC, &self.root)?)
849        } else {
850            None
851        };
852        if let Some(h) = &header {
853            ensure(h.len() as u64 <= record::HEADER_BYTES, Error::Budget)?;
854        }
855        record::reserve(&self.tree.memory, self.root.cursor)?;
856        self.tree.poisoned.set(true);
857        if self.root.cursor < self.tree.root.borrow().cursor
858            || self.root.top != self.tree.root.borrow().top
859        {
860            self.tree.cache.borrow_mut().clear();
861        }
862        for (offset, bytes) in self.writes {
863            self.tree.cache.borrow_mut().invalidate(offset);
864            self.tree.memory.write(offset, &bytes);
865        }
866        if let Some(h) = header {
867            self.tree.memory.write(0, &h);
868        }
869        *self.tree.root.borrow_mut() = self.root;
870        self.tree.poisoned.set(false);
871        Ok(())
872    }
873}
874///
875/// `MapEntry`
876///
877/// An owned key and lazily decoded value from an ordered scan.
878///
879pub struct MapEntry<K, V> {
880    key: K,
881    bytes: Vec<u8>,
882    value: PhantomData<V>,
883}
884impl<K, V: Storable> MapEntry<K, V> {
885    /// Borrow the decoded key.
886    pub const fn key(&self) -> &K {
887        &self.key
888    }
889    /// Decode the value using its declared codec.
890    ///
891    /// # Errors
892    /// Returns the application codec or declared-bound error for this value.
893    pub fn value(&self) -> Result<V> {
894        decode(self.bytes.clone())
895    }
896    /// Consume the entry and decode its value.
897    ///
898    /// # Errors
899    /// Returns the application codec or declared-bound error for this value.
900    pub fn into_pair(self) -> Result<(K, V)> {
901        Ok((self.key, decode(self.bytes)?))
902    }
903    /// Consume the entry without decoding its value.
904    pub fn into_key(self) -> K {
905        self.key
906    }
907    /// Consume the entry and decode only its value.
908    ///
909    /// # Errors
910    /// Returns the application codec or declared-bound error for this value.
911    pub fn into_value(self) -> Result<V> {
912        decode(self.bytes)
913    }
914}
915
916///
917/// `TreeIter`
918///
919/// Lazy bidirectional traversal with bounded routing-node stacks. External values
920/// are read only for in-range entries and decoded when requested by the caller.
921///
922pub struct TreeIter<'a, K, V, M: Memory> {
923    tree: &'a Tree<K, V, M>,
924    lower: Bound<K>,
925    upper: Bound<K>,
926    front: Option<Vec<Frame>>,
927    back: Option<Vec<Frame>>,
928    exhausted: bool,
929    invalid: bool,
930    keys_only: bool,
931}
932impl<K: Storable + Ord + Clone, V: Storable, M: Memory> TreeIter<'_, K, V, M> {
933    fn step(&mut self, forward: bool) -> Result<Option<MapEntry<K, V>>> {
934        ensure(!self.invalid, Error::Bounds)?;
935        self.tree.healthy()?;
936        let empty = match (&self.lower, &self.upper) {
937            (Bound::Included(a), Bound::Included(b)) => a > b,
938            (Bound::Included(a) | Bound::Excluded(a), Bound::Included(b) | Bound::Excluded(b)) => {
939                a >= b
940            }
941            _ => false,
942        };
943        if empty {
944            return Ok(None);
945        }
946        let (path, bound) = if forward {
947            (&mut self.front, &self.lower)
948        } else {
949            (&mut self.back, &self.upper)
950        };
951        if path.is_none() {
952            *path = Some(self.tree.seek_path(bound, forward)?);
953        }
954        let entry = self.tree.advance_path(path.as_mut().unwrap(), forward)?;
955        if let Some(entry) = entry {
956            let key: K = decode(entry.key)?;
957            let after = match &self.lower {
958                Bound::Unbounded => true,
959                Bound::Included(k) => key >= *k,
960                Bound::Excluded(k) => key > *k,
961            };
962            let before = match &self.upper {
963                Bound::Unbounded => true,
964                Bound::Included(k) => key <= *k,
965                Bound::Excluded(k) => key < *k,
966            };
967            ensure(
968                if forward { after } else { before },
969                Error::Corrupt("map cursor key order"),
970            )?;
971            if !(after && before) {
972                return Ok(None);
973            }
974            if forward {
975                self.lower = Bound::Excluded(key.clone());
976            } else {
977                self.upper = Bound::Excluded(key.clone());
978            }
979            return Ok(Some(MapEntry {
980                key,
981                bytes: if self.keys_only {
982                    vec![]
983                } else {
984                    self.tree.value(entry.value)?
985                },
986                value: PhantomData,
987            }));
988        }
989        Ok(None)
990    }
991    fn advance(&mut self, forward: bool) -> Option<Result<MapEntry<K, V>>> {
992        if self.exhausted {
993            return None;
994        }
995        match self.step(forward) {
996            Ok(Some(entry)) => Some(Ok(entry)),
997            Ok(None) => {
998                self.exhausted = true;
999                None
1000            }
1001            Err(error) => {
1002                self.exhausted = true;
1003                Some(Err(error))
1004            }
1005        }
1006    }
1007}
1008impl<K: Storable + Ord + Clone, V: Storable, M: Memory> Iterator for TreeIter<'_, K, V, M> {
1009    type Item = Result<MapEntry<K, V>>;
1010    fn next(&mut self) -> Option<Self::Item> {
1011        self.advance(true)
1012    }
1013}
1014impl<K: Storable + Ord + Clone, V: Storable, M: Memory> DoubleEndedIterator
1015    for TreeIter<'_, K, V, M>
1016{
1017    fn next_back(&mut self) -> Option<Self::Item> {
1018        self.advance(false)
1019    }
1020}
1021impl<K: Storable + Ord + Clone, V: Storable, M: Memory> std::iter::FusedIterator
1022    for TreeIter<'_, K, V, M>
1023{
1024}
1025
1026#[cfg(test)]
1027mod tests;
ic_storage/adapter/tree.rs

ic_storage/adapter/
tree.rs
