ryo_analysis/query/
index_vec.rs

1//! IndexVec - Type-safe indexed vector (rustc-style)
2//!
3//! Provides a Vec wrapper that only accepts strongly-typed indices,
4//! preventing accidental index confusion between different ID types.
5
6use std::fmt::Debug;
7use std::marker::PhantomData;
8use std::ops::{Index, IndexMut};
9
10/// Trait for index types that can be used with IndexVec.
11pub trait Idx: Copy + Eq + Debug {
12    /// Construct the index type from its raw `u32` representation.
13    /// Implementations should not validate the value beyond their own
14    /// invariants.
15    fn new(raw: u32) -> Self;
16    /// Extract the raw `u32` representation of this index, suitable for
17    /// indexing into an [`IndexVec`].
18    fn index(self) -> u32;
19}
20
21/// A Vec that can only be indexed by a specific index type `I`.
22#[derive(Clone)]
23pub struct IndexVec<I: Idx, T> {
24    raw: Vec<T>,
25    _marker: PhantomData<fn(&I)>,
26}
27
28impl<I: Idx, T: Debug> Debug for IndexVec<I, T> {
29    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
30        f.debug_struct("IndexVec")
31            .field("len", &self.raw.len())
32            .field("data", &self.raw)
33            .finish()
34    }
35}
36
37impl<I: Idx, T> Default for IndexVec<I, T> {
38    fn default() -> Self {
39        Self::new()
40    }
41}
42
43impl<I: Idx, T> IndexVec<I, T> {
44    /// Create a new empty IndexVec.
45    #[inline]
46    pub const fn new() -> Self {
47        Self {
48            raw: Vec::new(),
49            _marker: PhantomData,
50        }
51    }
52
53    /// Create with pre-allocated capacity.
54    #[inline]
55    pub fn with_capacity(capacity: usize) -> Self {
56        Self {
57            raw: Vec::with_capacity(capacity),
58            _marker: PhantomData,
59        }
60    }
61
62    /// Push a value and return its index.
63    #[inline]
64    pub fn push(&mut self, val: T) -> I {
65        let idx = self.raw.len() as u32;
66        self.raw.push(val);
67        I::new(idx)
68    }
69
70    /// Get a reference to the value at the given index.
71    #[inline]
72    pub fn get(&self, idx: I) -> Option<&T> {
73        self.raw.get(idx.index() as usize)
74    }
75
76    /// Get a mutable reference to the value at the given index.
77    #[inline]
78    pub fn get_mut(&mut self, idx: I) -> Option<&mut T> {
79        self.raw.get_mut(idx.index() as usize)
80    }
81
82    /// Returns the number of elements.
83    #[inline]
84    pub fn len(&self) -> usize {
85        self.raw.len()
86    }
87
88    /// Returns true if empty.
89    #[inline]
90    pub fn is_empty(&self) -> bool {
91        self.raw.is_empty()
92    }
93
94    /// Returns an iterator over the values.
95    #[inline]
96    pub fn iter(&self) -> impl Iterator<Item = &T> {
97        self.raw.iter()
98    }
99
100    /// Returns a mutable iterator over the values.
101    #[inline]
102    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut T> {
103        self.raw.iter_mut()
104    }
105
106    /// Returns an iterator over (index, value) pairs.
107    #[inline]
108    pub fn iter_enumerated(&self) -> impl Iterator<Item = (I, &T)> {
109        self.raw
110            .iter()
111            .enumerate()
112            .map(|(i, v)| (I::new(i as u32), v))
113    }
114
115    /// Returns an iterator over indices.
116    #[inline]
117    pub fn indices(&self) -> impl Iterator<Item = I> {
118        (0..self.raw.len() as u32).map(I::new)
119    }
120
121    /// Access the underlying raw Vec.
122    #[inline]
123    pub fn raw(&self) -> &Vec<T> {
124        &self.raw
125    }
126
127    /// Consume and return the underlying Vec.
128    #[inline]
129    pub fn into_raw(self) -> Vec<T> {
130        self.raw
131    }
132}
133
134impl<I: Idx, T> Index<I> for IndexVec<I, T> {
135    type Output = T;
136
137    #[inline]
138    fn index(&self, idx: I) -> &Self::Output {
139        &self.raw[idx.index() as usize]
140    }
141}
142
143impl<I: Idx, T> IndexMut<I> for IndexVec<I, T> {
144    #[inline]
145    fn index_mut(&mut self, idx: I) -> &mut Self::Output {
146        &mut self.raw[idx.index() as usize]
147    }
148}
149
150impl<I: Idx, T> FromIterator<T> for IndexVec<I, T> {
151    fn from_iter<It: IntoIterator<Item = T>>(iter: It) -> Self {
152        Self {
153            raw: iter.into_iter().collect(),
154            _marker: PhantomData,
155        }
156    }
157}
158
159/// Macro to define a newtype index.
160#[macro_export]
161macro_rules! define_index {
162    ($(#[$attr:meta])* $vis:vis struct $name:ident;) => {
163        $(#[$attr])*
164        #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, serde::Serialize, serde::Deserialize)]
165        #[repr(transparent)]
166        $vis struct $name(u32);
167
168        impl $crate::query::index_vec::Idx for $name {
169            #[inline]
170            fn new(raw: u32) -> Self {
171                Self(raw)
172            }
173
174            #[inline]
175            fn index(self) -> u32 {
176                self.0
177            }
178        }
179
180        impl $name {
181            /// Create a new index from a raw u32.
182            #[inline]
183            pub const fn from_raw(raw: u32) -> Self {
184                Self(raw)
185            }
186
187            /// Get the raw u32 value.
188            #[inline]
189            pub const fn as_u32(self) -> u32 {
190                self.0
191            }
192
193            /// Get as usize for array indexing.
194            #[inline]
195            pub const fn as_usize(self) -> usize {
196                self.0 as usize
197            }
198        }
199    };
200}
201
202#[cfg(test)]
203mod tests {
204    use super::*;
205
206    define_index! {
207        /// Test index type
208        struct TestId;
209    }
210
211    #[test]
212    fn test_push_and_get() {
213        let mut vec: IndexVec<TestId, &str> = IndexVec::new();
214
215        let a = vec.push("hello");
216        let b = vec.push("world");
217
218        assert_eq!(vec[a], "hello");
219        assert_eq!(vec[b], "world");
220        assert_eq!(vec.len(), 2);
221    }
222
223    #[test]
224    fn test_iter_enumerated() {
225        let mut vec: IndexVec<TestId, i32> = IndexVec::new();
226        vec.push(10);
227        vec.push(20);
228        vec.push(30);
229
230        let pairs: Vec<_> = vec.iter_enumerated().collect();
231        assert_eq!(pairs.len(), 3);
232        assert_eq!(pairs[0].0.as_u32(), 0);
233        assert_eq!(*pairs[0].1, 10);
234        assert_eq!(pairs[2].0.as_u32(), 2);
235        assert_eq!(*pairs[2].1, 30);
236    }
237
238    #[test]
239    fn test_indices() {
240        let mut vec: IndexVec<TestId, ()> = IndexVec::new();
241        vec.push(());
242        vec.push(());
243        vec.push(());
244
245        let indices: Vec<_> = vec.indices().collect();
246        assert_eq!(indices.len(), 3);
247        assert_eq!(indices[0].as_u32(), 0);
248        assert_eq!(indices[1].as_u32(), 1);
249        assert_eq!(indices[2].as_u32(), 2);
250        assert_eq!(indices[0].as_usize(), 0);
251        assert_eq!(TestId::from_raw(5).as_u32(), 5);
252    }
253}
ryo_analysis/query/index_vec.rs

ryo_analysis/query/
index_vec.rs
