ic_storage/btree/
mod.rs

1//! Fixed-size tree pages with separately allocated key/value chains.
2
3mod cursor;
4mod page;
5mod transaction;
6
7pub use cursor::{Entry, Range};
8use page::{Header, Node, Record};
9use std::{cmp::Ordering, fmt, ops::Bound};
10use transaction::Transaction;
11
12use crate::Memory;
13
14/// A rejected operation. Returned errors never publish a partial mutation.
15#[derive(Clone, Copy, Debug, Eq, PartialEq)]
16pub enum Error {
17    /// Creation requires a fresh, zero-sized memory.
18    NotEmpty,
19    /// Memory does not contain this tree's format.
20    Format,
21    /// Stored configuration or comparator identity differs from the caller.
22    Configuration,
23    /// Invalid page, pointer, length, checksum, or tree geometry.
24    Corrupt,
25    /// An input exceeds the configured key or value bound.
26    TooLarge,
27    /// The backing memory refused growth or arithmetic capacity was exhausted.
28    Capacity,
29    /// The comparator rejected malformed input.
30    InvalidKey,
31    /// The lower bound follows the upper bound.
32    InvalidRange,
33}
34
35impl fmt::Display for Error {
36    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
37        write!(f, "stable B-tree error: {self:?}")
38    }
39}
40impl std::error::Error for Error {}
41
42/// Immutable, bounded decoding limits stored in the tree header.
43#[derive(Clone, Copy, Debug, Eq, PartialEq)]
44pub struct Config {
45    /// Largest encoded key accepted by this tree, in bytes.
46    pub max_key_bytes: u32,
47    /// Largest encoded value accepted by this tree, in bytes.
48    pub max_value_bytes: u32,
49}
50impl Default for Config {
51    fn default() -> Self {
52        Self {
53            max_key_bytes: 128 * 1024,
54            max_value_bytes: 4 * 1024 * 1024,
55        }
56    }
57}
58impl Config {
59    fn check(self) -> Result<(), Error> {
60        if self.max_key_bytes == 0
61            || self.max_key_bytes > 1024 * 1024
62            || self.max_value_bytes == 0
63            || self.max_value_bytes > 64 * 1024 * 1024
64        {
65            return Err(Error::Configuration);
66        }
67        Ok(())
68    }
69}
70
71/// A deterministic total order over encoded keys.
72///
73/// `IDENTITY` must change when ordering or accepted key encoding changes.
74/// Equal keys must have identical encodings. The comparator must be immutable
75/// and return `InvalidKey` for malformed keys, including `compare(key, key)`.
76/// Persisted comparator code itself remains an application responsibility.
77pub trait Comparator {
78    /// Stable identity of the encoding and ordering contract.
79    const IDENTITY: [u8; 16];
80    /// Compare two encoded keys, rejecting invalid encodings.
81    ///
82    /// # Errors
83    /// Returns [`Error::InvalidKey`] when an encoded key is invalid.
84    fn compare(&self, left: &[u8], right: &[u8]) -> Result<Ordering, Error>;
85}
86
87/// Ordinary unsigned bytewise ordering for already order-preserving encodings.
88#[derive(Clone, Copy, Debug, Default)]
89pub struct Lexicographic;
90impl Comparator for Lexicographic {
91    const IDENTITY: [u8; 16] = *b"ic-storage-lex01";
92    fn compare(&self, left: &[u8], right: &[u8]) -> Result<Ordering, Error> {
93        Ok(left.cmp(right))
94    }
95}
96
97/// A stable-memory B-tree, independent of DFINITY's collection formats.
98///
99/// Exactly one live tree handle may own a memory region. The memory must not be
100/// mutated through aliases. Mutations stage touched pages, reserve capacity,
101/// then write them. Backing writes must either complete or trap in an atomic IC
102/// message; this is not a crash-safe filesystem database. See the package guide.
103pub struct BTreeMap<M, C = Lexicographic> {
104    memory: M,
105    order: C,
106    header: Header,
107}
108
109impl<M: Memory, C: Comparator> BTreeMap<M, C> {
110    /// Create a tree in an empty memory without overwriting an existing format.
111    ///
112    /// # Errors
113    /// Rejects nonempty memory, invalid limits, and refused growth.
114    pub fn create(memory: M, config: Config, order: C) -> Result<Self, Error> {
115        config.check()?;
116        if memory.size() != 0 {
117            return Err(Error::NotEmpty);
118        }
119        let mut tree = Self {
120            memory,
121            order,
122            header: Header::new(config, C::IDENTITY),
123        };
124        let mut tx = Transaction::new(&tree);
125        let root = tx.allocate()?;
126        tx.put_node(root, &Node::leaf())?;
127        tx.header.root = root;
128        tree.header = tx.commit()?;
129        Ok(tree)
130    }
131
132    /// Open a tree with its exact configuration and comparator identity.
133    ///
134    /// Opening checks the header and root, not every page. Reads validate pages
135    /// as visited; [`Self::validate`] performs a complete structural audit.
136    /// # Errors
137    /// Rejects unknown formats, incompatible limits/order, and observed corruption.
138    pub fn open(memory: M, config: Config, order: C) -> Result<Self, Error> {
139        config.check()?;
140        let header = Header::read(&memory)?;
141        if header.config != config || header.identity != C::IDENTITY {
142            return Err(Error::Configuration);
143        }
144        let tree = Self {
145            memory,
146            order,
147            header,
148        };
149        tree.node(tree.header.root)?;
150        Ok(tree)
151    }
152
153    /// Number of live entries, persisted with each completed mutation.
154    #[must_use]
155    pub const fn len(&self) -> u64 {
156        self.header.len
157    }
158    /// Whether the tree contains no entries.
159    #[must_use]
160    pub const fn is_empty(&self) -> bool {
161        self.header.len == 0
162    }
163    /// Bytes reserved in the backing memory, including reusable free pages.
164    #[must_use]
165    pub fn memory_bytes(&self) -> u64 {
166        self.memory.size().saturating_mul(65_536)
167    }
168    /// Release the sole handle and return its backing memory for later reopening.
169    #[must_use]
170    pub fn into_memory(self) -> M {
171        self.memory
172    }
173
174    /// Read one value by key.
175    /// # Errors
176    /// Rejects malformed keys, oversized keys, and observed storage corruption.
177    pub fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Error> {
178        self.find(key)?.map(|record| self.value(record)).transpose()
179    }
180
181    /// Test key membership without reading the value payload.
182    /// # Errors
183    /// Rejects invalid keys and observed key or tree-page corruption.
184    pub fn contains_key(&self, key: &[u8]) -> Result<bool, Error> {
185        Ok(self.find(key)?.is_some())
186    }
187
188    fn find(&self, key: &[u8]) -> Result<Option<Record>, Error> {
189        self.check_key(key)?;
190        let mut id = self.header.root;
191        for _ in 0..page::MAX_HEIGHT {
192            let node = self.node(id)?;
193            let (at, found) = self.locate(&node, key)?;
194            if found {
195                return Ok(Some(node.entries[at]));
196            }
197            if node.children.is_empty() {
198                return Ok(None);
199            }
200            id = node.children[at];
201        }
202        Err(Error::Corrupt)
203    }
204
205    /// Insert or replace a value, returning the previous value if present.
206    /// # Errors
207    /// Rejects invalid keys, oversized inputs, corruption, or insufficient capacity
208    /// before publishing changes. Memory traps follow the backing contract.
209    pub fn insert(&mut self, key: &[u8], value: &[u8]) -> Result<Option<Vec<u8>>, Error> {
210        self.check_key(key)?;
211        if value.len() > self.header.config.max_value_bytes as usize {
212            return Err(Error::TooLarge);
213        }
214        let previous = self.get(key)?;
215        let mut tx = Transaction::new(self);
216        tx.insert(key, value)?;
217        if previous.is_none() {
218            tx.header.len = tx.header.len.checked_add(1).ok_or(Error::Capacity)?;
219        }
220        self.header = tx.commit()?;
221        Ok(previous)
222    }
223
224    /// Remove a key, returning its previous value.
225    /// # Errors
226    /// Rejects malformed keys and observed corruption without publishing changes.
227    pub fn remove(&mut self, key: &[u8]) -> Result<Option<Vec<u8>>, Error> {
228        let previous = self.get(key)?;
229        if previous.is_none() {
230            return Ok(None);
231        }
232        let mut tx = Transaction::new(self);
233        tx.remove(key)?;
234        tx.header.len = tx.header.len.checked_sub(1).ok_or(Error::Corrupt)?;
235        self.header = tx.commit()?;
236        Ok(previous)
237    }
238
239    /// Iterate a bounded range in either direction without reading value payloads.
240    ///
241    /// Bounds use the supplied comparator. Equal exclusive bounds yield an empty
242    /// range. `rev()` and interleaved `next`/`next_back` are supported.
243    /// # Errors
244    /// Rejects invalid bounds; traversal errors are yielded by the iterator.
245    pub fn range(
246        &self,
247        lower: Bound<&[u8]>,
248        upper: Bound<&[u8]>,
249    ) -> Result<Range<'_, M, C>, Error> {
250        Range::new(self, lower, upper)
251    }
252
253    /// Traverse all entries in comparator order, with lazy value reads.
254    #[must_use]
255    pub fn iter(&self) -> Range<'_, M, C> {
256        Range::all(self)
257    }
258
259    /// Check every reachable page, key order, depth, occupancy, entry count and
260    /// free-list ownership. This explicit audit is linear in allocated storage.
261    /// # Errors
262    /// Returns an error on corruption, invalid encoded keys or invalid ordering.
263    pub fn validate(&self) -> Result<(), Error> {
264        transaction::validate(self)
265    }
266
267    fn check_key(&self, key: &[u8]) -> Result<(), Error> {
268        if key.len() > self.header.config.max_key_bytes as usize {
269            return Err(Error::TooLarge);
270        }
271        if self.order.compare(key, key)? != Ordering::Equal {
272            return Err(Error::InvalidKey);
273        }
274        Ok(())
275    }
276    fn node(&self, id: u64) -> Result<Node, Error> {
277        Node::decode(
278            &page::read(&self.memory, self.header.next, id)?,
279            self.header.next,
280            self.header.config,
281        )
282    }
283    fn key(&self, record: Record) -> Result<Vec<u8>, Error> {
284        if record.key == 0 {
285            return Ok(record.key_inline
286                [..usize::try_from(record.key_len).map_err(|_| Error::Corrupt)?]
287                .to_vec());
288        }
289        page::blob(&self.memory, self.header.next, record.key, record.key_len)
290    }
291    fn value(&self, record: Record) -> Result<Vec<u8>, Error> {
292        if record.value == 0 {
293            return Ok(record.value_inline
294                [..usize::try_from(record.value_len).map_err(|_| Error::Corrupt)?]
295                .to_vec());
296        }
297        page::blob(
298            &self.memory,
299            self.header.next,
300            record.value,
301            record.value_len,
302        )
303    }
304    fn locate(&self, node: &Node, key: &[u8]) -> Result<(usize, bool), Error> {
305        let (mut lo, mut hi) = (0, node.entries.len());
306        while lo < hi {
307            let mid = lo + (hi - lo) / 2;
308            match self.order.compare(&self.key(node.entries[mid])?, key)? {
309                Ordering::Less => lo = mid + 1,
310                Ordering::Equal => return Ok((mid, true)),
311                Ordering::Greater => hi = mid,
312            }
313        }
314        Ok((lo, false))
315    }
316}
317
318impl<'a, M: Memory, C: Comparator> IntoIterator for &'a BTreeMap<M, C> {
319    type Item = Result<Entry<'a, M, C>, Error>;
320    type IntoIter = Range<'a, M, C>;
321    fn into_iter(self) -> Self::IntoIter {
322        self.iter()
323    }
324}
ic_storage/btree/mod.rs

ic_storage/btree/
mod.rs
