tribles/patch/
leaf.rs

1use core::sync::atomic;
2use core::sync::atomic::Ordering::{Acquire, Relaxed, Release};
3use siphasher::sip128::SipHasher24;
4use std::alloc::*;
5use std::ptr::addr_of;
6
7use super::*;
8
9#[derive(Debug)]
10#[repr(C)]
11pub(crate) struct Leaf<const KEY_LEN: usize> {
12    pub key: [u8; KEY_LEN],
13    pub hash: u128,
14    rc: atomic::AtomicU32,
15}
16
17impl<const KEY_LEN: usize> Body for Leaf<KEY_LEN> {
18    fn tag(_body: NonNull<Self>) -> HeadTag {
19        HeadTag::Leaf
20    }
21}
22
23impl<const KEY_LEN: usize> Leaf<KEY_LEN> {
24    pub(super) unsafe fn new(key: &[u8; KEY_LEN]) -> NonNull<Self> {
25        unsafe {
26            let layout = Layout::new::<Self>();
27            if let Some(ptr) = NonNull::new(alloc(layout) as *mut Self) {
28                let hash = SipHasher24::new_with_key(&*addr_of!(SIP_KEY))
29                    .hash(&key[..])
30                    .into();
31
32                ptr.write(Self {
33                    key: *key,
34                    hash,
35                    rc: atomic::AtomicU32::new(1),
36                });
37
38                ptr
39            } else {
40                panic!("Allocation failed!");
41            }
42        }
43    }
44
45    pub(crate) unsafe fn rc_inc(leaf: NonNull<Self>) -> NonNull<Self> {
46        unsafe {
47            let leaf = leaf.as_ptr();
48            let mut current = (*leaf).rc.load(Relaxed);
49            loop {
50                if current == u32::MAX {
51                    panic!("max refcount exceeded");
52                }
53                match (*leaf)
54                    .rc
55                    .compare_exchange(current, current + 1, Relaxed, Relaxed)
56                {
57                    Ok(_) => return NonNull::new_unchecked(leaf),
58                    Err(v) => current = v,
59                }
60            }
61        }
62    }
63
64    pub(crate) unsafe fn rc_dec(leaf: NonNull<Self>) {
65        unsafe {
66            let ptr = leaf.as_ptr();
67            let rc = (*ptr).rc.fetch_sub(1, Release);
68            if rc != 1 {
69                return;
70            }
71            (*ptr).rc.load(Acquire);
72
73            std::ptr::drop_in_place(ptr);
74
75            let layout = Layout::new::<Self>();
76            let ptr = ptr as *mut u8;
77            dealloc(ptr, layout);
78        }
79    }
80
81    pub(crate) fn infixes<
82        const PREFIX_LEN: usize,
83        const INFIX_LEN: usize,
84        O: KeyOrdering<KEY_LEN>,
85        S: KeySegmentation<KEY_LEN>,
86        F,
87    >(
88        leaf: NonNull<Self>,
89        prefix: &[u8; PREFIX_LEN],
90        at_depth: usize,
91        f: &mut F,
92    ) where
93        F: FnMut(&[u8; INFIX_LEN]),
94    {
95        let leaf = unsafe { leaf.as_ref() };
96        let leaf_key = (*leaf).key;
97        for depth in at_depth..PREFIX_LEN {
98            if leaf_key[O::key_index(depth)] != prefix[depth] {
99                return;
100            }
101        }
102
103        let infix: [u8; INFIX_LEN] =
104            core::array::from_fn(|i| (*leaf).key[O::key_index(PREFIX_LEN + i)]);
105        f(&infix);
106    }
107
108    pub(crate) fn has_prefix<O: KeyOrdering<KEY_LEN>, const PREFIX_LEN: usize>(
109        leaf: NonNull<Self>,
110        at_depth: usize,
111        prefix: &[u8; PREFIX_LEN],
112    ) -> bool {
113        let leaf_key: &[u8; KEY_LEN] = unsafe { &(*leaf.as_ptr()).key };
114        for depth in at_depth..PREFIX_LEN {
115            if leaf_key[O::key_index(depth)] != prefix[depth] {
116                return false;
117            }
118        }
119        return true;
120    }
121
122    pub(crate) fn segmented_len<O: KeyOrdering<KEY_LEN>, const PREFIX_LEN: usize>(
123        leaf: NonNull<Self>,
124        at_depth: usize,
125        prefix: &[u8; PREFIX_LEN],
126    ) -> u64 {
127        let leaf_key: &[u8; KEY_LEN] = unsafe { &(*leaf.as_ptr()).key };
128        for depth in at_depth..PREFIX_LEN {
129            let key_depth = O::key_index(depth);
130            if leaf_key[key_depth] != prefix[depth] {
131                return 0;
132            }
133        }
134        return 1;
135    }
136}
tribles/patch/leaf.rs

tribles/patch/
leaf.rs
