ryo-analysis 0.1.0

//! IndexVec - Type-safe indexed vector (rustc-style)
//!
//! Provides a Vec wrapper that only accepts strongly-typed indices,
//! preventing accidental index confusion between different ID types.

use std::fmt::Debug;
use std::marker::PhantomData;
use std::ops::{Index, IndexMut};

/// Trait for index types that can be used with IndexVec.
pub trait Idx: Copy + Eq + Debug {
    /// Construct the index type from its raw `u32` representation.
    /// Implementations should not validate the value beyond their own
    /// invariants.
    fn new(raw: u32) -> Self;
    /// Extract the raw `u32` representation of this index, suitable for
    /// indexing into an [`IndexVec`].
    fn index(self) -> u32;
}

/// A Vec that can only be indexed by a specific index type `I`.
#[derive(Clone)]
pub struct IndexVec<I: Idx, T> {
    raw: Vec<T>,
    _marker: PhantomData<fn(&I)>,
}

impl<I: Idx, T: Debug> Debug for IndexVec<I, T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("IndexVec")
            .field("len", &self.raw.len())
            .field("data", &self.raw)
            .finish()
    }
}

impl<I: Idx, T> Default for IndexVec<I, T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<I: Idx, T> IndexVec<I, T> {
    /// Create a new empty IndexVec.
    #[inline]
    pub const fn new() -> Self {
        Self {
            raw: Vec::new(),
            _marker: PhantomData,
        }
    }

    /// Create with pre-allocated capacity.
    #[inline]
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            raw: Vec::with_capacity(capacity),
            _marker: PhantomData,
        }
    }

    /// Push a value and return its index.
    #[inline]
    pub fn push(&mut self, val: T) -> I {
        let idx = self.raw.len() as u32;
        self.raw.push(val);
        I::new(idx)
    }

    /// Get a reference to the value at the given index.
    #[inline]
    pub fn get(&self, idx: I) -> Option<&T> {
        self.raw.get(idx.index() as usize)
    }

    /// Get a mutable reference to the value at the given index.
    #[inline]
    pub fn get_mut(&mut self, idx: I) -> Option<&mut T> {
        self.raw.get_mut(idx.index() as usize)
    }

    /// Returns the number of elements.
    #[inline]
    pub fn len(&self) -> usize {
        self.raw.len()
    }

    /// Returns true if empty.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.raw.is_empty()
    }

    /// Returns an iterator over the values.
    #[inline]
    pub fn iter(&self) -> impl Iterator<Item = &T> {
        self.raw.iter()
    }

    /// Returns a mutable iterator over the values.
    #[inline]
    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut T> {
        self.raw.iter_mut()
    }

    /// Returns an iterator over (index, value) pairs.
    #[inline]
    pub fn iter_enumerated(&self) -> impl Iterator<Item = (I, &T)> {
        self.raw
            .iter()
            .enumerate()
            .map(|(i, v)| (I::new(i as u32), v))
    }

    /// Returns an iterator over indices.
    #[inline]
    pub fn indices(&self) -> impl Iterator<Item = I> {
        (0..self.raw.len() as u32).map(I::new)
    }

    /// Access the underlying raw Vec.
    #[inline]
    pub fn raw(&self) -> &Vec<T> {
        &self.raw
    }

    /// Consume and return the underlying Vec.
    #[inline]
    pub fn into_raw(self) -> Vec<T> {
        self.raw
    }
}

impl<I: Idx, T> Index<I> for IndexVec<I, T> {
    type Output = T;

    #[inline]
    fn index(&self, idx: I) -> &Self::Output {
        &self.raw[idx.index() as usize]
    }
}

impl<I: Idx, T> IndexMut<I> for IndexVec<I, T> {
    #[inline]
    fn index_mut(&mut self, idx: I) -> &mut Self::Output {
        &mut self.raw[idx.index() as usize]
    }
}

impl<I: Idx, T> FromIterator<T> for IndexVec<I, T> {
    fn from_iter<It: IntoIterator<Item = T>>(iter: It) -> Self {
        Self {
            raw: iter.into_iter().collect(),
            _marker: PhantomData,
        }
    }
}

/// Macro to define a newtype index.
#[macro_export]
macro_rules! define_index {
    ($(#[$attr:meta])* $vis:vis struct $name:ident;) => {
        $(#[$attr])*
        #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, serde::Serialize, serde::Deserialize)]
        #[repr(transparent)]
        $vis struct $name(u32);

        impl $crate::query::index_vec::Idx for $name {
            #[inline]
            fn new(raw: u32) -> Self {
                Self(raw)
            }

            #[inline]
            fn index(self) -> u32 {
                self.0
            }
        }

        impl $name {
            /// Create a new index from a raw u32.
            #[inline]
            pub const fn from_raw(raw: u32) -> Self {
                Self(raw)
            }

            /// Get the raw u32 value.
            #[inline]
            pub const fn as_u32(self) -> u32 {
                self.0
            }

            /// Get as usize for array indexing.
            #[inline]
            pub const fn as_usize(self) -> usize {
                self.0 as usize
            }
        }
    };
}

#[cfg(test)]
mod tests {
    use super::*;

    define_index! {
        /// Test index type
        struct TestId;
    }

    #[test]
    fn test_push_and_get() {
        let mut vec: IndexVec<TestId, &str> = IndexVec::new();

        let a = vec.push("hello");
        let b = vec.push("world");

        assert_eq!(vec[a], "hello");
        assert_eq!(vec[b], "world");
        assert_eq!(vec.len(), 2);
    }

    #[test]
    fn test_iter_enumerated() {
        let mut vec: IndexVec<TestId, i32> = IndexVec::new();
        vec.push(10);
        vec.push(20);
        vec.push(30);

        let pairs: Vec<_> = vec.iter_enumerated().collect();
        assert_eq!(pairs.len(), 3);
        assert_eq!(pairs[0].0.as_u32(), 0);
        assert_eq!(*pairs[0].1, 10);
        assert_eq!(pairs[2].0.as_u32(), 2);
        assert_eq!(*pairs[2].1, 30);
    }

    #[test]
    fn test_indices() {
        let mut vec: IndexVec<TestId, ()> = IndexVec::new();
        vec.push(());
        vec.push(());
        vec.push(());

        let indices: Vec<_> = vec.indices().collect();
        assert_eq!(indices.len(), 3);
        assert_eq!(indices[0].as_u32(), 0);
        assert_eq!(indices[1].as_u32(), 1);
        assert_eq!(indices[2].as_u32(), 2);
        assert_eq!(indices[0].as_usize(), 0);
        assert_eq!(TestId::from_raw(5).as_u32(), 5);
    }
}