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    const 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        record::encode_payload_with(NODE, capacity, |bytes| {
111            bytes.extend_from_slice(&[
112                u8::try_from(self.entries.len()).unwrap(),
113                u8::from(!self.children.is_empty()),
114            ]);
115            for p in &self.children {
116                p.put(bytes);
117            }
118            for entry in &self.entries {
119                bytes.extend_from_slice(
120                    &u32::try_from(entry.key.len())
121                        .map_err(|_| Error::Budget)?
122                        .to_le_bytes(),
123                );
124                bytes.extend_from_slice(&entry.key);
125                match &entry.value {
126                    Value::Inline(value) => {
127                        bytes.push(u8::try_from(value.len()).unwrap());
128                        bytes.extend_from_slice(value);
129                    }
130                    Value::External(p) => {
131                        bytes.push(255);
132                        p.put(bytes);
133                    }
134                }
135            }
136            Ok(())
137        })
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    pub fn set_cache_budget(&mut self, bytes: usize) {
250        self.cache.get_mut().configure(bytes);
251    }
252    fn read_node(&self, p: Pointer, depth: usize) -> Result<Rc<ReadNode>> {
253        let cursor = self.root.borrow().cursor;
254        p.check(cursor)?;
255        if let Some(node) = self.cache.borrow_mut().get(p, depth) {
256            return Ok(node);
257        }
258        let node = Rc::new(ReadNode::decode::<K, V>(
259            record::read_bytes(&self.memory, p.offset, p.capacity(), NODE)?,
260            cursor,
261        )?);
262        self.cache.borrow_mut().put(p, Rc::clone(&node), depth);
263        Ok(node)
264    }
265    fn value(&self, v: Value) -> Result<Vec<u8>> {
266        let bytes = match v {
267            Value::Inline(bytes) => bytes,
268            Value::External(p) => {
269                p.check(self.root.borrow().cursor)?;
270                record::read_payload(&self.memory, p.offset, p.capacity(), VALUE)?
271            }
272        };
273        V::check_bounds(&bytes)?;
274        Ok(bytes)
275    }
276    fn find(&self, key: &K) -> Result<Option<(Rc<ReadNode>, usize)>> {
277        self.healthy()?;
278        let mut p = self.root.borrow().top;
279        for depth in 0..MAX_DEPTH {
280            let Some(current) = p else { return Ok(None) };
281            let node = self.read_node(current, depth)?;
282            match node.search(key)? {
283                Ok(i) => return Ok(Some((node, i))),
284                Err(i) => p = node.children.get(i).copied(),
285            }
286        }
287        Err(Error::Corrupt("map path exceeds height bound"))
288    }
289    pub fn get(&self, key: &K) -> Result<Option<V>> {
290        let Some((node, index)) = self.find(key)? else {
291            return Ok(None);
292        };
293        if let Some(bytes) = node.inline_value(index) {
294            V::check_bounds(bytes)?;
295            Ok(Some(V::from_bytes(Cow::Borrowed(bytes))?))
296        } else {
297            let value = node.value(index);
298            // External decoding does not borrow the routing buffer. Release it
299            // before allocating the payload, as in the original lookup path.
300            drop(node);
301            decode(self.value(value)?).map(Some)
302        }
303    }
304    pub fn contains_key(&self, key: &K) -> Result<bool> {
305        Ok(self.find(key)?.is_some())
306    }
307    #[expect(
308        clippy::needless_pass_by_ref_mut,
309        reason = "Mutations require exclusive collection access despite internal transactional borrowing."
310    )]
311    pub fn insert(&mut self, key: K, value: V) -> Result<Option<V>> {
312        self.healthy()?;
313        let key_bytes = key.to_bytes_checked()?.into_owned();
314        let bytes = value.into_bytes_checked()?;
315        let mut tx = Transaction::new(self);
316        let (top, split, old) = tx.insert(tx.root.top, &key, key_bytes, bytes, 0)?;
317        tx.root.top = Some(if let Some((entry, right)) = split {
318            tx.save(
319                None,
320                &Node {
321                    entries: vec![entry],
322                    children: vec![top, right],
323                },
324            )?
325        } else {
326            top
327        });
328        drop(key);
329        let old = old.map(decode).transpose()?;
330        tx.commit()?;
331        Ok(old)
332    }
333    #[expect(
334        clippy::needless_pass_by_ref_mut,
335        reason = "Mutations require exclusive collection access despite internal transactional borrowing."
336    )]
337    pub fn remove(&mut self, key: &K) -> Result<Option<V>> {
338        self.healthy()?;
339        let mut tx = Transaction::new(self);
340        let Some(top) = tx.root.top else {
341            return Ok(None);
342        };
343        let (changed, old) = tx.remove(top, key, 0)?;
344        let Some(old) = old else { return Ok(None) };
345        let value = decode(old)?;
346        if let Some(node) = changed {
347            tx.root.top = if node.entries.is_empty() {
348                tx.free(top)?;
349                node.children.first().copied()
350            } else {
351                Some(tx.save(Some(top), &node)?)
352            };
353        }
354        tx.root.length -= 1;
355        tx.commit()?;
356        Ok(Some(value))
357    }
358    #[expect(
359        clippy::needless_pass_by_ref_mut,
360        reason = "Mutations require exclusive collection access despite internal transactional borrowing."
361    )]
362    pub fn clear(&mut self) -> Result<()> {
363        self.healthy()?;
364        let mut tx = Transaction::new(self);
365        tx.root.top = None;
366        tx.root.length = 0;
367        tx.root.cursor = record::HEADER_BYTES;
368        tx.root.free.fill(None);
369        tx.commit()
370    }
371    fn seek_path(&self, bound: &Bound<K>, forward: bool) -> Result<Vec<Frame>> {
372        self.healthy()?;
373        let mut path = Vec::new();
374        let mut pointer = self.root.borrow().top;
375        while let Some(p) = pointer {
376            ensure(
377                path.len() < MAX_DEPTH,
378                Error::Corrupt("map cursor exceeds height bound"),
379            )?;
380            let node = self.read_node(p, path.len())?;
381            let index = match bound {
382                Bound::Unbounded => {
383                    if forward {
384                        0
385                    } else {
386                        node.len()
387                    }
388                }
389                Bound::Included(k) => match node.search(k)? {
390                    Ok(i) => {
391                        if forward {
392                            i
393                        } else {
394                            i + 1
395                        }
396                    }
397                    Err(i) => i,
398                },
399                Bound::Excluded(k) => match node.search(k)? {
400                    Ok(i) => {
401                        if forward {
402                            i + 1
403                        } else {
404                            i
405                        }
406                    }
407                    Err(i) => i,
408                },
409            };
410            pointer = node.children.get(index).copied();
411            path.push(Frame {
412                node,
413                index,
414                descend: false,
415            });
416        }
417        Ok(path)
418    }
419    fn advance_path(
420        &self,
421        path: &mut Vec<Frame>,
422        forward: bool,
423        keys_only: bool,
424    ) -> Result<Option<(K, Option<Value>)>> {
425        self.healthy()?;
426        while let Some(frame) = path.last_mut() {
427            if frame.descend {
428                frame.descend = false;
429                let mut child = frame.node.children.get(frame.index).copied();
430                while let Some(p) = child {
431                    ensure(
432                        path.len() < MAX_DEPTH,
433                        Error::Corrupt("map cursor exceeds height bound"),
434                    )?;
435                    let node = self.read_node(p, path.len())?;
436                    let index = if forward { 0 } else { node.len() };
437                    child = node.children.get(index).copied();
438                    path.push(Frame {
439                        node,
440                        index,
441                        descend: false,
442                    });
443                }
444                continue;
445            }
446            let index = if forward {
447                frame.index
448            } else {
449                if frame.index == 0 {
450                    path.pop();
451                    continue;
452                }
453                frame.index - 1
454            };
455            if index >= frame.node.len() {
456                path.pop();
457                continue;
458            }
459            let bytes = frame.node.key_bytes(index);
460            K::check_bounds(bytes)?;
461            let key = K::from_bytes(Cow::Borrowed(bytes))?;
462            let value = (!keys_only).then(|| frame.node.value(index));
463            frame.index = if forward { index + 1 } else { index };
464            // A caller may stop at this entry. Read the following subtree only
465            // when another step in this direction is actually requested.
466            frame.descend = !frame.node.children.is_empty();
467            return Ok(Some((key, value)));
468        }
469        Ok(None)
470    }
471    pub fn range(&self, bounds: impl RangeBounds<K>) -> TreeIter<'_, K, V, M> {
472        let lower = bounds.start_bound().cloned();
473        let upper = bounds.end_bound().cloned();
474        let invalid = match (&lower, &upper) {
475            (Bound::Included(a) | Bound::Excluded(a), Bound::Included(b) | Bound::Excluded(b)) => {
476                a > b
477                    || (a == b
478                        && matches!((&lower, &upper), (Bound::Excluded(_), Bound::Excluded(_))))
479            }
480            _ => false,
481        };
482        TreeIter {
483            tree: self,
484            lower,
485            upper,
486            front: None,
487            back: None,
488            exhausted: false,
489            invalid,
490            keys_only: false,
491        }
492    }
493    pub fn keys_range(
494        &self,
495        bounds: impl RangeBounds<K>,
496    ) -> impl DoubleEndedIterator<Item = Result<K>> + '_ {
497        let mut iter = self.range(bounds);
498        iter.keys_only = true;
499        iter.map(|e| e.map(MapEntry::into_key))
500    }
501}
502struct Frame {
503    node: Rc<ReadNode>,
504    index: usize,
505    descend: bool,
506}
507fn decode<T: Storable>(bytes: Vec<u8>) -> Result<T> {
508    T::check_bounds(&bytes)?;
509    Ok(T::from_bytes(Cow::Owned(bytes))?)
510}
511
512type Split = Option<(Entry, Pointer)>;
513struct Transaction<'a, K, V, M: Memory> {
514    tree: &'a Tree<K, V, M>,
515    root: Root,
516    writes: BTreeMap<u64, Vec<u8>>,
517}
518impl<'a, K: Storable + Ord + Clone, V: Storable, M: Memory> Transaction<'a, K, V, M> {
519    fn new(tree: &'a Tree<K, V, M>) -> Self {
520        Self {
521            tree,
522            root: tree.root.borrow().clone(),
523            writes: BTreeMap::new(),
524        }
525    }
526    fn bytes(&self, p: Pointer, magic: [u8; 8]) -> Result<Vec<u8>> {
527        p.check(self.root.cursor)?;
528        self.writes.get(&p.offset).map_or_else(
529            || record::read_bytes(&self.tree.memory, p.offset, p.capacity(), magic),
530            |bytes| Ok(bytes.clone()),
531        )
532    }
533    fn node(&self, p: Pointer, depth: usize) -> Result<Node> {
534        Ok(self.read_node(p, depth)?.owned())
535    }
536    fn read_node(&self, p: Pointer, depth: usize) -> Result<Rc<ReadNode>> {
537        if let Some(bytes) = self.writes.get(&p.offset) {
538            Ok(Rc::new(ReadNode::decode::<K, V>(
539                bytes.clone(),
540                self.root.cursor,
541            )?))
542        } else {
543            self.tree.read_node(p, depth)
544        }
545    }
546    fn value(&self, v: &Value) -> Result<Vec<u8>> {
547        let bytes = match v {
548            Value::Inline(b) => b.clone(),
549            Value::External(p) => {
550                p.check(self.root.cursor)?;
551                if let Some(bytes) = self.writes.get(&p.offset) {
552                    record::decode_payload(VALUE, bytes)?.to_vec()
553                } else {
554                    record::read_payload(&self.tree.memory, p.offset, p.capacity(), VALUE)?
555                }
556            }
557        };
558        V::check_bounds(&bytes)?;
559        Ok(bytes)
560    }
561    fn free(&mut self, p: Pointer) -> Result<()> {
562        let index = usize::from(p.order);
563        self.writes
564            .insert(p.offset, record::encode(FREE, &self.root.free[index])?);
565        self.root.free[index] = Some(p);
566        Ok(())
567    }
568    fn allocate(&mut self, len: usize) -> Result<Pointer> {
569        let size = len
570            .max(256)
571            .checked_next_power_of_two()
572            .ok_or(Error::Budget)?;
573        let order = size.trailing_zeros() as usize;
574        ensure(order < BINS, Error::Budget)?;
575        for bin in order..BINS {
576            if let Some(mut p) = self.root.free[bin] {
577                let next: Option<Pointer> = record::decode(FREE, &self.bytes(p, FREE)?)?;
578                if let Some(next) = next {
579                    next.check(self.root.cursor)?;
580                    ensure(
581                        next.order == p.order && next != p,
582                        Error::Corrupt("map free-list link"),
583                    )?;
584                }
585                self.root.free[bin] = next;
586                // Return only the requested capacity. Each unused half remains
587                // an ordinary free-list block, staged with the same mutation.
588                while usize::from(p.order) > order {
589                    p.order -= 1;
590                    let sibling = Pointer {
591                        offset: p.offset + p.capacity(),
592                        order: p.order,
593                    };
594                    self.free(sibling)?;
595                }
596                return Ok(p);
597            }
598        }
599        let p = Pointer {
600            offset: self.root.cursor,
601            order: u8::try_from(order).unwrap(),
602        };
603        self.root.cursor = self
604            .root
605            .cursor
606            .checked_add(p.capacity())
607            .ok_or(Error::Budget)?;
608        Ok(p)
609    }
610    fn save_bytes(&mut self, previous: Option<Pointer>, bytes: Vec<u8>) -> Result<Pointer> {
611        let p = if let Some(p) = previous.filter(|p| p.capacity() >= bytes.len() as u64) {
612            let needed = bytes.len().max(256).next_power_of_two() as u64;
613            // Reclaim only into already reusable smaller capacity. Splitting the
614            // live block itself loses its large size class and causes unbounded
615            // growth under repeated large/small replacements without coalescing.
616            let smaller_available = (needed.trailing_zeros() as usize..usize::from(p.order))
617                .any(|order| self.root.free[order].is_some());
618            if p.capacity() / needed >= 4 && smaller_available {
619                let smaller = self.allocate(bytes.len())?;
620                self.free(p)?;
621                smaller
622            } else {
623                p
624            }
625        } else {
626            if let Some(p) = previous {
627                self.free(p)?;
628            }
629            self.allocate(bytes.len())?
630        };
631        self.writes.insert(p.offset, bytes);
632        Ok(p)
633    }
634    fn save(&mut self, p: Option<Pointer>, node: &Node) -> Result<Pointer> {
635        self.save_bytes(p, node.encode()?)
636    }
637    fn save_value(&mut self, previous: Option<&Value>, bytes: Vec<u8>) -> Result<Value> {
638        let old = match previous {
639            Some(Value::External(p)) => Some(*p),
640            _ => None,
641        };
642        if bytes.len() <= INLINE {
643            if let Some(p) = old {
644                self.free(p)?;
645            }
646            Ok(Value::Inline(bytes))
647        } else {
648            Ok(Value::External(
649                self.save_bytes(old, record::encode_payload(VALUE, &bytes)?)?,
650            ))
651        }
652    }
653    fn insert(
654        &mut self,
655        p: Option<Pointer>,
656        key: &K,
657        key_bytes: Vec<u8>,
658        bytes: Vec<u8>,
659        depth: usize,
660    ) -> Result<(Pointer, Split, Option<Vec<u8>>)> {
661        ensure(
662            depth < MAX_DEPTH,
663            Error::Corrupt("map path exceeds height bound"),
664        )?;
665        let read = p.map(|p| self.read_node(p, depth)).transpose()?;
666        let search = read.as_ref().map_or(Ok(Err(0)), |node| node.search(key))?;
667        let (mut node, old) = match search {
668            Ok(i) => {
669                let read = read.as_ref().unwrap();
670                let previous = read.value(i);
671                let old = self.value(&previous)?;
672                let value = self.save_value(Some(&previous), bytes)?;
673                let same_pointer =
674                    matches!((&previous,&value),(Value::External(a),Value::External(b)) if a==b);
675                if same_pointer {
676                    return Ok((p.unwrap(), None, Some(old)));
677                }
678                if let (Value::Inline(_), Value::Inline(bytes)) = (&previous, &value) {
679                    let next = self.save_bytes(p, read.replace_inline(i, bytes)?)?;
680                    return Ok((next, None, Some(old)));
681                }
682                let mut node = read.owned();
683                node.entries[i].value = value;
684                (node, Some(old))
685            }
686            Err(i) => {
687                if let Some(child) = read.as_ref().and_then(|node| node.children.get(i)).copied() {
688                    let (next, split, old) =
689                        self.insert(Some(child), key, key_bytes, bytes, depth + 1)?;
690                    if split.is_none() && next == child {
691                        return Ok((p.unwrap(), None, old));
692                    }
693                    let mut node = read.as_ref().unwrap().owned();
694                    node.children[i] = next;
695                    if let Some((entry, right)) = split {
696                        node.entries.insert(i, entry);
697                        node.children.insert(i + 1, right);
698                    }
699                    (node, old)
700                } else {
701                    let value = self.save_value(None, bytes)?;
702                    let mut node = read
703                        .as_ref()
704                        .map_or_else(Node::default, |node| node.owned());
705                    node.entries.insert(
706                        i,
707                        Entry {
708                            key: key_bytes,
709                            value,
710                        },
711                    );
712                    self.root.length = self.root.length.checked_add(1).ok_or(Error::Budget)?;
713                    (node, None)
714                }
715            }
716        };
717        let split = if node.entries.len() > MAX_ENTRIES {
718            let right_entries = node.entries.split_off(DEGREE + 1);
719            let median = node.entries.pop().unwrap();
720            let children = if node.children.is_empty() {
721                vec![]
722            } else {
723                node.children.split_off(DEGREE + 1)
724            };
725            let right = self.save(
726                None,
727                &Node {
728                    entries: right_entries,
729                    children,
730                },
731            )?;
732            Some((median, right))
733        } else {
734            None
735        };
736        Ok((self.save(p, &node)?, split, old))
737    }
738    // Return a modified node without encoding it: its parent may still repair
739    // an underflow. None means its persisted routing record remains unchanged.
740    fn remove(
741        &mut self,
742        p: Pointer,
743        key: &K,
744        depth: usize,
745    ) -> Result<(Option<Node>, Option<Vec<u8>>)> {
746        ensure(
747            depth < MAX_DEPTH,
748            Error::Corrupt("map path exceeds height bound"),
749        )?;
750        let read = self.read_node(p, depth)?;
751        match read.search(key)? {
752            Ok(i) => {
753                let mut node = read.owned();
754                let old = self.value(&node.entries[i].value)?;
755                if let Value::External(p) = node.entries[i].value {
756                    self.free(p)?;
757                }
758                if node.children.is_empty() {
759                    node.entries.remove(i);
760                } else {
761                    let child_pointer = node.children[i + 1];
762                    let (child, entry) = self.extract_min(child_pointer, depth + 1)?;
763                    node.entries[i] = entry;
764                    if let Some(child) = child {
765                        self.finish_child(&mut node, i + 1, child, depth)?;
766                    }
767                }
768                Ok((Some(node), Some(old)))
769            }
770            Err(i) => {
771                let Some(child_pointer) = read.children.get(i).copied() else {
772                    return Ok((None, None));
773                };
774                let (child, old) = self.remove(child_pointer, key, depth + 1)?;
775                let node = child
776                    .map(|child| self.changed_parent(&read, i, child, depth))
777                    .transpose()?
778                    .flatten();
779                Ok((node, old))
780            }
781        }
782    }
783    fn extract_min(&mut self, p: Pointer, depth: usize) -> Result<(Option<Node>, Entry)> {
784        ensure(
785            depth < MAX_DEPTH,
786            Error::Corrupt("map path exceeds height bound"),
787        )?;
788        let read = self.read_node(p, depth)?;
789        ensure(read.len() != 0, Error::Corrupt("empty successor node"))?;
790        if read.children.is_empty() {
791            let mut node = read.owned();
792            let entry = node.entries.remove(0);
793            Ok((Some(node), entry))
794        } else {
795            let (child, entry) = self.extract_min(read.children[0], depth + 1)?;
796            let node = child
797                .map(|child| self.changed_parent(&read, 0, child, depth))
798                .transpose()?
799                .flatten();
800            Ok((node, entry))
801        }
802    }
803    // Keep unchanged ancestors borrowed, including when the child was rewritten
804    // in place. Only a moved child or an underflow requires owning the parent.
805    fn changed_parent(
806        &mut self,
807        read: &ReadNode,
808        i: usize,
809        child: Node,
810        depth: usize,
811    ) -> Result<Option<Node>> {
812        if child.entries.len() >= MIN_ENTRIES {
813            let previous = read.children[i];
814            let next = self.save(Some(previous), &child)?;
815            if next == previous {
816                return Ok(None);
817            }
818            let mut parent = read.owned();
819            parent.children[i] = next;
820            Ok(Some(parent))
821        } else {
822            let mut parent = read.owned();
823            self.rebalance(&mut parent, i, child, depth)?;
824            Ok(Some(parent))
825        }
826    }
827    fn finish_child(
828        &mut self,
829        parent: &mut Node,
830        i: usize,
831        child: Node,
832        depth: usize,
833    ) -> Result<()> {
834        if child.entries.len() >= MIN_ENTRIES {
835            parent.children[i] = self.save(Some(parent.children[i]), &child)?;
836            Ok(())
837        } else {
838            self.rebalance(parent, i, child, depth)
839        }
840    }
841    fn rebalance(
842        &mut self,
843        parent: &mut Node,
844        i: usize,
845        mut child: Node,
846        parent_depth: usize,
847    ) -> Result<()> {
848        let depth = parent_depth + 1;
849        let p = parent.children[i];
850        let mut left = if i > 0 {
851            Some(self.node(parent.children[i - 1], depth)?)
852        } else {
853            None
854        };
855        if let Some(left) = &mut left {
856            ensure(
857                left.children.is_empty() == child.children.is_empty(),
858                Error::Corrupt("B-tree sibling levels"),
859            )?;
860            if left.entries.len() > MIN_ENTRIES {
861                child.entries.insert(
862                    0,
863                    std::mem::replace(&mut parent.entries[i - 1], left.entries.pop().unwrap()),
864                );
865                if let Some(p) = left.children.pop() {
866                    child.children.insert(0, p);
867                }
868                parent.children[i - 1] = self.save(Some(parent.children[i - 1]), left)?;
869                parent.children[i] = self.save(Some(p), &child)?;
870                return Ok(());
871            }
872        }
873        let mut right = if i + 1 < parent.children.len() {
874            Some(self.node(parent.children[i + 1], depth)?)
875        } else {
876            None
877        };
878        if let Some(right) = &mut right {
879            ensure(
880                right.children.is_empty() == child.children.is_empty(),
881                Error::Corrupt("B-tree sibling levels"),
882            )?;
883            if right.entries.len() > MIN_ENTRIES {
884                child.entries.push(std::mem::replace(
885                    &mut parent.entries[i],
886                    right.entries.remove(0),
887                ));
888                if !right.children.is_empty() {
889                    child.children.push(right.children.remove(0));
890                }
891                parent.children[i + 1] = self.save(Some(parent.children[i + 1]), right)?;
892                parent.children[i] = self.save(Some(p), &child)?;
893                return Ok(());
894            }
895        }
896        let (left_index, mut left, right) = if let Some(left) = left {
897            (i - 1, left, child)
898        } else {
899            (
900                i,
901                child,
902                right.ok_or(Error::Corrupt("B-tree missing separator"))?,
903            )
904        };
905        ensure(
906            left_index < parent.entries.len(),
907            Error::Corrupt("B-tree missing separator"),
908        )?;
909        let left_p = parent.children[left_index];
910        let right_p = parent.children[left_index + 1];
911        left.entries.push(parent.entries.remove(left_index));
912        left.entries.extend(right.entries);
913        left.children.extend(right.children);
914        self.free(right_p)?;
915        parent.children[left_index] = self.save(Some(left_p), &left)?;
916        parent.children.remove(left_index + 1);
917        Ok(())
918    }
919    fn commit(self) -> Result<()> {
920        let changed = self.root != *self.tree.root.borrow();
921        let header = if changed {
922            Some(record::encode(MAGIC, &self.root)?)
923        } else {
924            None
925        };
926        if let Some(h) = &header {
927            ensure(h.len() as u64 <= record::HEADER_BYTES, Error::Budget)?;
928        }
929        record::reserve(&self.tree.memory, self.root.cursor)?;
930        self.tree.poisoned.set(true);
931        if self.root.cursor < self.tree.root.borrow().cursor
932            || self.root.top != self.tree.root.borrow().top
933        {
934            self.tree.cache.borrow_mut().clear();
935        }
936        for (offset, bytes) in self.writes {
937            self.tree.cache.borrow_mut().invalidate(offset);
938            self.tree.memory.write(offset, &bytes);
939        }
940        if let Some(h) = header {
941            self.tree.memory.write(0, &h);
942        }
943        *self.tree.root.borrow_mut() = self.root;
944        self.tree.poisoned.set(false);
945        Ok(())
946    }
947}
948///
949/// `MapEntry`
950///
951/// An owned key and lazily decoded value from an ordered scan.
952///
953pub struct MapEntry<K, V> {
954    key: K,
955    bytes: Vec<u8>,
956    value: PhantomData<V>,
957}
958impl<K, V: Storable> MapEntry<K, V> {
959    /// Borrow the decoded key.
960    pub const fn key(&self) -> &K {
961        &self.key
962    }
963    /// Decode the value using its declared codec.
964    ///
965    /// # Errors
966    /// Returns the application codec or declared-bound error for this value.
967    pub fn value(&self) -> Result<V> {
968        V::check_bounds(&self.bytes)?;
969        Ok(V::from_bytes(Cow::Borrowed(&self.bytes))?)
970    }
971    /// Consume the entry and decode its value.
972    ///
973    /// # Errors
974    /// Returns the application codec or declared-bound error for this value.
975    pub fn into_pair(self) -> Result<(K, V)> {
976        Ok((self.key, decode(self.bytes)?))
977    }
978    /// Consume the entry without decoding its value.
979    pub fn into_key(self) -> K {
980        self.key
981    }
982    /// Consume the entry and decode only its value.
983    ///
984    /// # Errors
985    /// Returns the application codec or declared-bound error for this value.
986    pub fn into_value(self) -> Result<V> {
987        decode(self.bytes)
988    }
989}
990
991///
992/// `TreeIter`
993///
994/// Lazy bidirectional traversal with bounded routing-node stacks. External values
995/// are read only for in-range entries and decoded when requested by the caller.
996///
997pub struct TreeIter<'a, K, V, M: Memory> {
998    tree: &'a Tree<K, V, M>,
999    lower: Bound<K>,
1000    upper: Bound<K>,
1001    front: Option<Vec<Frame>>,
1002    back: Option<Vec<Frame>>,
1003    exhausted: bool,
1004    invalid: bool,
1005    keys_only: bool,
1006}
1007impl<K: Storable + Ord + Clone, V: Storable, M: Memory> TreeIter<'_, K, V, M> {
1008    fn step(&mut self, forward: bool) -> Result<Option<MapEntry<K, V>>> {
1009        ensure(!self.invalid, Error::Bounds)?;
1010        self.tree.healthy()?;
1011        let empty = match (&self.lower, &self.upper) {
1012            (Bound::Included(a), Bound::Included(b)) => a > b,
1013            (Bound::Included(a) | Bound::Excluded(a), Bound::Included(b) | Bound::Excluded(b)) => {
1014                a >= b
1015            }
1016            _ => false,
1017        };
1018        if empty {
1019            return Ok(None);
1020        }
1021        let (path, bound) = if forward {
1022            (&mut self.front, &self.lower)
1023        } else {
1024            (&mut self.back, &self.upper)
1025        };
1026        if path.is_none() {
1027            *path = Some(self.tree.seek_path(bound, forward)?);
1028        }
1029        let entry = self
1030            .tree
1031            .advance_path(path.as_mut().unwrap(), forward, self.keys_only)?;
1032        if let Some((key, value)) = entry {
1033            let after = match &self.lower {
1034                Bound::Unbounded => true,
1035                Bound::Included(k) => key >= *k,
1036                Bound::Excluded(k) => key > *k,
1037            };
1038            let before = match &self.upper {
1039                Bound::Unbounded => true,
1040                Bound::Included(k) => key <= *k,
1041                Bound::Excluded(k) => key < *k,
1042            };
1043            ensure(
1044                if forward { after } else { before },
1045                Error::Corrupt("map cursor key order"),
1046            )?;
1047            if !(after && before) {
1048                return Ok(None);
1049            }
1050            if forward {
1051                self.lower = Bound::Excluded(key.clone());
1052            } else {
1053                self.upper = Bound::Excluded(key.clone());
1054            }
1055            return Ok(Some(MapEntry {
1056                key,
1057                bytes: value
1058                    .map(|value| self.tree.value(value))
1059                    .transpose()?
1060                    .unwrap_or_default(),
1061                value: PhantomData,
1062            }));
1063        }
1064        Ok(None)
1065    }
1066    fn advance(&mut self, forward: bool) -> Option<Result<MapEntry<K, V>>> {
1067        if self.exhausted {
1068            return None;
1069        }
1070        match self.step(forward) {
1071            Ok(Some(entry)) => Some(Ok(entry)),
1072            Ok(None) => {
1073                self.exhausted = true;
1074                None
1075            }
1076            Err(error) => {
1077                self.exhausted = true;
1078                Some(Err(error))
1079            }
1080        }
1081    }
1082}
1083impl<K: Storable + Ord + Clone, V: Storable, M: Memory> Iterator for TreeIter<'_, K, V, M> {
1084    type Item = Result<MapEntry<K, V>>;
1085    fn next(&mut self) -> Option<Self::Item> {
1086        self.advance(true)
1087    }
1088}
1089impl<K: Storable + Ord + Clone, V: Storable, M: Memory> DoubleEndedIterator
1090    for TreeIter<'_, K, V, M>
1091{
1092    fn next_back(&mut self) -> Option<Self::Item> {
1093        self.advance(false)
1094    }
1095}
1096impl<K: Storable + Ord + Clone, V: Storable, M: Memory> std::iter::FusedIterator
1097    for TreeIter<'_, K, V, M>
1098{
1099}
1100
1101#[cfg(test)]
1102mod tests;
ic_storage/adapter/tree.rs

ic_storage/adapter/
tree.rs
