ic_storage/btree/
mod.rs

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

ic_storage/btree/
mod.rs
