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

miden_utils_indexing/
lib.rs
