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        let tree = UntypedBtree::new(
387            self.root,
388            self.mem.clone(),
389            K::fixed_width(),
390            V::fixed_width(),
391        );
392        tree.visit_all_pages(visitor)
393    }
394
395    pub(crate) fn get_root(&self) -> Option<BtreeHeader> {
396        self.root
397    }
398
399    pub(crate) fn relocate(
400        &mut self,
401        relocation_map: &HashMap<PageNumber, PageNumber>,
402    ) -> Result<bool> {
403        let mut tree = UntypedBtreeMut::new(
404            self.get_root(),
405            self.mem.clone(),
406            self.freed_pages.clone(),
407            K::fixed_width(),
408            V::fixed_width(),
409        );
410        if tree.relocate(relocation_map)? {
411            self.root = tree.get_root();
412            Ok(true)
413        } else {
414            Ok(false)
415        }
416    }
417
418    pub(crate) fn insert(
419        &mut self,
420        key: &K::SelfType<'_>,
421        value: &V::SelfType<'_>,
422    ) -> Result<Option<AccessGuard<V>>> {
423        #[cfg(feature = "logging")]
424        trace!(
425            "Btree(root={:?}): Inserting {:?} with value of length {}",
426            &self.root,
427            key,
428            V::as_bytes(value).as_ref().len()
429        );
430        let mut freed_pages = self.freed_pages.lock().unwrap();
431        let mut operation: MutateHelper<'_, '_, K, V> =
432            MutateHelper::new(&mut self.root, self.mem.clone(), freed_pages.as_mut());
433        let (old_value, _) = operation.insert(key, value)?;
434        Ok(old_value)
435    }
436
437    // Insert without allocating or freeing any pages. This requires that you've previously
438    // inserted the same key, with a value of at least the same serialized length, earlier
439    // in the same transaction. If those preconditions aren't satisfied, insert_inplace()
440    // will panic; it won't allocate under any circumstances
441    pub(crate) fn insert_inplace(
442        &mut self,
443        key: &K::SelfType<'_>,
444        value: &V::SelfType<'_>,
445    ) -> Result<()> {
446        let mut fake_freed_pages = vec![];
447        let mut operation =
448            MutateHelper::<K, V>::new(&mut self.root, self.mem.clone(), fake_freed_pages.as_mut());
449        operation.insert_inplace(key, value)?;
450        assert!(fake_freed_pages.is_empty());
451        Ok(())
452    }
453
454    pub(crate) fn remove(&mut self, key: &K::SelfType<'_>) -> Result<Option<AccessGuard<V>>> {
455        #[cfg(feature = "logging")]
456        trace!("Btree(root={:?}): Deleting {:?}", &self.root, key);
457        let mut freed_pages = self.freed_pages.lock().unwrap();
458        let mut operation: MutateHelper<'_, '_, K, V> =
459            MutateHelper::new(&mut self.root, self.mem.clone(), freed_pages.as_mut());
460        let result = operation.delete(key)?;
461        Ok(result)
462    }
463
464    #[allow(dead_code)]
465    pub(crate) fn print_debug(&self, include_values: bool) -> Result {
466        self.read_tree()?.print_debug(include_values)
467    }
468
469    pub(crate) fn stats(&self) -> Result<BtreeStats> {
470        btree_stats(
471            self.get_root().map(|x| x.root),
472            &self.mem,
473            K::fixed_width(),
474            V::fixed_width(),
475        )
476    }
477
478    fn read_tree(&self) -> Result<Btree<K, V>> {
479        Btree::new(
480            self.get_root(),
481            PageHint::None,
482            self.transaction_guard.clone(),
483            self.mem.clone(),
484        )
485    }
486
487    pub(crate) fn get(&self, key: &K::SelfType<'_>) -> Result<Option<AccessGuard<'_, V>>> {
488        self.read_tree()?.get(key)
489    }
490
491    pub(crate) fn first(
492        &self,
493    ) -> Result<Option<(AccessGuard<'static, K>, AccessGuard<'static, V>)>> {
494        self.read_tree()?.first()
495    }
496
497    pub(crate) fn last(
498        &self,
499    ) -> Result<Option<(AccessGuard<'static, K>, AccessGuard<'static, V>)>> {
500        self.read_tree()?.last()
501    }
502
503    pub(crate) fn range<'a0, T: RangeBounds<KR> + 'a0, KR: Borrow<K::SelfType<'a0>> + 'a0>(
504        &self,
505        range: &'_ T,
506    ) -> Result<BtreeRangeIter<K, V>>
507    where
508        K: 'a0,
509    {
510        self.read_tree()?.range(range)
511    }
512
513    pub(crate) fn extract_from_if<
514        'a,
515        'a0,
516        T: RangeBounds<KR> + 'a0,
517        KR: Borrow<K::SelfType<'a0>> + 'a0,
518        F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool,
519    >(
520        &'a mut self,
521        range: &'_ T,
522        predicate: F,
523    ) -> Result<BtreeExtractIf<'a, K, V, F>>
524    where
525        K: 'a0,
526    {
527        let iter = self.range(range)?;
528
529        let result = BtreeExtractIf::new(
530            &mut self.root,
531            iter,
532            predicate,
533            self.freed_pages.clone(),
534            self.mem.clone(),
535        );
536
537        Ok(result)
538    }
539
540    pub(crate) fn retain_in<'a, KR, F: for<'f> FnMut(K::SelfType<'f>, V::SelfType<'f>) -> bool>(
541        &mut self,
542        mut predicate: F,
543        range: impl RangeBounds<KR> + 'a,
544    ) -> Result
545    where
546        KR: Borrow<K::SelfType<'a>> + 'a,
547    {
548        let iter = self.range(&range)?;
549        let mut freed = vec![];
550        let mut operation: MutateHelper<'_, '_, K, V> =
551            MutateHelper::new_do_not_modify(&mut self.root, self.mem.clone(), &mut freed);
552        for entry in iter {
553            let entry = entry?;
554            if !predicate(entry.key(), entry.value()) {
555                assert!(operation.delete(&entry.key())?.is_some());
556            }
557        }
558        self.freed_pages.lock().unwrap().extend_from_slice(&freed);
559
560        Ok(())
561    }
562
563    pub(crate) fn len(&self) -> Result<u64> {
564        self.read_tree()?.len()
565    }
566}
567
568impl<'a, K: Key + 'a, V: MutInPlaceValue + 'a> BtreeMut<'a, K, V> {
569    /// Reserve space to insert a key-value pair
570    /// The returned reference will have length equal to `value_length`
571    // Return type has the same lifetime as &self, because the tree must not be modified until the mutable guard is dropped
572    pub(crate) fn insert_reserve(
573        &mut self,
574        key: &K::SelfType<'_>,
575        value_length: u32,
576    ) -> Result<AccessGuardMut<V>> {
577        #[cfg(feature = "logging")]
578        trace!(
579            "Btree(root={:?}): Inserting {:?} with {} reserved bytes for the value",
580            &self.root,
581            key,
582            value_length
583        );
584        let mut freed_pages = self.freed_pages.lock().unwrap();
585        let mut value = vec![0u8; value_length as usize];
586        V::initialize(&mut value);
587        let mut operation =
588            MutateHelper::<K, V>::new(&mut self.root, self.mem.clone(), freed_pages.as_mut());
589        let (_, guard) = operation.insert(key, &V::from_bytes(&value))?;
590        Ok(guard)
591    }
592}
593
594pub(crate) struct RawBtree {
595    mem: Arc<TransactionalMemory>,
596    root: Option<BtreeHeader>,
597    fixed_key_size: Option<usize>,
598    fixed_value_size: Option<usize>,
599}
600
601impl RawBtree {
602    pub(crate) fn new(
603        root: Option<BtreeHeader>,
604        fixed_key_size: Option<usize>,
605        fixed_value_size: Option<usize>,
606        mem: Arc<TransactionalMemory>,
607    ) -> Self {
608        Self {
609            mem,
610            root,
611            fixed_key_size,
612            fixed_value_size,
613        }
614    }
615
616    pub(crate) fn get_root(&self) -> Option<BtreeHeader> {
617        self.root
618    }
619
620    pub(crate) fn stats(&self) -> Result<BtreeStats> {
621        btree_stats(
622            self.root.map(|x| x.root),
623            &self.mem,
624            self.fixed_key_size,
625            self.fixed_value_size,
626        )
627    }
628
629    pub(crate) fn len(&self) -> Result<u64> {
630        Ok(self.root.map_or(0, |x| x.length))
631    }
632
633    pub(crate) fn verify_checksum(&self) -> Result<bool> {
634        if let Some(header) = self.root {
635            self.verify_checksum_helper(header.root, header.checksum)
636        } else {
637            Ok(true)
638        }
639    }
640
641    fn verify_checksum_helper(
642        &self,
643        page_number: PageNumber,
644        expected_checksum: Checksum,
645    ) -> Result<bool> {
646        let page = self.mem.get_page(page_number)?;
647        let node_mem = page.memory();
648        Ok(match node_mem[0] {
649            LEAF => {
650                if let Ok(computed) =
651                    leaf_checksum(&page, self.fixed_key_size, self.fixed_value_size)
652                {
653                    expected_checksum == computed
654                } else {
655                    false
656                }
657            }
658            BRANCH => {
659                if let Ok(computed) = branch_checksum(&page, self.fixed_key_size) {
660                    if expected_checksum != computed {
661                        return Ok(false);
662                    }
663                } else {
664                    return Ok(false);
665                }
666                let accessor = BranchAccessor::new(&page, self.fixed_key_size);
667                for i in 0..accessor.count_children() {
668                    if !self.verify_checksum_helper(
669                        accessor.child_page(i).unwrap(),
670                        accessor.child_checksum(i).unwrap(),
671                    )? {
672                        return Ok(false);
673                    }
674                }
675                true
676            }
677            _ => false,
678        })
679    }
680}
681
682pub(crate) struct Btree<K: Key + 'static, V: Value + 'static> {
683    mem: Arc<TransactionalMemory>,
684    transaction_guard: Arc<TransactionGuard>,
685    // Cache of the root page to avoid repeated lookups
686    cached_root: Option<PageImpl>,
687    root: Option<BtreeHeader>,
688    hint: PageHint,
689    _key_type: PhantomData<K>,
690    _value_type: PhantomData<V>,
691}
692
693impl<K: Key, V: Value> Btree<K, V> {
694    pub(crate) fn new(
695        root: Option<BtreeHeader>,
696        hint: PageHint,
697        guard: Arc<TransactionGuard>,
698        mem: Arc<TransactionalMemory>,
699    ) -> Result<Self> {
700        let cached_root = if let Some(header) = root {
701            Some(mem.get_page_extended(header.root, hint)?)
702        } else {
703            None
704        };
705        Ok(Self {
706            mem,
707            transaction_guard: guard,
708            cached_root,
709            root,
710            hint,
711            _key_type: Default::default(),
712            _value_type: Default::default(),
713        })
714    }
715
716    pub(crate) fn transaction_guard(&self) -> &Arc<TransactionGuard> {
717        &self.transaction_guard
718    }
719
720    pub(crate) fn get_root(&self) -> Option<BtreeHeader> {
721        self.root
722    }
723
724    pub(crate) fn get(&self, key: &K::SelfType<'_>) -> Result<Option<AccessGuard<'static, V>>> {
725        if let Some(ref root_page) = self.cached_root {
726            self.get_helper(root_page.clone(), K::as_bytes(key).as_ref())
727        } else {
728            Ok(None)
729        }
730    }
731
732    // Returns the value for the queried key, if present
733    fn get_helper(&self, page: PageImpl, query: &[u8]) -> Result<Option<AccessGuard<'static, V>>> {
734        let node_mem = page.memory();
735        match node_mem[0] {
736            LEAF => {
737                let accessor = LeafAccessor::new(page.memory(), K::fixed_width(), V::fixed_width());
738                if let Some(entry_index) = accessor.find_key::<K>(query) {
739                    let (start, end) = accessor.value_range(entry_index).unwrap();
740                    let guard = AccessGuard::with_page(page, start..end);
741                    Ok(Some(guard))
742                } else {
743                    Ok(None)
744                }
745            }
746            BRANCH => {
747                let accessor = BranchAccessor::new(&page, K::fixed_width());
748                let (_, child_page) = accessor.child_for_key::<K>(query);
749                self.get_helper(self.mem.get_page_extended(child_page, self.hint)?, query)
750            }
751            _ => unreachable!(),
752        }
753    }
754
755    pub(crate) fn first(
756        &self,
757    ) -> Result<Option<(AccessGuard<'static, K>, AccessGuard<'static, V>)>> {
758        if let Some(ref root) = self.cached_root {
759            self.first_helper(root.clone())
760        } else {
761            Ok(None)
762        }
763    }
764
765    fn first_helper(
766        &self,
767        page: PageImpl,
768    ) -> Result<Option<(AccessGuard<'static, K>, AccessGuard<'static, V>)>> {
769        let node_mem = page.memory();
770        match node_mem[0] {
771            LEAF => {
772                let accessor = LeafAccessor::new(page.memory(), K::fixed_width(), V::fixed_width());
773                let (key_range, value_range) = accessor.entry_ranges(0).unwrap();
774                let key_guard = AccessGuard::with_page(page.clone(), key_range);
775                let value_guard = AccessGuard::with_page(page, value_range);
776                Ok(Some((key_guard, value_guard)))
777            }
778            BRANCH => {
779                let accessor = BranchAccessor::new(&page, K::fixed_width());
780                let child_page = accessor.child_page(0).unwrap();
781                self.first_helper(self.mem.get_page_extended(child_page, self.hint)?)
782            }
783            _ => unreachable!(),
784        }
785    }
786
787    pub(crate) fn last(
788        &self,
789    ) -> Result<Option<(AccessGuard<'static, K>, AccessGuard<'static, V>)>> {
790        if let Some(ref root) = self.cached_root {
791            self.last_helper(root.clone())
792        } else {
793            Ok(None)
794        }
795    }
796
797    fn last_helper(
798        &self,
799        page: PageImpl,
800    ) -> Result<Option<(AccessGuard<'static, K>, AccessGuard<'static, V>)>> {
801        let node_mem = page.memory();
802        match node_mem[0] {
803            LEAF => {
804                let accessor = LeafAccessor::new(page.memory(), K::fixed_width(), V::fixed_width());
805                let (key_range, value_range) =
806                    accessor.entry_ranges(accessor.num_pairs() - 1).unwrap();
807                let key_guard = AccessGuard::with_page(page.clone(), key_range);
808                let value_guard = AccessGuard::with_page(page, value_range);
809                Ok(Some((key_guard, value_guard)))
810            }
811            BRANCH => {
812                let accessor = BranchAccessor::new(&page, K::fixed_width());
813                let child_page = accessor.child_page(accessor.count_children() - 1).unwrap();
814                self.last_helper(self.mem.get_page_extended(child_page, self.hint)?)
815            }
816            _ => unreachable!(),
817        }
818    }
819
820    pub(crate) fn range<'a0, T: RangeBounds<KR>, KR: Borrow<K::SelfType<'a0>>>(
821        &self,
822        range: &'_ T,
823    ) -> Result<BtreeRangeIter<K, V>> {
824        BtreeRangeIter::new(range, self.root.map(|x| x.root), self.mem.clone())
825    }
826
827    pub(crate) fn len(&self) -> Result<u64> {
828        Ok(self.root.map_or(0, |x| x.length))
829    }
830
831    pub(crate) fn stats(&self) -> Result<BtreeStats> {
832        btree_stats(
833            self.root.map(|x| x.root),
834            &self.mem,
835            K::fixed_width(),
836            V::fixed_width(),
837        )
838    }
839
840    #[allow(dead_code)]
841    pub(crate) fn print_debug(&self, include_values: bool) -> Result {
842        if let Some(p) = self.root.map(|x| x.root) {
843            let mut pages = vec![self.mem.get_page(p)?];
844            while !pages.is_empty() {
845                let mut next_children = vec![];
846                for page in pages.drain(..) {
847                    let node_mem = page.memory();
848                    match node_mem[0] {
849                        LEAF => {
850                            eprint!("Leaf[ (page={:?})", page.get_page_number());
851                            LeafAccessor::new(page.memory(), K::fixed_width(), V::fixed_width())
852                                .print_node::<K, V>(include_values);
853                            eprint!("]");
854                        }
855                        BRANCH => {
856                            let accessor = BranchAccessor::new(&page, K::fixed_width());
857                            for i in 0..accessor.count_children() {
858                                let child = accessor.child_page(i).unwrap();
859                                next_children.push(self.mem.get_page(child)?);
860                            }
861                            accessor.print_node::<K>();
862                        }
863                        _ => unreachable!(),
864                    }
865                    eprint!("  ");
866                }
867                eprintln!();
868
869                pages = next_children;
870            }
871        }
872
873        Ok(())
874    }
875}
876
877pub(crate) fn btree_stats(
878    root: Option<PageNumber>,
879    mem: &TransactionalMemory,
880    fixed_key_size: Option<usize>,
881    fixed_value_size: Option<usize>,
882) -> Result<BtreeStats> {
883    if let Some(root) = root {
884        stats_helper(root, mem, fixed_key_size, fixed_value_size)
885    } else {
886        Ok(BtreeStats {
887            tree_height: 0,
888            leaf_pages: 0,
889            branch_pages: 0,
890            stored_leaf_bytes: 0,
891            metadata_bytes: 0,
892            fragmented_bytes: 0,
893        })
894    }
895}
896
897fn stats_helper(
898    page_number: PageNumber,
899    mem: &TransactionalMemory,
900    fixed_key_size: Option<usize>,
901    fixed_value_size: Option<usize>,
902) -> Result<BtreeStats> {
903    let page = mem.get_page(page_number)?;
904    let node_mem = page.memory();
905    match node_mem[0] {
906        LEAF => {
907            let accessor = LeafAccessor::new(page.memory(), fixed_key_size, fixed_value_size);
908            let leaf_bytes = accessor.length_of_pairs(0, accessor.num_pairs());
909            let overhead_bytes = accessor.total_length() - leaf_bytes;
910            let fragmented_bytes = (page.memory().len() - accessor.total_length()) as u64;
911            Ok(BtreeStats {
912                tree_height: 1,
913                leaf_pages: 1,
914                branch_pages: 0,
915                stored_leaf_bytes: leaf_bytes.try_into().unwrap(),
916                metadata_bytes: overhead_bytes.try_into().unwrap(),
917                fragmented_bytes,
918            })
919        }
920        BRANCH => {
921            let accessor = BranchAccessor::new(&page, fixed_key_size);
922            let mut max_child_height = 0;
923            let mut leaf_pages = 0;
924            let mut branch_pages = 1;
925            let mut stored_leaf_bytes = 0;
926            let mut metadata_bytes = accessor.total_length() as u64;
927            let mut fragmented_bytes = (page.memory().len() - accessor.total_length()) as u64;
928            for i in 0..accessor.count_children() {
929                if let Some(child) = accessor.child_page(i) {
930                    let stats = stats_helper(child, mem, fixed_key_size, fixed_value_size)?;
931                    max_child_height = max(max_child_height, stats.tree_height);
932                    leaf_pages += stats.leaf_pages;
933                    branch_pages += stats.branch_pages;
934                    stored_leaf_bytes += stats.stored_leaf_bytes;
935                    metadata_bytes += stats.metadata_bytes;
936                    fragmented_bytes += stats.fragmented_bytes;
937                }
938            }
939
940            Ok(BtreeStats {
941                tree_height: max_child_height + 1,
942                leaf_pages,
943                branch_pages,
944                stored_leaf_bytes,
945                metadata_bytes,
946                fragmented_bytes,
947            })
948        }
949        _ => unreachable!(),
950    }
951}
redb/tree_store/btree.rs

redb/tree_store/
btree.rs
