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    /// Insert `value`, returning its freshly allocated key.
56    pub fn insert(&mut self, value: T) -> K {
57        self.insert_with(|_| value)
58    }
59
60    /// Insert a value that needs to know its own key (e.g. a node that stores
61    /// its own [`NodeId`](crate::NodeId)).
62    pub fn insert_with(&mut self, make: impl FnOnce(K) -> T) -> K {
63        let key = match self.free.pop() {
64            Some(idx) => K::from_raw(idx),
65            None => {
66                let idx = self.slots.len() as u32;
67                self.slots.push(None);
68                K::from_raw(idx)
69            }
70        };
71        self.slots[key.to_raw() as usize] = Some(make(key));
72        self.len += 1;
73        key
74    }
75
76    /// Whether `key` refers to a live entry.
77    pub fn contains(&self, key: K) -> bool {
78        self.get(key).is_some()
79    }
80
81    /// Borrow the entry for `key`, if live.
82    pub fn get(&self, key: K) -> Option<&T> {
83        self.slots.get(key.to_raw() as usize).and_then(Option::as_ref)
84    }
85
86    /// Mutably borrow the entry for `key`, if live.
87    pub fn get_mut(&mut self, key: K) -> Option<&mut T> {
88        self.slots
89            .get_mut(key.to_raw() as usize)
90            .and_then(Option::as_mut)
91    }
92
93    /// Remove and return the entry for `key`, recycling the slot.
94    pub fn remove(&mut self, key: K) -> Option<T> {
95        let slot = self.slots.get_mut(key.to_raw() as usize)?;
96        let taken = slot.take();
97        if taken.is_some() {
98            self.free.push(key.to_raw());
99            self.len -= 1;
100        }
101        taken
102    }
103
104    /// Iterate over `(key, &value)` for all live entries, in ascending key order.
105    pub fn iter(&self) -> impl Iterator<Item = (K, &T)> {
106        self.slots
107            .iter()
108            .enumerate()
109            .filter_map(|(i, slot)| slot.as_ref().map(|v| (K::from_raw(i as u32), v)))
110    }
111
112    /// Iterate over the keys of all live entries.
113    pub fn keys(&self) -> impl Iterator<Item = K> + '_ {
114        self.slots
115            .iter()
116            .enumerate()
117            .filter_map(|(i, slot)| slot.as_ref().map(|_| K::from_raw(i as u32)))
118    }
119
120    /// Iterate over `&value` for all live entries.
121    pub fn values(&self) -> impl Iterator<Item = &T> {
122        self.slots.iter().filter_map(Option::as_ref)
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129
130    #[derive(Clone, Copy, PartialEq, Eq, Debug)]
131    struct K(u32);
132    impl ArenaKey for K {
133        fn from_raw(raw: u32) -> Self {
134            K(raw)
135        }
136        fn to_raw(self) -> u32 {
137            self.0
138        }
139    }
140
141    #[test]
142    fn insert_get_remove() {
143        let mut a: Arena<K, &str> = Arena::new();
144        assert!(a.is_empty());
145        let x = a.insert("x");
146        let y = a.insert("y");
147        assert_eq!(a.len(), 2);
148        assert_eq!(a.get(x), Some(&"x"));
149        assert_eq!(a.get(y), Some(&"y"));
150        assert_eq!(a.remove(x), Some("x"));
151        assert!(!a.contains(x));
152        assert_eq!(a.len(), 1);
153    }
154
155    #[test]
156    fn slots_are_recycled() {
157        let mut a: Arena<K, u32> = Arena::new();
158        let x = a.insert(1);
159        a.remove(x);
160        let z = a.insert(2);
161        // freed slot reused
162        assert_eq!(x.to_raw(), z.to_raw());
163        assert_eq!(a.get(z), Some(&2));
164    }
165
166    #[test]
167    fn insert_with_sees_own_key() {
168        let mut a: Arena<K, K> = Arena::new();
169        let k = a.insert_with(|self_key| self_key);
170        assert_eq!(a.get(k), Some(&k));
171    }
172
173    #[test]
174    fn iter_skips_tombstones() {
175        let mut a: Arena<K, u32> = Arena::new();
176        let x = a.insert(10);
177        let _y = a.insert(20);
178        a.remove(x);
179        let collected: Vec<u32> = a.values().copied().collect();
180        assert_eq!(collected, vec![20]);
181    }
182}
onnx_runtime_ir/arena.rs

onnx_runtime_ir/
arena.rs
