1 use std::marker::PhantomData; 2 3 use crate::Idx; 4 5 /// A map from arena indexes to some other type. 6 /// Space requirement is O(highest index). 7 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] 8 pub struct ArenaMap<IDX, V> { 9 v: Vec<Option<V>>, 10 _ty: PhantomData<IDX>, 11 } 12 13 impl<T, V> ArenaMap<Idx<T>, V> { 14 /// Inserts a value associated with a given arena index into the map. insert(&mut self, idx: Idx<T>, t: V)15 pub fn insert(&mut self, idx: Idx<T>, t: V) { 16 let idx = Self::to_idx(idx); 17 18 self.v.resize_with((idx + 1).max(self.v.len()), || None); 19 self.v[idx] = Some(t); 20 } 21 22 /// Returns a reference to the value associated with the provided index 23 /// if it is present. get(&self, idx: Idx<T>) -> Option<&V>24 pub fn get(&self, idx: Idx<T>) -> Option<&V> { 25 self.v.get(Self::to_idx(idx)).and_then(|it| it.as_ref()) 26 } 27 28 /// Returns a mutable reference to the value associated with the provided index 29 /// if it is present. get_mut(&mut self, idx: Idx<T>) -> Option<&mut V>30 pub fn get_mut(&mut self, idx: Idx<T>) -> Option<&mut V> { 31 self.v.get_mut(Self::to_idx(idx)).and_then(|it| it.as_mut()) 32 } 33 34 /// Returns an iterator over the values in the map. values(&self) -> impl Iterator<Item = &V>35 pub fn values(&self) -> impl Iterator<Item = &V> { 36 self.v.iter().filter_map(|o| o.as_ref()) 37 } 38 39 /// Returns an iterator over mutable references to the values in the map. values_mut(&mut self) -> impl Iterator<Item = &mut V>40 pub fn values_mut(&mut self) -> impl Iterator<Item = &mut V> { 41 self.v.iter_mut().filter_map(|o| o.as_mut()) 42 } 43 44 /// Returns an iterator over the arena indexes and values in the map. iter(&self) -> impl Iterator<Item = (Idx<T>, &V)>45 pub fn iter(&self) -> impl Iterator<Item = (Idx<T>, &V)> { 46 self.v.iter().enumerate().filter_map(|(idx, o)| Some((Self::from_idx(idx), o.as_ref()?))) 47 } 48 to_idx(idx: Idx<T>) -> usize49 fn to_idx(idx: Idx<T>) -> usize { 50 u32::from(idx.into_raw()) as usize 51 } 52 from_idx(idx: usize) -> Idx<T>53 fn from_idx(idx: usize) -> Idx<T> { 54 Idx::from_raw((idx as u32).into()) 55 } 56 } 57 58 impl<T, V> std::ops::Index<Idx<V>> for ArenaMap<Idx<V>, T> { 59 type Output = T; index(&self, idx: Idx<V>) -> &T60 fn index(&self, idx: Idx<V>) -> &T { 61 self.v[Self::to_idx(idx)].as_ref().unwrap() 62 } 63 } 64 65 impl<T, V> std::ops::IndexMut<Idx<V>> for ArenaMap<Idx<V>, T> { index_mut(&mut self, idx: Idx<V>) -> &mut T66 fn index_mut(&mut self, idx: Idx<V>) -> &mut T { 67 self.v[Self::to_idx(idx)].as_mut().unwrap() 68 } 69 } 70 71 impl<T, V> Default for ArenaMap<Idx<V>, T> { default() -> Self72 fn default() -> Self { 73 ArenaMap { v: Vec::new(), _ty: PhantomData } 74 } 75 } 76