brie_tree/
int.rs

1use core::{cmp::Ordering, fmt::Debug, mem, ops::IndexMut};
2
3use crate::{
4    BTreeKey,
5    node::{NodePos, NodeRef},
6    simd::SimdSearch,
7    stack::{Height, Stack, max_height},
8};
9
10/// Helper function to convert a key directly to a raw integer.
11#[inline]
12pub(crate) fn int_from_key<K: BTreeKey>(key: K) -> <K::Int as BTreeInteger>::Raw {
13    key.to_int().to_raw()
14}
15
16/// Helper function to convert a raw integer to a key.
17#[inline]
18pub(crate) fn key_from_int<K: BTreeKey>(int: <K::Int as BTreeInteger>::Raw) -> Option<K> {
19    K::Int::from_raw(int).map(K::from_int)
20}
21
22/// B is selected so that all the keys fit in 128 bytes (2 cache lines).
23pub(crate) const KEYS_BYTES: usize = 128;
24
25/// Nodes are aligned to 128 bytes so they fit exactly in cache lines.
26#[repr(C, align(128))]
27pub(crate) struct AlignedKeys<T>(pub(crate) T);
28
29/// This trait covers all operations that are specific to the integer type used
30/// as a key.
31///
32/// # Safety
33///
34/// All items must be implemented as documented.
35pub(crate) unsafe trait BTreeInteger: Copy + Eq + Debug {
36    /// Number of elements per node, which must be at least 4.
37    ///
38    /// The number of elements may vary depending on the integer size to fit in
39    /// cache lines or to make optimal use of SIMD instructions.
40    const B: usize;
41
42    /// Maximum integer value.
43    ///
44    /// `search` and `cmp` must compare this as larger than any other integer
45    /// value.
46    const MAX: Self::Raw;
47
48    /// Raw integer type that is actually stored in the tree.
49    type Raw: Copy + Eq + Debug + SimdSearch;
50
51    /// Conversion from a `Self` to a raw integer.
52    fn to_raw(self) -> Self::Raw;
53
54    /// Conversion from a raw integer to a `Self`.
55    fn from_raw(int: Self::Raw) -> Option<Self>;
56
57    /// Compares 2 integers. We don't just use the `Ord` trait here because some
58    /// implementations add a bias to the integer values.
59    fn cmp(a: Self::Raw, b: Self::Raw) -> Ordering;
60
61    /// Increments a raw integer by 1.
62    fn increment(int: Self::Raw) -> Self::Raw;
63
64    /// Array of keys used for SIMD comparison in `rank`.
65    ///
66    /// This must have the same layout as `[Self; Self::B]`.
67    type Keys;
68
69    /// Returns the index of the first key greater than or equal to `search`.
70    ///
71    /// Because this assumes that keys are sorted, it can be implemented either
72    /// as a binary search or by counting the number of keys less than `search`.
73    ///
74    ///  # Safety
75    ///
76    /// The last key must be `Self::MAX`, which guarantees that the returned
77    /// position is less than `Self::B`.
78    unsafe fn search(keys: &Self::Keys, search: Self::Raw) -> NodePos<Self>;
79
80    /// Array of `(NodeRef, NodePos)` pairs which can be indexed by a `Height`.
81    type Stack: IndexMut<Height<Self>, Output = (NodeRef, NodePos<Self>)> + Default + Clone;
82}
83
84macro_rules! impl_int {
85    ($($int:ident $nonmax:ident,)*) => {
86        $(
87            unsafe impl BTreeInteger for nonmax::$nonmax {
88                const B: usize = KEYS_BYTES / mem::size_of::<Self>();
89
90                const MAX: Self::Raw = $int::MAX.wrapping_add(Self::Raw::BIAS);
91
92                type Raw = $int;
93
94                #[inline]
95                fn to_raw(self) -> Self::Raw {
96                    self.get().wrapping_add(Self::Raw::BIAS)
97                }
98
99                #[inline]
100                fn from_raw(int: Self::Raw) -> Option<Self> {
101                    Self::new(int.wrapping_sub(Self::Raw::BIAS))
102                }
103
104                #[inline]
105                fn cmp(a: Self::Raw, b: Self::Raw) -> Ordering {
106                    Self::Raw::bias_cmp(a, b)
107                }
108
109                #[inline]
110                fn increment(int: Self::Raw) -> Self::Raw {
111                    int.wrapping_add(1)
112                }
113
114                type Keys = AlignedKeys<[Self::Raw; Self::B]>;
115
116                #[inline]
117                unsafe fn search(keys: &Self::Keys, search: Self::Raw) -> NodePos<Self> {
118                    unsafe { NodePos::new_unchecked(Self::Raw::search(&keys.0, search)) }
119                }
120
121                type Stack = Stack<Self, { max_height::<Self>() }>;
122            }
123
124            impl BTreeKey for nonmax::$nonmax {
125                type Int = Self;
126
127                #[inline]
128                fn to_int(self) -> Self::Int {
129                    self
130                }
131
132                #[inline]
133                fn from_int(int: Self::Int) -> Self {
134                    int
135                }
136            }
137        )*
138    };
139}
140
141impl_int! {
142    u8 NonMaxU8,
143    u16 NonMaxU16,
144    u32 NonMaxU32,
145    u64 NonMaxU64,
146    u128 NonMaxU128,
147    i8 NonMaxI8,
148    i16 NonMaxI16,
149    i32 NonMaxI32,
150    i64 NonMaxI64,
151    i128 NonMaxI128,
152}
brie_tree/int.rs

brie_tree/
int.rs
