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

egglog_numeric_id/
lib.rs
