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.get(key.to_raw() as usize).and_then(Option::as_ref)
89    }
90
91    /// Mutably borrow the entry for `key`, if live.
92    pub fn get_mut(&mut self, key: K) -> Option<&mut T> {
93        self.slots
94            .get_mut(key.to_raw() as usize)
95            .and_then(Option::as_mut)
96    }
97
98    /// Remove and return the entry for `key`, recycling the slot.
99    pub fn remove(&mut self, key: K) -> Option<T> {
100        let slot = self.slots.get_mut(key.to_raw() as usize)?;
101        let taken = slot.take();
102        if taken.is_some() {
103            self.free.push(key.to_raw());
104            self.len -= 1;
105        }
106        taken
107    }
108
109    /// Iterate over `(key, &value)` for all live entries, in ascending key order.
110    pub fn iter(&self) -> impl Iterator<Item = (K, &T)> {
111        self.slots
112            .iter()
113            .enumerate()
114            .filter_map(|(i, slot)| slot.as_ref().map(|v| (K::from_raw(i as u32), v)))
115    }
116
117    /// Iterate over the keys of all live entries.
118    pub fn keys(&self) -> impl Iterator<Item = K> + '_ {
119        self.slots
120            .iter()
121            .enumerate()
122            .filter_map(|(i, slot)| slot.as_ref().map(|_| K::from_raw(i as u32)))
123    }
124
125    /// Iterate over `&value` for all live entries.
126    pub fn values(&self) -> impl Iterator<Item = &T> {
127        self.slots.iter().filter_map(Option::as_ref)
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use super::*;
134
135    #[derive(Clone, Copy, PartialEq, Eq, Debug)]
136    struct K(u32);
137    impl ArenaKey for K {
138        fn from_raw(raw: u32) -> Self {
139            K(raw)
140        }
141        fn to_raw(self) -> u32 {
142            self.0
143        }
144    }
145
146    #[test]
147    fn insert_get_remove() {
148        let mut a: Arena<K, &str> = Arena::new();
149        assert!(a.is_empty());
150        let x = a.insert("x");
151        let y = a.insert("y");
152        assert_eq!(a.len(), 2);
153        assert_eq!(a.get(x), Some(&"x"));
154        assert_eq!(a.get(y), Some(&"y"));
155        assert_eq!(a.remove(x), Some("x"));
156        assert!(!a.contains(x));
157        assert_eq!(a.len(), 1);
158    }
159
160    #[test]
161    fn slots_are_recycled() {
162        let mut a: Arena<K, u32> = Arena::new();
163        let x = a.insert(1);
164        a.remove(x);
165        let z = a.insert(2);
166        // freed slot reused
167        assert_eq!(x.to_raw(), z.to_raw());
168        assert_eq!(a.get(z), Some(&2));
169    }
170
171    #[test]
172    fn insert_with_sees_own_key() {
173        let mut a: Arena<K, K> = Arena::new();
174        let k = a.insert_with(|self_key| self_key);
175        assert_eq!(a.get(k), Some(&k));
176    }
177
178    #[test]
179    fn iter_skips_tombstones() {
180        let mut a: Arena<K, u32> = Arena::new();
181        let x = a.insert(10);
182        let _y = a.insert(20);
183        a.remove(x);
184        let collected: Vec<u32> = a.values().copied().collect();
185        assert_eq!(collected, vec![20]);
186    }
187}
onnx_runtime_ir/arena.rs

onnx_runtime_ir/
arena.rs
