cranelift_codegen/
ctxhash.rs1use hashbrown::raw::RawTable;
8use std::hash::{Hash, Hasher};
9
10pub trait CtxEq<V1: ?Sized, V2: ?Sized> {
18 fn ctx_eq(&self, a: &V1, b: &V2) -> bool;
21}
22
23pub trait CtxHash<Value: ?Sized>: CtxEq<Value, Value> {
25 fn ctx_hash<H: Hasher>(&self, state: &mut H, value: &Value);
28}
29
30#[derive(Default)]
33pub struct NullCtx;
34
35impl<V: Eq + Hash> CtxEq<V, V> for NullCtx {
36 fn ctx_eq(&self, a: &V, b: &V) -> bool {
37 a.eq(b)
38 }
39}
40impl<V: Eq + Hash> CtxHash<V> for NullCtx {
41 fn ctx_hash<H: Hasher>(&self, state: &mut H, value: &V) {
42 value.hash(state);
43 }
44}
45
46struct BucketData<K, V> {
55 hash: u32,
56 k: K,
57 v: V,
58}
59
60pub struct CtxHashMap<K, V> {
62 raw: RawTable<BucketData<K, V>>,
63}
64
65impl<K, V> CtxHashMap<K, V> {
66 pub fn with_capacity(capacity: usize) -> Self {
69 Self {
70 raw: RawTable::with_capacity(capacity),
71 }
72 }
73}
74
75fn compute_hash<Ctx, K>(ctx: &Ctx, k: &K) -> u32
76where
77 Ctx: CtxHash<K>,
78{
79 let mut hasher = crate::fx::FxHasher::default();
80 ctx.ctx_hash(&mut hasher, k);
81 hasher.finish() as u32
82}
83
84impl<K, V> CtxHashMap<K, V> {
85 pub fn insert<Ctx>(&mut self, k: K, v: V, ctx: &Ctx) -> Option<V>
88 where
89 Ctx: CtxEq<K, K> + CtxHash<K>,
90 {
91 let hash = compute_hash(ctx, &k);
92 match self.raw.find(hash as u64, |bucket| {
93 hash == bucket.hash && ctx.ctx_eq(&bucket.k, &k)
94 }) {
95 Some(bucket) => {
96 let data = unsafe { bucket.as_mut() };
97 Some(std::mem::replace(&mut data.v, v))
98 }
99 None => {
100 let data = BucketData { hash, k, v };
101 self.raw
102 .insert_entry(hash as u64, data, |bucket| bucket.hash as u64);
103 None
104 }
105 }
106 }
107
108 pub fn get<'a, Q, Ctx>(&'a self, k: &Q, ctx: &Ctx) -> Option<&'a V>
110 where
111 Ctx: CtxEq<K, Q> + CtxHash<Q> + CtxHash<K>,
112 {
113 let hash = compute_hash(ctx, k);
114 self.raw
115 .find(hash as u64, |bucket| {
116 hash == bucket.hash && ctx.ctx_eq(&bucket.k, k)
117 })
118 .map(|bucket| {
119 let data = unsafe { bucket.as_ref() };
120 &data.v
121 })
122 }
123}
124
125#[cfg(test)]
126mod test {
127 use super::*;
128 use std::hash::Hash;
129
130 #[derive(Clone, Copy, Debug)]
131 struct Key {
132 index: u32,
133 }
134 struct Ctx {
135 vals: &'static [&'static str],
136 }
137 impl CtxEq<Key, Key> for Ctx {
138 fn ctx_eq(&self, a: &Key, b: &Key) -> bool {
139 self.vals[a.index as usize].eq(self.vals[b.index as usize])
140 }
141 }
142 impl CtxHash<Key> for Ctx {
143 fn ctx_hash<H: Hasher>(&self, state: &mut H, value: &Key) {
144 self.vals[value.index as usize].hash(state);
145 }
146 }
147
148 #[test]
149 fn test_basic() {
150 let ctx = Ctx {
151 vals: &["a", "b", "a"],
152 };
153
154 let k0 = Key { index: 0 };
155 let k1 = Key { index: 1 };
156 let k2 = Key { index: 2 };
157
158 assert!(ctx.ctx_eq(&k0, &k2));
159 assert!(!ctx.ctx_eq(&k0, &k1));
160 assert!(!ctx.ctx_eq(&k2, &k1));
161
162 let mut map: CtxHashMap<Key, u64> = CtxHashMap::with_capacity(4);
163 assert_eq!(map.insert(k0, 42, &ctx), None);
164 assert_eq!(map.insert(k2, 84, &ctx), Some(42));
165 assert_eq!(map.get(&k1, &ctx), None);
166 assert_eq!(*map.get(&k0, &ctx).unwrap(), 84);
167 }
168}