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

miden_utils_indexing/
lib.rs
