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

ic_storage/adapter/
tree.rs
