Skip to main content

jsonld_core/context/
memo.rs

1use hashbrown::{Equivalent, HashTable, hash_table::Entry};
2use once_cell::sync::OnceCell;
3use parking_lot::Mutex;
4use std::hash::{BuildHasher, Hash};
5
6const SHARDS: usize = 16;
7
8/// One shard of a [`ShardedMemo`], aligned so that neighbouring shards' locks
9/// never share a cache line. 128 bytes is the line size on Apple silicon and
10/// the span of the adjacent-line prefetcher on x86.
11#[repr(align(128))]
12struct Shard<K, V>(Mutex<HashTable<(K, V)>>);
13
14/// Memo map split into independently locked shards.
15///
16/// A processed context is shared across threads — a server caches one and
17/// hands a clone to every request — so its memos see concurrent lookups. With
18/// a single lock every lookup serialises; here a key's hash picks one of
19/// sixteen shards, so lookups of different keys rarely meet on the same lock.
20///
21/// The key is hashed once per operation: the same hash selects the shard and
22/// probes that shard's table.
23pub struct ShardedMemo<K, V> {
24    // Allocated on the first insert. Expansion and context processing create a
25    // fresh context, and with it a fresh memo, for every document and scoped
26    // context, many of which are dropped after a lookup or two; building 2 KB
27    // of padded shards for each of those costs more than the lookups save.
28    shards: OnceCell<Box<[Shard<K, V>; SHARDS]>>,
29    hasher: crate::DefaultBuildHasher,
30}
31
32impl<K, V> Default for ShardedMemo<K, V> {
33    fn default() -> Self {
34        Self {
35            shards: OnceCell::new(),
36            hasher: crate::DefaultBuildHasher::default(),
37        }
38    }
39}
40
41impl<K: Hash + Eq, V: Clone> ShardedMemo<K, V> {
42    // The table indexes buckets with the low bits of the hash and tags them
43    // with the top seven, so the shard comes from bits in between, masked to
44    // four bits before the cast.
45    #[allow(clippy::cast_possible_truncation)]
46    fn shard(shards: &[Shard<K, V>; SHARDS], hash: u64) -> &Mutex<HashTable<(K, V)>> {
47        &shards[((hash >> 48) & (SHARDS as u64 - 1)) as usize].0
48    }
49
50    /// Returns a clone of the value memoised for `key`, if any.
51    pub fn get<Q: Hash + Equivalent<K> + ?Sized>(&self, key: &Q) -> Option<V> {
52        let shards = self.shards.get()?;
53        let hash = self.hasher.hash_one(key);
54        Self::shard(shards, hash).lock().find(hash, |(k, _)| key.equivalent(k)).map(|(_, v)| v.clone())
55    }
56
57    /// Memoises `value` for `key`, replacing any value already there.
58    pub fn insert(&self, key: K, value: V) {
59        let shards = self
60            .shards
61            .get_or_init(|| Box::new(std::array::from_fn(|_| Shard(Mutex::new(HashTable::new())))));
62        let hash = self.hasher.hash_one(&key);
63        let mut shard = Self::shard(shards, hash).lock();
64        match shard.entry(hash, |(k, _)| *k == key, |(k, _)| self.hasher.hash_one(k)) {
65            Entry::Occupied(mut entry) => entry.get_mut().1 = value,
66            Entry::Vacant(entry) => {
67                entry.insert((key, value));
68            }
69        }
70    }
71}