Skip to main content

cranpose_core/collections/
perfect_map.rs

1//! A map from strings to values whose keys are all known when the program
2//! is built. [`PerfectMap::new`] runs in a `const` context and finds a hash
3//! that gives every key a slot of its own, so a lookup hashes the key once,
4//! reads one slot and compares one key. The search follows PTHash: the keys
5//! fall into buckets, and each bucket, the fullest first, takes the first
6//! pilot that moves all its keys to free slots. The hash folds the key in
7//! 8-byte lanes, one multiply a lane.
8
9// The buckets the keys fall into, and the slots they are placed in.
10const BUCKETS: usize = 32;
11const SLOTS: usize = 128;
12// A slot no key took.
13const EMPTY: u8 = u8::MAX;
14// Seeds tried before a set of keys is refused.
15const SEEDS: u64 = 64;
16
17/// The most keys a [`PerfectMap`] holds: three of every four slots.
18pub const MAX_KEYS: usize = SLOTS * 3 / 4;
19
20/// A map from the strings of a fixed set to values, built at compile time.
21///
22/// ```
23/// use cranpose_core::collections::perfect_map::PerfectMap;
24///
25/// const DIGITS: PerfectMap<u8, 3> = PerfectMap::new([("zero", 0), ("one", 1), ("two", 2)]);
26/// assert_eq!(DIGITS.get("one"), Some(&1));
27/// assert_eq!(DIGITS.get("three"), None);
28/// ```
29pub struct PerfectMap<V: Copy + 'static, const N: usize> {
30    entries: [(&'static str, V); N],
31    seed: u64,
32    pilots: [u8; BUCKETS],
33    slots: [u8; SLOTS],
34}
35
36impl<V: Copy + 'static, const N: usize> PerfectMap<V, N> {
37    /// The map of `entries`. Called in a `const`, it fails the build when two
38    /// keys are equal, when there are more than [`MAX_KEYS`], or when no seed
39    /// places every key.
40    pub const fn new(entries: [(&'static str, V); N]) -> Self {
41        assert!(N <= MAX_KEYS, "a PerfectMap holds at most MAX_KEYS keys");
42        let mut seed = 0;
43        while seed < SEEDS {
44            seed += 1;
45            if let Some((pilots, slots)) = place(&entries, seed) {
46                return Self {
47                    entries,
48                    seed,
49                    pilots,
50                    slots,
51                };
52            }
53        }
54        panic!("no seed places these keys: is one of them there twice?");
55    }
56
57    /// The value of `key`, if it is one of the map's keys.
58    #[inline]
59    pub fn get(&self, key: &str) -> Option<&V> {
60        let hash = lane_hash(self.seed, key.as_bytes());
61        let slot = self.slots[slot_of(hash, self.pilots[bucket_of(hash)])];
62        let (name, value) = self.entries.get(usize::from(slot))?;
63        (*name == key).then_some(value)
64    }
65}
66
67// Folds `bytes` into the seed 8 bytes at a time, each lane with one
68// multiply, then mixes the high bits into the low ones the bucket reads.
69const fn lane_hash(seed: u64, bytes: &[u8]) -> u64 {
70    let mut hash = seed ^ (bytes.len() as u64).wrapping_mul(0x9e37_79b9_7f4a_7c15);
71    let mut at = 0;
72    while at < bytes.len() {
73        let mut word = 0u64;
74        let mut index = 0;
75        while index < 8 && at + index < bytes.len() {
76            word |= (bytes[at + index] as u64) << (8 * index);
77            index += 1;
78        }
79        hash = (hash ^ word).wrapping_mul(0xbf58_476d_1ce4_e5b9);
80        hash ^= hash >> 31;
81        at += 8;
82    }
83    // A multiply carries a difference only to higher bits: keys that differ
84    // in their last byte, as "key 94" and "key 95", left the low bits alike
85    // and 96 such keys found no seed. One more multiply between two shifts
86    // spreads them.
87    hash = (hash ^ (hash >> 32)).wrapping_mul(0x94d0_49bb_1331_11eb);
88    hash ^ (hash >> 29)
89}
90
91// The hash modulo BUCKETS, read from its low byte: BUCKETS divides 256.
92const fn bucket_of(hash: u64) -> usize {
93    hash.to_le_bytes()[0] as usize % BUCKETS
94}
95const _: () = assert!(256 % BUCKETS == 0, "a hash's low byte holds its bucket");
96
97// The slot a key of `hash` takes with `pilot`: the top seven bits of the
98// hash mixed with the pilot.
99const fn slot_of(hash: u64, pilot: u8) -> usize {
100    let mixed = (hash ^ (pilot as u64).wrapping_mul(0x94d0_49bb_1331_11eb))
101        .wrapping_mul(0x9e37_79b9_7f4a_7c15);
102    // The top bits, read from the top byte: SLOTS is at most 256.
103    (mixed.to_be_bytes()[0] >> (8 - SLOTS.trailing_zeros())) as usize
104}
105const _: () = assert!(
106    SLOTS.is_power_of_two() && SLOTS <= 256,
107    "a hash's top byte holds its slot"
108);
109const _: () = assert!(MAX_KEYS < 256, "a key's index fits its slot's byte");
110
111// Every key's slot under `seed`: each bucket, the fullest first, takes the
112// first pilot that gives its keys free slots. `None` when a bucket finds
113// none.
114const fn place<V: Copy, const N: usize>(
115    entries: &[(&'static str, V); N],
116    seed: u64,
117) -> Option<([u8; BUCKETS], [u8; SLOTS])> {
118    let mut hashes = [0u64; N];
119    let mut sizes = [0usize; BUCKETS];
120    let mut key = 0;
121    while key < N {
122        hashes[key] = lane_hash(seed, entries[key].0.as_bytes());
123        sizes[bucket_of(hashes[key])] += 1;
124        key += 1;
125    }
126    let mut pilots = [0u8; BUCKETS];
127    let mut slots = [EMPTY; SLOTS];
128    let mut size = N;
129    while size > 0 {
130        let mut bucket = 0;
131        while bucket < BUCKETS {
132            if sizes[bucket] == size {
133                match place_bucket(&hashes, bucket, &mut slots) {
134                    Some(pilot) => pilots[bucket] = pilot,
135                    None => return None,
136                }
137            }
138            bucket += 1;
139        }
140        size -= 1;
141    }
142    Some((pilots, slots))
143}
144
145// Puts the keys of `bucket` in `slots` with the first pilot that finds each
146// a free slot, and returns that pilot.
147const fn place_bucket<const N: usize>(
148    hashes: &[u64; N],
149    bucket: usize,
150    slots: &mut [u8; SLOTS],
151) -> Option<u8> {
152    let mut pilot = 0u8;
153    loop {
154        if bucket_fits(hashes, bucket, slots, pilot) {
155            let mut key = 0;
156            while key < N {
157                if bucket_of(hashes[key]) == bucket {
158                    // A key's index is its low byte: N is at most MAX_KEYS, below 256.
159                    slots[slot_of(hashes[key], pilot)] = key.to_le_bytes()[0];
160                }
161                key += 1;
162            }
163            return Some(pilot);
164        }
165        if pilot == u8::MAX {
166            return None;
167        }
168        pilot += 1;
169    }
170}
171
172// Whether every key of `bucket` finds a free slot with `pilot`, no two the
173// same.
174const fn bucket_fits<const N: usize>(
175    hashes: &[u64; N],
176    bucket: usize,
177    slots: &[u8; SLOTS],
178    pilot: u8,
179) -> bool {
180    let mut taken = [false; SLOTS];
181    let mut key = 0;
182    while key < N {
183        if bucket_of(hashes[key]) == bucket {
184            let slot = slot_of(hashes[key], pilot);
185            if slots[slot] != EMPTY || taken[slot] {
186                return false;
187            }
188            taken[slot] = true;
189        }
190        key += 1;
191    }
192    true
193}
194
195#[cfg(test)]
196#[path = "tests/perfect_map_tests.rs"]
197mod tests;