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#[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
44mod 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 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 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 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 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 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 if Arc::strong_count(arc) != 2 {
110 return Err(f);
111 }
112 {
113 let mut write_guard = INTERNED.write().unwrap();
115 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 return Err(f);
124 }
125 }
128 let ret = match Arc::get_mut(arc) {
130 Some(val) => Ok(f(val)),
131 None => Err(f),
132 };
133 *x = HashConsed::from_arc(arc.clone());
136 ret
137 }
138}
139
140impl<T> HashConsed<T>
141where
142 T: HashConsable,
143{
144 pub fn new(inner: T) -> Self {
147 intern_table::intern(inner)
148 }
149 pub fn from_arc(inner: Arc<T>) -> Self {
151 intern_table::intern(inner)
152 }
153
154 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 let mut value = self.inner().clone();
161 let ret = f(&mut value);
162 *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 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}
194impl<'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
205mod 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 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 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 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}