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

brie_tree/
int.rs
