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, mem::size_of, ops};
13
14pub use csr::{CsrMatrix, CsrValidationError};
15#[doc(hidden)]
16pub use miden_serde_utils;
17#[cfg(feature = "arbitrary")]
18use proptest::prelude::*;
19#[cfg(feature = "serde")]
20use serde::{Deserialize, Serialize};
21use thiserror::Error;
22
23/// Error returned when too many items are added to an IndexedVec.
24#[derive(Debug, Clone, PartialEq, Eq, Error)]
25pub enum IndexedVecError {
26    /// The number of items exceeds the maximum supported by ID type.
27    #[error("IndexedVec contains maximum number of items")]
28    TooManyItems,
29}
30
31#[cfg(feature = "arbitrary")]
32impl Arbitrary for IndexedVecError {
33    type Parameters = ();
34    type Strategy = BoxedStrategy<Self>;
35
36    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
37        Just(Self::TooManyItems).boxed()
38    }
39}
40
41/// A trait for u32-backed, 0-based IDs.
42pub trait Idx: Copy + Eq + Ord + Debug + From<u32> + Into<u32> {
43    /// Convert from this ID type to usize.
44    #[inline]
45    fn to_usize(self) -> usize {
46        self.into() as usize
47    }
48}
49
50/// Macro to create a newtyped ID that implements Idx.
51#[macro_export]
52macro_rules! newtype_id {
53    (
54        $(#[$a:meta])*
55        $vis:vis struct $name:ident;
56    ) => {
57        $(#[$a])*
58        #[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
59        #[repr(transparent)]
60        $vis struct $name(u32);
61
62        impl core::fmt::Debug for $name {
63            fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
64                write!(f, "{}({})", stringify!($name), self.0)
65            }
66        }
67        impl core::fmt::Display for $name {
68            fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
69                core::fmt::Display::fmt(&self.0, f)
70            }
71        }
72        impl From<u32> for $name {
73            fn from(v: u32) -> Self {
74                Self(v)
75            }
76        }
77        impl From<$name> for u32 {
78            fn from(v: $name) -> Self {
79                v.0
80            }
81        }
82        impl $crate::Idx for $name {}
83
84        impl $crate::miden_serde_utils::Serializable for $name {
85            fn write_into<W: $crate::miden_serde_utils::ByteWriter>(&self, target: &mut W) {
86                target.write_u32(self.0);
87            }
88        }
89
90        impl $crate::miden_serde_utils::Deserializable for $name {
91            fn read_from<R: $crate::miden_serde_utils::ByteReader>(source: &mut R) -> Result<Self, $crate::miden_serde_utils::DeserializationError> {
92                Ok(Self(source.read_u32()?))
93            }
94
95            fn min_serialized_size() -> usize {
96                4
97            }
98        }
99    };
100
101    ($name:ident) => {
102        $crate::newtype_id!(pub struct $name;);
103    };
104}
105
106#[cfg(test)]
107#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
108#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
109#[repr(transparent)]
110pub struct SerdeTestId(u32);
111
112#[cfg(test)]
113impl From<u32> for SerdeTestId {
114    fn from(v: u32) -> Self {
115        Self(v)
116    }
117}
118
119#[cfg(test)]
120impl From<SerdeTestId> for u32 {
121    fn from(v: SerdeTestId) -> Self {
122        v.0
123    }
124}
125
126#[cfg(test)]
127impl Idx for SerdeTestId {}
128
129/// A dense vector indexed by ID types.
130///
131/// This provides O(1) access and storage for dense ID-indexed data.
132#[derive(Clone, Debug, PartialEq, Eq)]
133#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
134#[cfg_attr(
135    all(feature = "arbitrary", test),
136    miden_test_serde_macros::serde_test(binary_serde(true), types(SerdeTestId, u32))
137)]
138pub struct IndexVec<I: Idx, T> {
139    raw: Vec<T>,
140    _m: PhantomData<I>,
141}
142
143#[cfg(feature = "arbitrary")]
144impl<I, T> Arbitrary for IndexVec<I, T>
145where
146    I: Idx + 'static,
147    T: Arbitrary + 'static,
148    T::Strategy: 'static,
149{
150    type Parameters = T::Parameters;
151    type Strategy = BoxedStrategy<Self>;
152
153    fn arbitrary_with(args: Self::Parameters) -> Self::Strategy {
154        proptest::collection::vec(any_with::<T>(args), 0..32)
155            .prop_map(|raw| Self::try_from(raw).expect("generated vector length fits in u32"))
156            .boxed()
157    }
158}
159
160impl<I: Idx, T> Default for IndexVec<I, T> {
161    fn default() -> Self {
162        Self { raw: Vec::new(), _m: PhantomData }
163    }
164}
165
166impl<I: Idx, T> IndexVec<I, T> {
167    /// Create a new empty IndexVec.
168    #[inline]
169    pub fn new() -> Self {
170        Self { raw: Vec::new(), _m: PhantomData }
171    }
172
173    /// Create a new IndexVec with pre-allocated capacity.
174    #[inline]
175    pub fn with_capacity(n: usize) -> Self {
176        Self {
177            raw: Vec::with_capacity(n),
178            _m: PhantomData,
179        }
180    }
181
182    /// Get the number of elements in the IndexVec.
183    #[inline]
184    pub fn len(&self) -> usize {
185        self.raw.len()
186    }
187
188    /// Check if the IndexVec is empty.
189    #[inline]
190    pub fn is_empty(&self) -> bool {
191        self.raw.is_empty()
192    }
193
194    /// Push an element and return its ID.
195    ///
196    /// Returns an error if the length would exceed the maximum representable by the ID type.
197    #[inline]
198    pub fn push(&mut self, v: T) -> Result<I, IndexedVecError> {
199        if self.raw.len() >= u32::MAX as usize {
200            return Err(IndexedVecError::TooManyItems);
201        }
202        let id = I::from(self.raw.len() as u32);
203        self.raw.push(v);
204        Ok(id)
205    }
206
207    /// Insert an element at the specified ID.
208    ///
209    /// This sets the value at the given index. It does **not** insert or shift elements.
210    /// If you need to append elements, use `push()` instead.
211    ///
212    /// # Panics
213    /// - If the ID is out of bounds.
214    #[inline]
215    pub(crate) fn insert_at(&mut self, idx: I, v: T) {
216        self.raw[idx.to_usize()] = v;
217    }
218
219    /// Get an element by ID, returning None if the ID is out of bounds.
220    #[inline]
221    pub fn get(&self, idx: I) -> Option<&T> {
222        self.raw.get(idx.to_usize())
223    }
224
225    /// Get a slice of all elements.
226    #[inline]
227    pub fn as_slice(&self) -> &[T] {
228        &self.raw
229    }
230
231    /// Consume this IndexVec and return the underlying Vec.
232    #[inline]
233    pub fn into_inner(self) -> Vec<T> {
234        self.raw
235    }
236
237    /// Remove an element at the specified index and return it.
238    pub fn swap_remove(&mut self, index: usize) -> T {
239        self.raw.swap_remove(index)
240    }
241
242    /// Shortens the vector, keeping the first `new_len` elements and dropping the rest
243    pub fn truncate(&mut self, new_len: usize) {
244        self.raw.truncate(new_len);
245    }
246
247    /// Check if this IndexVec contains a specific element.
248    pub fn contains(&self, item: &T) -> bool
249    where
250        T: PartialEq,
251    {
252        self.raw.contains(item)
253    }
254
255    /// Get an iterator over the elements in this IndexVec.
256    pub fn iter(&self) -> core::slice::Iter<'_, T> {
257        self.raw.iter()
258    }
259
260    /// Get a mutable iterator over the elements in this IndexVec.
261    pub fn iter_mut(&mut self) -> core::slice::IterMut<'_, T> {
262        self.raw.iter_mut()
263    }
264}
265
266impl<I: Idx, T> ops::Index<I> for IndexVec<I, T> {
267    type Output = T;
268    #[inline]
269    fn index(&self, index: I) -> &Self::Output {
270        &self.raw[index.to_usize()]
271    }
272}
273
274impl<I: Idx, T> ops::IndexMut<I> for IndexVec<I, T> {
275    #[inline]
276    fn index_mut(&mut self, index: I) -> &mut Self::Output {
277        &mut self.raw[index.to_usize()]
278    }
279}
280
281/// A dense mapping from ID to ID.
282///
283/// This is equivalent to `IndexVec<From, Option<To>>` and provides
284/// efficient dense ID remapping.
285#[derive(Clone)]
286pub struct DenseIdMap<From: Idx, To: Idx> {
287    inner: IndexVec<From, Option<To>>,
288}
289
290impl<From: Idx, To: Idx> DenseIdMap<From, To> {
291    /// Create a new dense ID mapping with the specified length.
292    #[inline]
293    pub fn with_len(length: usize) -> Self {
294        Self {
295            inner: IndexVec { raw: vec![None; length], _m: PhantomData },
296        }
297    }
298
299    /// Insert a mapping from source ID to target ID.
300    ///
301    /// # Panics
302    ///
303    /// Panics if the source ID is beyond the length of this DenseIdMap.
304    /// This DenseIdMap should be created with sufficient length to accommodate
305    /// all expected source IDs.
306    #[inline]
307    pub fn insert(&mut self, k: From, v: To) {
308        let idx = k.to_usize();
309        let len = self.len();
310
311        assert!(idx < len, "source ID {idx} exceeds DenseIdMap length {len}");
312        self.inner.insert_at(k, Some(v));
313    }
314
315    /// Get the target ID for the given source ID.
316    #[inline]
317    pub fn get(&self, k: From) -> Option<To> {
318        *self.inner.get(k)?
319    }
320
321    /// Get the number of source IDs in this mapping.
322    #[inline]
323    pub fn len(&self) -> usize {
324        self.inner.len()
325    }
326
327    /// Check if the mapping is empty.
328    #[inline]
329    pub fn is_empty(&self) -> bool {
330        self.inner.is_empty()
331    }
332}
333
334/// A trait for looking up values by ID.
335pub trait LookupByIdx<ID, V>
336where
337    ID: Idx,
338{
339    /// Get the value for the given ID.
340    fn get(&self, id: ID) -> Option<&V>;
341}
342
343/// A trait for looking up values by key that doesn't need to implement Idx.
344pub trait LookupByKey<K, V> {
345    /// Get the value for the given key.
346    fn get(&self, key: &K) -> Option<&V>;
347}
348
349impl<I, T> LookupByIdx<I, T> for IndexVec<I, T>
350where
351    I: Idx,
352{
353    fn get(&self, id: I) -> Option<&T> {
354        IndexVec::get(self, id)
355    }
356}
357
358impl<K, V> LookupByKey<K, V> for BTreeMap<K, V>
359where
360    K: Ord,
361{
362    fn get(&self, key: &K) -> Option<&V> {
363        BTreeMap::get(self, key)
364    }
365}
366
367impl<K, V> LookupByIdx<K, V> for BTreeMap<K, V>
368where
369    K: Idx,
370{
371    fn get(&self, id: K) -> Option<&V> {
372        BTreeMap::get(self, &id)
373    }
374}
375
376impl<I, T> LookupByIdx<I, T> for DenseIdMap<I, T>
377where
378    I: Idx,
379    T: Idx,
380{
381    fn get(&self, id: I) -> Option<&T> {
382        IndexVec::get(&self.inner, id).and_then(Option::as_ref)
383    }
384}
385
386impl<I: Idx, T> IntoIterator for IndexVec<I, T> {
387    type Item = T;
388    type IntoIter = vec::IntoIter<T>;
389
390    fn into_iter(self) -> Self::IntoIter {
391        self.raw.into_iter()
392    }
393}
394
395impl<'a, I: Idx, T> IntoIterator for &'a IndexVec<I, T> {
396    type Item = &'a T;
397    type IntoIter = core::slice::Iter<'a, T>;
398
399    fn into_iter(self) -> Self::IntoIter {
400        self.iter()
401    }
402}
403
404impl<I: Idx, T> TryFrom<Vec<T>> for IndexVec<I, T> {
405    type Error = IndexedVecError;
406
407    /// Create an IndexVec from a Vec.
408    ///
409    /// Returns an error if the Vec length exceeds u32::MAX.
410    fn try_from(raw: Vec<T>) -> Result<Self, Self::Error> {
411        if raw.len() > u32::MAX as usize {
412            return Err(IndexedVecError::TooManyItems);
413        }
414        Ok(Self { raw, _m: PhantomData })
415    }
416}
417
418// SERIALIZATION
419// ================================================================================================
420
421use miden_serde_utils::{
422    ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
423};
424
425impl<I, T> Serializable for IndexVec<I, T>
426where
427    I: Idx,
428    T: Serializable,
429{
430    fn write_into<W: ByteWriter>(&self, target: &mut W) {
431        self.as_slice().write_into(target);
432    }
433}
434
435impl<I, T> Deserializable for IndexVec<I, T>
436where
437    I: Idx,
438    T: Deserializable,
439{
440    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
441        let vec: Vec<T> = Deserializable::read_from(source)?;
442        IndexVec::try_from(vec).map_err(|_| {
443            DeserializationError::InvalidValue("IndexVec length exceeds u32::MAX".into())
444        })
445    }
446}
447
448impl<I, T> IndexVec<I, T>
449where
450    I: Idx,
451    T: Deserializable,
452{
453    /// Reads and validates a serialized length before it is used for allocation.
454    pub fn read_from_bounded<R: ByteReader>(
455        source: &mut R,
456        label: &str,
457    ) -> Result<Self, DeserializationError> {
458        let len = read_bounded_len(source, label, <T as Deserializable>::min_serialized_size())?;
459        if len > u32::MAX as usize {
460            return Err(DeserializationError::InvalidValue(
461                "IndexVec length exceeds u32::MAX".into(),
462            ));
463        }
464
465        let mut vec = Vec::<T>::with_capacity(bounded_initial_capacity::<T, _>(source, len));
466        for element in source.read_many_iter(len)? {
467            vec.push(element?);
468        }
469
470        Ok(Self { raw: vec, _m: PhantomData })
471    }
472}
473
474impl<I, T> IndexVec<I, T>
475where
476    I: Idx,
477{
478    /// Reads and validates a serialized length before it is used for allocation, using the provided
479    /// function to deserializing each element
480    pub fn read_from_bounded_with<R: ByteReader>(
481        source: &mut R,
482        label: &str,
483        min_element_size: usize,
484        deserializer: impl Fn(&mut R) -> Result<T, DeserializationError>,
485    ) -> Result<Self, DeserializationError> {
486        let len = read_bounded_len(source, label, min_element_size)?;
487        if len > u32::MAX as usize {
488            return Err(DeserializationError::InvalidValue(
489                "IndexVec length exceeds u32::MAX".into(),
490            ));
491        }
492
493        let mut vec = Vec::<T>::with_capacity(bounded_initial_capacity::<T, _>(source, len));
494        for _ in 0..len {
495            vec.push(deserializer(source)?);
496        }
497
498        Ok(Self { raw: vec, _m: PhantomData })
499    }
500}
501
502/// Reads and validates a serialized length before it is used for allocation.
503fn read_bounded_len<R: ByteReader>(
504    source: &mut R,
505    label: &str,
506    min_element_size: usize,
507) -> Result<usize, DeserializationError> {
508    let len = source.read_usize()?;
509    validate_bounded_len(source, label, len, min_element_size)?;
510    Ok(len)
511}
512
513/// Validates that a serialized length fits both the reader budget and remaining input.
514fn validate_bounded_len<R: ByteReader>(
515    source: &R,
516    label: &str,
517    len: usize,
518    min_element_size: usize,
519) -> Result<(), DeserializationError> {
520    let max_len = source.max_alloc(min_element_size);
521    if len > max_len {
522        return Err(DeserializationError::InvalidValue(alloc::format!(
523            "{label} count {len} exceeds budget {max_len}"
524        )));
525    }
526
527    let min_bytes = len.checked_mul(min_element_size).ok_or_else(|| {
528        DeserializationError::InvalidValue(alloc::format!(
529            "{label} count {len} overflows minimum serialized size {min_element_size}"
530        ))
531    })?;
532    source.check_eor(min_bytes).map_err(|err| match err {
533        DeserializationError::UnexpectedEOF => DeserializationError::InvalidValue(alloc::format!(
534            "{label} count {len} exceeds remaining input"
535        )),
536        err => err,
537    })
538}
539
540/// Bounds speculative collection capacity by both the declared length and the reader's remaining
541/// budget expressed in bytes of the in-memory element type.
542///
543/// Variable-width values can have a much smaller minimum serialized size than their in-memory
544/// representation. Reserving their full declared length before decoding the first value would
545/// amplify a compact malformed payload into a much larger allocation. A valid input can still grow
546/// the vector as each element is successfully decoded.
547fn bounded_initial_capacity<T, R: ByteReader>(source: &R, len: usize) -> usize {
548    let element_size = size_of::<T>();
549    if element_size == 0 {
550        len
551    } else {
552        len.min(source.max_alloc(element_size))
553    }
554}
555
556#[cfg(test)]
557mod tests {
558    use alloc::string::{String, ToString};
559
560    use miden_serde_utils::{BudgetedReader, SliceReader};
561
562    use super::*;
563
564    // Test ID types
565    newtype_id!(TestId);
566    newtype_id!(TestId2);
567
568    #[test]
569    fn bounded_initial_capacity_uses_in_memory_element_size() {
570        let reader = BudgetedReader::new(SliceReader::new(&[]), 256);
571
572        assert_eq!(bounded_initial_capacity::<[u8; 64], _>(&reader, 100), 4);
573        assert_eq!(bounded_initial_capacity::<[u8; 64], _>(&reader, 2), 2);
574        assert_eq!(bounded_initial_capacity::<(), _>(&reader, 100), 100);
575    }
576
577    #[test]
578    fn test_indexvec_basic() {
579        let mut vec = IndexVec::<TestId, String>::new();
580        let id1 = vec.push("hello".to_string()).unwrap();
581        let id2 = vec.push("world".to_string()).unwrap();
582
583        assert_eq!(vec.len(), 2);
584        assert_eq!(&vec[id1], "hello");
585        assert_eq!(&vec[id2], "world");
586        assert_eq!(vec.get(TestId::from(0)), Some(&"hello".to_string()));
587        assert_eq!(vec.get(TestId::from(2)), None);
588    }
589
590    #[test]
591    fn test_dense_id_map() {
592        let mut map = DenseIdMap::<TestId, TestId2>::with_len(2);
593        map.insert(TestId::from(0), TestId2::from(10));
594        map.insert(TestId::from(1), TestId2::from(11));
595
596        assert_eq!(map.len(), 2);
597        assert_eq!(map.get(TestId::from(0)), Some(TestId2::from(10)));
598        assert_eq!(map.get(TestId::from(1)), Some(TestId2::from(11)));
599        assert_eq!(map.get(TestId::from(2)), None);
600    }
601}
miden_utils_indexing/lib.rs

miden_utils_indexing/
lib.rs
