redb/tree_store/
btree.rs

1use crate::db::TransactionGuard;
2use crate::tree_store::btree_base::{
3    branch_checksum, leaf_checksum, BranchAccessor, BranchMutator, BtreeHeader, Checksum,
4    LeafAccessor, BRANCH, DEFERRED, LEAF,
5};
6use crate::tree_store::btree_iters::BtreeExtractIf;
7use crate::tree_store::btree_mutator::MutateHelper;
8use crate::tree_store::page_store::{Page, PageImpl, PageMut, TransactionalMemory};
9use crate::tree_store::{
10    AccessGuardMut, AllPageNumbersBtreeIter, BtreeRangeIter, PageHint, PageNumber,
11};
12use crate::types::{Key, MutInPlaceValue, Value};
13use crate::{AccessGuard, Result};
14#[cfg(feature = "logging")]
15use log::trace;
16use std::borrow::Borrow;
17use std::cmp::max;
18use std::collections::HashMap;
19use std::marker::PhantomData;
20use std::ops::RangeBounds;
21use std::sync::{Arc, Mutex};
22
23pub(crate) struct BtreeStats {
24    pub(crate) tree_height: u32,
25    pub(crate) leaf_pages: u64,
26    pub(crate) branch_pages: u64,
27    pub(crate) stored_leaf_bytes: u64,
28    pub(crate) metadata_bytes: u64,
29    pub(crate) fragmented_bytes: u64,
30}
31
32#[derive(Clone)]
33pub(crate) struct PagePath {
34    path: Vec<PageNumber>,
35}
36
37impl PagePath {
38    pub(crate) fn new_root(page_number: PageNumber) -> Self {
39        Self {
40            path: vec![page_number],
41        }
42    }
43
44    pub(crate) fn with_child(&self, page_number: PageNumber) -> Self {
45        let mut path = self.path.clone();
46        path.push(page_number);
47        Self { path }
48    }
49
50    pub(crate) fn with_subpath(&self, other: &Self) -> Self {
51        let mut path = self.path.clone();
52        path.extend(&other.path);
53        Self { path }
54    }
55
56    pub(crate) fn parents(&self) -> &[PageNumber] {
57        &self.path[..self.path.len() - 1]
58    }
59
60    pub(crate) fn page_number(&self) -> PageNumber {
61        self.path[self.path.len() - 1]
62    }
63}
64
65pub(crate) struct UntypedBtree {
66    mem: Arc<TransactionalMemory>,
67    root: Option<BtreeHeader>,
68    key_width: Option<usize>,
69    _value_width: Option<usize>,
70}
71
72impl UntypedBtree {
73    pub(crate) fn new(
74        root: Option<BtreeHeader>,
75        mem: Arc<TransactionalMemory>,
76        key_width: Option<usize>,
77        value_width: Option<usize>,
78    ) -> Self {
79        Self {
80            mem,
81            root,
82            key_width,
83            _value_width: value_width,
84        }
85    }
86
87    // Applies visitor to pages in the tree
88    pub(crate) fn visit_all_pages<F>(&self, mut visitor: F) -> Result
89    where
90        F: FnMut(&PagePath) -> Result,
91    {
92        if let Some(page_number) = self.root.map(|x| x.root) {
93            self.visit_pages_helper(PagePath::new_root(page_number), &mut visitor)?;
94        }
95
96        Ok(())
97    }
98
99    fn visit_pages_helper<F>(&self, path: PagePath, visitor: &mut F) -> Result
100    where
101        F: FnMut(&PagePath) -> Result,
102    {
103        visitor(&path)?;
104        let page = self.mem.get_page(path.page_number())?;
105
106        match page.memory()[0] {
107            LEAF => {
108                // No-op
109            }
110            BRANCH => {
111                let accessor = BranchAccessor::new(&page, self.key_width);
112                for i in 0..accessor.count_children() {
113                    let child_page = accessor.child_page(i).unwrap();
114                    let child_path = path.with_child(child_page);
115                    self.visit_pages_helper(child_path, visitor)?;
116                }
117            }
118            _ => unreachable!(),
119        }
120
121        Ok(())
122    }
123}
124
125pub(crate) struct UntypedBtreeMut {
126    mem: Arc<TransactionalMemory>,
127    root: Option<BtreeHeader>,
128    freed_pages: Arc<Mutex<Vec<PageNumber>>>,
129    key_width: Option<usize>,
130    value_width: Option<usize>,
131}
132
133impl UntypedBtreeMut {
134    pub(crate) fn new(
135        root: Option<BtreeHeader>,
136        mem: Arc<TransactionalMemory>,
137        freed_pages: Arc<Mutex<Vec<PageNumber>>>,
138        key_width: Option<usize>,
139        value_width: Option<usize>,
140    ) -> Self {
141        Self {
142            mem,
143            root,
144            freed_pages,
145            key_width,
146            value_width,
147        }
148    }
149
150    pub(crate) fn get_root(&self) -> Option<BtreeHeader> {
151        self.root
152    }
153
154    // Recomputes the checksum for all pages that are uncommitted
155    pub(crate) fn finalize_dirty_checksums(&mut self) -> Result<Option<BtreeHeader>> {
156        let mut root = self.root;
157        if let Some(BtreeHeader {
158            root: ref p,
159            ref mut checksum,
160            length: _,
161        }) = root
162        {
163            if !self.mem.uncommitted(*p) {
164                // root page is clean
165                return Ok(root);
166            }
167
168            *checksum = self.finalize_dirty_checksums_helper(*p)?;
169            self.root = root;
170        }
171
172        Ok(root)
173    }
174
175    fn finalize_dirty_checksums_helper(&mut self, page_number: PageNumber) -> Result<Checksum> {
176        assert!(self.mem.uncommitted(page_number));
177        let mut page = self.mem.get_page_mut(page_number)?;
178
179        match page.memory()[0] {
180            LEAF => leaf_checksum(&page, self.key_width, self.value_width),
181            BRANCH => {
182                let accessor = BranchAccessor::new(&page, self.key_width);
183                let mut new_children = vec![];
184                for i in 0..accessor.count_children() {
185                    let child_page = accessor.child_page(i).unwrap();
186                    if self.mem.uncommitted(child_page) {
187                        let new_checksum = self.finalize_dirty_checksums_helper(child_page)?;
188                        new_children.push(Some((i, child_page, new_checksum)));
189                    } else {
190                        // Child is clean, skip it
191                        new_children.push(None);
192                    }
193                }
194
195                let mut mutator = BranchMutator::new(&mut page);
196                for (child_index, child_page, child_checksum) in new_children.into_iter().flatten()
197                {
198                    mutator.write_child_page(child_index, child_page, child_checksum);
199                }
200
201                branch_checksum(&page, self.key_width)
202            }
203            _ => unreachable!(),
204        }
205    }
206
207    // Applies visitor to all dirty leaf pages in the tree
208    pub(crate) fn dirty_leaf_visitor<F>(&mut self, visitor: F) -> Result
209    where
210        F: Fn(PageMut) -> Result,
211    {
212        if let Some(page_number) = self.root.map(|x| x.root) {
213            if !self.mem.uncommitted(page_number) {
214                // root page is clean
215                return Ok(());
216            }
217
218            let page = self.mem.get_page_mut(page_number)?;
219            match page.memory()[0] {
220                LEAF => {
221                    visitor(page)?;
222                }
223                BRANCH => {
224                    drop(page);
225                    self.dirty_leaf_visitor_helper(page_number, &visitor)?;
226                }
227                _ => unreachable!(),
228            }
229        }
230
231        Ok(())
232    }
233
234    fn dirty_leaf_visitor_helper<F>(&mut self, page_number: PageNumber, visitor: &F) -> Result
235    where
236        F: Fn(PageMut) -> Result,
237    {
238        assert!(self.mem.uncommitted(page_number));
239        let page = self.mem.get_page_mut(page_number)?;
240
241        match page.memory()[0] {
242            LEAF => {
243                visitor(page)?;
244            }
245            BRANCH => {
246                let accessor = BranchAccessor::new(&page, self.key_width);
247                for i in 0..accessor.count_children() {
248                    let child_page = accessor.child_page(i).unwrap();
249                    if self.mem.uncommitted(child_page) {
250                        self.dirty_leaf_visitor_helper(child_page, visitor)?;
251                    }
252                }
253            }
254            _ => unreachable!(),
255        }
256
257        Ok(())
258    }
259
260    pub(crate) fn relocate(
261        &mut self,
262        relocation_map: &HashMap<PageNumber, PageNumber>,
263    ) -> Result<bool> {
264        if let Some(root) = self.get_root() {
265            if let Some((new_root, new_checksum)) =
266                self.relocate_helper(root.root, relocation_map)?
267            {
268                self.root = Some(BtreeHeader::new(new_root, new_checksum, root.length));
269                return Ok(true);
270            }
271        }
272        Ok(false)
273    }
274
275    // Relocates the given subtree to the pages specified in relocation_map
276    fn relocate_helper(
277        &mut self,
278        page_number: PageNumber,
279        relocation_map: &HashMap<PageNumber, PageNumber>,
280    ) -> Result<Option<(PageNumber, Checksum)>> {
281        let old_page = self.mem.get_page(page_number)?;
282        let mut new_page = if let Some(new_page_number) = relocation_map.get(&page_number) {
283            self.mem.get_page_mut(*new_page_number)?
284        } else {
285            return Ok(None);
286        };
287        new_page.memory_mut().copy_from_slice(old_page.memory());
288
289        let node_mem = old_page.memory();
290        match node_mem[0] {
291            LEAF => {
292                // No-op
293            }
294            BRANCH => {
295                let accessor = BranchAccessor::new(&old_page, self.key_width);
296                let mut mutator = BranchMutator::new(&mut new_page);
297                for i in 0..accessor.count_children() {
298                    let child = accessor.child_page(i).unwrap();
299                    if let Some((new_child, new_checksum)) =
300                        self.relocate_helper(child, relocation_map)?
301                    {
302                        mutator.write_child_page(i, new_child, new_checksum);
303                    }
304                }
305            }
306            _ => unreachable!(),
307        }
308
309        let mut freed_pages = self.freed_pages.lock().unwrap();
310        if !self.mem.free_if_uncommitted(page_number) {
311            freed_pages.push(page_number);
312        }
313
314        Ok(Some((new_page.get_page_number(), DEFERRED)))
315    }
316}
317
318pub(crate) struct BtreeMut<'a, K: Key + 'static, V: Value + 'static> {
319    mem: Arc<TransactionalMemory>,
320    transaction_guard: Arc<TransactionGuard>,
321    root: Option<BtreeHeader>,
322    freed_pages: Arc<Mutex<Vec<PageNumber>>>,
323    _key_type: PhantomData<K>,
324    _value_type: PhantomData<V>,
325    _lifetime: PhantomData<&'a ()>,
326}
327
328impl<K: Key + 'static, V: Value + 'static> BtreeMut<'_, K, V> {
329    pub(crate) fn new(
330        root: Option<BtreeHeader>,
331        guard: Arc<TransactionGuard>,
332        mem: Arc<TransactionalMemory>,
333        freed_pages: Arc<Mutex<Vec<PageNumber>>>,
334    ) -> Self {
335        Self {
336            mem,
337            transaction_guard: guard,
338            root,
339            freed_pages,
340            _key_type: Default::default(),
341            _value_type: Default::default(),
342            _lifetime: Default::default(),
343        }
344    }
345
346    pub(crate) fn verify_checksum(&self) -> Result<bool> {
347        RawBtree::new(
348            self.get_root(),
349            K::fixed_width(),
350            V::fixed_width(),
351            self.mem.clone(),
352        )
353        .verify_checksum()
354    }
355
356    pub(crate) fn finalize_dirty_checksums(&mut self) -> Result<Option<BtreeHeader>> {
357        let mut tree = UntypedBtreeMut::new(
358            self.get_root(),
359            self.mem.clone(),
360            self.freed_pages.clone(),
361            K::fixed_width(),
362            V::fixed_width(),
363        );
364        self.root = tree.finalize_dirty_checksums()?;
365        Ok(self.root)
366    }
367
368    #[allow(dead_code)]
369    pub(crate) fn all_pages_iter(&self) -> Result<Option<AllPageNumbersBtreeIter>> {
370        if let Some(root) = self.root.map(|x| x.root) {
371            Ok(Some(AllPageNumbersBtreeIter::new(
372                root,
373                K::fixed_width(),
374                V::fixed_width(),
375                self.mem.clone(),
376            )?))
377        } else {
378            Ok(None)
379        }
380    }
381
382    pub(crate) fn visit_all_pages<F>(&self, visitor: F) -> Result
383    where
384        F: FnMut(&PagePath) -> Result,
385    {
386        self.read_tree()?.visit_all_pages(visitor)
387    }
388
389    pub(crate) fn get_root(&self) -> Option<BtreeHeader> {
390        self.root
391    }
392
393    pub(crate) fn relocate(
394        &mut self,
395        relocation_map: &HashMap<PageNumber, PageNumber>,
396    ) -> Result<bool> {
397        let mut tree = UntypedBtreeMut::new(
398            self.get_root(),
399            self.mem.clone(),
400            self.freed_pages.clone(),
401            K::fixed_width(),
402            V::fixed_width(),
403        );
404        if tree.relocate(relocation_map)? {
405            self.root = tree.get_root();
406            Ok(true)
407        } else {
408            Ok(false)
409        }
410    }
411
412    pub(crate) fn insert(
413        &mut self,
414        key: &K::SelfType<'_>,
415        value: &V::SelfType<'_>,
416    ) -> Result<Option<AccessGuard<V>>> {
417        #[cfg(feature = "logging")]
418        trace!(
419            "Btree(root={:?}): Inserting {:?} with value of length {}",
420            &self.root,
421            key,
422            V::as_bytes(value).as_ref().len()
423        );
424        let mut freed_pages = self.freed_pages.lock().unwrap();
425        let mut operation: MutateHelper<'_, '_, K, V> =
426            MutateHelper::new(&mut self.root, self.mem.clone(), freed_pages.as_mut());
427        let (old_value, _) = operation.insert(key, value)?;
428        Ok(old_value)
429    }
430
431    // Insert without allocating or freeing any pages. This requires that you've previously
432    // inserted the same key, with a value of at least the same serialized length, earlier
433    // in the same transaction. If those preconditions aren't satisfied, insert_inplace()
434    // will panic; it won't allocate under any circumstances
435    pub(crate) fn insert_inplace(
436        &mut self,
437        key: &K::SelfType<'_>,
438        value: &V::SelfType<'_>,
439    ) -> Result<()> {
440        let mut fake_freed_pages = vec![];
441        let mut operation =
442            MutateHelper::<K, V>::new(&mut self.root, self.mem.clone(), fake_freed_pages.as_mut());
443        operation.insert_inplace(key, value)?;
444        assert!(fake_freed_pages.is_empty());
445        Ok(())
446    }
447
448    pub(crate) fn remove(&mut self, key: &K::SelfType<'_>) -> Result<Option<AccessGuard<V>>> {
449        #[cfg(feature = "logging")]
450        trace!("Btree(root={:?}): Deleting {:?}", &self.root, key);
451        let mut freed_pages = self.freed_pages.lock().unwrap();
452        let mut operation: MutateHelper<'_, '_, K, V> =
453            MutateHelper::new(&mut self.root, self.mem.clone(), freed_pages.as_mut());
454        let result = operation.delete(key)?;
455        Ok(result)
456    }
457
458    #[allow(dead_code)]
459    pub(crate) fn print_debug(&self, include_values: bool) -> Result {
460        self.read_tree()?.print_debug(include_values)
461    }
462
463    pub(crate) fn stats(&self) -> Result<BtreeStats> {
464        btree_stats(
465            self.get_root().map(|x| x.root),
466            &self.mem,
467            K::fixed_width(),
468            V::fixed_width(),
469        )
470    }
471
472    fn read_tree(&self) -> Result<Btree<K, V>> {
473        Btree::new(
474            self.get_root(),
475            PageHint::None,
476            self.transaction_guard.clone(),
477            self.mem.clone(),
478        )
479    }
480
481    pub(crate) fn get(&self, key: &K::SelfType<'_>) -> Result<Option<AccessGuard<'_, V>>> {
482        self.read_tree()?.get(key)
483    }
484
485    pub(crate) fn first(
486        &self,
487    ) -> Result<Option<(AccessGuard<'static, K>, AccessGuard<'static, V>)>> {
488        self.read_tree()?.first()
489    }
490
491    pub(crate) fn last(
492        &self,
493    ) -> Result<Option<(AccessGuard<'static, K>, AccessGuard<'static, V>)>> {
494        self.read_tree()?.last()
495    }
496
497    pub(crate) fn range<'a0, T: RangeBounds<KR> + 'a0, KR: Borrow<K::SelfType<'a0>> + 'a0>(
498        &self,
499        range: &'_ T,
500    ) -> Result<BtreeRangeIter<K, V>>
501    where
502        K: 'a0,
503    {
504        self.read_tree()?.range(range)
505    }
506
507    pub(crate) fn extract_from_if<
508        'a,
509        'a0,
510        T: RangeBounds<KR> + 'a0,
511        KR: Borrow<K::SelfType<'a0>> + 'a0,
512        F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool,
513    >(
514        &'a mut self,
515        range: &'_ T,
516        predicate: F,
517    ) -> Result<BtreeExtractIf<'a, K, V, F>>
518    where
519        K: 'a0,
520    {
521        let iter = self.range(range)?;
522
523        let result = BtreeExtractIf::new(
524            &mut self.root,
525            iter,
526            predicate,
527            self.freed_pages.clone(),
528            self.mem.clone(),
529        );
530
531        Ok(result)
532    }
533
534    pub(crate) fn retain_in<'a, KR, F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool>(
535        &mut self,
536        mut predicate: F,
537        range: impl RangeBounds<KR> + 'a,
538    ) -> Result
539    where
540        KR: Borrow<K::SelfType<'a>> + 'a,
541    {
542        let iter = self.range(&range)?;
543        let mut freed = vec![];
544        let mut operation: MutateHelper<'_, '_, K, V> =
545            MutateHelper::new_do_not_modify(&mut self.root, self.mem.clone(), &mut freed);
546        for entry in iter {
547            let entry = entry?;
548            if !predicate(entry.key(), entry.value()) {
549                assert!(operation.delete(&entry.key())?.is_some());
550            }
551        }
552        self.freed_pages.lock().unwrap().extend_from_slice(&freed);
553
554        Ok(())
555    }
556
557    pub(crate) fn len(&self) -> Result<u64> {
558        self.read_tree()?.len()
559    }
560}
561
562impl<'a, K: Key + 'a, V: MutInPlaceValue + 'a> BtreeMut<'a, K, V> {
563    /// Reserve space to insert a key-value pair
564    /// The returned reference will have length equal to `value_length`
565    // Return type has the same lifetime as &self, because the tree must not be modified until the mutable guard is dropped
566    pub(crate) fn insert_reserve(
567        &mut self,
568        key: &K::SelfType<'_>,
569        value_length: u32,
570    ) -> Result<AccessGuardMut<V>> {
571        #[cfg(feature = "logging")]
572        trace!(
573            "Btree(root={:?}): Inserting {:?} with {} reserved bytes for the value",
574            &self.root,
575            key,
576            value_length
577        );
578        let mut freed_pages = self.freed_pages.lock().unwrap();
579        let mut value = vec![0u8; value_length as usize];
580        V::initialize(&mut value);
581        let mut operation =
582            MutateHelper::<K, V>::new(&mut self.root, self.mem.clone(), freed_pages.as_mut());
583        let (_, guard) = operation.insert(key, &V::from_bytes(&value))?;
584        Ok(guard)
585    }
586}
587
588pub(crate) struct RawBtree {
589    mem: Arc<TransactionalMemory>,
590    root: Option<BtreeHeader>,
591    fixed_key_size: Option<usize>,
592    fixed_value_size: Option<usize>,
593}
594
595impl RawBtree {
596    pub(crate) fn new(
597        root: Option<BtreeHeader>,
598        fixed_key_size: Option<usize>,
599        fixed_value_size: Option<usize>,
600        mem: Arc<TransactionalMemory>,
601    ) -> Self {
602        Self {
603            mem,
604            root,
605            fixed_key_size,
606            fixed_value_size,
607        }
608    }
609
610    pub(crate) fn get_root(&self) -> Option<BtreeHeader> {
611        self.root
612    }
613
614    pub(crate) fn stats(&self) -> Result<BtreeStats> {
615        btree_stats(
616            self.root.map(|x| x.root),
617            &self.mem,
618            self.fixed_key_size,
619            self.fixed_value_size,
620        )
621    }
622
623    pub(crate) fn len(&self) -> Result<u64> {
624        Ok(self.root.map_or(0, |x| x.length))
625    }
626
627    pub(crate) fn verify_checksum(&self) -> Result<bool> {
628        if let Some(header) = self.root {
629            self.verify_checksum_helper(header.root, header.checksum)
630        } else {
631            Ok(true)
632        }
633    }
634
635    fn verify_checksum_helper(
636        &self,
637        page_number: PageNumber,
638        expected_checksum: Checksum,
639    ) -> Result<bool> {
640        let page = self.mem.get_page(page_number)?;
641        let node_mem = page.memory();
642        Ok(match node_mem[0] {
643            LEAF => {
644                if let Ok(computed) =
645                    leaf_checksum(&page, self.fixed_key_size, self.fixed_value_size)
646                {
647                    expected_checksum == computed
648                } else {
649                    false
650                }
651            }
652            BRANCH => {
653                if let Ok(computed) = branch_checksum(&page, self.fixed_key_size) {
654                    if expected_checksum != computed {
655                        return Ok(false);
656                    }
657                } else {
658                    return Ok(false);
659                }
660                let accessor = BranchAccessor::new(&page, self.fixed_key_size);
661                for i in 0..accessor.count_children() {
662                    if !self.verify_checksum_helper(
663                        accessor.child_page(i).unwrap(),
664                        accessor.child_checksum(i).unwrap(),
665                    )? {
666                        return Ok(false);
667                    }
668                }
669                true
670            }
671            _ => false,
672        })
673    }
674}
675
676pub(crate) struct Btree<K: Key + 'static, V: Value + 'static> {
677    mem: Arc<TransactionalMemory>,
678    transaction_guard: Arc<TransactionGuard>,
679    // Cache of the root page to avoid repeated lookups
680    cached_root: Option<PageImpl>,
681    root: Option<BtreeHeader>,
682    hint: PageHint,
683    _key_type: PhantomData<K>,
684    _value_type: PhantomData<V>,
685}
686
687impl<K: Key, V: Value> Btree<K, V> {
688    pub(crate) fn new(
689        root: Option<BtreeHeader>,
690        hint: PageHint,
691        guard: Arc<TransactionGuard>,
692        mem: Arc<TransactionalMemory>,
693    ) -> Result<Self> {
694        let cached_root = if let Some(header) = root {
695            Some(mem.get_page_extended(header.root, hint)?)
696        } else {
697            None
698        };
699        Ok(Self {
700            mem,
701            transaction_guard: guard,
702            cached_root,
703            root,
704            hint,
705            _key_type: Default::default(),
706            _value_type: Default::default(),
707        })
708    }
709
710    pub(crate) fn transaction_guard(&self) -> &Arc<TransactionGuard> {
711        &self.transaction_guard
712    }
713
714    pub(crate) fn get_root(&self) -> Option<BtreeHeader> {
715        self.root
716    }
717
718    pub(crate) fn visit_all_pages<F>(&self, visitor: F) -> Result
719    where
720        F: FnMut(&PagePath) -> Result,
721    {
722        let tree = UntypedBtree::new(
723            self.root,
724            self.mem.clone(),
725            K::fixed_width(),
726            V::fixed_width(),
727        );
728        tree.visit_all_pages(visitor)
729    }
730
731    pub(crate) fn get(&self, key: &K::SelfType<'_>) -> Result<Option<AccessGuard<'static, V>>> {
732        if let Some(ref root_page) = self.cached_root {
733            self.get_helper(root_page.clone(), K::as_bytes(key).as_ref())
734        } else {
735            Ok(None)
736        }
737    }
738
739    // Returns the value for the queried key, if present
740    fn get_helper(&self, page: PageImpl, query: &[u8]) -> Result<Option<AccessGuard<'static, V>>> {
741        let node_mem = page.memory();
742        match node_mem[0] {
743            LEAF => {
744                let accessor = LeafAccessor::new(page.memory(), K::fixed_width(), V::fixed_width());
745                if let Some(entry_index) = accessor.find_key::<K>(query) {
746                    let (start, end) = accessor.value_range(entry_index).unwrap();
747                    let guard = AccessGuard::with_page(page, start..end);
748                    Ok(Some(guard))
749                } else {
750                    Ok(None)
751                }
752            }
753            BRANCH => {
754                let accessor = BranchAccessor::new(&page, K::fixed_width());
755                let (_, child_page) = accessor.child_for_key::<K>(query);
756                self.get_helper(self.mem.get_page_extended(child_page, self.hint)?, query)
757            }
758            _ => unreachable!(),
759        }
760    }
761
762    pub(crate) fn first(
763        &self,
764    ) -> Result<Option<(AccessGuard<'static, K>, AccessGuard<'static, V>)>> {
765        if let Some(ref root) = self.cached_root {
766            self.first_helper(root.clone())
767        } else {
768            Ok(None)
769        }
770    }
771
772    fn first_helper(
773        &self,
774        page: PageImpl,
775    ) -> Result<Option<(AccessGuard<'static, K>, AccessGuard<'static, V>)>> {
776        let node_mem = page.memory();
777        match node_mem[0] {
778            LEAF => {
779                let accessor = LeafAccessor::new(page.memory(), K::fixed_width(), V::fixed_width());
780                let (key_range, value_range) = accessor.entry_ranges(0).unwrap();
781                let key_guard = AccessGuard::with_page(page.clone(), key_range);
782                let value_guard = AccessGuard::with_page(page, value_range);
783                Ok(Some((key_guard, value_guard)))
784            }
785            BRANCH => {
786                let accessor = BranchAccessor::new(&page, K::fixed_width());
787                let child_page = accessor.child_page(0).unwrap();
788                self.first_helper(self.mem.get_page_extended(child_page, self.hint)?)
789            }
790            _ => unreachable!(),
791        }
792    }
793
794    pub(crate) fn last(
795        &self,
796    ) -> Result<Option<(AccessGuard<'static, K>, AccessGuard<'static, V>)>> {
797        if let Some(ref root) = self.cached_root {
798            self.last_helper(root.clone())
799        } else {
800            Ok(None)
801        }
802    }
803
804    fn last_helper(
805        &self,
806        page: PageImpl,
807    ) -> Result<Option<(AccessGuard<'static, K>, AccessGuard<'static, V>)>> {
808        let node_mem = page.memory();
809        match node_mem[0] {
810            LEAF => {
811                let accessor = LeafAccessor::new(page.memory(), K::fixed_width(), V::fixed_width());
812                let (key_range, value_range) =
813                    accessor.entry_ranges(accessor.num_pairs() - 1).unwrap();
814                let key_guard = AccessGuard::with_page(page.clone(), key_range);
815                let value_guard = AccessGuard::with_page(page, value_range);
816                Ok(Some((key_guard, value_guard)))
817            }
818            BRANCH => {
819                let accessor = BranchAccessor::new(&page, K::fixed_width());
820                let child_page = accessor.child_page(accessor.count_children() - 1).unwrap();
821                self.last_helper(self.mem.get_page_extended(child_page, self.hint)?)
822            }
823            _ => unreachable!(),
824        }
825    }
826
827    pub(crate) fn range<'a0, T: RangeBounds<KR>, KR: Borrow<K::SelfType<'a0>>>(
828        &self,
829        range: &'_ T,
830    ) -> Result<BtreeRangeIter<K, V>> {
831        BtreeRangeIter::new(range, self.root.map(|x| x.root), self.mem.clone())
832    }
833
834    pub(crate) fn len(&self) -> Result<u64> {
835        Ok(self.root.map_or(0, |x| x.length))
836    }
837
838    pub(crate) fn stats(&self) -> Result<BtreeStats> {
839        btree_stats(
840            self.root.map(|x| x.root),
841            &self.mem,
842            K::fixed_width(),
843            V::fixed_width(),
844        )
845    }
846
847    #[allow(dead_code)]
848    pub(crate) fn print_debug(&self, include_values: bool) -> Result {
849        if let Some(p) = self.root.map(|x| x.root) {
850            let mut pages = vec![self.mem.get_page(p)?];
851            while !pages.is_empty() {
852                let mut next_children = vec![];
853                for page in pages.drain(..) {
854                    let node_mem = page.memory();
855                    match node_mem[0] {
856                        LEAF => {
857                            eprint!("Leaf[ (page={:?})", page.get_page_number());
858                            LeafAccessor::new(page.memory(), K::fixed_width(), V::fixed_width())
859                                .print_node::<K, V>(include_values);
860                            eprint!("]");
861                        }
862                        BRANCH => {
863                            let accessor = BranchAccessor::new(&page, K::fixed_width());
864                            for i in 0..accessor.count_children() {
865                                let child = accessor.child_page(i).unwrap();
866                                next_children.push(self.mem.get_page(child)?);
867                            }
868                            accessor.print_node::<K>();
869                        }
870                        _ => unreachable!(),
871                    }
872                    eprint!("  ");
873                }
874                eprintln!();
875
876                pages = next_children;
877            }
878        }
879
880        Ok(())
881    }
882}
883
884pub(crate) fn btree_stats(
885    root: Option<PageNumber>,
886    mem: &TransactionalMemory,
887    fixed_key_size: Option<usize>,
888    fixed_value_size: Option<usize>,
889) -> Result<BtreeStats> {
890    if let Some(root) = root {
891        stats_helper(root, mem, fixed_key_size, fixed_value_size)
892    } else {
893        Ok(BtreeStats {
894            tree_height: 0,
895            leaf_pages: 0,
896            branch_pages: 0,
897            stored_leaf_bytes: 0,
898            metadata_bytes: 0,
899            fragmented_bytes: 0,
900        })
901    }
902}
903
904fn stats_helper(
905    page_number: PageNumber,
906    mem: &TransactionalMemory,
907    fixed_key_size: Option<usize>,
908    fixed_value_size: Option<usize>,
909) -> Result<BtreeStats> {
910    let page = mem.get_page(page_number)?;
911    let node_mem = page.memory();
912    match node_mem[0] {
913        LEAF => {
914            let accessor = LeafAccessor::new(page.memory(), fixed_key_size, fixed_value_size);
915            let leaf_bytes = accessor.length_of_pairs(0, accessor.num_pairs());
916            let overhead_bytes = accessor.total_length() - leaf_bytes;
917            let fragmented_bytes = (page.memory().len() - accessor.total_length()) as u64;
918            Ok(BtreeStats {
919                tree_height: 1,
920                leaf_pages: 1,
921                branch_pages: 0,
922                stored_leaf_bytes: leaf_bytes.try_into().unwrap(),
923                metadata_bytes: overhead_bytes.try_into().unwrap(),
924                fragmented_bytes,
925            })
926        }
927        BRANCH => {
928            let accessor = BranchAccessor::new(&page, fixed_key_size);
929            let mut max_child_height = 0;
930            let mut leaf_pages = 0;
931            let mut branch_pages = 1;
932            let mut stored_leaf_bytes = 0;
933            let mut metadata_bytes = accessor.total_length() as u64;
934            let mut fragmented_bytes = (page.memory().len() - accessor.total_length()) as u64;
935            for i in 0..accessor.count_children() {
936                if let Some(child) = accessor.child_page(i) {
937                    let stats = stats_helper(child, mem, fixed_key_size, fixed_value_size)?;
938                    max_child_height = max(max_child_height, stats.tree_height);
939                    leaf_pages += stats.leaf_pages;
940                    branch_pages += stats.branch_pages;
941                    stored_leaf_bytes += stats.stored_leaf_bytes;
942                    metadata_bytes += stats.metadata_bytes;
943                    fragmented_bytes += stats.fragmented_bytes;
944                }
945            }
946
947            Ok(BtreeStats {
948                tree_height: max_child_height + 1,
949                leaf_pages,
950                branch_pages,
951                stored_leaf_bytes,
952                metadata_bytes,
953                fragmented_bytes,
954            })
955        }
956        _ => unreachable!(),
957    }
958}
redb/tree_store/btree.rs

redb/tree_store/
btree.rs
