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

miden_utils_indexing/
lib.rs
