egglog_numeric_id/
lib.rs

1//! A crate with utilities for working with numeric Ids.
2use std::{
3    fmt::{self, Debug},
4    hash::Hash,
5    marker::PhantomData,
6    ops,
7};
8
9#[cfg(test)]
10mod tests;
11
12/// A trait describing "newtypes" that wrap an integer.
13pub trait NumericId: Copy + Clone + PartialEq + Eq + PartialOrd + Ord + Hash + Send + Sync {
14    type Rep;
15    type Atomic;
16    fn new(val: Self::Rep) -> Self;
17    fn from_usize(index: usize) -> Self;
18    fn index(self) -> usize;
19    fn rep(self) -> Self::Rep;
20    fn inc(self) -> Self {
21        Self::from_usize(self.index() + 1)
22    }
23}
24
25impl NumericId for usize {
26    type Rep = usize;
27    type Atomic = std::sync::atomic::AtomicUsize;
28    fn new(val: usize) -> Self {
29        val
30    }
31    fn from_usize(index: usize) -> Self {
32        index
33    }
34
35    fn rep(self) -> usize {
36        self
37    }
38
39    fn index(self) -> usize {
40        self
41    }
42}
43
44/// A mapping from a [`NumericId`] to some value.
45///
46/// This mapping is _dense_: it stores a flat array indexed by `K::index()`,
47/// with no hashing. For sparse mappings, use a HashMap.
48#[derive(Clone, PartialEq, Eq, Hash)]
49pub struct DenseIdMap<K, V> {
50    data: Vec<Option<V>>,
51    _marker: PhantomData<K>,
52}
53
54impl<K: NumericId + Debug, V: Debug> Debug for DenseIdMap<K, V> {
55    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
56        let mut map = f.debug_map();
57        for (k, v) in self.iter() {
58            map.entry(&k, v);
59        }
60        map.finish()
61    }
62}
63
64impl<K, V> Default for DenseIdMap<K, V> {
65    fn default() -> Self {
66        Self {
67            data: Vec::new(),
68            _marker: PhantomData,
69        }
70    }
71}
72
73impl<K: NumericId, V> DenseIdMap<K, V> {
74    /// Create an empty map with space for `n` entries pre-allocated.
75    pub fn with_capacity(n: usize) -> Self {
76        let mut res = Self::new();
77        res.reserve_space(K::from_usize(n.saturating_sub(1)));
78        res
79    }
80
81    /// Create an empty map.
82    pub fn new() -> Self {
83        Self::default()
84    }
85
86    /// Clear the table's contents.
87    pub fn clear(&mut self) {
88        self.data.clear();
89    }
90
91    /// Get the current capacity for the table.
92    pub fn capacity(&self) -> usize {
93        self.data.capacity()
94    }
95
96    /// Get the number of ids currently indexed by the table (including "null"
97    /// entries). This is a less useful version of "length" in other containers.
98    pub fn n_ids(&self) -> usize {
99        self.data.len()
100    }
101
102    /// Insert the given mapping into the table.
103    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
104        self.reserve_space(key);
105        self.data[key.index()].replace(value)
106    }
107
108    /// Get the key that would be returned by the next call to [`DenseIdMap::push`].
109    pub fn next_id(&self) -> K {
110        K::from_usize(self.data.len())
111    }
112
113    /// Add the given mapping to the table, returning the key corresponding to
114    /// [`DenseIdMap::n_ids`].
115    pub fn push(&mut self, val: V) -> K {
116        let res = self.next_id();
117        self.data.push(Some(val));
118        res
119    }
120
121    /// Test whether `key` is set in this map.
122    pub fn contains_key(&self, key: K) -> bool {
123        self.data.get(key.index()).is_some_and(Option::is_some)
124    }
125
126    /// Get the current mapping for `key` in the table.
127    pub fn get(&self, key: K) -> Option<&V> {
128        self.data.get(key.index())?.as_ref()
129    }
130
131    /// Get a mutable reference to the current mapping for `key` in the table.
132    pub fn get_mut(&mut self, key: K) -> Option<&mut V> {
133        self.reserve_space(key);
134        self.data.get_mut(key.index())?.as_mut()
135    }
136
137    /// Extract the value mapped to by `key` from the table.
138    ///
139    /// # Panics
140    /// This method panics if `key` is not in the table.
141    pub fn unwrap_val(&mut self, key: K) -> V {
142        self.reserve_space(key);
143        self.data.get_mut(key.index()).unwrap().take().unwrap()
144    }
145
146    /// Extract the value mapped to by `key` from the table, if it is present.
147    pub fn take(&mut self, key: K) -> Option<V> {
148        self.reserve_space(key);
149        self.data.get_mut(key.index()).unwrap().take()
150    }
151
152    /// Get the current mapping for `key` in the table, or insert the value
153    /// returned by `f` and return a mutable reference to it.
154    pub fn get_or_insert(&mut self, key: K, f: impl FnOnce() -> V) -> &mut V {
155        self.reserve_space(key);
156        self.data[key.index()].get_or_insert_with(f)
157    }
158
159    pub fn raw(&self) -> &[Option<V>] {
160        &self.data
161    }
162
163    pub fn iter(&self) -> impl Iterator<Item = (K, &V)> {
164        self.data
165            .iter()
166            .enumerate()
167            .filter_map(|(i, v)| Some((K::from_usize(i), v.as_ref()?)))
168    }
169
170    pub fn iter_mut(&mut self) -> impl Iterator<Item = (K, &mut V)> {
171        self.data
172            .iter_mut()
173            .enumerate()
174            .filter_map(|(i, v)| Some((K::from_usize(i), v.as_mut()?)))
175    }
176
177    #[allow(clippy::should_implement_trait)]
178    pub fn into_iter(self) -> impl Iterator<Item = (K, V)> {
179        self.data
180            .into_iter()
181            .enumerate()
182            .filter_map(|(i, v)| Some((K::from_usize(i), v?)))
183    }
184
185    /// Reserve space up to the given key in the table.
186    pub fn reserve_space(&mut self, key: K) {
187        let index = key.index();
188        if index >= self.data.len() {
189            self.data.resize_with(index + 1, || None);
190        }
191    }
192
193    pub fn drain(&mut self) -> impl Iterator<Item = (K, V)> + '_ {
194        // To avoid the need to write down the return type.
195        self.data
196            .drain(..)
197            .enumerate()
198            .filter_map(|(i, v)| Some((K::from_usize(i), v?)))
199    }
200
201    pub fn retain(&mut self, mut f: impl FnMut(K, &V) -> bool) {
202        for (i, v) in self.data.iter_mut().enumerate() {
203            if let Some(inner) = v
204                && !f(K::from_usize(i), inner)
205            {
206                *v = None;
207            }
208        }
209    }
210
211    pub fn len(&self) -> usize {
212        self.data.iter().filter(|v| v.is_some()).count()
213    }
214
215    pub fn is_empty(&self) -> bool {
216        self.data.iter().all(|v| v.is_none())
217    }
218}
219
220impl<K: NumericId, V: Send + Sync> DenseIdMap<K, V> {
221    /// Get a parallel iterator over the entries in the table.
222    pub fn par_iter(&self) -> impl ParallelIterator<Item = (K, &V)> {
223        self.data
224            .par_iter()
225            .enumerate()
226            .filter_map(|(i, v)| Some((K::from_usize(i), v.as_ref()?)))
227    }
228
229    /// Get a parallel iterator over mutable references to the entries in the table.
230    pub fn par_iter_mut(&mut self) -> impl ParallelIterator<Item = (K, &mut V)> {
231        self.data
232            .par_iter_mut()
233            .enumerate()
234            .filter_map(|(i, v)| Some((K::from_usize(i), v.as_mut()?)))
235    }
236}
237
238impl<K: NumericId, V> ops::Index<K> for DenseIdMap<K, V> {
239    type Output = V;
240
241    fn index(&self, key: K) -> &Self::Output {
242        self.get(key).unwrap()
243    }
244}
245
246impl<K: NumericId, V> ops::IndexMut<K> for DenseIdMap<K, V> {
247    fn index_mut(&mut self, key: K) -> &mut Self::Output {
248        self.get_mut(key).unwrap()
249    }
250}
251
252impl<K: NumericId, V: Default> DenseIdMap<K, V> {
253    pub fn get_or_default(&mut self, key: K) -> &mut V {
254        self.get_or_insert(key, V::default)
255    }
256}
257
258impl<K: NumericId, V: Clone> FromIterator<(K, V)> for DenseIdMap<K, V> {
259    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
260        let mut res = DenseIdMap::new();
261        for (k, v) in iter {
262            res.insert(k, v);
263        }
264        res
265    }
266}
267
268#[derive(Debug)]
269pub struct IdVec<K, V> {
270    data: Vec<V>,
271    _marker: std::marker::PhantomData<K>,
272}
273
274impl<K, V> IdVec<K, V> {
275    pub fn clear(&mut self) {
276        self.data.clear();
277    }
278    pub fn len(&self) -> usize {
279        self.data.len()
280    }
281    pub fn capacity(&self) -> usize {
282        self.data.capacity()
283    }
284}
285
286impl<K, V> Default for IdVec<K, V> {
287    fn default() -> IdVec<K, V> {
288        IdVec {
289            data: Default::default(),
290            _marker: std::marker::PhantomData,
291        }
292    }
293}
294
295impl<K, V: Clone> Clone for IdVec<K, V> {
296    fn clone(&self) -> Self {
297        IdVec {
298            data: self.data.clone(),
299            _marker: std::marker::PhantomData,
300        }
301    }
302}
303
304/// Like a [`DenseIdMap`], but supports freeing (and reusing) slots.
305#[derive(Clone)]
306pub struct DenseIdMapWithReuse<K, V> {
307    data: DenseIdMap<K, V>,
308    free: Vec<K>,
309}
310
311impl<K, V> Default for DenseIdMapWithReuse<K, V> {
312    fn default() -> Self {
313        Self {
314            data: Default::default(),
315            free: Default::default(),
316        }
317    }
318}
319
320impl<K: NumericId, V> DenseIdMapWithReuse<K, V> {
321    /// Reserve a slot in the map for use later with [`DenseIdMapWithReuse::insert`].
322    pub fn reserve_slot(&mut self) -> K {
323        match self.free.pop() {
324            Some(res) => res,
325            None => {
326                let res = self.data.next_id();
327                self.data.reserve_space(res);
328                res
329            }
330        }
331    }
332
333    /// Insert the given mapping into the table. You probably
334    /// want to use [`DenseIdMapWithReuse::push`] instead, unless you need to use
335    /// the key to build the value, in which case you can
336    /// use [`DenseIdMapWithReuse::reserve_slot`] to get the key for this method.
337    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
338        self.data.insert(key, value)
339    }
340
341    /// Add the given value to the table.
342    pub fn push(&mut self, value: V) -> K {
343        let res = self.reserve_slot();
344        self.insert(res, value);
345        res
346    }
347
348    /// Remove the given key from the table, if it is present.
349    pub fn take(&mut self, id: K) -> Option<V> {
350        let res = self.data.take(id);
351        if res.is_some() {
352            self.free.push(id);
353        }
354        res
355    }
356}
357
358impl<K: NumericId, V> std::ops::Index<K> for DenseIdMapWithReuse<K, V> {
359    type Output = V;
360    fn index(&self, key: K) -> &V {
361        &self.data[key]
362    }
363}
364
365impl<K: NumericId, V> std::ops::IndexMut<K> for DenseIdMapWithReuse<K, V> {
366    fn index_mut(&mut self, key: K) -> &mut V {
367        &mut self.data[key]
368    }
369}
370
371impl<K: NumericId, V> IdVec<K, V> {
372    pub fn with_capacity(cap: usize) -> IdVec<K, V> {
373        IdVec {
374            data: Vec::with_capacity(cap),
375            _marker: std::marker::PhantomData,
376        }
377    }
378
379    pub fn push(&mut self, elt: V) -> K {
380        let res = K::from_usize(self.data.len());
381        self.data.push(elt);
382        res
383    }
384
385    pub fn resize_with(&mut self, size: usize, init: impl FnMut() -> V) {
386        self.data.resize_with(size, init)
387    }
388
389    pub fn is_empty(&self) -> bool {
390        self.data.is_empty()
391    }
392
393    pub fn values(&self) -> impl Iterator<Item = &V> {
394        self.data.iter()
395    }
396
397    pub fn iter(&self) -> impl Iterator<Item = (K, &V)> {
398        self.data
399            .iter()
400            .enumerate()
401            .map(|(i, v)| (K::from_usize(i), v))
402    }
403    pub fn iter_mut(&mut self) -> impl Iterator<Item = (K, &mut V)> {
404        self.data
405            .iter_mut()
406            .enumerate()
407            .map(|(i, v)| (K::from_usize(i), v))
408    }
409    pub fn drain(&mut self) -> impl Iterator<Item = (K, V)> + '_ {
410        self.data
411            .drain(..)
412            .enumerate()
413            .map(|(i, v)| (K::from_usize(i), v))
414    }
415    pub fn get(&self, key: K) -> Option<&V> {
416        self.data.get(key.index())
417    }
418}
419
420impl<K: NumericId, V: Send + Sync> IdVec<K, V> {
421    pub fn par_iter_mut(&mut self) -> impl IndexedParallelIterator<Item = (K, &mut V)> {
422        self.data
423            .par_iter_mut()
424            .with_max_len(1)
425            .enumerate()
426            .map(|(i, v)| (K::from_usize(i), v))
427    }
428}
429
430impl<K: NumericId, V> ops::Index<K> for IdVec<K, V> {
431    type Output = V;
432
433    fn index(&self, key: K) -> &Self::Output {
434        &self.data[key.index()]
435    }
436}
437
438impl<K: NumericId, V> ops::IndexMut<K> for IdVec<K, V> {
439    fn index_mut(&mut self, key: K) -> &mut Self::Output {
440        &mut self.data[key.index()]
441    }
442}
443
444use rayon::iter::{
445    IndexedParallelIterator, IntoParallelRefIterator, IntoParallelRefMutIterator, ParallelIterator,
446};
447
448#[macro_export]
449#[doc(hidden)]
450macro_rules! atomic_of {
451    (usize) => {
452        std::sync::atomic::AtomicUsize
453    };
454    (u8) => {
455        std::sync::atomic::AtomicU8
456    };
457    (u16) => {
458        std::sync::atomic::AtomicU16
459    };
460    (u32) => {
461        std::sync::atomic::AtomicU32
462    };
463    (u64) => {
464        std::sync::atomic::AtomicU64
465    };
466}
467
468#[macro_export]
469macro_rules! define_id {
470    ($v:vis $name:ident, $repr:tt) => { define_id!($v $name, $repr, "", pretty ""); };
471    ($v:vis $name:ident, $repr:tt, $doc:tt) => { define_id!($v $name, $repr, $doc, pretty ""); };
472    ($v:vis $name:ident, $repr:tt, pretty $pretty_name:expr) => { define_id!($v $name, $repr, "", pretty $pretty_name); };
473    ($v:vis $name:ident, $repr:tt, $doc:tt, pretty $pretty_name:tt) => {
474        #[derive(Copy, Clone)]
475        #[doc = $doc]
476        $v struct $name {
477            rep: $repr,
478        }
479
480        impl PartialEq for $name {
481            fn eq(&self, other: &Self) -> bool {
482                self.rep == other.rep
483            }
484        }
485
486        impl Eq for $name {}
487
488        impl PartialOrd for $name {
489            fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
490                Some(self.cmp(other))
491            }
492        }
493
494        impl Ord for $name {
495            fn cmp(&self, other: &Self) -> std::cmp::Ordering {
496                self.rep.cmp(&other.rep)
497            }
498        }
499
500        impl std::hash::Hash for $name {
501            fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
502                self.rep.hash(state);
503            }
504        }
505
506        impl $name {
507            #[allow(unused)]
508            $v const fn new_const(id: $repr) -> Self {
509                $name {
510                    rep: id,
511                }
512            }
513
514            #[allow(unused)]
515            $v fn range(low: Self, high: Self) -> impl Iterator<Item = Self> {
516                use $crate::NumericId;
517                (low.rep..high.rep).map(|i| $name::new(i))
518            }
519
520        }
521
522        impl $crate::NumericId for $name {
523            type Rep = $repr;
524            type Atomic = $crate::atomic_of!($repr);
525            fn new(id: $repr) -> Self {
526                Self::new_const(id)
527            }
528            fn from_usize(index: usize) -> Self {
529                assert!(<$repr>::MAX as usize >= index,
530                    "overflowing id type {} (represented as {}) with index {}", stringify!($name), stringify!($repr), index);
531                $name::new(index as $repr)
532            }
533            /// return the inner representation of id as usize
534            fn index(self) -> usize {
535                self.rep as usize
536            }
537            /// return the inner representation of id.
538            fn rep(self) -> $repr {
539                self.rep
540            }
541        }
542
543        impl std::fmt::Debug for $name {
544            fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
545                let name = if $pretty_name.is_empty() {
546                    stringify!($name).to_string()
547                } else {
548                    $pretty_name.to_string()
549                };
550                write!(fmt, "{}({:?})", name, self.rep)
551            }
552        }
553    };
554}
egglog_numeric_id/lib.rs

egglog_numeric_id/
lib.rs
