Skip to main content

scattered_collect/map/
mod.rs

1//! A swiss-table-style lookup table initialized with link-time data.
2#![doc = concat!("```rust\n", include_str!("../../examples/map.rs"), "\n```\n")]
3
4use link_section::{TypedMutableSection, TypedSection};
5use std::{
6    mem::MaybeUninit,
7    ptr,
8    sync::atomic::{AtomicU8, Ordering},
9};
10
11use crate::hash::ConstHash;
12pub use crate::map::table::ScatteredMapTable;
13#[cfg(feature = "__internal")]
14pub use build::{initialize_scattered_map, safe_byte_count_for_capacity};
15
16mod build;
17mod probe;
18mod table;
19
20/// One gathered map entry.
21///
22/// Scatter sites store the hash next to the key and value so map initialization
23/// only needs to place rows into metadata slots.
24#[repr(C)]
25pub struct MapRecord<K, V> {
26    /// The key of the record.
27    pub key: K,
28    /// The value of the record.
29    pub value: V,
30    /// The hash of the record.
31    pub hash: u64,
32}
33
34impl<K, V> MapRecord<K, V> {
35    /// Create a new map record with a key, value and known `const` hash.
36    pub const fn new(key: K, value: V, hash: u64) -> Self {
37        Self { key, value, hash }
38    }
39}
40
41impl<K, V> ::core::fmt::Debug for MapRecord<K, V>
42where
43    K: ::core::fmt::Debug,
44    V: ::core::fmt::Debug,
45{
46    fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result {
47        f.debug_struct("MapRecord")
48            .field("key", &self.key)
49            .field("value", &self.value)
50            .finish()
51    }
52}
53
54/// One zeroed aux padding block per scatter site.
55#[doc(hidden)]
56pub type MapMetadataChunk = [u8; build::PER_ITEM_CAPACITY];
57
58/// Zero-filled padding chunk for the map metadata aux section.
59#[doc(hidden)]
60pub const MAP_METADATA_CHUNK_ZERO: MapMetadataChunk = [0; _];
61
62/// One zeroed aux padding block per scatter site.
63#[doc(hidden)]
64pub type MapMetadataChunkBase = [u8; build::BASE_CAPACITY];
65
66/// Zero-filled padding chunk for the map metadata aux section.
67#[doc(hidden)]
68pub const MAP_METADATA_CHUNK_BASE_ZERO: MapMetadataChunkBase = [0; _];
69
70impl<K: ConstHash + PartialEq + 'static, V: 'static> ScatteredMap<K, V> {
71    #[doc(hidden)]
72    pub const fn __new(state: &'static __ScatteredMapState<K, V>) -> Self {
73        Self { state }
74    }
75
76    /// Lookup a value by key.
77    #[inline]
78    pub fn find(&self, key: &K) -> Option<&V> {
79        let this = self.state;
80        let table = this.ensure_initialized();
81        let hash = ConstHash::hash(key);
82        let offset = (table.lookup_fn)(table, hash)?;
83        let record = &this.records[offset as usize];
84        if record.key == *key {
85            Some(&record.value)
86        } else {
87            None
88        }
89    }
90
91    /// Iterate over the entries in the map.
92    pub fn entries(&self) -> impl Iterator<Item = (&K, &V)> {
93        self.state
94            .records
95            .iter()
96            .map(|record| (&record.key, &record.value))
97    }
98
99    /// Iterate over the keys in the map.
100    pub fn keys(&self) -> impl Iterator<Item = &K> {
101        self.state.records.iter().map(|record| &record.key)
102    }
103
104    /// Iterate over the values in the map.
105    pub fn values(&self) -> impl Iterator<Item = &V> {
106        self.state.records.iter().map(|record| &record.value)
107    }
108
109    /// True when a key is present in the map.
110    #[inline]
111    pub fn contains_key(&self, key: &K) -> bool {
112        self.find(key).is_some()
113    }
114
115    /// The number of records in the map.
116    #[inline]
117    pub fn len(&self) -> usize {
118        self.state.records.len()
119    }
120
121    /// True if the map has no records.
122    #[inline]
123    pub fn is_empty(&self) -> bool {
124        self.state.records.is_empty()
125    }
126}
127
128impl<K: 'static, V: 'static> IntoIterator for &'static ScatteredMap<K, V> {
129    type Item = (&'static K, &'static V);
130    type IntoIter = ::std::iter::Map<
131        ::std::slice::Iter<'static, MapRecord<K, V>>,
132        fn(&MapRecord<K, V>) -> (&K, &V),
133    >;
134    fn into_iter(self) -> Self::IntoIter {
135        self.state
136            .records
137            .iter()
138            .map(|record| (&record.key, &record.value))
139    }
140}
141
142/// A swiss-table-style lookup table initialized with link-time data. Each item
143/// in the map must have a unique hash.
144///
145/// For a flat list in arbitrary link order, use [`crate::ScatteredSlice`]. For a sorted
146/// list without per-item handles, use [`crate::ScatteredSortedSlice`]. For `static`
147/// handles at scatter sites, use [`crate::ScatteredReferencedSlice`]; when sorted order
148/// is required as well, use [`crate::ScatteredSortedReferencedSlice`].
149///
150/// ## Performance notes
151///
152/// The map's metadata section is only ever written to once, and then becomes
153/// read-only. This means that we can avoid tombstone logic.
154///
155/// Metadata is arranged in 16-slot SIMD groups for optimal performance.
156/// ```rust
157#[doc = include_str!("../../examples/map.rs")]
158/// ```
159pub struct ScatteredMap<K: 'static, V: 'static> {
160    state: &'static __ScatteredMapState<K, V>,
161}
162
163#[doc(hidden)]
164pub struct __ScatteredMapState<K: 'static, V: 'static> {
165    state: AtomicU8,
166    records: &'static TypedSection<MapRecord<K, V>>,
167    refs: &'static TypedMutableSection<u8>,
168    table: ::core::mem::MaybeUninit<ScatteredMapTable>,
169}
170
171impl<K: 'static, V: 'static> __ScatteredMapState<K, V> {
172    #[doc(hidden)]
173    pub const fn new(
174        records: &'static TypedSection<MapRecord<K, V>>,
175        refs: &'static TypedMutableSection<u8>,
176    ) -> Self {
177        Self {
178            state: AtomicU8::new(0),
179            records,
180            refs,
181            table: MaybeUninit::uninit(),
182        }
183    }
184
185    #[doc(hidden)]
186    pub fn __initialize(&self) {
187        self.ensure_initialized();
188    }
189
190    #[allow(unsafe_code)]
191    fn ensure_initialized(&self) -> &ScatteredMapTable {
192        match self
193            .state
194            .compare_exchange(0, 1, Ordering::Relaxed, Ordering::Relaxed)
195        {
196            Ok(_) => {
197                let table = build::initialize_scattered_map(self.records, unsafe {
198                    self.refs.as_mut_slice()
199                });
200                unsafe { ptr::write(self.table.as_ptr() as _, table) };
201                self.state.store(2, Ordering::Relaxed);
202                unsafe { self.table.assume_init_ref() }
203            }
204            Err(2) => unsafe { self.table.assume_init_ref() },
205            Err(_) => panic!("Recursive or overlapping initialization of static variable"),
206        }
207    }
208}
209
210#[macro_export]
211#[doc(hidden)]
212macro_rules! __map {
213    (@gather $(#[$meta:meta])* $vis:vis static $name:ident: $map:ident < $key:ty, $value:ty >;) => {
214        $crate::__map!(@declare_scatter_macro $name, $key, $value, $vis);
215
216        $(#[$meta])*
217        $vis static $name: $map<$key, $value> = {
218            $crate::__support::link_section::declarative::section!(
219                #[section(typed, no_macro)]
220                static $name: $crate::__support::link_section::TypedSection<
221                    $crate::map::MapRecord<$key, $value>
222                >;
223            );
224
225            $crate::__support::link_section::declarative::section!(
226                #[section(mutable, no_macro, aux(main = $name))]
227                static MAP_META: $crate::__support::link_section::TypedMutableSection<u8>;
228            );
229
230            $crate::__support::link_section::declarative::in_section!(
231                #[in_section(unsafe, name = MAP_META, aux(main = $name), type = mutable)]
232                const _: $crate::map::MapMetadataChunkBase = $crate::map::MAP_METADATA_CHUNK_BASE_ZERO;
233            );
234
235            static __MAP_STATE: $crate::map::__ScatteredMapState<$key, $value> = $crate::map::__ScatteredMapState::new($name.const_deref(), MAP_META.const_deref());
236
237            $crate::__support::ctor::declarative::ctor!(
238                #[ctor(unsafe, anonymous, priority = 0)]
239                fn __map_init() {
240                    __MAP_STATE.__initialize();
241                }
242            );
243
244            $crate::map::ScatteredMap::__new(&__MAP_STATE)
245        };
246    };
247    (@declare_scatter_macro $name:ident, $key:ty, $value:ty, $vis:vis) => {
248        $crate::__support::ident_concat!((#[doc(hidden)] #[macro_export] macro_rules!) (
249            __ $name __map_scatter__
250        ) ({
251            ($passthru:tt) => {
252                $crate::__map!(@scatter [$name] [$key] [$value] $passthru);
253            };
254        }));
255
256        $crate::__support::ident_concat!(
257            (#[doc(hidden)] $vis use)
258            (__ $name __map_scatter__)
259            (as $name;)
260        );
261    };
262    (scatter [$collection:ident] => [$key:ty] [$value:ty] $vis:vis $name:ident: $ty:ty = ($key_expr:expr, $value_expr:expr)) => {
263        $collection ! (( [$collection] => $vis static $name: $ty = ($key_expr, $value_expr); ));
264    };
265    (@scatter [$collection_name:ident] [$key:ty] [$value:ty] ([$($meta:tt)*] => $(#[$imeta:meta])* $vis:vis $kind:ident $name:tt: $ty:ty = ($key_expr:expr, $value_expr:expr);)) => {
266        $crate::__support::link_section::declarative::in_section!(
267            #[in_section(unsafe, name = $collection_name, type = typed)]
268            $(#[$imeta])*
269            $vis $kind $name: $crate::map::MapRecord<$key, $value> = $crate::map::MapRecord::new(
270                $key_expr,
271                $value_expr,
272                $crate::const_hash!($key_expr)
273            );
274        );
275        $crate::__support::link_section::declarative::in_section!(
276            #[in_section(unsafe, name = MAP_META, aux(main = $collection_name), type = mutable)]
277            const _: $crate::map::MapMetadataChunk = $crate::map::MAP_METADATA_CHUNK_ZERO;
278        );
279    };
280}
281
282#[cfg(all(test, not(miri)))]
283mod link_tests {
284    use crate::ScatteredMap;
285
286    __map!(@gather pub static TEST_MAP: ScatteredMap<&'static str, u32>;);
287    __map!(scatter [TEST_MAP] => [&'static str] [u32] APPLE: (&'static str, u32) = ("apple", 1));
288    __map!(scatter [TEST_MAP] => [&'static str] [u32] BANANA: (&'static str, u32) = ("banana", 2));
289
290    #[test]
291    fn scattered_map_gather_scatter_find() {
292        assert_eq!(TEST_MAP.len(), 2);
293        assert_eq!(TEST_MAP.find(&"apple"), Some(&1));
294        assert_eq!(TEST_MAP.find(&"banana"), Some(&2));
295        assert_eq!(TEST_MAP.find(&"orange"), None);
296        assert!(TEST_MAP.contains_key(&"apple"));
297    }
298
299    struct Record {
300        key: &'static str,
301        f: fn(&'static str),
302    }
303
304    impl Record {
305        pub const fn new(key: &'static str, f: fn(&'static str)) -> Self {
306            Self { key, f }
307        }
308
309        pub fn call(&self) {
310            (self.f)(self.key);
311        }
312    }
313
314    __map!(@gather static TEST_MAP_2: ScatteredMap<&'static str, Record>;);
315
316    macro_rules! make_test {
317        ($($name:ident)*) => {
318            $(
319            __map!(scatter [TEST_MAP_2] => [&'static str] [Record]
320                $name: (&'static str, Record) = (
321                    stringify!($name),
322                    Record::new(stringify!($name), |_key| println!(stringify!($name)))
323                )
324            );
325            )*
326        }
327    }
328
329    make_test!(
330        A B C D E F G H I J K L M N O P Q R S T U V W X Y Z
331        A1 B1 C1 D1 E1 F1 G1 H1 I1 J1 K1 L1 M1 N1 O1 P1 Q1 R1 S1 T1 U1 V1 W1 X1 Y1 Z1
332        A2 B2 C2 D2 E2 F2 G2 H2 I2 J2 K2 L2 M2 N2 O2 P2 Q2 R2 S2 T2 U2 V2 W2 X2 Y2 Z2
333        A3 B3 C3 D3 E3 F3 G3 H3 I3 J3 K3 L3 M3 N3 O3 P3 Q3 R3 S3 T3 U3 V3 W3 X3 Y3 Z3
334        A4 B4 C4 D4 E4 F4 G4 H4 I4 J4 K4 L4 M4 N4 O4 P4 Q4 R4 S4 T4 U4 V4 W4 X4 Y4 Z4
335        A5 B5 C5 D5 E5 F5 G5 H5 I5 J5 K5 L5 M5 N5 O5 P5 Q5 R5 S5 T5 U5 V5 W5 X5 Y5 Z5
336    );
337
338    #[test]
339    fn test_scattered_map_2() {
340        TEST_MAP_2.find(&"A").unwrap().call();
341        TEST_MAP_2.find(&"B").unwrap().call();
342        TEST_MAP_2.find(&"C").unwrap().call();
343    }
344}