Skip to main content

embed_btree/
compact.rs

1use crate::CACHE_LINE_SIZE;
2use crate::bptree::{
3    self, BTree, entry::EntryInner, helper::PathBuffer, inter::*, node::*, tree::BTreeInner, *,
4};
5use alloc::alloc::{Layout, alloc, dealloc, handle_alloc_error, realloc};
6use core::borrow::BorrowMut;
7use core::fmt::{self, Debug};
8use core::marker::PhantomData;
9use core::mem::{MaybeUninit, align_of, size_of, transmute};
10use core::num::{NonZeroU8, NonZeroU64};
11use core::ptr::null_mut;
12
13pub type BTreeMap<K, V> = bptree::BTree<K, V, Compact>;
14pub use bptree::cursor::Cursor;
15pub use bptree::iter::{Iter, IterMut, Keys, Range, RangeMut, Values, ValuesMut};
16pub type IntoIter<K, V> = bptree::iter::IntoIter<K, V, Compact>;
17pub type Entry<'a, K, V> = bptree::entry::Entry<'a, K, V, Compact>;
18pub type OccupiedEntry<'a, K, V> = bptree::entry::OccupiedEntry<'a, K, V, Compact>;
19pub type VacantEntry<'a, K, V> = bptree::entry::VacantEntry<'a, K, V, Compact>;
20
21/// Compact metadata for `BTreeMap`.
22///
23/// It's either CompactBufHeader cast as u64, or ptr for meta on heap.
24///
25/// NOTE: use NonZeroU64 to optimize `Option<BTreeMap>`
26pub struct Compact(NonZeroU64);
27
28const STACK_CAP: usize = 3;
29const ENTRY_CAP: usize = 1;
30
31unsafe impl Send for Compact {}
32unsafe impl Sync for Compact {}
33
34impl Default for Compact {
35    fn default() -> Self {
36        Self(unsafe {
37            NonZeroU64::new_unchecked(transmute::<CompactBufHeader, u64>(
38                CompactBufHeader::default(),
39            ))
40        })
41    }
42}
43
44impl Drop for Compact {
45    #[inline]
46    fn drop(&mut self) {
47        let p = self.try_heap_ptr();
48        if !p.is_null() {
49            unsafe {
50                let header = p as *mut CompactBufHeader;
51                let size = Self::buf_size_from_cap((*header).cap.get());
52                // InterNode does not have drop
53                dealloc(p, Self::get_layout(size));
54            }
55        }
56    }
57}
58
59// for log
60impl Debug for Compact {
61    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
62        write!(f, "Compact")
63    }
64}
65
66impl Compact {
67    const ITEM_SIZE: usize = size_of::<(*mut u8, u8)>();
68
69    // Offset at which items start (header size rounded up to item alignment).
70    const ITEMS_OFFSET: usize = {
71        let hs = size_of::<CompactBufHeader>();
72        let ia = align_of::<(*mut u8, u8)>();
73        let offset = (hs + ia - 1) & !(ia - 1);
74        if offset != hs {
75            panic!("TreeInfoHeader is not aligned");
76        }
77        offset
78    };
79
80    #[inline(always)]
81    fn try_heap_ptr(&self) -> *mut u8 {
82        let this = self.0.get();
83        // the cap field
84        if this & (STACK_FLAG as u64) == 0 { this as *mut u8 } else { null_mut() }
85    }
86
87    #[inline(always)]
88    fn header_ptr(&self) -> *const CompactBufHeader {
89        let this = self.0.get();
90        // the cap field
91        if this & (STACK_FLAG as u64) > 0 {
92            &self.0 as *const NonZeroU64 as *const CompactBufHeader
93        } else {
94            this as *const CompactBufHeader
95        }
96    }
97
98    #[inline(always)]
99    fn header_ptr_mut(&mut self) -> *mut CompactBufHeader {
100        let this = self.0.get();
101        // the cap field
102        if this & (STACK_FLAG as u64) > 0 {
103            &mut self.0 as *mut NonZeroU64 as *mut CompactBufHeader
104        } else {
105            this as *mut CompactBufHeader
106        }
107    }
108
109    #[inline(always)]
110    fn get_header(&self) -> &CompactBufHeader {
111        unsafe { &*self.header_ptr() }
112    }
113
114    #[inline(always)]
115    fn get_header_mut(&mut self) -> &mut CompactBufHeader {
116        unsafe { &mut *self.header_ptr_mut() }
117    }
118
119    /// the input cap is not including root, but for actual items need to store on heap.
120    ///
121    /// return (new_buf_size, cap_field)
122    ///
123    /// cap_field = actual heap cap + 1 (root)
124    #[inline(always)]
125    const fn cal_new_buf_size(cap: u8) -> (usize, u8) {
126        let need = cap as usize * Self::ITEM_SIZE + Self::ITEMS_OFFSET;
127        // aligned multiple of CACHE_LINE_SIZE
128        let new_size = (need + CACHE_LINE_SIZE - 1) & !(CACHE_LINE_SIZE - 1);
129        let aligned_cap = ((new_size - Self::ITEMS_OFFSET) / Self::ITEM_SIZE) as u8;
130        // one cap for root not stored on the heap
131        (new_size, aligned_cap + 1)
132    }
133
134    #[inline(always)]
135    const fn get_layout(buf_size: usize) -> Layout {
136        let align = align_of::<usize>();
137        #[cfg(debug_assertions)]
138        {
139            if align_of::<(*mut u8, u8)>() != align {
140                panic!("TreeInfoHeader is not aligned");
141            }
142        }
143        unsafe { Layout::from_size_align_unchecked(buf_size, align) }
144    }
145
146    /// the input cap is not including root, but for actual items can store on heap.
147    #[inline(always)]
148    fn buf_size_from_cap(aligned_cap: u8) -> usize {
149        // one cap for root not stored on the heap
150        let size = (aligned_cap - 1) as usize * Self::ITEM_SIZE + Self::ITEMS_OFFSET;
151        (size + CACHE_LINE_SIZE - 1) & !(CACHE_LINE_SIZE - 1)
152    }
153
154    /// height is the root.height (tree height - 1)
155    #[inline]
156    fn ensure_cap<const N: usize>(&mut self, root_height: u8) -> *mut CompactBufHeader {
157        assert!(root_height < u8::MAX);
158        let mut header = self.header_ptr_mut();
159        // cap include root although it's not stored
160        let cur_cap = unsafe { (*header).cap.get() };
161        let new_size: usize;
162        let new_cap: u8;
163        let new_p;
164        // STACK_FLAG == 1 can represent root, and on heap is 1 + real number
165        if root_height <= (N as u8) + cur_cap {
166            return header;
167        }
168        let need_cap = root_height - 1 - N as u8;
169        if cur_cap == STACK_FLAG {
170            (new_size, new_cap) = Self::cal_new_buf_size(need_cap);
171            crate::trace_log!("alloc pathbuf cap {need_cap} actual {new_cap} cap");
172            let layout = Self::get_layout(new_size);
173            new_p = unsafe { alloc(layout) };
174            if !new_p.is_null() {
175                header = new_p as *mut CompactBufHeader;
176                unsafe {
177                    (*header).entry_idx = 0;
178                    (*header).len = 0;
179                    (*header).buffer_pos = 0;
180                    (*header).cap = NonZeroU8::new_unchecked(new_cap);
181                }
182            } else {
183                handle_alloc_error(Self::get_layout(new_size));
184            }
185        } else {
186            (new_size, new_cap) = Self::cal_new_buf_size(need_cap);
187            crate::trace_log!("grow pathbuf cap {cur_cap}->{need_cap}");
188            let old_layout = Self::get_layout(Self::buf_size_from_cap(cur_cap));
189            new_p = unsafe { realloc(self.0.get() as *mut u8, old_layout, new_size) };
190            if !new_p.is_null() {
191                header = new_p as *mut CompactBufHeader;
192                unsafe {
193                    (*header).cap = NonZeroU8::new_unchecked(new_cap);
194                }
195            } else {
196                handle_alloc_error(Self::get_layout(new_size));
197            }
198        }
199        self.0 = unsafe { NonZeroU64::new_unchecked(new_p as u64) };
200        unsafe { debug_assert!((*header).cap.get() != STACK_FLAG, "{}", (*header).cap) };
201        header
202    }
203}
204
205impl Stats for Compact {}
206
207impl StatsPriv for Compact {
208    type EntryInner<'a, K, V>
209        = CompactEntry<'a, K, V>
210    where
211        K: Key + 'a,
212        V: Value + 'a;
213
214    type PathBufferRef<'a>
215        = CompactBufRef<'a, STACK_CAP, &'a mut CompactBufStack<STACK_CAP>>
216    where
217        Self: 'a;
218
219    type PathBuffer = CompactBufIntoIter<STACK_CAP>;
220
221    type BufferStack = CompactBufStack<STACK_CAP>;
222
223    #[inline]
224    fn get_cache<'a, K>(
225        &'a mut self, stack: &'a mut Self::BufferStack, root: Option<InterNode<K>>,
226    ) -> Self::PathBufferRef<'a> {
227        let _header = if let Some(root) = root.as_ref() {
228            self.ensure_cap::<STACK_CAP>(root.height())
229        } else {
230            self.header_ptr_mut()
231        };
232        unsafe { (*_header).clear() };
233        CompactBufRef { root: root.map(|n| n.into()), _header, stack, _phan: Default::default() }
234    }
235
236    #[inline]
237    fn take_cache<K>(self, root: Option<InterNode<K>>) -> Self::PathBuffer {
238        let heap_ptr: *mut CompactBufHeader;
239        let mut header: CompactBufHeader;
240        // take_cache will not increase height, it's always enough
241        unsafe {
242            let this = self.0.get();
243            // the cap field
244            if this & (STACK_FLAG as u64) > 0 {
245                heap_ptr = null_mut();
246                let _header: &CompactBufHeader = transmute(&self.0);
247                header = *_header;
248            } else {
249                heap_ptr = this as *mut CompactBufHeader;
250                header = *heap_ptr;
251            }
252        }
253        header.clear();
254        core::mem::forget(self);
255        CompactBufIntoIter {
256            root: root.map(|n| n.into()),
257            // because it's moving, getting the address now is unsafe
258            header,
259            _heap: heap_ptr as *mut u8,
260            stack: CompactBufStack::default(),
261        }
262    }
263
264    #[inline]
265    fn search_entry<'a, K: Key, V: Value>(
266        tree: &'a mut BTree<K, V, Self>, key: K,
267    ) -> Entry<'a, K, V> {
268        let mut inner = CompactEntry { tree, stack: Default::default() };
269        let (_tree, mut cache) = inner.init_cache();
270        if let Some(root) = _tree.get_root() {
271            let mut is_seq = true;
272            let leaf = match root {
273                Node::Inter(inter) => {
274                    inter.find_leaf_with_cache_smart::<K, V, _>(&mut cache, &key, &mut is_seq)
275                }
276                Node::Leaf(leaf) => leaf,
277            };
278            let (idx, is_equal) = leaf.search_smart(&key, is_seq);
279            // use cache' header ptr to accelerate
280            cache.set_entry_idx(idx);
281            if is_equal {
282                Entry::Occupied(OccupiedEntry { inner, leaf })
283            } else {
284                Entry::Vacant(VacantEntry { inner, key, leaf: Some(leaf) })
285            }
286        } else {
287            // Safety: entry_idx already reset 0 when init_cache
288            Entry::Vacant(VacantEntry { inner, key, leaf: None })
289        }
290    }
291
292    /// seek first or last entry
293    #[inline]
294    fn seek_entry<'a, K: Key, V: Value, const FIRST: bool>(
295        tree: &'a mut BTree<K, V, Self>,
296    ) -> Option<OccupiedEntry<'a, K, V>> {
297        let mut inner = CompactEntry { tree, stack: Default::default() };
298        let (_tree, mut cache) = inner.init_cache();
299        if let Some(root) = _tree.get_root() {
300            let leaf = match root {
301                Node::Inter(inter) => {
302                    if FIRST {
303                        inter.find_first_leaf_with_cache::<V, _>(&mut cache)
304                    } else {
305                        inter.find_last_leaf_with_cache::<V, _>(&mut cache)
306                    }
307                }
308                Node::Leaf(leaf) => leaf,
309            };
310            let count = leaf.key_count();
311            if count > 0 {
312                let idx = if FIRST { 0 } else { count - 1 };
313                // use cache' header ptr to accelerate
314                cache.set_entry_idx(idx);
315                return Some(OccupiedEntry { inner, leaf });
316            }
317        }
318        None
319    }
320
321    #[cfg(test)]
322    fn init_count(&mut self, _leaf_count: usize, _inter_count: u32) {}
323
324    #[cfg(test)]
325    fn assert_leaf_count(&self, _leaf_count: usize) {}
326
327    #[cfg(test)]
328    fn assert_inter_count(&self, _inter_count: usize) {}
329}
330
331#[cfg(target_endian = "little")]
332#[repr(C)]
333#[derive(Clone, Copy)]
334pub(crate) struct CompactBufHeader {
335    /// when on stack, we mark cap == 1, the lowest bit on u64 (we convert CompactBufHeader as u64,
336    /// we can determine it's not a pointer
337    pub cap: NonZeroU8,
338    /// left < 0, or right > 0, or center = 0
339    pub buffer_pos: i8,
340    /// the items stored in PathBuffer (not including the root)
341    pub len: u8,
342    /// tempoary storage for Entry.idx
343    pub entry_idx: u8,
344    pub idxs: MaybeUninit<[u8; STACK_CAP + 1]>,
345}
346
347#[cfg(target_endian = "big")]
348#[repr(C)]
349#[derive(Clone, Copy)]
350pub(crate) struct CompactBufHeader {
351    pub idxs: MaybeUninit<[u8; STACK_CAP + 1]>,
352    /// tempoary storage for Entry.idx
353    pub entry_idx: u8,
354    /// the items stored in PathBuffer (not including the root)
355    pub len: u8,
356    /// left < 0, or right > 0, or center = 0
357    pub buffer_pos: i8,
358    /// when on stack, we mark cap == 1, the lowest bit on u64 (we convert CompactBufHeader as u64,
359    /// we can determine it's not a pointer
360    pub cap: NonZeroU8,
361}
362
363const STACK_FLAG: u8 = 1;
364
365impl Default for CompactBufHeader {
366    #[inline]
367    fn default() -> Self {
368        Self {
369            buffer_pos: 0,
370            len: 0,
371            entry_idx: 0,
372            cap: unsafe { NonZeroU8::new_unchecked(STACK_FLAG) },
373            idxs: MaybeUninit::zeroed(),
374        }
375    }
376}
377
378impl CompactBufHeader {
379    #[inline(always)]
380    fn get_idxs(&self, i: u8) -> u8 {
381        unsafe { self.idxs.assume_init_ref()[i as usize] }
382    }
383
384    fn put_idxs(&mut self, i: u8, value: u8) {
385        unsafe { self.idxs.assume_init_mut()[i as usize] = value };
386    }
387
388    /// Reset the stack without freeing the buffer.
389    #[inline]
390    fn clear(&mut self) {
391        self.len = 0;
392        self.buffer_pos = 0;
393        self.entry_idx = 0;
394    }
395}
396
397pub(crate) struct CompactBufStack<const N: usize>(MaybeUninit<[NodeBase; N]>);
398
399impl<const N: usize> Default for CompactBufStack<N> {
400    #[inline]
401    fn default() -> Self {
402        Self(MaybeUninit::uninit())
403    }
404}
405
406impl<const N: usize> CompactBufStack<N> {
407    #[inline(always)]
408    fn _get(&self, i: u8) -> NodeBase {
409        debug_assert!((i as usize) < N);
410        unsafe { self.0.assume_init_ref()[i as usize].clone() }
411    }
412
413    #[inline(always)]
414    fn _put(&mut self, i: u8, inter: NodeBase) {
415        debug_assert!((i as usize) < N);
416        unsafe { self.0.assume_init_mut()[i as usize] = inter };
417    }
418}
419
420trait CompactBufCommon {
421    const STACK_N: usize;
422
423    fn header(&self) -> &CompactBufHeader;
424
425    fn header_mut(&mut self) -> &mut CompactBufHeader;
426
427    #[inline(always)]
428    fn set_entry_idx(&mut self, entry_idx: u8) {
429        self.header_mut().entry_idx = entry_idx;
430    }
431
432    unsafe fn heap_item_ptr(&self, idx: u8) -> *mut (NodeBase, u8);
433
434    fn stack_get(&self, i: u8) -> NodeBase;
435
436    fn stack_put(&mut self, i: u8, inter: NodeBase);
437
438    fn get_root(&self) -> NodeBase;
439}
440
441pub(crate) struct CompactBufRef<'a, const N: usize, S: BorrowMut<CompactBufStack<N>>> {
442    root: Option<NodeBase>,
443    // for fast access
444    _header: *mut CompactBufHeader,
445    stack: S,
446    _phan: PhantomData<&'a ()>,
447}
448
449impl<'a, const N: usize, S: BorrowMut<CompactBufStack<N>>> CompactBufCommon
450    for CompactBufRef<'a, N, S>
451{
452    const STACK_N: usize = N;
453
454    #[inline(always)]
455    fn header(&self) -> &CompactBufHeader {
456        unsafe { &*self._header }
457    }
458
459    #[inline(always)]
460    fn header_mut(&mut self) -> &mut CompactBufHeader {
461        unsafe { &mut *self._header }
462    }
463
464    #[inline(always)]
465    unsafe fn heap_item_ptr(&self, i: u8) -> *mut (NodeBase, u8) {
466        let ptr = self._header;
467        #[cfg(debug_assertions)]
468        {
469            let cap = unsafe { (*ptr).cap.get() };
470            debug_assert!(cap > STACK_FLAG, "cap {cap}");
471            debug_assert!(i < cap, "i {i} >= cap {cap}");
472        }
473        unsafe {
474            let p = (ptr as *mut u8).add(Compact::ITEMS_OFFSET) as *mut (NodeBase, u8);
475            p.add(i as usize)
476        }
477    }
478
479    #[inline(always)]
480    fn stack_get(&self, i: u8) -> NodeBase {
481        self.stack.borrow()._get(i)
482    }
483
484    #[inline(always)]
485    fn stack_put(&mut self, i: u8, inter: NodeBase) {
486        self.stack.borrow_mut()._put(i, inter);
487    }
488
489    #[inline(always)]
490    fn get_root(&self) -> NodeBase {
491        self.root.as_ref().unwrap().clone()
492    }
493}
494
495unsafe impl<'a, const N: usize, S: BorrowMut<CompactBufStack<N>>> Send for CompactBufRef<'a, N, S> {}
496
497/// Once we are calling into iter the cap will not increase.
498///
499/// Because CompactBufIntoIter can be moving, we extract header here,
500pub(crate) struct CompactBufIntoIter<const N: usize> {
501    root: Option<NodeBase>,
502    header: CompactBufHeader,
503    // just for dropping when we are done
504    _heap: *mut u8,
505    stack: CompactBufStack<N>,
506}
507
508unsafe impl<const N: usize> Send for CompactBufIntoIter<N> {}
509
510impl<const N: usize> CompactBufCommon for CompactBufIntoIter<N> {
511    const STACK_N: usize = N;
512
513    #[inline(always)]
514    fn header(&self) -> &CompactBufHeader {
515        &self.header
516    }
517
518    #[inline(always)]
519    fn header_mut(&mut self) -> &mut CompactBufHeader {
520        &mut self.header
521    }
522
523    #[inline(always)]
524    unsafe fn heap_item_ptr(&self, i: u8) -> *mut (NodeBase, u8) {
525        let ptr = self._heap;
526        #[cfg(debug_assertions)]
527        {
528            debug_assert!(!ptr.is_null());
529            let cap = self.header.cap.get();
530            debug_assert!(cap > STACK_FLAG, "cap {cap}");
531            debug_assert!(i < cap, "i {i} >= cap {cap}");
532        }
533        unsafe {
534            let p = ptr.add(Compact::ITEMS_OFFSET) as *mut (NodeBase, u8);
535            p.add(i as usize)
536        }
537    }
538
539    #[inline(always)]
540    fn stack_get(&self, i: u8) -> NodeBase {
541        self.stack._get(i)
542    }
543
544    #[inline(always)]
545    fn stack_put(&mut self, i: u8, inter: NodeBase) {
546        self.stack._put(i, inter);
547    }
548
549    #[inline(always)]
550    fn get_root(&self) -> NodeBase {
551        self.root.as_ref().unwrap().clone()
552    }
553}
554
555impl<const N: usize> Drop for CompactBufIntoIter<N> {
556    #[inline]
557    fn drop(&mut self) {
558        if !self._heap.is_null() {
559            unsafe {
560                let size = Compact::buf_size_from_cap(self.header.cap.get());
561                // InterNode does not have drop
562                dealloc(self._heap, Compact::get_layout(size));
563            }
564        }
565    }
566}
567
568impl<T: CompactBufCommon> PathBuffer for T {
569    // --- stats method begins ---
570
571    #[inline(always)]
572    fn inc_leaf_count(&mut self) {}
573
574    #[inline(always)]
575    fn dec_leaf_count(&mut self) {}
576    #[inline(always)]
577    fn inc_inter_count(&mut self) {}
578
579    #[inline(always)]
580    fn dec_inter_count(&mut self) {}
581
582    // --- stats method ends ---
583
584    /// The count of current level items contains by the PathBuffer
585    #[inline]
586    fn buffer_len(&self) -> u8 {
587        self.header().len
588    }
589
590    /// The delta position of current entry to the PathBuffer.
591    /// < 0 for left, > 0 for right, ==0 for center
592    #[inline]
593    fn buffer_pos(&self) -> i8 {
594        self.header().buffer_pos
595    }
596
597    #[inline]
598    fn _set_pos(&mut self, buffer_pos: i8) {
599        self.header_mut().buffer_pos = buffer_pos;
600    }
601
602    /// Push one entry onto the cache stack, growing the buffer if needed.
603    #[inline]
604    fn _push(&mut self, inter: NodeBase, idx: u8) {
605        let header = self.header_mut();
606        let i = header.len;
607        header.len = i + 1;
608        if i < Self::STACK_N as u8 + 1 {
609            debug_assert!((i as usize) < Self::STACK_N + 1, "{i} < N {} + 1", Self::STACK_N);
610            header.put_idxs(i, idx);
611            if i > 0 {
612                self.stack_put(i - 1, inter);
613            }
614        } else {
615            unsafe { self.heap_item_ptr(i - Self::STACK_N as u8 - 1).write((inter, idx)) };
616        }
617    }
618
619    #[inline]
620    unsafe fn _get_unchecked(&self, i: u8) -> (NodeBase, u8) {
621        let header = self.header();
622        if i < Self::STACK_N as u8 + 1 {
623            let idx = header.get_idxs(i);
624            if i == 0 { (self.get_root(), idx) } else { (self.stack_get(i - 1), idx) }
625        } else {
626            let p = unsafe { &*self.heap_item_ptr(i - 1 - Self::STACK_N as u8) };
627            (p.0.clone(), p.1)
628        }
629    }
630
631    /// Pop the top entry from the cache stack.
632    #[inline]
633    fn _pop(&mut self) -> Option<(NodeBase, u8)> {
634        let header = self.header_mut();
635        let mut i = header.len;
636        if i > 0 {
637            i -= 1;
638            header.len = i;
639            if i < Self::STACK_N as u8 + 1 {
640                let idx = header.get_idxs(i);
641                if i == 0 {
642                    Some((self.get_root(), idx))
643                } else {
644                    let node = self.stack_get(i - 1);
645                    Some((node, idx))
646                }
647            } else {
648                let p = unsafe { &*self.heap_item_ptr(i - 1 - Self::STACK_N as u8) };
649                Some((p.0.clone(), p.1))
650            }
651        } else {
652            None
653        }
654    }
655
656    /// Peek the top entry from the cache stack.
657    #[inline]
658    fn _last(&self) -> Option<(NodeBase, u8)> {
659        let header = self.header();
660        let mut i = header.len;
661        if i > 0 {
662            i -= 1;
663            if i < Self::STACK_N as u8 + 1 {
664                let idx = header.get_idxs(i);
665                if i == 0 {
666                    Some((self.get_root(), idx))
667                } else {
668                    let node = self.stack_get(i - 1);
669                    Some((node, idx))
670                }
671            } else {
672                let p = unsafe { &*self.heap_item_ptr(i - 1 - Self::STACK_N as u8) };
673                Some((p.0.clone(), p.1))
674            }
675        } else {
676            None
677        }
678    }
679}
680
681type CompactBufEntry<'a> = CompactBufRef<'a, ENTRY_CAP, &'a mut CompactBufStack<ENTRY_CAP>>;
682
683pub(crate) struct CompactEntry<'a, K: Key, V: Value> {
684    tree: &'a mut BTree<K, V, Compact>,
685    stack: CompactBufStack<ENTRY_CAP>,
686}
687
688impl<'a, K: Key, V: Value> CompactEntry<'a, K, V> {
689    #[inline]
690    fn init_cache(&mut self) -> (&BTreeInner<K, V>, CompactBufEntry<'_>) {
691        let (tree, stats, stack) = (&self.tree.inner, &mut self.tree.stats, &mut self.stack);
692        let _header;
693        let root = if let Some(inter) = tree.root_as_inter() {
694            _header = stats.ensure_cap::<ENTRY_CAP>(inter.height());
695            Some(inter)
696        } else {
697            _header = stats.header_ptr_mut();
698            None
699        };
700        unsafe { (*_header).clear() };
701        (
702            tree,
703            CompactBufEntry {
704                _header,
705                root: root.map(|n| n.into()),
706                stack,
707                _phan: Default::default(),
708            },
709        )
710    }
711}
712
713impl<'a, K: Key, V: Value> EntryInner<K, V> for CompactEntry<'a, K, V> {
714    #[inline(always)]
715    fn get_cache(&mut self) -> impl PathBuffer {
716        let root = self.tree.inner.root_as_inter().map(|_node| _node.into());
717        let _header = self.tree.stats.header_ptr_mut();
718        CompactBufEntry { root, _header, stack: &mut self.stack, _phan: Default::default() }
719    }
720
721    #[cfg(test)]
722    #[inline(always)]
723    fn get_tree(&self) -> &BTreeInner<K, V> {
724        &self.tree.inner
725    }
726
727    #[inline(always)]
728    fn get_tree_cache(&mut self) -> (&mut BTreeInner<K, V>, impl PathBuffer) {
729        let (tree, stats) = (&mut self.tree.inner, &mut self.tree.stats);
730        let root = tree.root_as_inter();
731        (
732            tree,
733            CompactBufRef {
734                root: root.map(|n| n.into()),
735                _header: stats.header_ptr_mut(),
736                stack: &mut self.stack,
737                _phan: Default::default(),
738            },
739        )
740    }
741
742    #[inline(always)]
743    fn set_idx(&mut self, idx: u8) {
744        self.tree.stats.get_header_mut().entry_idx = idx;
745    }
746
747    #[inline(always)]
748    fn get_idx(&self) -> u8 {
749        self.tree.stats.get_header().entry_idx
750    }
751}
752
753#[cfg(test)]
754mod tests {
755    use super::*;
756    use crate::bptree::tests::*;
757    use captains_log::logfn;
758    use core::mem::size_of;
759    use rstest::*;
760    use std::println;
761    use std::vec::Vec;
762
763    #[test]
764    fn test_compact_buf_size() {
765        type StackBuf<'a> = CompactBufRef<'a, STACK_CAP, &'a mut CompactBufStack<STACK_CAP>>;
766        type OwnedBuf<'a> = CompactBufIntoIter<STACK_CAP>;
767        println!("stack buf size: {}", size_of::<StackBuf>());
768        println!("owned buf size: {}", size_of::<OwnedBuf>());
769        //        println!("CompactBufferStack {}", size_of::<CompactBufferStack>());
770        println!("CompactEntry {}", size_of::<CompactEntry::<u32, u32>>());
771    }
772
773    #[logfn]
774    #[rstest]
775    fn test_compact_buf_path_buffer(setup_log: ()) {
776        fn test_push_pop(info: &mut Compact, height: u8) {
777            let mut nodes: Vec<InterNode<u32>> = Vec::new();
778            let root = unsafe { InterNode::<u32>::alloc(height) };
779            let mut stack = CompactBufStack::<STACK_CAP>::default();
780            let mut cache = info.get_cache(&mut stack, Some(root.clone()));
781            cache.push_path(root.clone(), height);
782            nodes.push(root);
783            for i in 1..height {
784                let node = unsafe { InterNode::<u32>::alloc(height - i) };
785                cache.push_path(node.clone(), height - i);
786                nodes.push(node);
787            }
788            let mut _nodes = Vec::new();
789            let mut i = 0;
790            while let Some((node, idx)) = cache.pop_path() {
791                i += 1;
792                assert_eq!(idx, i);
793                _nodes.insert(0, node);
794            }
795            assert_eq!(i, height);
796            assert_eq!(nodes, _nodes);
797            for node in nodes {
798                node.dealloc::<false>();
799            }
800        }
801        {
802            let mut info = Compact::default();
803            assert!(info.try_heap_ptr().is_null());
804            let inner0 = info.0.get();
805            assert_eq!(inner0 & (STACK_FLAG as u64), STACK_FLAG as u64);
806            assert_eq!(info.get_header().cap.get(), STACK_FLAG);
807            assert_eq!(STACK_CAP, 3);
808
809            info.ensure_cap::<STACK_CAP>(4);
810            let inner1 = info.0.get();
811            assert_eq!(inner0 & (STACK_FLAG as u64), STACK_FLAG as u64);
812            assert_eq!(info.get_header().cap.get(), STACK_FLAG);
813            info.ensure_cap::<STACK_CAP>(4);
814            assert_eq!(info.0.get(), inner1); // not changed
815            test_push_pop(&mut info, 4);
816
817            info.ensure_cap::<STACK_CAP>(7);
818            assert_eq!(info.get_header().cap.get(), 3 + 1); // heap + 1
819            assert!(info.0.get() != inner1); // re-alloc
820            test_push_pop(&mut info, 7);
821            info.ensure_cap::<STACK_CAP>(9);
822            // round up
823            assert_eq!(info.get_header().cap.get(), 3 + 4 + 1);
824            test_push_pop(&mut info, 10);
825            test_push_pop(&mut info, 3 + 4 + 1 + STACK_CAP as u8);
826        }
827    }
828}