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

ic_storage/adapter/
tree.rs
