onnx_runtime_ir/
arena.rs

1//! A generational-free slot arena keyed by a typed index.
2//!
3//! [`Graph`](crate::Graph) stores nodes and values in `Arena`s. Removal leaves
4//! a tombstone and recycles the slot on the next insert, so live ids stay
5//! stable across unrelated mutations. Ids are **not** generational: reusing a
6//! removed id after its slot has been recycled will alias the new occupant.
7//! Optimization passes are expected to drop stale ids when they remove nodes.
8
9use std::marker::PhantomData;
10
11/// A type usable as an arena key: a newtype around a `u32` index.
12pub trait ArenaKey: Copy {
13    /// Build a key from a raw slot index.
14    fn from_raw(raw: u32) -> Self;
15    /// The raw slot index of this key.
16    fn to_raw(self) -> u32;
17}
18
19/// A slot map: dense `Vec` storage with `O(1)` insert/remove/lookup by key.
20#[derive(Clone, Debug)]
21pub struct Arena<K: ArenaKey, T> {
22    slots: Vec<Option<T>>,
23    free: Vec<u32>,
24    len: usize,
25    _marker: PhantomData<K>,
26}
27
28impl<K: ArenaKey, T> Default for Arena<K, T> {
29    fn default() -> Self {
30        Self {
31            slots: Vec::new(),
32            free: Vec::new(),
33            len: 0,
34            _marker: PhantomData,
35        }
36    }
37}
38
39impl<K: ArenaKey, T> Arena<K, T> {
40    /// An empty arena.
41    pub fn new() -> Self {
42        Self::default()
43    }
44
45    /// Number of live entries.
46    pub fn len(&self) -> usize {
47        self.len
48    }
49
50    /// Whether there are no live entries.
51    pub fn is_empty(&self) -> bool {
52        self.len == 0
53    }
54
55    /// Number of allocated slots, including tombstones.
56    pub fn capacity(&self) -> usize {
57        self.slots.len()
58    }
59
60    /// Insert `value`, returning its freshly allocated key.
61    pub fn insert(&mut self, value: T) -> K {
62        self.insert_with(|_| value)
63    }
64
65    /// Insert a value that needs to know its own key (e.g. a node that stores
66    /// its own [`NodeId`](crate::NodeId)).
67    pub fn insert_with(&mut self, make: impl FnOnce(K) -> T) -> K {
68        let key = match self.free.pop() {
69            Some(idx) => K::from_raw(idx),
70            None => {
71                let idx = self.slots.len() as u32;
72                self.slots.push(None);
73                K::from_raw(idx)
74            }
75        };
76        self.slots[key.to_raw() as usize] = Some(make(key));
77        self.len += 1;
78        key
79    }
80
81    /// Whether `key` refers to a live entry.
82    pub fn contains(&self, key: K) -> bool {
83        self.get(key).is_some()
84    }
85
86    /// Borrow the entry for `key`, if live.
87    pub fn get(&self, key: K) -> Option<&T> {
88        self.slots
89            .get(key.to_raw() as usize)
90            .and_then(Option::as_ref)
91    }
92
93    /// Mutably borrow the entry for `key`, if live.
94    pub fn get_mut(&mut self, key: K) -> Option<&mut T> {
95        self.slots
96            .get_mut(key.to_raw() as usize)
97            .and_then(Option::as_mut)
98    }
99
100    /// Remove and return the entry for `key`, recycling the slot.
101    pub fn remove(&mut self, key: K) -> Option<T> {
102        let slot = self.slots.get_mut(key.to_raw() as usize)?;
103        let taken = slot.take();
104        if taken.is_some() {
105            self.free.push(key.to_raw());
106            self.len -= 1;
107        }
108        taken
109    }
110
111    /// Iterate over `(key, &value)` for all live entries, in ascending key order.
112    pub fn iter(&self) -> impl Iterator<Item = (K, &T)> {
113        self.slots
114            .iter()
115            .enumerate()
116            .filter_map(|(i, slot)| slot.as_ref().map(|v| (K::from_raw(i as u32), v)))
117    }
118
119    /// Iterate over the keys of all live entries.
120    pub fn keys(&self) -> impl Iterator<Item = K> + '_ {
121        self.slots
122            .iter()
123            .enumerate()
124            .filter_map(|(i, slot)| slot.as_ref().map(|_| K::from_raw(i as u32)))
125    }
126
127    /// Iterate over `&value` for all live entries.
128    pub fn values(&self) -> impl Iterator<Item = &T> {
129        self.slots.iter().filter_map(Option::as_ref)
130    }
131
132    /// Iterate over `&mut value` for all live entries.
133    pub fn values_mut(&mut self) -> impl Iterator<Item = &mut T> {
134        self.slots.iter_mut().filter_map(Option::as_mut)
135    }
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141
142    #[derive(Clone, Copy, PartialEq, Eq, Debug)]
143    struct K(u32);
144    impl ArenaKey for K {
145        fn from_raw(raw: u32) -> Self {
146            K(raw)
147        }
148        fn to_raw(self) -> u32 {
149            self.0
150        }
151    }
152
153    #[test]
154    fn insert_get_remove() {
155        let mut a: Arena<K, &str> = Arena::new();
156        assert!(a.is_empty());
157        let x = a.insert("x");
158        let y = a.insert("y");
159        assert_eq!(a.len(), 2);
160        assert_eq!(a.get(x), Some(&"x"));
161        assert_eq!(a.get(y), Some(&"y"));
162        assert_eq!(a.remove(x), Some("x"));
163        assert!(!a.contains(x));
164        assert_eq!(a.len(), 1);
165    }
166
167    #[test]
168    fn slots_are_recycled() {
169        let mut a: Arena<K, u32> = Arena::new();
170        let x = a.insert(1);
171        a.remove(x);
172        let z = a.insert(2);
173        // freed slot reused
174        assert_eq!(x.to_raw(), z.to_raw());
175        assert_eq!(a.get(z), Some(&2));
176    }
177
178    #[test]
179    fn insert_with_sees_own_key() {
180        let mut a: Arena<K, K> = Arena::new();
181        let k = a.insert_with(|self_key| self_key);
182        assert_eq!(a.get(k), Some(&k));
183    }
184
185    #[test]
186    fn iter_skips_tombstones() {
187        let mut a: Arena<K, u32> = Arena::new();
188        let x = a.insert(10);
189        let _y = a.insert(20);
190        a.remove(x);
191        let collected: Vec<u32> = a.values().copied().collect();
192        assert_eq!(collected, vec![20]);
193    }
194}
onnx_runtime_ir/arena.rs

onnx_runtime_ir/
arena.rs
