Skip to main content

wasmtime_environ/collections/
primary_map.rs

1use crate::{collections::TryExtend, error::OutOfMemory};
2use core::{
3    fmt,
4    ops::{Index, IndexMut},
5};
6use cranelift_entity::EntityRef;
7use serde::{Serialize, ser::SerializeSeq};
8
9/// Like [`cranelift_entity::PrimaryMap`] but enforces fallible allocation for
10/// all methods that allocate.
11#[derive(Hash, PartialEq, Eq)]
12pub struct PrimaryMap<K, V>
13where
14    K: EntityRef,
15{
16    inner: cranelift_entity::PrimaryMap<K, V>,
17}
18
19impl<K, V> Default for PrimaryMap<K, V>
20where
21    K: EntityRef,
22{
23    fn default() -> Self {
24        Self {
25            inner: cranelift_entity::PrimaryMap::default(),
26        }
27    }
28}
29
30impl<K, V> fmt::Debug for PrimaryMap<K, V>
31where
32    K: EntityRef + fmt::Debug,
33    V: fmt::Debug,
34{
35    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
36        fmt::Debug::fmt(&self.inner, f)
37    }
38}
39
40impl<K, V> From<crate::collections::Vec<V>> for PrimaryMap<K, V>
41where
42    K: EntityRef,
43{
44    fn from(values: crate::collections::Vec<V>) -> Self {
45        let values: ::alloc::vec::Vec<V> = values.into();
46        let inner = cranelift_entity::PrimaryMap::from(values);
47        Self { inner }
48    }
49}
50
51impl<K, V> serde::ser::Serialize for PrimaryMap<K, V>
52where
53    K: EntityRef,
54    V: Serialize,
55{
56    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
57    where
58        S: serde::Serializer,
59    {
60        let mut seq = serializer.serialize_seq(Some(self.len()))?;
61        for val in self.values() {
62            seq.serialize_element(val)?;
63        }
64        seq.end()
65    }
66}
67
68impl<'de, K, V> serde::de::Deserialize<'de> for PrimaryMap<K, V>
69where
70    K: EntityRef,
71    V: serde::de::Deserialize<'de>,
72{
73    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
74    where
75        D: serde::Deserializer<'de>,
76    {
77        let v: crate::collections::Vec<V> = serde::de::Deserialize::deserialize(deserializer)?;
78        Ok(Self::from(v))
79    }
80}
81
82impl<K, V> PrimaryMap<K, V>
83where
84    K: EntityRef,
85{
86    /// Same as [`cranelift_entity::PrimaryMap::new`].
87    pub fn new() -> Self {
88        Self {
89            inner: cranelift_entity::PrimaryMap::new(),
90        }
91    }
92
93    /// Same as [`cranelift_entity::PrimaryMap::try_with_capacity`].
94    pub fn with_capacity(capacity: usize) -> Result<Self, OutOfMemory> {
95        let mut map = Self::new();
96        map.reserve(capacity)?;
97        Ok(map)
98    }
99
100    /// Same as [`cranelift_entity::PrimaryMap::is_valid`].
101    pub fn is_valid(&self, k: K) -> bool {
102        self.inner.is_valid(k)
103    }
104
105    /// Same as [`cranelift_entity::PrimaryMap::get`].
106    pub fn get(&self, k: K) -> Option<&V> {
107        self.inner.get(k)
108    }
109
110    /// Same as [`cranelift_entity::PrimaryMap::get_range`].
111    pub fn get_range(&self, range: core::ops::Range<K>) -> Option<&[V]> {
112        self.inner.get_range(range)
113    }
114
115    /// Same as [`cranelift_entity::PrimaryMap::get_mut`].
116    pub fn get_mut(&mut self, k: K) -> Option<&mut V> {
117        self.inner.get_mut(k)
118    }
119
120    /// Same as [`cranelift_entity::PrimaryMap::is_empty`].
121    pub fn is_empty(&self) -> bool {
122        self.inner.is_empty()
123    }
124
125    /// Same as [`cranelift_entity::PrimaryMap::len`].
126    pub fn len(&self) -> usize {
127        self.inner.len()
128    }
129
130    /// Same as [`cranelift_entity::PrimaryMap::keys`].
131    pub fn keys(&self) -> cranelift_entity::Keys<K> {
132        self.inner.keys()
133    }
134
135    /// Same as [`cranelift_entity::PrimaryMap::values`].
136    pub fn values(&self) -> core::slice::Iter<'_, V> {
137        self.inner.values()
138    }
139
140    /// Same as [`cranelift_entity::PrimaryMap::values_mut`].
141    pub fn values_mut(&mut self) -> core::slice::IterMut<'_, V> {
142        self.inner.values_mut()
143    }
144
145    /// Same as [`cranelift_entity::PrimaryMap::as_values_slice`].
146    pub fn as_values_slice(&self) -> &[V] {
147        self.inner.as_values_slice()
148    }
149
150    /// Same as [`cranelift_entity::PrimaryMap::iter`].
151    pub fn iter(&self) -> cranelift_entity::Iter<'_, K, V> {
152        self.inner.iter()
153    }
154
155    /// Same as [`cranelift_entity::PrimaryMap::iter_mut`].
156    pub fn iter_mut(&mut self) -> cranelift_entity::IterMut<'_, K, V> {
157        self.inner.iter_mut()
158    }
159
160    /// Same as [`cranelift_entity::PrimaryMap::clear`].
161    pub fn clear(&mut self) {
162        self.inner.clear()
163    }
164
165    /// Same as [`cranelift_entity::PrimaryMap::next_key`].
166    pub fn next_key(&self) -> K {
167        self.inner.next_key()
168    }
169
170    /// Same as [`cranelift_entity::PrimaryMap::push`] but returns an error on
171    /// allocation failure.
172    pub fn push(&mut self, v: V) -> Result<K, OutOfMemory> {
173        self.reserve(1)?;
174        Ok(self.inner.push(v))
175    }
176
177    /// Same as [`cranelift_entity::PrimaryMap::last`].
178    pub fn last(&self) -> Option<(K, &V)> {
179        self.inner.last()
180    }
181
182    /// Same as [`cranelift_entity::PrimaryMap::last_mut`].
183    pub fn last_mut(&mut self) -> Option<(K, &mut V)> {
184        self.inner.last_mut()
185    }
186
187    /// Same as [`cranelift_entity::PrimaryMap::try_reserve`].
188    pub fn reserve(&mut self, additional: usize) -> Result<(), OutOfMemory> {
189        self.inner.try_reserve(additional)
190    }
191
192    /// Same as [`cranelift_entity::PrimaryMap::try_reserve_exact`].
193    pub fn reserve_exact(&mut self, additional: usize) -> Result<(), OutOfMemory> {
194        self.inner.try_reserve_exact(additional)
195    }
196
197    /// Same as [`cranelift_entity::PrimaryMap::get_disjoint_mut`].
198    pub fn get_disjoint_mut<const N: usize>(
199        &mut self,
200        indices: [K; N],
201    ) -> Result<[&mut V; N], core::slice::GetDisjointMutError> {
202        self.inner.get_disjoint_mut(indices)
203    }
204
205    /// Same as [`cranelift_entity::PrimaryMap::binary_search_values_by_key`].
206    pub fn binary_search_values_by_key<'a, B, F>(&'a self, b: &B, f: F) -> Result<K, K>
207    where
208        F: FnMut(&'a V) -> B,
209        B: Ord,
210    {
211        self.inner.binary_search_values_by_key(b, f)
212    }
213
214    /// Same as [`cranelift_entity::PrimaryMap::get_raw_mut`].
215    pub fn get_raw_mut(&mut self, k: K) -> Option<*mut V> {
216        self.inner.get_raw_mut(k)
217    }
218}
219
220impl<K, V> TryExtend<V> for PrimaryMap<K, V>
221where
222    K: EntityRef,
223{
224    fn try_extend<I>(&mut self, iter: I) -> Result<(), OutOfMemory>
225    where
226        I: IntoIterator<Item = V>,
227    {
228        let iter = iter.into_iter();
229        let (min, max) = iter.size_hint();
230        let cap = max.unwrap_or(min);
231        self.reserve(cap)?;
232        for v in iter {
233            self.push(v)?;
234        }
235        Ok(())
236    }
237}
238
239impl<K, V> Index<K> for PrimaryMap<K, V>
240where
241    K: EntityRef,
242{
243    type Output = V;
244
245    fn index(&self, k: K) -> &V {
246        &self.inner[k]
247    }
248}
249
250impl<K, V> IndexMut<K> for PrimaryMap<K, V>
251where
252    K: EntityRef,
253{
254    fn index_mut(&mut self, k: K) -> &mut V {
255        &mut self.inner[k]
256    }
257}
258
259impl<K, V> IntoIterator for PrimaryMap<K, V>
260where
261    K: EntityRef,
262{
263    type Item = (K, V);
264    type IntoIter = cranelift_entity::IntoIter<K, V>;
265    fn into_iter(self) -> Self::IntoIter {
266        self.inner.into_iter()
267    }
268}
269
270impl<'a, K, V> IntoIterator for &'a PrimaryMap<K, V>
271where
272    K: EntityRef,
273{
274    type Item = (K, &'a V);
275    type IntoIter = cranelift_entity::Iter<'a, K, V>;
276    fn into_iter(self) -> Self::IntoIter {
277        self.iter()
278    }
279}
280
281impl<'a, K, V> IntoIterator for &'a mut PrimaryMap<K, V>
282where
283    K: EntityRef,
284{
285    type Item = (K, &'a mut V);
286    type IntoIter = cranelift_entity::IterMut<'a, K, V>;
287    fn into_iter(self) -> Self::IntoIter {
288        self.iter_mut()
289    }
290}