Skip to main content

charon_lib/utils/
hash_cons.rs

1use derive_generic_visitor::{Drive, DriveMut, DriveTwo, Visit, VisitMut, VisitTwo};
2use std::hash::Hash;
3use std::ops::{ControlFlow, Deref};
4use std::sync::Arc;
5
6use crate::utils::hash_by_addr::HashByAddr;
7use crate::utils::type_map::Mappable;
8
9/// Hash-consed data structure: a reference-counted wrapper that guarantees that two equal
10/// value will be stored at the same address. This makes it possible to use the pointer address
11/// as a hash value.
12// Warning: a `derive` should not introduce a way to create a new `HashConsed` value without
13// going through the interning table.
14#[derive(PartialEq, Eq, Hash)]
15pub struct HashConsed<T>(HashByAddr<Arc<T>>);
16
17impl<T> Clone for HashConsed<T> {
18    fn clone(&self) -> Self {
19        Self(self.0.clone())
20    }
21}
22
23impl<T> HashConsed<T> {
24    pub fn inner(&self) -> &T {
25        self.0.0.as_ref()
26    }
27}
28
29impl<T: PartialOrd> PartialOrd for HashConsed<T> {
30    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
31        self.inner().partial_cmp(other.inner())
32    }
33}
34
35impl<T: Ord> Ord for HashConsed<T> {
36    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
37        self.inner().cmp(other.inner())
38    }
39}
40
41pub trait HashConsable: Hash + PartialEq + Eq + Clone + Mappable {}
42impl<T> HashConsable for T where T: Hash + PartialEq + Eq + Clone + Mappable {}
43
44// Private module that contains the static we'll use as interning map. A value of type
45// `HashCons` MUST NOT be created in any other way than this table, else hashing and euqality
46// on it will be broken. Note that this likely means that if a crate uses charon both as a
47// direct dependency and as a dylib, then the static will be duplicated, causing hashing and
48// equality on `HashCons` to be broken.
49mod intern_table {
50    use rustc_hash::FxBuildHasher;
51    use std::borrow::Borrow;
52    use std::sync::{Arc, LazyLock, RwLock};
53
54    use super::{HashConsable, HashConsed};
55    use crate::utils::hash_by_addr::HashByAddr;
56    use crate::utils::type_map::{Mappable, Mapper, TypeMap};
57
58    type SeqHashSet<T> = indexmap::IndexSet<T, FxBuildHasher>;
59
60    // This is a static mutable `SeqHashSet<Arc<T>>` that records for each `T` value a unique
61    // `Arc<T>` that contains the same value. Values inside the set are hashed/compared
62    // as is normal for `T`.
63    // Once we've gotten an `Arc` out of the set however, we're sure that "T-equality"
64    // implies address-equality, hence the `HashByAddr` wrapper preserves correct equality
65    // and hashing behavior.
66    struct InternMapper;
67    impl Mapper for InternMapper {
68        type Value<T: Mappable> = SeqHashSet<Arc<T>>;
69    }
70    static INTERNED: LazyLock<RwLock<TypeMap<InternMapper>>> = LazyLock::new(Default::default);
71
72    // The excessive generality is to make it work for both `U = T` and `U = Arc<T>`.
73    pub fn intern<T: HashConsable, U>(inner: U) -> HashConsed<T>
74    where
75        Arc<T>: Borrow<U>,
76        U: Into<Arc<T>> + std::hash::Hash,
77        U: indexmap::Equivalent<Arc<T>>,
78    {
79        // Fast read-only check.
80        let arc = if let read_guard = INTERNED.read().unwrap()
81            && let Some(set) = read_guard.get::<T>()
82            && let Some(arc) = set.get(&inner)
83        {
84            arc.clone()
85        } else {
86            // Concurrent access is possible right here, so we have to check everything again.
87            let mut write_guard = INTERNED.write().unwrap();
88            let set: &mut SeqHashSet<Arc<T>> = write_guard.or_default::<T>();
89            if let Some(arc) = set.get(&inner) {
90                arc.clone()
91            } else {
92                let arc: Arc<T> = inner.into();
93                set.insert(arc.clone());
94                arc
95            }
96        };
97        HashConsed(HashByAddr(arc))
98    }
99
100    /// Mutate the contents in-place if possible.
101    pub fn mutate_in_place<T: HashConsable, R, F: FnOnce(&mut T) -> R>(
102        x: &mut HashConsed<T>,
103        f: F,
104    ) -> Result<R, F> {
105        let arc = &mut x.0.0;
106        // Every value has at least two pointers: the current value and the one stored in the
107        // global map. If there are exactly two, we may mutate directly by discarding the one in
108        // the global map temporarily.
109        if Arc::strong_count(arc) != 2 {
110            return Err(f);
111        }
112        {
113            // Take the write guard just long enough to drop the other `Arc` to this value.
114            let mut write_guard = INTERNED.write().unwrap();
115            // Check the count again, it could have changed concurrently.
116            if Arc::strong_count(arc) != 2 {
117                return Err(f);
118            }
119            if let Some(other_arc) = write_guard.or_default::<T>().swap_take(&*arc) {
120                drop(other_arc);
121            } else {
122                // Nothing was removed, early return.
123                return Err(f);
124            }
125            // The Arc was removed from the map; `x` is invalid as interning the same value would
126            // result in a different pointer. NO MORE EARLY RETURN until we fix that.
127        }
128        // If we are still the sole owner, we can now mutate in-place.
129        let ret = match Arc::get_mut(arc) {
130            Some(val) => Ok(f(val)),
131            None => Err(f),
132        };
133        // Re-establish the interning invariant. If the same value was added to the map in the
134        // meantime, we'll get a pointer to that.
135        *x = HashConsed::from_arc(arc.clone());
136        ret
137    }
138}
139
140impl<T> HashConsed<T>
141where
142    T: HashConsable,
143{
144    /// Deduplicate the values by hashing them. This deduplication is crucial for the hashing
145    /// function to be correct. This is the only function allowed to create `Self` values.
146    pub fn new(inner: T) -> Self {
147        intern_table::intern(inner)
148    }
149    /// Rarely used: in case we already have an `Arc`, may avoid an allocation.
150    pub fn from_arc(inner: Arc<T>) -> Self {
151        intern_table::intern(inner)
152    }
153
154    /// Clones if needed to get mutable access to the inner value.
155    pub fn with_inner_mut<R>(&mut self, f: impl FnOnce(&mut T) -> R) -> R {
156        match intern_table::mutate_in_place(self, f) {
157            Ok(r) => r,
158            Err(f) => {
159                // The value is behind a shared `Arc`, we clone it in order to mutate it.
160                let mut value = self.inner().clone();
161                let ret = f(&mut value);
162                // Re-intern the new value.
163                *self = Self::new(value);
164                ret
165            }
166        }
167    }
168}
169
170impl<T> Deref for HashConsed<T> {
171    type Target = T;
172    fn deref(&self) -> &Self::Target {
173        self.inner()
174    }
175}
176
177impl<T: std::fmt::Debug> std::fmt::Debug for HashConsed<T> {
178    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
179        // Hide the `HashByAddr` wrapper.
180        f.debug_tuple("HashConsed").field(self.inner()).finish()
181    }
182}
183
184impl<'s, T, V: Visit<'s, T>> Drive<'s, V> for HashConsed<T> {
185    fn drive_inner(&'s self, v: &mut V) -> ControlFlow<V::Break> {
186        v.visit(self.inner())
187    }
188}
189impl<'s, T, V: VisitTwo<'s, T>> DriveTwo<'s, V> for HashConsed<T> {
190    fn drive_two_inner(&'s self, other: &'s Self, v: &mut V) -> ControlFlow<V::Break> {
191        v.visit(self.inner(), other.inner())
192    }
193}
194/// Note: this explores the inner value mutably by cloning and re-hashing afterwards.
195impl<'s, T, V> DriveMut<'s, V> for HashConsed<T>
196where
197    T: HashConsable,
198    V: for<'a> VisitMut<'a, T>,
199{
200    fn drive_inner_mut(&'s mut self, v: &mut V) -> ControlFlow<V::Break> {
201        self.with_inner_mut(|inner| v.visit(inner))
202    }
203}
204
205/// `HashCons` values are deduplicated in the serialized output: see [`crate::utils::dedup`].
206mod serialize {
207    use serde::{Deserialize, Serialize};
208    use serde_state::{DeserializeState, SerializeState};
209
210    use super::{HashConsable, HashConsed};
211    use crate::utils::dedup::*;
212
213    impl<T> Serialize for HashConsed<T>
214    where
215        T: Serialize + HashConsable,
216    {
217        fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
218        where
219            S: serde::Serializer,
220        {
221            SerDedup::Untagged(self.inner()).serialize(serializer)
222        }
223    }
224    /// Options for the state are `()` to serialize values normally and `DedupSerializer`
225    /// to deduplicate identical values in the serialized output.
226    impl<T, State> SerializeState<State> for HashConsed<T>
227    where
228        T: SerializeState<State> + HashConsable,
229        State: DedupSerializerState,
230    {
231        fn serialize_state<S>(&self, state: &State, serializer: S) -> Result<S::Ok, S::Error>
232        where
233            S: serde::Serializer,
234        {
235            serialize_dedup(self, self.inner(), state, serializer)
236        }
237    }
238
239    impl<'de, T> Deserialize<'de> for HashConsed<T>
240    where
241        T: Deserialize<'de> + HashConsable,
242    {
243        fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
244        where
245            D: serde::Deserializer<'de>,
246        {
247            use serde::de::Error;
248            let repr: SerDedup<T> = SerDedup::deserialize(deserializer)?;
249            match repr {
250                SerDedup::Value { .. } | SerDedup::Deduplicated { .. } => {
251                    Err(D::Error::custom(stateless_deserialize_error::<T>()))
252                }
253                SerDedup::Untagged(val) => Ok(HashConsed::new(val)),
254            }
255        }
256    }
257    impl<'de, T, State> DeserializeState<'de, State> for HashConsed<T>
258    where
259        T: DeserializeState<'de, State> + HashConsable,
260        State: DedupSerializerState,
261    {
262        fn deserialize_state<D>(state: &State, deserializer: D) -> Result<Self, D::Error>
263        where
264            D: serde::Deserializer<'de>,
265        {
266            deserialize_dedup(state, deserializer, HashConsed::new)
267        }
268    }
269}
270
271#[test]
272fn test_hash_cons() {
273    let x = HashConsed::new(42u32);
274    let y = HashConsed::new(42u32);
275    assert_eq!(x, y);
276    // Test a serialization round-trip.
277    let z = serde_json::from_value(serde_json::to_value(x.clone()).unwrap()).unwrap();
278    assert_eq!(x, z);
279}
280
281#[test]
282fn test_hash_cons_concurrent() {
283    use itertools::Itertools;
284    let handles = (0..10)
285        .map(|_| std::thread::spawn(|| std::hint::black_box(HashConsed::new(42u32))))
286        .collect_vec();
287    let values = handles.into_iter().map(|h| h.join().unwrap()).collect_vec();
288    assert!(values.iter().all_equal())
289}
290
291#[test]
292fn test_hash_cons_dedup() {
293    use crate::utils::dedup::DedupSerializer;
294    use serde_state::{DeserializeState, SerializeState};
295    type Ty = HashConsed<TyKind>;
296    #[derive(Debug, Clone, PartialEq, Eq, Hash, SerializeState, DeserializeState)]
297    #[serde_state(state = DedupSerializer)]
298    enum TyKind {
299        Bool,
300        Pair(Ty, Ty),
301    }
302
303    // Build a value with some redundancy.
304    let bool1 = HashConsed::new(TyKind::Bool);
305    let bool2 = HashConsed::new(TyKind::Bool);
306    let pair = HashConsed::new(TyKind::Pair(bool1.clone(), bool2));
307    let triple = HashConsed::new(TyKind::Pair(bool1, pair));
308
309    let state = DedupSerializer::default();
310    let json_val = triple
311        .serialize_state(&state, serde_json::value::Serializer)
312        .unwrap();
313    let state = DedupSerializer::default();
314    let round_tripped = Ty::deserialize_state(&state, json_val).unwrap();
315
316    assert_eq!(triple, round_tripped);
317}