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

miden_utils_indexing/
lib.rs
