Skip to main content

lattices/
map_union_with_tombstones.rs

1//! Module containing the [`MapUnionWithTombstones`] lattice and aliases for different datastructures.
2//!
3//! See [`crate::tombstone`] for documentation on choosing a tombstone implementation.
4
5use core::cmp::Ordering::{self, *};
6use core::fmt::Debug;
7#[cfg(feature = "std")]
8use std::collections::{HashMap, HashSet};
9#[cfg(feature = "std")]
10use std::string::String;
11
12#[cfg(feature = "std")]
13use cc_traits::Remove;
14use cc_traits::{Get, Iter, Len};
15
16#[cfg(feature = "std")]
17use crate::Merge;
18#[cfg(feature = "std")]
19use crate::cc_traits::GetMut;
20use crate::cc_traits::{Keyed, Map, MapIter, SimpleKeyedRef};
21use crate::collections::{EmptyMap, EmptySet, SingletonMap, SingletonSet};
22#[cfg(feature = "std")]
23use crate::tombstone::{FstTombstoneSet, RoaringTombstoneSet, TombstoneSet};
24use crate::{IsBot, IsTop, LatticeFrom, LatticeOrd};
25
26/// Map-union-with-tombstones compound lattice.
27///
28/// When a key is deleted from the map-union-with-tombstones lattice, it is removed from the underlying `map` and placed into
29/// the `tombstones` set.
30///
31/// This forms the first invariant for this data structure. A key should appear either nowhere, in `map` or in `tombstones`.
32/// but never in `map` and `tombstones` at the same time.
33///
34/// merging is done by merging the underlying `map` and then merging the `tombstones` set, then doing `map` = `map` - `tombstones`.
35///
36/// The implementation of `tombstones` can be any set-like thing. This allows a user to plug in their own set-like implementation.
37/// For example, if the user knows that keys will be created and deleted strictly sequentially, then they could create a highly optimized set implementation
38/// which would just be a single integer, correpsonding to the current key value that the set is up to. Queries for keys below that integer would return true,
39/// queries for keys above it would return false.
40#[derive(Copy, Clone, Debug, Default)]
41#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
42pub struct MapUnionWithTombstones<Map, TombstoneSet> {
43    map: Map,
44    tombstones: TombstoneSet,
45}
46
47impl<Map, TombstoneSet> MapUnionWithTombstones<Map, TombstoneSet> {
48    /// Create a new `MapUnionWithTombstones` from a `Map` and a `TombstoneSet`.
49    pub fn new(map: Map, tombstones: TombstoneSet) -> Self {
50        Self { map, tombstones }
51    }
52
53    /// Create a new `MapUnionWithTombstones` from an `Into<Map>` and an `Into<TombstoneSet>`.
54    pub fn new_from(map: impl Into<Map>, tombstones: impl Into<TombstoneSet>) -> Self {
55        Self::new(map.into(), tombstones.into())
56    }
57
58    /// Reveal the inner value as a shared reference.
59    pub fn as_reveal_ref(&self) -> (&Map, &TombstoneSet) {
60        (&self.map, &self.tombstones)
61    }
62
63    /// Reveal the inner value as an exclusive reference.
64    pub fn as_reveal_mut(&mut self) -> (&mut Map, &mut TombstoneSet) {
65        (&mut self.map, &mut self.tombstones)
66    }
67
68    /// Gets the inner by value, consuming self.
69    pub fn into_reveal(self) -> (Map, TombstoneSet) {
70        (self.map, self.tombstones)
71    }
72}
73
74/// Merge implementation using TombstoneSet trait for optimized union operations
75#[cfg(feature = "std")]
76impl<MapSelf, MapOther, K, ValSelf, ValOther, TombstoneSetSelf, TombstoneSetOther>
77    Merge<MapUnionWithTombstones<MapOther, TombstoneSetOther>>
78    for MapUnionWithTombstones<MapSelf, TombstoneSetSelf>
79where
80    MapSelf: Keyed<Key = K, Item = ValSelf>
81        + Extend<(K, ValSelf)>
82        + for<'a> GetMut<&'a K, Item = ValSelf>
83        + for<'b> Remove<&'b K>,
84    MapOther: IntoIterator<Item = (K, ValOther)>,
85    ValSelf: Merge<ValOther> + LatticeFrom<ValOther>,
86    ValOther: IsBot,
87    TombstoneSetSelf: TombstoneSet<K>,
88    TombstoneSetOther: IntoIterator<Item = K>,
89{
90    fn merge(&mut self, other: MapUnionWithTombstones<MapOther, TombstoneSetOther>) -> bool {
91        use alloc::vec::Vec;
92
93        let mut changed = false;
94
95        // Collect other.tombstones into a vec to avoid borrowing issues
96        let other_tombstones: Vec<_> = other.tombstones.into_iter().collect();
97        let old_tombstones_len = self.tombstones.len();
98
99        // This vec collect is needed to prevent simultaneous mut references `self.0.extend` and
100        // `self.0.get_mut`.
101        // TODO(mingwei): This could be fixed with a different structure, maybe some sort of
102        // `Collection` entry API.
103        let iter: Vec<_> = other
104            .map
105            .into_iter()
106            .filter(|(k_other, val_other)| {
107                !val_other.is_bot() && !self.tombstones.contains(k_other)
108            })
109            .filter_map(|(k_other, val_other)| {
110                if let Some(mut val_self) = self.map.get_mut(&k_other) {
111                    // Key collision, merge into `self`.
112                    changed |= val_self.merge(val_other);
113                    None
114                } else {
115                    // New value, convert for extending.
116                    changed = true;
117                    Some((k_other, ValSelf::lattice_from(val_other)))
118                }
119            })
120            .collect();
121        self.map.extend(iter);
122
123        // Extend tombstones and remove tombstoned keys from the map
124        self.tombstones
125            .extend(other_tombstones.into_iter().inspect(|k| {
126                self.map.remove(k);
127            }));
128
129        if old_tombstones_len != self.tombstones.len() {
130            changed = true;
131        }
132
133        changed
134    }
135}
136
137impl<MapSelf, MapOther, K, ValSelf, ValOther, TombstoneSetSelf, TombstoneSetOther>
138    LatticeFrom<MapUnionWithTombstones<MapOther, TombstoneSetOther>>
139    for MapUnionWithTombstones<MapSelf, TombstoneSetSelf>
140where
141    MapSelf: Keyed<Key = K, Item = ValSelf> + FromIterator<(K, ValSelf)>,
142    MapOther: IntoIterator<Item = (K, ValOther)>,
143    ValSelf: LatticeFrom<ValOther>,
144    TombstoneSetSelf: FromIterator<K>,
145    TombstoneSetOther: IntoIterator<Item = K>,
146{
147    fn lattice_from(other: MapUnionWithTombstones<MapOther, TombstoneSetOther>) -> Self {
148        Self {
149            map: other
150                .map
151                .into_iter()
152                .map(|(k_other, val_other)| (k_other, LatticeFrom::lattice_from(val_other)))
153                .collect(),
154            tombstones: other.tombstones.into_iter().collect(),
155        }
156    }
157}
158
159impl<MapSelf, MapOther, K, ValSelf, ValOther, TombstoneSetSelf, TombstoneSetOther>
160    PartialOrd<MapUnionWithTombstones<MapOther, TombstoneSetOther>>
161    for MapUnionWithTombstones<MapSelf, TombstoneSetSelf>
162where
163    MapSelf: Map<K, ValSelf, Key = K, Item = ValSelf> + MapIter + SimpleKeyedRef,
164    MapOther: Map<K, ValOther, Key = K, Item = ValOther> + MapIter + SimpleKeyedRef,
165    ValSelf: PartialOrd<ValOther> + IsBot,
166    ValOther: IsBot,
167    TombstoneSetSelf: Len + Iter<Item = K> + for<'a> Get<&'a K>,
168    TombstoneSetOther: Len + Iter<Item = K> + for<'a> Get<&'a K>,
169{
170    fn partial_cmp(
171        &self,
172        other: &MapUnionWithTombstones<MapOther, TombstoneSetOther>,
173    ) -> Option<Ordering> {
174        let self_tombstones_greater = self
175            .tombstones
176            .iter()
177            .any(|k| !other.tombstones.contains(&*k));
178
179        let other_tombstones_greater = other
180            .tombstones
181            .iter()
182            .any(|k| !self.tombstones.contains(&*k));
183
184        if self_tombstones_greater && other_tombstones_greater {
185            return None;
186        }
187
188        let mut self_any_greater = false;
189        let mut other_any_greater = false;
190        let self_keys = self
191            .map
192            .iter()
193            .filter(|(k, v)| {
194                !v.is_bot() && !self.tombstones.contains(k) && !other.tombstones.contains(k)
195            })
196            .map(|(k, _v)| <MapSelf as SimpleKeyedRef>::into_ref(k));
197        let other_keys = other
198            .map
199            .iter()
200            .filter(|(k, v)| {
201                !v.is_bot() && !self.tombstones.contains(k) && !other.tombstones.contains(k)
202            })
203            .map(|(k, _v)| <MapOther as SimpleKeyedRef>::into_ref(k));
204
205        for k in self_keys.chain(other_keys) {
206            match (self.map.get(k), other.map.get(k)) {
207                (Some(self_value), Some(other_value)) => {
208                    match self_value.partial_cmp(&*other_value)? {
209                        Less => {
210                            other_any_greater = true;
211                        }
212                        Greater => {
213                            self_any_greater = true;
214                        }
215                        Equal => {}
216                    }
217                }
218                (Some(_), None) => {
219                    self_any_greater = true;
220                }
221                (None, Some(_)) => {
222                    other_any_greater = true;
223                }
224                (None, None) => unreachable!(),
225            }
226            if self_any_greater && other_any_greater {
227                return None;
228            }
229        }
230        match (
231            self_any_greater,
232            other_any_greater,
233            self_tombstones_greater,
234            other_tombstones_greater,
235        ) {
236            (false, false, false, false) => Some(Equal),
237
238            (false, false, true, false) => Some(Greater),
239            (false, false, false, true) => Some(Less),
240
241            (true, false, false, false) => Some(Greater),
242            (false, true, false, false) => Some(Less),
243
244            (true, false, true, false) => Some(Greater),
245            (false, true, false, true) => Some(Less),
246
247            (true, false, false, true) => None,
248            (false, true, true, false) => None,
249
250            (true, true, _, _) => unreachable!(),
251            (_, _, true, true) => unreachable!(),
252        }
253    }
254}
255impl<MapSelf, MapOther, TombstoneSetSelf, TombstoneSetOther>
256    LatticeOrd<MapUnionWithTombstones<MapOther, TombstoneSetOther>>
257    for MapUnionWithTombstones<MapSelf, TombstoneSetSelf>
258where
259    Self: PartialOrd<MapUnionWithTombstones<MapOther, TombstoneSetOther>>,
260{
261}
262
263impl<MapSelf, MapOther, K, ValSelf, ValOther, TombstoneSetSelf, TombstoneSetOther>
264    PartialEq<MapUnionWithTombstones<MapOther, TombstoneSetOther>>
265    for MapUnionWithTombstones<MapSelf, TombstoneSetSelf>
266where
267    MapSelf: Map<K, ValSelf, Key = K, Item = ValSelf> + MapIter + SimpleKeyedRef,
268    MapOther: Map<K, ValOther, Key = K, Item = ValOther> + MapIter + SimpleKeyedRef,
269    ValSelf: PartialEq<ValOther> + IsBot,
270    ValOther: IsBot,
271    TombstoneSetSelf: Len + Iter<Item = K> + for<'a> Get<&'a K>,
272    TombstoneSetOther: Len + Iter<Item = K> + for<'b> Get<&'b K>,
273{
274    fn eq(&self, other: &MapUnionWithTombstones<MapOther, TombstoneSetOther>) -> bool {
275        if self.tombstones.len() != other.tombstones.len() {
276            return false;
277        }
278
279        if self
280            .tombstones
281            .iter()
282            .any(|k| !other.tombstones.contains(&*k))
283        {
284            return false;
285        }
286
287        if other
288            .tombstones
289            .iter()
290            .any(|k| !self.tombstones.contains(&*k))
291        {
292            return false;
293        }
294
295        let self_keys = self
296            .map
297            .iter()
298            .filter(|(_k, v)| !v.is_bot())
299            .map(|(k, _v)| <MapSelf as SimpleKeyedRef>::into_ref(k));
300        let other_keys = other
301            .map
302            .iter()
303            .filter(|(_k, v)| !v.is_bot())
304            .map(|(k, _v)| <MapOther as SimpleKeyedRef>::into_ref(k));
305        for k in self_keys.chain(other_keys) {
306            match (self.map.get(k), other.map.get(k)) {
307                (Some(self_value), Some(other_value)) => {
308                    if *self_value != *other_value {
309                        return false;
310                    }
311                }
312                (None, None) => unreachable!(),
313                _ => {
314                    return false;
315                }
316            }
317        }
318
319        true
320    }
321}
322impl<MapSelf, TombstoneSetSelf> Eq for MapUnionWithTombstones<MapSelf, TombstoneSetSelf> where
323    Self: PartialEq
324{
325}
326
327impl<Map, TombstoneSet> IsBot for MapUnionWithTombstones<Map, TombstoneSet>
328where
329    Map: Iter,
330    Map::Item: IsBot,
331    TombstoneSet: Len,
332{
333    fn is_bot(&self) -> bool {
334        self.map.iter().all(|v| v.is_bot()) && self.tombstones.is_empty()
335    }
336}
337
338impl<Map, TombstoneSet> IsTop for MapUnionWithTombstones<Map, TombstoneSet> {
339    fn is_top(&self) -> bool {
340        false
341    }
342}
343
344/// [`std::collections::HashMap`]-backed [`MapUnionWithTombstones`] lattice.
345#[cfg(feature = "std")]
346pub type MapUnionHashMapWithTombstoneHashSet<K, Val> =
347    MapUnionWithTombstones<HashMap<K, Val>, HashSet<K>>;
348
349/// [`crate::collections::SingletonMap`]-backed [`MapUnionWithTombstones`] lattice.
350pub type MapUnionWithTombstonesSingletonMapOnly<K, Val> =
351    MapUnionWithTombstones<SingletonMap<K, Val>, EmptySet<K>>;
352
353/// [`crate::collections::SingletonSet`]-backed [`MapUnionWithTombstones`] lattice.
354pub type MapUnionWithTombstonesTombstoneSingletonSetOnly<K, Val> =
355    MapUnionWithTombstones<EmptyMap<K, Val>, SingletonSet<K>>;
356
357/// [`crate::tombstone::RoaringTombstoneSet`]-backed tombstone set with [`HashMap`] for the main map.
358/// Provides space-efficient tombstone storage for u64 integer keys.
359#[cfg(feature = "std")]
360pub type MapUnionWithTombstonesRoaring<Val> =
361    MapUnionWithTombstones<HashMap<u64, Val>, RoaringTombstoneSet>;
362
363/// FST-backed tombstone set with [`HashMap`] for the main map.
364/// Provides space-efficient, collision-free tombstone storage for String keys.
365#[cfg(feature = "std")]
366pub type MapUnionWithTombstonesFstString<Val> =
367    MapUnionWithTombstones<HashMap<String, Val>, FstTombstoneSet<String>>;
368
369#[cfg(test)]
370mod test {
371    use std::borrow::ToOwned;
372
373    use super::*;
374    use crate::NaiveLatticeOrd;
375    use crate::set_union::{SetUnion, SetUnionHashSet, SetUnionSingletonSet};
376    use crate::test::check_all;
377
378    #[test]
379    fn test_map_union() {
380        type K = &'static str;
381        type V = usize;
382
383        type M = MapUnionWithTombstones<HashMap<K, SetUnionHashSet<V>>, HashSet<K>>;
384        type S = MapUnionWithTombstones<SingletonMap<K, SetUnionSingletonSet<V>>, EmptySet<K>>;
385        type T = MapUnionWithTombstones<EmptyMap<K, SetUnion<EmptySet<V>>>, SingletonSet<K>>;
386
387        let mut my_map_a = M::default();
388        let my_map_b = S::new(
389            SingletonMap("hello", SetUnion::new(SingletonSet(100))),
390            Default::default(),
391        );
392
393        let my_map_c = T::new(Default::default(), SingletonSet("hello"));
394
395        my_map_a.merge(my_map_b);
396        my_map_a.merge(my_map_c);
397
398        assert!(!my_map_a.as_reveal_ref().0.contains_key("hello"));
399    }
400
401    #[test]
402    fn contrain1() {
403        type T = MapUnionWithTombstones<HashMap<i32, SetUnion<HashSet<i32>>>, HashSet<i32>>;
404
405        let a = T::new_from([], HashSet::from_iter([0]));
406        let b = T::new_from(
407            [(0, SetUnionHashSet::new_from([0]))],
408            HashSet::from_iter([]),
409        );
410
411        assert_eq!(a.naive_cmp(&b), Some(Greater));
412        assert_eq!(a.partial_cmp(&b), Some(Greater));
413
414        let a = T::new_from([], HashSet::from_iter([1]));
415        let b = T::new_from([(0, SetUnionHashSet::new_from([0]))], HashSet::default());
416
417        assert_eq!(a.naive_cmp(&b), None);
418        assert_eq!(a.partial_cmp(&b), None);
419    }
420
421    #[cfg(feature = "alloc")]
422    #[test]
423    fn consistency() {
424        use alloc::vec::Vec;
425
426        type K = &'static str;
427        type V = SetUnion<HashSet<i32>>;
428
429        type M = MapUnionWithTombstones<HashMap<K, V>, HashSet<K>>;
430
431        let mut test_vec = Vec::new();
432
433        #[rustfmt::skip]
434        {
435            test_vec.push(M::new_from([], HashSet::from_iter([])));
436
437            test_vec.push(M::new_from([], HashSet::from_iter(["a"])));
438            test_vec.push(M::new_from([], HashSet::from_iter(["b"])));
439            test_vec.push(M::new_from([], HashSet::from_iter(["a", "b"])));
440
441            test_vec.push(M::new_from([("a", SetUnionHashSet::new_from([]))], HashSet::from_iter([])));
442            test_vec.push(M::new_from([("a", SetUnionHashSet::new_from([0]))], HashSet::from_iter([])));
443            test_vec.push(M::new_from([("a", SetUnionHashSet::new_from([1]))], HashSet::from_iter([])));
444            test_vec.push(M::new_from([("a", SetUnionHashSet::new_from([0, 1]))], HashSet::from_iter([])));
445
446            test_vec.push(M::new_from([("b", SetUnionHashSet::new_from([]))], HashSet::from_iter([])));
447            test_vec.push(M::new_from([("b", SetUnionHashSet::new_from([0]))], HashSet::from_iter([])));
448            test_vec.push(M::new_from([("b", SetUnionHashSet::new_from([1]))], HashSet::from_iter([])));
449            test_vec.push(M::new_from([("b", SetUnionHashSet::new_from([0, 1]))], HashSet::from_iter([])));
450        };
451
452        check_all(&test_vec);
453    }
454
455    /// Check that a key with a value of bottom is the same as an empty map, etc.
456    #[test]
457    fn test_collapses_bot() {
458        type K = &'static str;
459        type V = SetUnion<HashSet<i32>>;
460
461        type A = MapUnionWithTombstones<HashMap<K, V>, HashSet<K>>;
462        type B = MapUnionWithTombstones<SingletonMap<K, V>, HashSet<K>>;
463
464        let map_empty = A::default();
465
466        let map_a_bot = B::new(SingletonMap("a", Default::default()), Default::default());
467        let map_b_bot = B::new(SingletonMap("b", Default::default()), Default::default());
468
469        assert_eq!(map_empty, map_a_bot);
470        assert_eq!(map_empty, map_b_bot);
471        assert_eq!(map_a_bot, map_b_bot);
472    }
473
474    #[test]
475    fn roaring_u64_basic() {
476        let mut x = MapUnionWithTombstonesRoaring::new_from(
477            HashMap::from([
478                (1u64, SetUnionHashSet::new_from([10])),
479                (2, SetUnionHashSet::new_from([20])),
480            ]),
481            RoaringTombstoneSet::new(),
482        );
483        let mut y = MapUnionWithTombstonesRoaring::new_from(
484            HashMap::from([
485                (2u64, SetUnionHashSet::new_from([21])),
486                (3, SetUnionHashSet::new_from([30])),
487            ]),
488            RoaringTombstoneSet::new(),
489        );
490
491        // Add tombstone for key 2
492        y.as_reveal_mut().1.insert(2);
493
494        x.merge(y);
495
496        // Should have keys 1 and 3, but not 2 (tombstoned)
497        assert!(!x.as_reveal_ref().0.contains_key(&2));
498        assert!(x.as_reveal_ref().0.contains_key(&1));
499        assert!(x.as_reveal_ref().0.contains_key(&3));
500        assert!(x.as_reveal_ref().1.contains(&2));
501    }
502
503    #[test]
504    fn fst_string_basic() {
505        let mut x = MapUnionWithTombstonesFstString::new_from(
506            HashMap::from([
507                ("apple".to_owned(), SetUnionHashSet::new_from([1])),
508                ("banana".to_owned(), SetUnionHashSet::new_from([2])),
509            ]),
510            FstTombstoneSet::new(),
511        );
512        let mut y = MapUnionWithTombstonesFstString::new_from(
513            HashMap::from([
514                ("banana".to_owned(), SetUnionHashSet::new_from([3])),
515                ("cherry".to_owned(), SetUnionHashSet::new_from([4])),
516            ]),
517            FstTombstoneSet::new(),
518        );
519
520        // Add tombstone for "banana"
521        y.as_reveal_mut().1.extend(["banana".to_owned()]);
522
523        x.merge(y);
524
525        // Should have "apple" and "cherry", but not "banana" (tombstoned)
526        assert!(!x.as_reveal_ref().0.contains_key("banana"));
527        assert!(x.as_reveal_ref().0.contains_key("apple"));
528        assert!(x.as_reveal_ref().0.contains_key("cherry"));
529        assert!(x.as_reveal_ref().1.contains(b"banana"));
530    }
531
532    #[test]
533    fn roaring_merge_efficiency() {
534        // Test that merging roaring bitmaps works correctly
535        let mut x = MapUnionWithTombstonesRoaring::new_from(
536            HashMap::from([
537                (1u64, SetUnionHashSet::new_from([1])),
538                (2, SetUnionHashSet::new_from([2])),
539            ]),
540            RoaringTombstoneSet::from_iter([10u64, 20]),
541        );
542
543        let y = MapUnionWithTombstonesRoaring::new_from(
544            HashMap::from([(3u64, SetUnionHashSet::new_from([3]))]),
545            RoaringTombstoneSet::from_iter([30u64, 2]),
546        );
547
548        x.merge(y);
549
550        // Should have all tombstones
551        assert!(x.as_reveal_ref().1.contains(&10));
552        assert!(x.as_reveal_ref().1.contains(&20));
553        assert!(x.as_reveal_ref().1.contains(&30));
554        assert!(x.as_reveal_ref().1.contains(&2));
555
556        // Should not have key 2 in the map
557        assert!(!x.as_reveal_ref().0.contains_key(&2));
558
559        // Should have keys 1 and 3
560        assert!(x.as_reveal_ref().0.contains_key(&1));
561        assert!(x.as_reveal_ref().0.contains_key(&3));
562    }
563}