charon_lib/ast/
hash_cons.rs

1use derive_generic_visitor::{Drive, DriveMut, DriveTwo, Visit, VisitMut, VisitTwo};
2use serde::{Deserialize, Serialize};
3use std::hash::Hash;
4use std::ops::{ControlFlow, Deref};
5use std::sync::Arc;
6
7use crate::common::hash_by_addr::HashByAddr;
8use crate::common::type_map::Mappable;
9
10/// Hash-consed data structure: a reference-counted wrapper that guarantees that two equal
11/// value will be stored at the same address. This makes it possible to use the pointer address
12/// as a hash value.
13// Warning: a `derive` should not introduce a way to create a new `HashConsed` value without
14// going through the interning table.
15#[derive(PartialEq, Eq, Hash)]
16pub struct HashConsed<T>(HashByAddr<Arc<T>>);
17
18impl<T> Clone for HashConsed<T> {
19    fn clone(&self) -> Self {
20        Self(self.0.clone())
21    }
22}
23
24impl<T> HashConsed<T> {
25    pub fn inner(&self) -> &T {
26        self.0.0.as_ref()
27    }
28}
29
30impl<T: PartialOrd> PartialOrd for HashConsed<T> {
31    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
32        self.inner().partial_cmp(other.inner())
33    }
34}
35
36impl<T: Ord> Ord for HashConsed<T> {
37    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
38        self.inner().cmp(other.inner())
39    }
40}
41
42pub trait HashConsable: Hash + PartialEq + Eq + Clone + Mappable {}
43impl<T> HashConsable for T where T: Hash + PartialEq + Eq + Clone + Mappable {}
44
45/// Unique id identifying a hashconsed value amongst all hashconsed values.
46#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
47pub struct HashConsId(u64);
48
49// Private module that contains the static we'll use as interning map. A value of type
50// `HashCons` MUST NOT be created in any other way than this table, else hashing and euqality
51// on it will be broken. Note that this likely means that if a crate uses charon both as a
52// direct dependency and as a dylib, then the static will be duplicated, causing hashing and
53// equality on `HashCons` to be broken.
54mod intern_table {
55    use rustc_hash::FxBuildHasher;
56    use std::borrow::Borrow;
57    use std::sync::atomic::{AtomicU64, Ordering};
58    use std::sync::{Arc, LazyLock, RwLock};
59
60    use super::{HashConsId, HashConsable, HashConsed};
61    use crate::common::hash_by_addr::HashByAddr;
62    use crate::common::type_map::{Mappable, Mapper, TypeMap};
63
64    type SeqHashMap<K, V> = indexmap::IndexMap<K, V, FxBuildHasher>;
65
66    // Only way we create a `HashConsId`.
67    fn fresh_id() -> HashConsId {
68        static ID: AtomicU64 = AtomicU64::new(0);
69        HashConsId(ID.fetch_add(1, Ordering::Relaxed))
70    }
71
72    // This is a static mutable `SeqHashSet<Arc<T>>` that records for each `T` value a unique
73    // `Arc<T>` that contains the same value. Values inside the set are hashed/compared
74    // as is normal for `T`.
75    // Once we've gotten an `Arc` out of the set however, we're sure that "T-equality"
76    // implies address-equality, hence the `HashByAddr` wrapper preserves correct equality
77    // and hashing behavior.
78    // Note that we also store a `HashConsId` for each item so this is a map instead of a set, but
79    // what matters is really the map keys.
80    struct InternMapper;
81    impl Mapper for InternMapper {
82        type Value<T: Mappable> = SeqHashMap<Arc<T>, HashConsId>;
83    }
84    static INTERNED: LazyLock<RwLock<TypeMap<InternMapper>>> = LazyLock::new(Default::default);
85
86    // The excessive generality is to make it work for both `U = T` and `U = Arc<T>`.
87    pub fn intern<T: HashConsable, U>(inner: U) -> HashConsed<T>
88    where
89        Arc<T>: Borrow<U>,
90        U: Into<Arc<T>> + std::hash::Hash,
91        U: indexmap::Equivalent<Arc<T>>,
92    {
93        // Fast read-only check.
94        let arc = if let read_guard = INTERNED.read().unwrap()
95            && let Some(map) = read_guard.get::<T>()
96            && let Some((arc, _id)) = map.get_key_value(&inner)
97        {
98            arc.clone()
99        } else {
100            // Concurrent access is possible right here, so we have to check everything again.
101            let mut write_guard = INTERNED.write().unwrap();
102            let map: &mut SeqHashMap<Arc<T>, _> = write_guard.or_default::<T>();
103            if let Some((arc, _id)) = map.get_key_value(&inner) {
104                arc.clone()
105            } else {
106                let arc: Arc<T> = inner.into();
107                map.insert(arc.clone(), fresh_id());
108                arc
109            }
110        };
111        HashConsed(HashByAddr(arc))
112    }
113
114    /// Mutate the contents in-place if possible.
115    pub fn mutate_in_place<T: HashConsable, R, F: FnOnce(&mut T) -> R>(
116        x: &mut HashConsed<T>,
117        f: F,
118    ) -> Result<R, F> {
119        let arc = &mut x.0.0;
120        // Every value has at least two pointers: the current value and the one stored in the
121        // global map. If there are exactly two, we may mutate directly by discarding the one in
122        // the global map temporarily.
123        if Arc::strong_count(arc) != 2 {
124            return Err(f);
125        }
126        {
127            // Take the write guard just long enough to drop the other `Arc` to this value.
128            let mut write_guard = INTERNED.write().unwrap();
129            if let Some((other_arc, _)) = write_guard.or_default::<T>().swap_remove_entry(&*arc) {
130                drop(other_arc);
131            } else {
132                // Nothing was removed, early return.
133                return Err(f);
134            }
135            // The Arc was removed from the map; `x` is invalid as interning the same value would
136            // result in a different pointer. NO MORE EARLY RETURN until we fix that.
137        }
138        // If we are still the sole owner, we can now mutate in-place.
139        let ret = match Arc::get_mut(arc) {
140            Some(val) => Ok(f(val)),
141            None => Err(f),
142        };
143        // Re-establish the interning invariant. If the same value was added to the map in the
144        // meantime, we'll get a pointer to that.
145        *x = HashConsed::from_arc(arc.clone());
146        ret
147    }
148
149    /// Identify this value uniquely amongst values of its type. The id depends on insertion
150    /// order into the interning table which makes them in principle deterministic.
151    pub fn id<T: HashConsable>(x: &HashConsed<T>) -> HashConsId {
152        // `HashConsed` can only be constructed via `intern`, so we know this value exists in the
153        // table.
154        let read_guard = INTERNED.read().unwrap();
155        let map = read_guard.get::<T>().unwrap();
156        let (_arc, id) = map.get_key_value(&x.0.0).unwrap();
157        *id
158    }
159}
160
161impl<T> HashConsed<T>
162where
163    T: HashConsable,
164{
165    /// Deduplicate the values by hashing them. This deduplication is crucial for the hashing
166    /// function to be correct. This is the only function allowed to create `Self` values.
167    pub fn new(inner: T) -> Self {
168        intern_table::intern(inner)
169    }
170    /// Rarely used: in case we already have an `Arc`, may avoid an allocation.
171    pub fn from_arc(inner: Arc<T>) -> Self {
172        intern_table::intern(inner)
173    }
174
175    /// Clones if needed to get mutable access to the inner value.
176    pub fn with_inner_mut<R>(&mut self, f: impl FnOnce(&mut T) -> R) -> R {
177        match intern_table::mutate_in_place(self, f) {
178            Ok(r) => r,
179            Err(f) => {
180                // The value is behind a shared `Arc`, we clone it in order to mutate it.
181                let mut value = self.inner().clone();
182                let ret = f(&mut value);
183                // Re-intern the new value.
184                *self = Self::new(value);
185                ret
186            }
187        }
188    }
189
190    pub fn id(&self) -> HashConsId {
191        intern_table::id(self)
192    }
193}
194
195impl<T> Deref for HashConsed<T> {
196    type Target = T;
197    fn deref(&self) -> &Self::Target {
198        self.inner()
199    }
200}
201
202impl<T: std::fmt::Debug> std::fmt::Debug for HashConsed<T> {
203    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
204        // Hide the `HashByAddr` wrapper.
205        f.debug_tuple("HashConsed").field(self.inner()).finish()
206    }
207}
208
209impl<'s, T, V: Visit<'s, T>> Drive<'s, V> for HashConsed<T> {
210    fn drive_inner(&'s self, v: &mut V) -> ControlFlow<V::Break> {
211        v.visit(self.inner())
212    }
213}
214impl<'s, T, V: VisitTwo<'s, T>> DriveTwo<'s, V> for HashConsed<T> {
215    fn drive_two_inner(&'s self, other: &'s Self, v: &mut V) -> ControlFlow<V::Break> {
216        v.visit(self.inner(), other.inner())
217    }
218}
219/// Note: this explores the inner value mutably by cloning and re-hashing afterwards.
220impl<'s, T, V> DriveMut<'s, V> for HashConsed<T>
221where
222    T: HashConsable,
223    V: for<'a> VisitMut<'a, T>,
224{
225    fn drive_inner_mut(&'s mut self, v: &mut V) -> ControlFlow<V::Break> {
226        self.with_inner_mut(|inner| v.visit(inner))
227    }
228}
229
230/// `HashCons` supports serializing each value to a unique id in order to serialize
231/// highly-shared values without explosion.
232///
233/// Note that the deduplication scheme is highly order-dependent: we serialize the real value
234/// the first time it comes up, and use ids only subsequent times. This relies on the fact that
235/// `derive(Serialize, Deserialize)` traverse the value in the same order.
236pub use serialize::{HashConsDedupSerializer, HashConsSerializerState};
237mod serialize {
238    use indexmap::IndexMap as SeqHashMap;
239    use serde::{Deserialize, Serialize};
240    use serde_state::{DeserializeState, SerializeState};
241    use std::any::type_name;
242    use std::cell::RefCell;
243    use std::collections::HashSet;
244
245    use super::{HashConsId, HashConsable, HashConsed};
246    use crate::common::type_map::{Mappable, Mapper, TypeMap};
247
248    pub trait HashConsSerializerState: Sized {
249        /// Record that this type is being serialized. Return `None` if we're not deduplicating
250        /// values, otherwise return whether this item was newly recorded.
251        fn record_serialized<T: Mappable>(&self, id: HashConsId) -> Option<bool>;
252        /// Record that we deserialized this type.
253        fn record_deserialized<T: Mappable>(&self, id: HashConsId, value: HashConsed<T>);
254        /// Find the previously-deserialized type with that id.
255        fn get_deserialized_val<T: Mappable>(&self, id: HashConsId) -> Option<HashConsed<T>>;
256    }
257
258    impl HashConsSerializerState for () {
259        fn record_serialized<T: Mappable>(&self, _id: HashConsId) -> Option<bool> {
260            None
261        }
262        fn record_deserialized<T: Mappable>(&self, _id: HashConsId, _value: HashConsed<T>) {}
263        fn get_deserialized_val<T: Mappable>(&self, _id: HashConsId) -> Option<HashConsed<T>> {
264            None
265        }
266    }
267
268    struct SerializeTableMapper;
269    impl Mapper for SerializeTableMapper {
270        type Value<T: Mappable> = HashSet<HashConsId>;
271    }
272    struct DeserializeTableMapper;
273    impl Mapper for DeserializeTableMapper {
274        type Value<T: Mappable> = SeqHashMap<HashConsId, HashConsed<T>>;
275    }
276    #[derive(Default)]
277    pub struct HashConsDedupSerializer {
278        // Table used for serialization.
279        ser: RefCell<TypeMap<SerializeTableMapper>>,
280        // Table used for deserialization.
281        de: RefCell<TypeMap<DeserializeTableMapper>>,
282    }
283    impl HashConsSerializerState for HashConsDedupSerializer {
284        fn record_serialized<T: Mappable>(&self, id: HashConsId) -> Option<bool> {
285            Some(self.ser.borrow_mut().or_default::<T>().insert(id))
286        }
287        fn record_deserialized<T: Mappable>(&self, id: HashConsId, val: HashConsed<T>) {
288            self.de.borrow_mut().or_default::<T>().insert(id, val);
289        }
290        fn get_deserialized_val<T: Mappable>(&self, id: HashConsId) -> Option<HashConsed<T>> {
291            self.de
292                .borrow()
293                .get::<T>()
294                .and_then(|map| map.get(&id))
295                .cloned()
296        }
297    }
298
299    /// A dummy enum used when serializing/deserializing a `HashConsed<T>`.
300    #[derive(Serialize, Deserialize, SerializeState, DeserializeState)]
301    #[serde_state(state_implements = HashConsSerializerState)]
302    enum SerRepr<T> {
303        /// A value represented normally, accompanied by its id. This is emitted the first time
304        /// we serialize a given value: subsequent times will use `SerRepr::Deduplicate`
305        /// instead.
306        HashConsedValue(#[serde_state(stateless)] HashConsId, T),
307        /// A value represented by its id. The actual value must have been emitted as a
308        /// `SerRepr::Value` with that same id earlier.
309        #[serde_state(stateless)]
310        Deduplicated(HashConsId),
311        /// A plain value without an id.
312        Untagged(T),
313    }
314
315    impl<T> Serialize for HashConsed<T>
316    where
317        T: Serialize + HashConsable,
318    {
319        fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
320        where
321            S: serde::Serializer,
322        {
323            SerRepr::Untagged(self.inner()).serialize(serializer)
324        }
325    }
326    /// Options for the state are `()` to serialize values normally and `HashConsDedupSerializer`
327    /// to deduplicate identical values in the serialized output.
328    impl<T, State> SerializeState<State> for HashConsed<T>
329    where
330        T: SerializeState<State> + HashConsable,
331        State: HashConsSerializerState,
332    {
333        fn serialize_state<S>(&self, state: &State, serializer: S) -> Result<S::Ok, S::Error>
334        where
335            S: serde::Serializer,
336        {
337            let hash_cons_id = self.id();
338            let repr = match state.record_serialized::<T>(hash_cons_id) {
339                Some(true) => SerRepr::HashConsedValue(hash_cons_id, self.inner()),
340                Some(false) => SerRepr::Deduplicated(hash_cons_id),
341                None => SerRepr::Untagged(self.inner()),
342            };
343            repr.serialize_state(state, serializer)
344        }
345    }
346
347    impl<'de, T> Deserialize<'de> for HashConsed<T>
348    where
349        T: Deserialize<'de> + HashConsable,
350    {
351        fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
352        where
353            D: serde::Deserializer<'de>,
354        {
355            use serde::de::Error;
356            let repr: SerRepr<T> = SerRepr::deserialize(deserializer)?;
357            match repr {
358                SerRepr::HashConsedValue { .. } | SerRepr::Deduplicated { .. } => {
359                    let msg = format!(
360                        "trying to deserialize a deduplicated value using serde's `{ty}::deserialize` method. \
361                        This won't work, use serde_state's \
362                        `{ty}::deserialize_state(&HashConsDedupSerializer::default(), _)` instead",
363                        ty = type_name::<T>(),
364                    );
365                    Err(D::Error::custom(msg))
366                }
367                SerRepr::Untagged(val) => Ok(HashConsed::new(val)),
368            }
369        }
370    }
371    impl<'de, T, State> DeserializeState<'de, State> for HashConsed<T>
372    where
373        T: DeserializeState<'de, State> + HashConsable,
374        State: HashConsSerializerState,
375    {
376        fn deserialize_state<D>(state: &State, deserializer: D) -> Result<Self, D::Error>
377        where
378            D: serde::Deserializer<'de>,
379        {
380            use serde::de::Error;
381            let repr: SerRepr<T> = SerRepr::deserialize_state(state, deserializer)?;
382            Ok(match repr {
383                SerRepr::HashConsedValue(hash_cons_id, value) => {
384                    let val = HashConsed::new(value);
385                    state.record_deserialized(hash_cons_id, val.clone());
386                    val
387                }
388                SerRepr::Deduplicated(hash_cons_id) => {
389                    state.get_deserialized_val(hash_cons_id).ok_or_else(|| {
390                        let msg = format!(
391                            "can't deserialize deduplicated value of type {}; \
392                            were you careful with managing the deduplication state?",
393                            type_name::<T>()
394                        );
395                        D::Error::custom(msg)
396                    })?
397                }
398                SerRepr::Untagged(val) => HashConsed::new(val),
399            })
400        }
401    }
402}
403
404#[test]
405fn test_hash_cons() {
406    let x = HashConsed::new(42u32);
407    let y = HashConsed::new(42u32);
408    assert_eq!(x, y);
409    // Test a serialization round-trip.
410    let z = serde_json::from_value(serde_json::to_value(x.clone()).unwrap()).unwrap();
411    assert_eq!(x, z);
412}
413
414#[test]
415fn test_hash_cons_concurrent() {
416    use itertools::Itertools;
417    let handles = (0..10)
418        .map(|_| std::thread::spawn(|| std::hint::black_box(HashConsed::new(42u32))))
419        .collect_vec();
420    let values = handles.into_iter().map(|h| h.join().unwrap()).collect_vec();
421    assert!(values.iter().all_equal())
422}
423
424#[test]
425fn test_hash_cons_dedup() {
426    use serde_state::{DeserializeState, SerializeState};
427    type Ty = HashConsed<TyKind>;
428    #[derive(Debug, Clone, PartialEq, Eq, Hash, SerializeState, DeserializeState)]
429    #[serde_state(state = HashConsDedupSerializer)]
430    enum TyKind {
431        Bool,
432        Pair(Ty, Ty),
433    }
434
435    // Build a value with some redundancy.
436    let bool1 = HashConsed::new(TyKind::Bool);
437    let bool2 = HashConsed::new(TyKind::Bool);
438    let pair = HashConsed::new(TyKind::Pair(bool1.clone(), bool2));
439    let triple = HashConsed::new(TyKind::Pair(bool1, pair));
440
441    let state = HashConsDedupSerializer::default();
442    let json_val = triple
443        .serialize_state(&state, serde_json::value::Serializer)
444        .unwrap();
445    let state = HashConsDedupSerializer::default();
446    let round_tripped = Ty::deserialize_state(&state, json_val).unwrap();
447
448    assert_eq!(triple, round_tripped);
449}
charon_lib/ast/hash_cons.rs

charon_lib/ast/
hash_cons.rs
