miden_utils_indexing/
lib.rs

1//! Type-safe u32-indexed vector utilities for Miden
2//!
3//! This module provides utilities for working with u32-indexed vectors in a type-safe manner,
4//! including the `IndexVec` type and related functionality.
5#![no_std]
6
7extern crate alloc;
8
9#[doc = include_str!("../README.md")]
10use alloc::{collections::BTreeMap, vec, vec::Vec};
11use core::{fmt::Debug, marker::PhantomData, ops};
12
13#[cfg(feature = "serde")]
14use serde::{Deserialize, Serialize};
15use thiserror::Error;
16
17/// Error returned when too many items are added to an IndexedVec.
18#[derive(Debug, Clone, PartialEq, Eq, Error)]
19pub enum IndexedVecError {
20    /// The number of items exceeds the maximum supported by ID type.
21    #[error("IndexedVec contains maximum number of items")]
22    TooManyItems,
23}
24
25/// A trait for u32-backed, 0-based IDs.
26pub trait Idx: Copy + Eq + Ord + Debug + From<u32> + Into<u32> {
27    /// Convert from this ID type to usize.
28    #[inline]
29    fn to_usize(self) -> usize {
30        self.into() as usize
31    }
32}
33
34/// Macro to create a newtyped ID that implements Idx.
35#[macro_export]
36macro_rules! newtype_id {
37    ($name:ident) => {
38        #[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
39        #[repr(transparent)]
40        pub struct $name(u32);
41
42        impl core::fmt::Debug for $name {
43            fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
44                write!(f, "{}({})", stringify!($name), self.0)
45            }
46        }
47        impl From<u32> for $name {
48            fn from(v: u32) -> Self {
49                Self(v)
50            }
51        }
52        impl From<$name> for u32 {
53            fn from(v: $name) -> Self {
54                v.0
55            }
56        }
57        impl $crate::Idx for $name {}
58    };
59}
60
61/// A dense vector indexed by ID types.
62///
63/// This provides O(1) access and storage for dense ID-indexed data.
64#[derive(Clone, Debug, PartialEq, Eq)]
65#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
66pub struct IndexVec<I: Idx, T> {
67    raw: Vec<T>,
68    _m: PhantomData<I>,
69}
70
71impl<I: Idx, T> Default for IndexVec<I, T> {
72    fn default() -> Self {
73        Self { raw: Vec::new(), _m: PhantomData }
74    }
75}
76
77impl<I: Idx, T> IndexVec<I, T> {
78    /// Create a new empty IndexVec.
79    #[inline]
80    pub fn new() -> Self {
81        Self { raw: Vec::new(), _m: PhantomData }
82    }
83
84    /// Create a new IndexVec with pre-allocated capacity.
85    #[inline]
86    pub fn with_capacity(n: usize) -> Self {
87        Self {
88            raw: Vec::with_capacity(n),
89            _m: PhantomData,
90        }
91    }
92
93    /// Get the number of elements in the IndexVec.
94    #[inline]
95    pub fn len(&self) -> usize {
96        self.raw.len()
97    }
98
99    /// Check if the IndexVec is empty.
100    #[inline]
101    pub fn is_empty(&self) -> bool {
102        self.raw.is_empty()
103    }
104
105    /// Push an element and return its ID.
106    ///
107    /// Returns an error if the length would exceed the maximum representable by the ID type.
108    #[inline]
109    pub fn push(&mut self, v: T) -> Result<I, IndexedVecError> {
110        if self.raw.len() >= u32::MAX as usize {
111            return Err(IndexedVecError::TooManyItems);
112        }
113        let id = I::from(self.raw.len() as u32);
114        self.raw.push(v);
115        Ok(id)
116    }
117
118    /// Insert an element at the specified ID.
119    ///
120    /// This sets the value at the given index. It does **not** insert or shift elements.
121    /// If you need to append elements, use `push()` instead.
122    ///
123    /// # Panics
124    /// - If the ID is out of bounds.
125    #[inline]
126    pub(crate) fn insert_at(&mut self, idx: I, v: T) {
127        self.raw[idx.to_usize()] = v;
128    }
129
130    /// Get an element by ID, returning None if the ID is out of bounds.
131    #[inline]
132    pub fn get(&self, idx: I) -> Option<&T> {
133        self.raw.get(idx.to_usize())
134    }
135
136    /// Get a slice of all elements.
137    #[inline]
138    pub fn as_slice(&self) -> &[T] {
139        &self.raw
140    }
141
142    /// Consume this IndexVec and return the underlying Vec.
143    #[inline]
144    pub fn into_inner(self) -> Vec<T> {
145        self.raw
146    }
147
148    /// Remove an element at the specified index and return it.
149    pub fn swap_remove(&mut self, index: usize) -> T {
150        self.raw.swap_remove(index)
151    }
152
153    /// Check if this IndexVec contains a specific element.
154    pub fn contains(&self, item: &T) -> bool
155    where
156        T: PartialEq,
157    {
158        self.raw.contains(item)
159    }
160
161    /// Get an iterator over the elements in this IndexVec.
162    pub fn iter(&self) -> core::slice::Iter<'_, T> {
163        self.raw.iter()
164    }
165
166    /// Get a mutable iterator over the elements in this IndexVec.
167    pub fn iter_mut(&mut self) -> core::slice::IterMut<'_, T> {
168        self.raw.iter_mut()
169    }
170}
171
172impl<I: Idx, T> ops::Index<I> for IndexVec<I, T> {
173    type Output = T;
174    #[inline]
175    fn index(&self, index: I) -> &Self::Output {
176        &self.raw[index.to_usize()]
177    }
178}
179
180impl<I: Idx, T> ops::IndexMut<I> for IndexVec<I, T> {
181    #[inline]
182    fn index_mut(&mut self, index: I) -> &mut Self::Output {
183        &mut self.raw[index.to_usize()]
184    }
185}
186
187/// A dense mapping from ID to ID.
188///
189/// This is equivalent to `IndexVec<From, Option<To>>` and provides
190/// efficient dense ID remapping.
191#[derive(Clone)]
192pub struct DenseIdMap<From: Idx, To: Idx> {
193    inner: IndexVec<From, Option<To>>,
194}
195
196impl<From: Idx, To: Idx> DenseIdMap<From, To> {
197    /// Create a new dense ID mapping with the specified length.
198    #[inline]
199    pub fn with_len(length: usize) -> Self {
200        Self {
201            inner: IndexVec { raw: vec![None; length], _m: PhantomData },
202        }
203    }
204
205    /// Insert a mapping from source ID to target ID.
206    ///
207    /// # Panics
208    ///
209    /// Panics if the source ID is beyond the length of this DenseIdMap.
210    /// This DenseIdMap should be created with sufficient length to accommodate
211    /// all expected source IDs.
212    #[inline]
213    pub fn insert(&mut self, k: From, v: To) {
214        let idx = k.to_usize();
215        let len = self.len();
216
217        assert!(idx < len, "source ID {idx} exceeds DenseIdMap length {len}");
218        self.inner.insert_at(k, Some(v));
219    }
220
221    /// Get the target ID for the given source ID.
222    #[inline]
223    pub fn get(&self, k: From) -> Option<To> {
224        *self.inner.get(k)?
225    }
226
227    /// Get the number of source IDs in this mapping.
228    #[inline]
229    pub fn len(&self) -> usize {
230        self.inner.len()
231    }
232
233    /// Check if the mapping is empty.
234    #[inline]
235    pub fn is_empty(&self) -> bool {
236        self.inner.is_empty()
237    }
238}
239
240/// A trait for looking up values by ID.
241pub trait LookupByIdx<ID, V>
242where
243    ID: Idx,
244{
245    /// Get the value for the given ID.
246    fn get(&self, id: ID) -> Option<&V>;
247}
248
249/// A trait for looking up values by key that doesn't need to implement Idx.
250pub trait LookupByKey<K, V> {
251    /// Get the value for the given key.
252    fn get(&self, key: &K) -> Option<&V>;
253}
254
255impl<I, T> LookupByIdx<I, T> for IndexVec<I, T>
256where
257    I: Idx,
258{
259    fn get(&self, id: I) -> Option<&T> {
260        IndexVec::get(self, id)
261    }
262}
263
264impl<K, V> LookupByKey<K, V> for BTreeMap<K, V>
265where
266    K: Ord,
267{
268    fn get(&self, key: &K) -> Option<&V> {
269        BTreeMap::get(self, key)
270    }
271}
272
273impl<K, V> LookupByIdx<K, V> for BTreeMap<K, V>
274where
275    K: Idx,
276{
277    fn get(&self, id: K) -> Option<&V> {
278        BTreeMap::get(self, &id)
279    }
280}
281
282impl<I, T> LookupByIdx<I, T> for DenseIdMap<I, T>
283where
284    I: Idx,
285    T: Idx,
286{
287    fn get(&self, id: I) -> Option<&T> {
288        IndexVec::get(&self.inner, id).and_then(Option::as_ref)
289    }
290}
291
292impl<I: Idx, T> IntoIterator for IndexVec<I, T> {
293    type Item = T;
294    type IntoIter = alloc::vec::IntoIter<T>;
295
296    fn into_iter(self) -> Self::IntoIter {
297        self.raw.into_iter()
298    }
299}
300
301impl<'a, I: Idx, T> IntoIterator for &'a IndexVec<I, T> {
302    type Item = &'a T;
303    type IntoIter = core::slice::Iter<'a, T>;
304
305    fn into_iter(self) -> Self::IntoIter {
306        self.iter()
307    }
308}
309
310#[cfg(test)]
311mod tests {
312    use alloc::string::{String, ToString};
313
314    use super::*;
315
316    // Test ID types
317    newtype_id!(TestId);
318    newtype_id!(TestId2);
319
320    #[test]
321    fn test_indexvec_basic() {
322        let mut vec = IndexVec::<TestId, String>::new();
323        let id1 = vec.push("hello".to_string()).unwrap();
324        let id2 = vec.push("world".to_string()).unwrap();
325
326        assert_eq!(vec.len(), 2);
327        assert_eq!(&vec[id1], "hello");
328        assert_eq!(&vec[id2], "world");
329        assert_eq!(vec.get(TestId::from(0)), Some(&"hello".to_string()));
330        assert_eq!(vec.get(TestId::from(2)), None);
331    }
332
333    #[test]
334    fn test_dense_id_map() {
335        let mut map = DenseIdMap::<TestId, TestId2>::with_len(2);
336        map.insert(TestId::from(0), TestId2::from(10));
337        map.insert(TestId::from(1), TestId2::from(11));
338
339        assert_eq!(map.len(), 2);
340        assert_eq!(map.get(TestId::from(0)), Some(TestId2::from(10)));
341        assert_eq!(map.get(TestId::from(1)), Some(TestId2::from(11)));
342        assert_eq!(map.get(TestId::from(2)), None);
343    }
344}
miden_utils_indexing/lib.rs

miden_utils_indexing/
lib.rs
