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 raw_mut(&mut self) -> &mut [Option<V>] {
164        &mut self.data
165    }
166
167    pub fn iter(&self) -> impl Iterator<Item = (K, &V)> {
168        self.data
169            .iter()
170            .enumerate()
171            .filter_map(|(i, v)| Some((K::from_usize(i), v.as_ref()?)))
172    }
173
174    pub fn iter_mut(&mut self) -> impl Iterator<Item = (K, &mut V)> {
175        self.data
176            .iter_mut()
177            .enumerate()
178            .filter_map(|(i, v)| Some((K::from_usize(i), v.as_mut()?)))
179    }
180
181    #[allow(clippy::should_implement_trait)]
182    pub fn into_iter(self) -> impl Iterator<Item = (K, V)> {
183        self.data
184            .into_iter()
185            .enumerate()
186            .filter_map(|(i, v)| Some((K::from_usize(i), v?)))
187    }
188
189    /// Reserve space up to the given key in the table.
190    pub fn reserve_space(&mut self, key: K) {
191        let index = key.index();
192        if index >= self.data.len() {
193            self.data.resize_with(index + 1, || None);
194        }
195    }
196
197    pub fn drain(&mut self) -> impl Iterator<Item = (K, V)> + '_ {
198        // To avoid the need to write down the return type.
199        self.data
200            .drain(..)
201            .enumerate()
202            .filter_map(|(i, v)| Some((K::from_usize(i), v?)))
203    }
204
205    pub fn retain(&mut self, mut f: impl FnMut(K, &V) -> bool) {
206        for (i, v) in self.data.iter_mut().enumerate() {
207            if let Some(inner) = v
208                && !f(K::from_usize(i), inner)
209            {
210                *v = None;
211            }
212        }
213    }
214
215    pub fn len(&self) -> usize {
216        self.data.iter().filter(|v| v.is_some()).count()
217    }
218
219    pub fn is_empty(&self) -> bool {
220        self.data.iter().all(|v| v.is_none())
221    }
222}
223
224impl<K: NumericId, V> ops::Index<K> for DenseIdMap<K, V> {
225    type Output = V;
226
227    fn index(&self, key: K) -> &Self::Output {
228        self.get(key).unwrap()
229    }
230}
231
232impl<K: NumericId, V> ops::IndexMut<K> for DenseIdMap<K, V> {
233    fn index_mut(&mut self, key: K) -> &mut Self::Output {
234        self.get_mut(key).unwrap()
235    }
236}
237
238impl<K: NumericId, V: Default> DenseIdMap<K, V> {
239    pub fn get_or_default(&mut self, key: K) -> &mut V {
240        self.get_or_insert(key, V::default)
241    }
242}
243
244impl<K: NumericId, V: Clone> FromIterator<(K, V)> for DenseIdMap<K, V> {
245    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
246        let mut res = DenseIdMap::new();
247        for (k, v) in iter {
248            res.insert(k, v);
249        }
250        res
251    }
252}
253
254#[derive(Debug)]
255pub struct IdVec<K, V> {
256    data: Vec<V>,
257    _marker: std::marker::PhantomData<K>,
258}
259
260impl<K, V> IdVec<K, V> {
261    pub fn clear(&mut self) {
262        self.data.clear();
263    }
264    pub fn len(&self) -> usize {
265        self.data.len()
266    }
267    pub fn capacity(&self) -> usize {
268        self.data.capacity()
269    }
270}
271
272impl<K, V> Default for IdVec<K, V> {
273    fn default() -> IdVec<K, V> {
274        IdVec {
275            data: Default::default(),
276            _marker: std::marker::PhantomData,
277        }
278    }
279}
280
281impl<K, V: Clone> Clone for IdVec<K, V> {
282    fn clone(&self) -> Self {
283        IdVec {
284            data: self.data.clone(),
285            _marker: std::marker::PhantomData,
286        }
287    }
288}
289
290/// Like a [`DenseIdMap`], but supports freeing (and reusing) slots.
291#[derive(Clone)]
292pub struct DenseIdMapWithReuse<K, V> {
293    data: DenseIdMap<K, V>,
294    free: Vec<K>,
295}
296
297impl<K, V> Default for DenseIdMapWithReuse<K, V> {
298    fn default() -> Self {
299        Self {
300            data: Default::default(),
301            free: Default::default(),
302        }
303    }
304}
305
306impl<K: NumericId, V> DenseIdMapWithReuse<K, V> {
307    /// Reserve a slot in the map for use later with [`DenseIdMapWithReuse::insert`].
308    pub fn reserve_slot(&mut self) -> K {
309        match self.free.pop() {
310            Some(res) => res,
311            None => {
312                let res = self.data.next_id();
313                self.data.reserve_space(res);
314                res
315            }
316        }
317    }
318
319    /// Insert the given mapping into the table. You probably
320    /// want to use [`DenseIdMapWithReuse::push`] instead, unless you need to use
321    /// the key to build the value, in which case you can
322    /// use [`DenseIdMapWithReuse::reserve_slot`] to get the key for this method.
323    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
324        self.data.insert(key, value)
325    }
326
327    /// Add the given value to the table.
328    pub fn push(&mut self, value: V) -> K {
329        let res = self.reserve_slot();
330        self.insert(res, value);
331        res
332    }
333
334    /// Remove the given key from the table, if it is present.
335    pub fn take(&mut self, id: K) -> Option<V> {
336        let res = self.data.take(id);
337        if res.is_some() {
338            self.free.push(id);
339        }
340        res
341    }
342}
343
344impl<K: NumericId, V> std::ops::Index<K> for DenseIdMapWithReuse<K, V> {
345    type Output = V;
346    fn index(&self, key: K) -> &V {
347        &self.data[key]
348    }
349}
350
351impl<K: NumericId, V> std::ops::IndexMut<K> for DenseIdMapWithReuse<K, V> {
352    fn index_mut(&mut self, key: K) -> &mut V {
353        &mut self.data[key]
354    }
355}
356
357impl<K: NumericId, V> IdVec<K, V> {
358    pub fn with_capacity(cap: usize) -> IdVec<K, V> {
359        IdVec {
360            data: Vec::with_capacity(cap),
361            _marker: std::marker::PhantomData,
362        }
363    }
364
365    pub fn push(&mut self, elt: V) -> K {
366        let res = K::from_usize(self.data.len());
367        self.data.push(elt);
368        res
369    }
370
371    pub fn resize_with(&mut self, size: usize, init: impl FnMut() -> V) {
372        self.data.resize_with(size, init)
373    }
374
375    pub fn is_empty(&self) -> bool {
376        self.data.is_empty()
377    }
378
379    pub fn values(&self) -> impl Iterator<Item = &V> {
380        self.data.iter()
381    }
382
383    pub fn iter(&self) -> impl Iterator<Item = (K, &V)> {
384        self.data
385            .iter()
386            .enumerate()
387            .map(|(i, v)| (K::from_usize(i), v))
388    }
389    pub fn iter_mut(&mut self) -> impl Iterator<Item = (K, &mut V)> {
390        self.data
391            .iter_mut()
392            .enumerate()
393            .map(|(i, v)| (K::from_usize(i), v))
394    }
395    pub fn drain(&mut self) -> impl Iterator<Item = (K, V)> + '_ {
396        self.data
397            .drain(..)
398            .enumerate()
399            .map(|(i, v)| (K::from_usize(i), v))
400    }
401    pub fn get(&self, key: K) -> Option<&V> {
402        self.data.get(key.index())
403    }
404
405    pub fn as_mut_slice(&mut self) -> &mut [V] {
406        &mut self.data
407    }
408}
409
410impl<K: NumericId, V> ops::Index<K> for IdVec<K, V> {
411    type Output = V;
412
413    fn index(&self, key: K) -> &Self::Output {
414        &self.data[key.index()]
415    }
416}
417
418impl<K: NumericId, V> ops::IndexMut<K> for IdVec<K, V> {
419    fn index_mut(&mut self, key: K) -> &mut Self::Output {
420        &mut self.data[key.index()]
421    }
422}
423
424#[macro_export]
425#[doc(hidden)]
426macro_rules! atomic_of {
427    (usize) => {
428        std::sync::atomic::AtomicUsize
429    };
430    (u8) => {
431        std::sync::atomic::AtomicU8
432    };
433    (u16) => {
434        std::sync::atomic::AtomicU16
435    };
436    (u32) => {
437        std::sync::atomic::AtomicU32
438    };
439    (u64) => {
440        std::sync::atomic::AtomicU64
441    };
442}
443
444#[macro_export]
445macro_rules! define_id {
446    ($v:vis $name:ident, $repr:tt) => { define_id!($v $name, $repr, "", pretty ""); };
447    ($v:vis $name:ident, $repr:tt, $doc:tt) => { define_id!($v $name, $repr, $doc, pretty ""); };
448    ($v:vis $name:ident, $repr:tt, pretty $pretty_name:expr) => { define_id!($v $name, $repr, "", pretty $pretty_name); };
449    ($v:vis $name:ident, $repr:tt, $doc:tt, pretty $pretty_name:tt) => {
450        #[derive(Copy, Clone)]
451        #[doc = $doc]
452        $v struct $name {
453            rep: $repr,
454        }
455
456        impl PartialEq for $name {
457            fn eq(&self, other: &Self) -> bool {
458                self.rep == other.rep
459            }
460        }
461
462        impl Eq for $name {}
463
464        impl PartialOrd for $name {
465            fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
466                Some(self.cmp(other))
467            }
468        }
469
470        impl Ord for $name {
471            fn cmp(&self, other: &Self) -> std::cmp::Ordering {
472                self.rep.cmp(&other.rep)
473            }
474        }
475
476        impl std::hash::Hash for $name {
477            fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
478                self.rep.hash(state);
479            }
480        }
481
482        impl $name {
483            #[allow(unused)]
484            $v const fn new_const(id: $repr) -> Self {
485                $name {
486                    rep: id,
487                }
488            }
489
490            #[allow(unused)]
491            $v fn range(low: Self, high: Self) -> impl Iterator<Item = Self> {
492                use $crate::NumericId;
493                (low.rep..high.rep).map(|i| $name::new(i))
494            }
495
496        }
497
498        impl $crate::NumericId for $name {
499            type Rep = $repr;
500            type Atomic = $crate::atomic_of!($repr);
501            fn new(id: $repr) -> Self {
502                Self::new_const(id)
503            }
504            fn from_usize(index: usize) -> Self {
505                assert!(<$repr>::MAX as usize >= index,
506                    "overflowing id type {} (represented as {}) with index {}", stringify!($name), stringify!($repr), index);
507                $name::new(index as $repr)
508            }
509            /// return the inner representation of id as usize
510            fn index(self) -> usize {
511                self.rep as usize
512            }
513            /// return the inner representation of id.
514            fn rep(self) -> $repr {
515                self.rep
516            }
517        }
518
519        impl std::fmt::Debug for $name {
520            fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
521                let name = if $pretty_name.is_empty() {
522                    stringify!($name).to_string()
523                } else {
524                    $pretty_name.to_string()
525                };
526                write!(fmt, "{}({:?})", name, self.rep)
527            }
528        }
529    };
530}
egglog_numeric_id/lib.rs

egglog_numeric_id/
lib.rs
