1use std::hash::{Hash, Hasher};
2use std::ops::Index;
3use std::sync::Arc;
4
5use celox_design::VarAtomBase;
6use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet, FxHasher};
7
8use crate::RangeStore;
9
10type SymbolicRange<A, N> = RangeStore<Option<(N, HashSet<VarAtomBase<A>>)>>;
11type RangeEntry<A, N> = Arc<SymbolicRange<A, N>>;
12type StoreShard<A, N> = HashMap<A, RangeEntry<A, N>>;
13
14const STORE_SHARDS: usize = 64;
19type StorePages<A, N> = [Option<Arc<StoreShard<A, N>>>; STORE_SHARDS];
20
21fn shard_index(key: &(impl Hash + ?Sized)) -> usize {
22 let mut hasher = FxHasher::default();
23 key.hash(&mut hasher);
24 (hasher.finish() as usize) & (STORE_SHARDS - 1)
25}
26
27#[derive(Clone, Debug)]
34pub struct SymbolicStore<A, N> {
35 entries: Option<Arc<StorePages<A, N>>>,
36 len: usize,
37}
38
39impl<A, N> Default for SymbolicStore<A, N> {
40 fn default() -> Self {
41 Self {
42 entries: None,
43 len: 0,
44 }
45 }
46}
47
48impl<A, N> SymbolicStore<A, N> {
49 pub fn fork(&self) -> Self {
55 Self {
56 entries: self.entries.as_ref().map(Arc::clone),
57 len: self.len,
58 }
59 }
60}
61
62impl<A: Eq + Hash, N> SymbolicStore<A, N> {
63 fn entry_arc(&self, key: &A) -> Option<&RangeEntry<A, N>> {
64 self.entries.as_deref()?[shard_index(key)]
65 .as_deref()?
66 .get(key)
67 }
68
69 pub fn len(&self) -> usize {
70 self.len
71 }
72
73 pub fn is_empty(&self) -> bool {
74 self.len == 0
75 }
76
77 pub fn reserve(&mut self, additional: usize)
78 where
79 A: Clone,
80 N: Clone,
81 {
82 if additional == 0 {
83 return;
84 }
85 let per_shard = additional.div_ceil(STORE_SHARDS);
86 let pages = Arc::make_mut(
87 self.entries
88 .get_or_insert_with(|| Arc::new(std::array::from_fn(|_| None))),
89 );
90 for shard in pages {
91 Arc::make_mut(shard.get_or_insert_with(|| Arc::new(HashMap::default())))
92 .reserve(per_shard);
93 }
94 }
95
96 pub fn contains_key(&self, key: &A) -> bool {
97 self.entry_arc(key).is_some()
98 }
99
100 pub fn get(&self, key: &A) -> Option<&SymbolicRange<A, N>> {
101 self.entry_arc(key).map(AsRef::as_ref)
102 }
103
104 pub fn insert(&mut self, key: A, value: SymbolicRange<A, N>) -> Option<RangeEntry<A, N>>
105 where
106 A: Clone,
107 N: Clone,
108 {
109 let pages = Arc::make_mut(
110 self.entries
111 .get_or_insert_with(|| Arc::new(std::array::from_fn(|_| None))),
112 );
113 let shard = Arc::make_mut(
114 pages[shard_index(&key)].get_or_insert_with(|| Arc::new(HashMap::default())),
115 );
116 let previous = shard.insert(key, Arc::new(value));
117 if previous.is_none() {
118 self.len += 1;
119 }
120 previous
121 }
122
123 pub fn remove(&mut self, key: &A) -> Option<RangeEntry<A, N>>
124 where
125 A: Clone,
126 N: Clone,
127 {
128 let index = shard_index(key);
129 if !self.entries.as_deref()?[index]
130 .as_deref()?
131 .contains_key(key)
132 {
133 return None;
134 }
135 let pages = Arc::make_mut(self.entries.as_mut().expect("the page table exists"));
136 let shard = pages[index].as_mut().expect("the populated shard exists");
137 let removed = Arc::make_mut(shard).remove(key);
138 if removed.is_some() {
139 self.len -= 1;
140 }
141 if shard.is_empty() {
142 pages[index] = None;
143 }
144 if self.len == 0 {
145 self.entries = None;
146 }
147 removed
148 }
149
150 pub fn iter(&self) -> Iter<'_, A, N> {
151 Iter {
152 shards: self.entries.as_deref().map(|pages| pages.iter()),
153 current: None,
154 remaining: self.len,
155 }
156 }
157
158 pub fn values(&self) -> Values<'_, A, N> {
159 Values {
160 shards: self.entries.as_deref().map(|pages| pages.iter()),
161 current: None,
162 remaining: self.len,
163 }
164 }
165
166 pub fn keys(&self) -> Keys<'_, A, N> {
167 Keys {
168 shards: self.entries.as_deref().map(|pages| pages.iter()),
169 current: None,
170 remaining: self.len,
171 }
172 }
173
174 pub fn extend<I>(&mut self, values: I)
175 where
176 A: Clone,
177 N: Clone,
178 I: IntoIterator<Item = (A, SymbolicRange<A, N>)>,
179 {
180 for (key, value) in values {
181 self.insert(key, value);
182 }
183 }
184
185 pub fn shares_entry_with(&self, other: &Self, key: &A) -> bool {
190 match (self.entry_arc(key), other.entry_arc(key)) {
191 (Some(left), Some(right)) => Arc::ptr_eq(left, right),
192 (None, None) => true,
193 _ => false,
194 }
195 }
196
197 pub fn differing_keys<'a>(&'a self, other: &'a Self) -> impl Iterator<Item = &'a A> {
205 let versions_match = match (&self.entries, &other.entries) {
206 (None, None) => true,
207 (Some(left), Some(right)) => Arc::ptr_eq(left, right),
208 _ => false,
209 };
210 DifferingKeys {
211 left_shards: self.entries.as_deref(),
212 right_shards: other.entries.as_deref(),
213 next_shard: if versions_match { STORE_SHARDS } else { 0 },
214 left_shard: None,
215 right_shard: None,
216 left_entries: None,
217 right_keys: None,
218 }
219 }
220}
221
222struct DifferingKeys<'a, A, N> {
223 left_shards: Option<&'a StorePages<A, N>>,
224 right_shards: Option<&'a StorePages<A, N>>,
225 next_shard: usize,
226 left_shard: Option<&'a StoreShard<A, N>>,
227 right_shard: Option<&'a StoreShard<A, N>>,
228 left_entries: Option<std::collections::hash_map::Iter<'a, A, RangeEntry<A, N>>>,
229 right_keys: Option<std::collections::hash_map::Keys<'a, A, RangeEntry<A, N>>>,
230}
231
232impl<'a, A: Eq + Hash, N> Iterator for DifferingKeys<'a, A, N> {
233 type Item = &'a A;
234
235 fn next(&mut self) -> Option<Self::Item> {
236 loop {
237 while let Some((key, left)) = self.left_entries.as_mut().and_then(Iterator::next) {
238 if self
239 .right_shard
240 .and_then(|right| right.get(key))
241 .is_none_or(|right| !Arc::ptr_eq(left, right))
242 {
243 return Some(key);
244 }
245 }
246
247 while let Some(key) = self.right_keys.as_mut().and_then(Iterator::next) {
248 if self.left_shard.is_none_or(|left| !left.contains_key(key)) {
249 return Some(key);
250 }
251 }
252
253 if self.next_shard == STORE_SHARDS {
254 return None;
255 }
256 let left = self
257 .left_shards
258 .and_then(|shards| shards[self.next_shard].as_ref());
259 let right = self
260 .right_shards
261 .and_then(|shards| shards[self.next_shard].as_ref());
262 self.next_shard += 1;
263 if match (left, right) {
264 (None, None) => true,
265 (Some(left), Some(right)) => Arc::ptr_eq(left, right),
266 _ => false,
267 } {
268 continue;
269 }
270
271 self.left_shard = left.map(|shard| shard.as_ref());
272 self.right_shard = right.map(|shard| shard.as_ref());
273 self.left_entries = left.map(|shard| shard.iter());
274 self.right_keys = right.map(|shard| shard.keys());
275 }
276 }
277}
278
279impl<A: Clone + Eq + Hash, N: Clone> SymbolicStore<A, N> {
280 pub fn clone_entry_from(&mut self, key: &A, source: &Self) -> bool {
281 let Some(value) = source.entry_arc(key) else {
282 return false;
283 };
284 let pages = Arc::make_mut(
285 self.entries
286 .get_or_insert_with(|| Arc::new(std::array::from_fn(|_| None))),
287 );
288 let shard = Arc::make_mut(
289 pages[shard_index(key)].get_or_insert_with(|| Arc::new(HashMap::default())),
290 );
291 if shard.insert(key.clone(), Arc::clone(value)).is_none() {
292 self.len += 1;
293 }
294 true
295 }
296
297 pub fn entry(&mut self, key: A) -> Entry<'_, A, N> {
298 let shard_index = shard_index(&key);
299 let pages = Arc::make_mut(
300 self.entries
301 .get_or_insert_with(|| Arc::new(std::array::from_fn(|_| None))),
302 );
303 Entry {
304 inner: Arc::make_mut(
305 pages[shard_index].get_or_insert_with(|| Arc::new(HashMap::default())),
306 )
307 .entry(key),
308 len: &mut self.len,
309 }
310 }
311
312 pub fn get_mut(&mut self, key: &A) -> Option<&mut SymbolicRange<A, N>> {
313 let index = shard_index(key);
314 if !self.entries.as_deref()?[index]
315 .as_deref()?
316 .contains_key(key)
317 {
318 return None;
319 }
320 let pages = Arc::make_mut(self.entries.as_mut().expect("the page table exists"));
321 let shard = pages[index].as_mut().expect("the populated shard exists");
322 Arc::make_mut(shard).get_mut(key).map(Arc::make_mut)
323 }
324
325 pub fn values_mut(&mut self) -> ValuesMut<'_, A, N> {
326 ValuesMut {
327 shards: self
328 .entries
329 .as_mut()
330 .map(Arc::make_mut)
331 .map(|pages| pages.iter_mut()),
332 current: None,
333 remaining: self.len,
334 }
335 }
336}
337
338pub struct Entry<'a, A, N> {
339 inner: std::collections::hash_map::Entry<'a, A, RangeEntry<A, N>>,
340 len: &'a mut usize,
341}
342
343impl<'a, A: Clone, N: Clone> Entry<'a, A, N> {
344 pub fn or_insert_with(
345 self,
346 default: impl FnOnce() -> SymbolicRange<A, N>,
347 ) -> &'a mut SymbolicRange<A, N> {
348 let value = match self.inner {
349 std::collections::hash_map::Entry::Occupied(entry) => entry.into_mut(),
350 std::collections::hash_map::Entry::Vacant(entry) => {
351 *self.len += 1;
352 entry.insert(Arc::new(default()))
353 }
354 };
355 Arc::make_mut(value)
356 }
357}
358
359impl<A: Eq + Hash, N> Index<&A> for SymbolicStore<A, N> {
360 type Output = SymbolicRange<A, N>;
361
362 fn index(&self, key: &A) -> &Self::Output {
363 self.get(key).expect("symbolic store key is absent")
364 }
365}
366
367pub struct Iter<'a, A, N> {
368 shards: Option<std::slice::Iter<'a, Option<Arc<StoreShard<A, N>>>>>,
369 current: Option<std::collections::hash_map::Iter<'a, A, RangeEntry<A, N>>>,
370 remaining: usize,
371}
372
373impl<'a, A, N> Iterator for Iter<'a, A, N> {
374 type Item = (&'a A, &'a SymbolicRange<A, N>);
375
376 fn next(&mut self) -> Option<Self::Item> {
377 loop {
378 if let Some((key, value)) = self.current.as_mut().and_then(Iterator::next) {
379 self.remaining -= 1;
380 return Some((key, value.as_ref()));
381 }
382 let Some(shard) = self.shards.as_mut()?.next()?.as_deref() else {
383 continue;
384 };
385 self.current = Some(shard.iter());
386 }
387 }
388
389 fn size_hint(&self) -> (usize, Option<usize>) {
390 (self.remaining, Some(self.remaining))
391 }
392}
393
394impl<A, N> ExactSizeIterator for Iter<'_, A, N> {}
395
396pub struct Values<'a, A, N> {
397 shards: Option<std::slice::Iter<'a, Option<Arc<StoreShard<A, N>>>>>,
398 current: Option<std::collections::hash_map::Values<'a, A, RangeEntry<A, N>>>,
399 remaining: usize,
400}
401
402impl<'a, A, N> Iterator for Values<'a, A, N> {
403 type Item = &'a SymbolicRange<A, N>;
404
405 fn next(&mut self) -> Option<Self::Item> {
406 loop {
407 if let Some(value) = self.current.as_mut().and_then(Iterator::next) {
408 self.remaining -= 1;
409 return Some(value.as_ref());
410 }
411 let Some(shard) = self.shards.as_mut()?.next()?.as_deref() else {
412 continue;
413 };
414 self.current = Some(shard.values());
415 }
416 }
417
418 fn size_hint(&self) -> (usize, Option<usize>) {
419 (self.remaining, Some(self.remaining))
420 }
421}
422
423impl<A, N> ExactSizeIterator for Values<'_, A, N> {}
424
425pub struct Keys<'a, A, N> {
426 shards: Option<std::slice::Iter<'a, Option<Arc<StoreShard<A, N>>>>>,
427 current: Option<std::collections::hash_map::Keys<'a, A, RangeEntry<A, N>>>,
428 remaining: usize,
429}
430
431impl<'a, A, N> Iterator for Keys<'a, A, N> {
432 type Item = &'a A;
433
434 fn next(&mut self) -> Option<Self::Item> {
435 loop {
436 if let Some(key) = self.current.as_mut().and_then(Iterator::next) {
437 self.remaining -= 1;
438 return Some(key);
439 }
440 let Some(shard) = self.shards.as_mut()?.next()?.as_deref() else {
441 continue;
442 };
443 self.current = Some(shard.keys());
444 }
445 }
446
447 fn size_hint(&self) -> (usize, Option<usize>) {
448 (self.remaining, Some(self.remaining))
449 }
450}
451
452impl<A, N> ExactSizeIterator for Keys<'_, A, N> {}
453
454pub struct ValuesMut<'a, A: Clone, N: Clone> {
455 shards: Option<std::slice::IterMut<'a, Option<Arc<StoreShard<A, N>>>>>,
456 current: Option<std::collections::hash_map::ValuesMut<'a, A, RangeEntry<A, N>>>,
457 remaining: usize,
458}
459
460impl<'a, A: Clone, N: Clone> Iterator for ValuesMut<'a, A, N> {
461 type Item = &'a mut SymbolicRange<A, N>;
462
463 fn next(&mut self) -> Option<Self::Item> {
464 loop {
465 if let Some(value) = self.current.as_mut().and_then(Iterator::next) {
466 self.remaining -= 1;
467 return Some(Arc::make_mut(value));
468 }
469 let Some(shard) = self.shards.as_mut()?.next()?.as_mut() else {
470 continue;
471 };
472 self.current = Some(Arc::make_mut(shard).values_mut());
473 }
474 }
475
476 fn size_hint(&self) -> (usize, Option<usize>) {
477 (self.remaining, Some(self.remaining))
478 }
479}
480
481impl<A: Clone, N: Clone> ExactSizeIterator for ValuesMut<'_, A, N> {}
482
483pub struct IntoIter<A, N> {
484 inner: std::vec::IntoIter<(A, SymbolicRange<A, N>)>,
485}
486
487impl<A, N> Iterator for IntoIter<A, N> {
488 type Item = (A, SymbolicRange<A, N>);
489
490 fn next(&mut self) -> Option<Self::Item> {
491 self.inner.next()
492 }
493
494 fn size_hint(&self) -> (usize, Option<usize>) {
495 self.inner.size_hint()
496 }
497}
498
499impl<A, N> ExactSizeIterator for IntoIter<A, N> {}
500
501impl<'a, A: Eq + Hash, N> IntoIterator for &'a SymbolicStore<A, N> {
502 type Item = (&'a A, &'a SymbolicRange<A, N>);
503 type IntoIter = Iter<'a, A, N>;
504
505 fn into_iter(self) -> Self::IntoIter {
506 self.iter()
507 }
508}
509
510impl<A: Clone + Eq + Hash, N: Clone> IntoIterator for SymbolicStore<A, N> {
511 type Item = (A, SymbolicRange<A, N>);
512 type IntoIter = IntoIter<A, N>;
513
514 fn into_iter(self) -> Self::IntoIter {
515 let mut values = Vec::with_capacity(self.len);
516 if let Some(pages) = self.entries {
517 let pages = Arc::try_unwrap(pages).unwrap_or_else(|shared| shared.as_ref().clone());
518 for shard in pages {
519 let Some(shard) = shard else {
520 continue;
521 };
522 let shard = Arc::try_unwrap(shard).unwrap_or_else(|shared| shared.as_ref().clone());
523 values.extend(shard.into_iter().map(|(key, value)| {
524 let value =
525 Arc::try_unwrap(value).unwrap_or_else(|shared| shared.as_ref().clone());
526 (key, value)
527 }));
528 }
529 }
530 IntoIter {
531 inner: values.into_iter(),
532 }
533 }
534}
535
536impl<A: Clone + Eq + Hash, N: Clone> FromIterator<(A, SymbolicRange<A, N>)>
537 for SymbolicStore<A, N>
538{
539 fn from_iter<T: IntoIterator<Item = (A, SymbolicRange<A, N>)>>(iter: T) -> Self {
540 let mut store = Self::default();
541 store.extend(iter);
542 store
543 }
544}
545
546#[cfg(test)]
547mod tests {
548 use celox_design::BitAccess;
549
550 use super::*;
551
552 #[test]
553 fn empty_store_allocates_no_shard_pages() {
554 let mut store = SymbolicStore::<u32, u32>::default();
555 assert!(store.entries.is_none());
556 assert!(store.get_mut(&7).is_none());
557 assert!(store.entries.is_none());
558
559 let branch = store.fork();
560 assert!(branch.entries.is_none());
561 }
562
563 #[test]
564 fn removing_a_last_entry_releases_its_shard_page() {
565 let mut store = SymbolicStore::<u32, u32>::default();
566 store.insert(7, RangeStore::new(None, 8));
567 let index = shard_index(&7);
568 assert!(store.entries.as_deref().unwrap()[index].is_some());
569
570 assert!(store.remove(&7).is_some());
571 assert!(store.entries.is_none());
572 assert!(store.is_empty());
573 }
574
575 #[test]
576 fn cloned_store_copies_a_range_only_when_it_is_mutated() {
577 let mut original = SymbolicStore::<u32, u32>::default();
578 original.insert(7, RangeStore::new(None, 8));
579
580 let mut branch = original.fork();
581 branch
582 .get_mut(&7)
583 .unwrap()
584 .update(BitAccess::new(0, 3), Some((11, HashSet::default())))
585 .unwrap();
586
587 assert_eq!(
588 original[&7].get_parts(BitAccess::new(0, 7)).unwrap()[0].0,
589 None
590 );
591 assert_eq!(
592 branch[&7].get_parts(BitAccess::new(0, 7)).unwrap()[0].0,
593 Some((11, HashSet::default()))
594 );
595 assert!(!original.shares_entry_with(&branch, &7));
596 }
597
598 #[test]
599 fn branch_write_copies_only_one_sparse_page() {
600 let mut original = SymbolicStore::<u32, u32>::default();
601 for key in 0..4096 {
602 original.insert(key, RangeStore::new(None, 8));
603 }
604
605 let mut branch = original.fork();
606 assert!(Arc::ptr_eq(
607 original.entries.as_ref().unwrap(),
608 branch.entries.as_ref().unwrap()
609 ));
610 branch
611 .get_mut(&2048)
612 .unwrap()
613 .update(BitAccess::new(0, 0), Some((1, HashSet::default())))
614 .unwrap();
615 assert!(!Arc::ptr_eq(
616 original.entries.as_ref().unwrap(),
617 branch.entries.as_ref().unwrap()
618 ));
619
620 let original_pages = original.entries.as_deref().unwrap();
621 let branch_pages = branch.entries.as_deref().unwrap();
622 let shared_pages = original_pages
623 .iter()
624 .zip(branch_pages)
625 .filter(|(left, right)| {
626 left.as_ref()
627 .zip(right.as_ref())
628 .is_some_and(|(left, right)| Arc::ptr_eq(left, right))
629 })
630 .count();
631 assert_eq!(shared_pages, STORE_SHARDS - 1);
632 assert!(original.shares_entry_with(&branch, &7));
633 assert!(!original.shares_entry_with(&branch, &2048));
634 assert_eq!(
635 original
636 .differing_keys(&branch)
637 .copied()
638 .collect::<Vec<_>>(),
639 vec![2048]
640 );
641 }
642
643 #[test]
644 fn branch_versions_report_only_their_divergent_definitions() {
645 let mut entry = SymbolicStore::<u32, u32>::default();
646 for key in 0..1024 {
647 entry.insert(key, RangeStore::new(None, 8));
648 }
649
650 let mut then_version = entry.fork();
651 then_version
652 .get_mut(&17)
653 .unwrap()
654 .update(BitAccess::new(0, 3), Some((1, HashSet::default())))
655 .unwrap();
656 let mut else_version = entry.fork();
657 else_version
658 .get_mut(&900)
659 .unwrap()
660 .update(BitAccess::new(4, 7), Some((2, HashSet::default())))
661 .unwrap();
662
663 let mut differing = then_version
664 .differing_keys(&else_version)
665 .copied()
666 .collect::<Vec<_>>();
667 differing.sort_unstable();
668 assert_eq!(differing, vec![17, 900]);
669 assert_eq!(
670 entry
671 .differing_keys(&then_version)
672 .copied()
673 .collect::<Vec<_>>(),
674 vec![17]
675 );
676 assert_eq!(
677 entry
678 .differing_keys(&else_version)
679 .copied()
680 .collect::<Vec<_>>(),
681 vec![900]
682 );
683 }
684
685 #[test]
686 fn branch_diff_iterator_reports_insertions_and_removals_once() {
687 let mut left = SymbolicStore::<u32, u32>::default();
688 left.insert(1, RangeStore::new(None, 8));
689 left.insert(2, RangeStore::new(None, 8));
690
691 let mut right = left.fork();
692 right.remove(&1);
693 right.insert(3, RangeStore::new(None, 8));
694
695 let mut differing = left.differing_keys(&right).copied().collect::<Vec<_>>();
696 differing.sort_unstable();
697 assert_eq!(differing, vec![1, 3]);
698
699 let mut reverse = right.differing_keys(&left).copied().collect::<Vec<_>>();
700 reverse.sort_unstable();
701 assert_eq!(reverse, vec![1, 3]);
702 }
703
704 #[test]
705 fn sharded_iterators_and_entry_api_preserve_map_semantics() {
706 let mut store = SymbolicStore::<u32, u32>::default();
707 for key in 0..256 {
708 store.entry(key).or_insert_with(|| RangeStore::new(None, 1));
709 }
710 assert_eq!(store.len(), 256);
711 assert_eq!(store.keys().count(), 256);
712 assert_eq!(store.values().count(), 256);
713 assert_eq!(store.iter().count(), 256);
714
715 for value in store.values_mut() {
716 value
717 .update(BitAccess::new(0, 0), Some((9, HashSet::default())))
718 .unwrap();
719 }
720 assert!(
721 store
722 .values()
723 .all(|value| value.get_parts(BitAccess::new(0, 0)).unwrap()[0].0
724 == Some((9, HashSet::default())))
725 );
726
727 assert!(store.remove(&128).is_some());
728 assert_eq!(store.len(), 255);
729 assert_eq!(store.into_iter().count(), 255);
730 }
731}