binius_utils/
sparse_index.rs

1// Copyright 2024-2025 Irreducible Inc.
2
3/// An index mapping positive integer IDs to optional values.
4#[derive(Debug, Clone)]
5pub struct SparseIndex<T> {
6	entries: Vec<Option<T>>,
7}
8
9impl<T> SparseIndex<T> {
10	pub fn new() -> Self {
11		Self::default()
12	}
13
14	pub fn with_capacity(capacity: usize) -> Self {
15		let mut entries = Vec::with_capacity(capacity);
16		entries.resize_with(capacity, || None);
17		Self { entries }
18	}
19
20	pub fn entry(&mut self, id: usize) -> Entry<T> {
21		if self.entries.len() <= id {
22			self.entries.resize_with(id + 1, || None);
23		}
24		Entry(&mut self.entries[id])
25	}
26
27	pub fn get(&self, id: usize) -> Option<&T> {
28		self.entries.get(id)?.as_ref()
29	}
30
31	pub fn get_mut(&mut self, id: usize) -> Option<&mut T> {
32		self.entries.get_mut(id)?.as_mut()
33	}
34
35	pub fn set(&mut self, id: usize, val: T) {
36		if self.entries.len() <= id {
37			self.entries.resize_with(id + 1, || None);
38		}
39		self.entries[id] = Some(val);
40	}
41
42	pub fn contains_key(&self, id: usize) -> bool {
43		self.entries.get(id).map(|x| x.is_some()).unwrap_or(false)
44	}
45
46	pub fn is_empty(&self) -> bool {
47		self.entries.iter().all(|v| v.is_none())
48	}
49
50	pub fn len(&self) -> usize {
51		self.entries.iter().flatten().count()
52	}
53
54	pub fn iter(&self) -> impl DoubleEndedIterator<Item = (usize, &T)> {
55		self.entries
56			.iter()
57			.enumerate()
58			.filter_map(|(i, v)| v.as_ref().map(|v| (i, v)))
59	}
60
61	pub fn iter_mut(&mut self) -> impl DoubleEndedIterator<Item = (usize, &mut T)> {
62		self.entries
63			.iter_mut()
64			.enumerate()
65			.filter_map(|(i, v)| v.as_mut().map(|v| (i, v)))
66	}
67
68	pub fn keys(&self) -> impl DoubleEndedIterator<Item = usize> + '_ {
69		self.entries
70			.iter()
71			.enumerate()
72			.filter_map(|(i, v)| v.as_ref().map(|_| i))
73	}
74
75	pub fn values(&self) -> impl DoubleEndedIterator<Item = &T> {
76		self.entries.iter().flatten()
77	}
78
79	pub fn values_mut(&mut self) -> impl DoubleEndedIterator<Item = &mut T> {
80		self.entries.iter_mut().flatten()
81	}
82}
83
84impl<T> Default for SparseIndex<T> {
85	fn default() -> Self {
86		Self {
87			entries: Vec::new(),
88		}
89	}
90}
91
92impl<T> IntoIterator for SparseIndex<T> {
93	type Item = (usize, T);
94	type IntoIter = std::iter::FilterMap<
95		std::iter::Enumerate<std::vec::IntoIter<Option<T>>>,
96		fn((usize, Option<T>)) -> Option<(usize, T)>,
97	>;
98
99	fn into_iter(self) -> Self::IntoIter {
100		self.entries
101			.into_iter()
102			.enumerate()
103			.filter_map(|(i, v)| v.map(|v| (i, v)))
104	}
105}
106
107impl<T> std::iter::FromIterator<(usize, T)> for SparseIndex<T> {
108	fn from_iter<I: IntoIterator<Item = (usize, T)>>(iter: I) -> Self {
109		let mut index = Self::new();
110		for (i, v) in iter {
111			index.set(i, v);
112		}
113		index
114	}
115}
116
117impl<T> std::ops::Index<usize> for SparseIndex<T> {
118	type Output = T;
119
120	fn index(&self, index: usize) -> &Self::Output {
121		self.get(index).unwrap()
122	}
123}
124
125impl<T> std::ops::IndexMut<usize> for SparseIndex<T> {
126	fn index_mut(&mut self, index: usize) -> &mut Self::Output {
127		self.get_mut(index).unwrap()
128	}
129}
130
131pub struct Entry<'a, V: 'a>(&'a mut Option<V>);
132
133impl<'a, V: 'a> Entry<'a, V> {
134	pub fn or_insert(self, default: V) -> &'a mut V {
135		self.0.get_or_insert(default)
136	}
137
138	pub fn or_insert_with(self, f: impl FnOnce() -> V) -> &'a mut V {
139		self.0.get_or_insert_with(f)
140	}
141}