Skip to main content

piece_tree/
cranelift_entity_vendor.rs

1// Portions of this file are copied from the `cranelift_entity` crate.
2// Original code Copyright (c) The Cranelift Project Developers.
3// Licensed under the Apache License, Version 2.0 with LLVM Exception.
4// See: https://llvm.org/LICENSE.txt
5
6extern crate alloc;
7
8use alloc::vec::Vec;
9use core::fmt;
10use core::marker::PhantomData;
11use core::ops::{Index, IndexMut};
12use core::slice;
13use alloc::vec;
14use core::iter::Enumerate;
15
16/// A type wrapping a small integer index should implement `EntityRef` so it can be used as the key
17/// of an `SecondaryMap` or `SparseMap`.
18pub trait EntityRef: Copy + Eq {
19    /// Create a new entity reference from a small integer.
20    /// This should crash if the requested index is not representable.
21    fn new(_: usize) -> Self;
22
23    /// Get the index that was used to create this entity reference.
24    fn index(self) -> usize;
25}
26
27/// Macro which provides the common implementation of a 32-bit entity reference.
28#[macro_export]
29macro_rules! entity_impl {
30    // Basic traits.
31    ($entity:ident) => {
32        impl $crate::EntityRef for $entity {
33            #[inline]
34            fn new(index: usize) -> Self {
35                debug_assert!(index < (core::u32::MAX as usize));
36                $entity(index as u32)
37            }
38
39            #[inline]
40            fn index(self) -> usize {
41                self.0 as usize
42            }
43        }
44
45        impl $entity {
46            /// Create a new instance from a `u32`.
47            #[allow(dead_code, reason = "macro-generated code")]
48            #[inline]
49            pub fn from_u32(x: u32) -> Self {
50                debug_assert!(x < core::u32::MAX);
51                $entity(x)
52            }
53
54            /// Return the underlying index value as a `u32`.
55            #[allow(dead_code, reason = "macro-generated code")]
56            #[inline]
57            pub fn as_u32(self) -> u32 {
58                self.0
59            }
60
61            /// Return the raw bit encoding for this instance.
62            ///
63            /// __Warning__: the raw bit encoding is opaque and has no
64            /// guaranteed correspondence to the entity's index. It encodes the
65            /// entire state of this index value: either a valid index or an
66            /// invalid-index sentinel. The value returned by this method should
67            /// only be passed to `from_bits`.
68            #[allow(dead_code, reason = "macro-generated code")]
69            #[inline]
70            pub fn as_bits(self) -> u32 {
71                self.0
72            }
73
74            /// Create a new instance from the raw bit encoding.
75            ///
76            /// __Warning__: the raw bit encoding is opaque and has no
77            /// guaranteed correspondence to the entity's index. It encodes the
78            /// entire state of this index value: either a valid index or an
79            /// invalid-index sentinel. The value returned by this method should
80            /// only be given bits from `as_bits`.
81            #[allow(dead_code, reason = "macro-generated code")]
82            #[inline]
83            pub fn from_bits(x: u32) -> Self {
84                $entity(x)
85            }
86        }
87    };
88
89    // Include basic `Display` impl using the given display prefix.
90    // Display a `Block` reference as "block12".
91    ($entity:ident, $display_prefix:expr) => {
92        $crate::entity_impl!($entity);
93
94        impl core::fmt::Display for $entity {
95            fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
96                write!(f, concat!($display_prefix, "{}"), self.0)
97            }
98        }
99
100        impl core::fmt::Debug for $entity {
101            fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
102                (self as &dyn core::fmt::Display).fmt(f)
103            }
104        }
105    };
106
107    // Alternate form for tuples we can't directly construct; providing "to" and "from" expressions
108    // to turn an index *into* an entity, or get an index *from* an entity.
109    ($entity:ident, $display_prefix:expr, $arg:ident, $to_expr:expr, $from_expr:expr) => {
110        impl $crate::EntityRef for $entity {
111            #[inline]
112            fn new(index: usize) -> Self {
113                debug_assert!(index < (core::u32::MAX as usize));
114                let $arg = index as u32;
115                $to_expr
116            }
117
118            #[inline]
119            fn index(self) -> usize {
120                let $arg = self;
121                $from_expr as usize
122            }
123        }
124
125        impl $crate::packed_option::ReservedValue for $entity {
126            #[inline]
127            fn reserved_value() -> $entity {
128                $entity::from_u32(core::u32::MAX)
129            }
130
131            #[inline]
132            fn is_reserved_value(&self) -> bool {
133                self.as_u32() == core::u32::MAX
134            }
135        }
136
137        impl $entity {
138            /// Create a new instance from a `u32`.
139            #[allow(dead_code, reason = "macro-generated code")]
140            #[inline]
141            pub fn from_u32(x: u32) -> Self {
142                debug_assert!(x < core::u32::MAX);
143                let $arg = x;
144                $to_expr
145            }
146
147            /// Return the underlying index value as a `u32`.
148            #[allow(dead_code, reason = "macro-generated code")]
149            #[inline]
150            pub fn as_u32(self) -> u32 {
151                let $arg = self;
152                $from_expr
153            }
154        }
155
156        impl core::fmt::Display for $entity {
157            fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
158                write!(f, concat!($display_prefix, "{}"), self.as_u32())
159            }
160        }
161
162        impl core::fmt::Debug for $entity {
163            fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
164                (self as &dyn core::fmt::Display).fmt(f)
165            }
166        }
167    };
168}
169
170/// Densely numbered entity references as mapping keys.
171
172/// A primary mapping `K -> V` allocating dense entity references.
173///
174/// The `PrimaryMap` data structure uses the dense index space to implement a map with a vector.
175///
176/// A primary map contains the main definition of an entity, and it can be used to allocate new
177/// entity references with the `push` method.
178///
179/// There should only be a single `PrimaryMap` instance for a given `EntityRef` type, otherwise
180/// conflicting references will be created. Using unknown keys for indexing will cause a panic.
181///
182/// Note that `PrimaryMap` doesn't implement `Deref` or `DerefMut`, which would allow
183/// `&PrimaryMap<K, V>` to convert to `&[V]`. One of the main advantages of `PrimaryMap` is
184/// that it only allows indexing with the distinct `EntityRef` key type, so converting to a
185/// plain slice would make it easier to use incorrectly.
186#[derive(Clone, Hash, PartialEq, Eq)]
187pub struct PrimaryMap<K, V>
188where
189    K: EntityRef,
190{
191    elems: Vec<V>,
192    unused: PhantomData<K>,
193}
194
195impl<K, V> PrimaryMap<K, V>
196where
197    K: EntityRef,
198{
199    /// Create a new empty map.
200    pub fn new() -> Self {
201        Self {
202            elems: Vec::new(),
203            unused: PhantomData,
204        }
205    }
206
207    /// Create a new empty map with the given capacity.
208    pub fn with_capacity(capacity: usize) -> Self {
209        Self {
210            elems: Vec::with_capacity(capacity),
211            unused: PhantomData,
212        }
213    }
214
215    /// Check if `k` is a valid key in the map.
216    pub fn is_valid(&self, k: K) -> bool {
217        k.index() < self.elems.len()
218    }
219
220    /// Get the element at `k` if it exists.
221    pub fn get(&self, k: K) -> Option<&V> {
222        self.elems.get(k.index())
223    }
224
225    /// Get the slice of values associated with the given range of keys, if any.
226    pub fn get_range(&self, range: core::ops::Range<K>) -> Option<&[V]> {
227        self.elems.get(range.start.index()..range.end.index())
228    }
229
230    /// Get the element at `k` if it exists, mutable version.
231    pub fn get_mut(&mut self, k: K) -> Option<&mut V> {
232        self.elems.get_mut(k.index())
233    }
234
235    /// Is this map completely empty?
236    pub fn is_empty(&self) -> bool {
237        self.elems.is_empty()
238    }
239
240    /// Get the total number of entity references created.
241    pub fn len(&self) -> usize {
242        self.elems.len()
243    }
244
245    /// Iterate over all the keys in this map.
246    pub fn keys(&self) -> Keys<K> {
247        Keys::with_len(self.elems.len())
248    }
249
250    /// Iterate over all the values in this map.
251    pub fn values(&self) -> slice::Iter<'_, V> {
252        self.elems.iter()
253    }
254
255    /// Iterate over all the values in this map, mutable edition.
256    pub fn values_mut(&mut self) -> slice::IterMut<'_, V> {
257        self.elems.iter_mut()
258    }
259
260    /// Get this map's underlying values as a slice.
261    pub fn as_values_slice(&self) -> &[V] {
262        &self.elems
263    }
264
265    /// Iterate over all the keys and values in this map.
266    pub fn iter(&self) -> Iter<'_, K, V> {
267        Iter::new(self.elems.iter())
268    }
269
270    /// Iterate over all the keys and values in this map, mutable edition.
271    pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
272        IterMut::new(self.elems.iter_mut())
273    }
274
275    /// Remove all entries from this map.
276    pub fn clear(&mut self) {
277        self.elems.clear()
278    }
279
280    /// Get the key that will be assigned to the next pushed value.
281    pub fn next_key(&self) -> K {
282        K::new(self.elems.len())
283    }
284
285    /// Append `v` to the mapping, assigning a new key which is returned.
286    pub fn push(&mut self, v: V) -> K {
287        let k = self.next_key();
288        self.elems.push(v);
289        k
290    }
291
292    /// Returns the last element that was inserted in the map.
293    pub fn last(&self) -> Option<(K, &V)> {
294        let len = self.elems.len();
295        let last = self.elems.last()?;
296        Some((K::new(len - 1), last))
297    }
298
299    /// Returns the last element that was inserted in the map.
300    pub fn last_mut(&mut self) -> Option<(K, &mut V)> {
301        let len = self.elems.len();
302        let last = self.elems.last_mut()?;
303        Some((K::new(len - 1), last))
304    }
305
306    /// Reserves capacity for at least `additional` more elements to be inserted.
307    pub fn reserve(&mut self, additional: usize) {
308        self.elems.reserve(additional)
309    }
310
311    /// Reserves the minimum capacity for exactly `additional` more elements to be inserted.
312    pub fn reserve_exact(&mut self, additional: usize) {
313        self.elems.reserve_exact(additional)
314    }
315
316    /// Shrinks the capacity of the `PrimaryMap` as much as possible.
317    pub fn shrink_to_fit(&mut self) {
318        self.elems.shrink_to_fit()
319    }
320
321    /// Returns mutable references to many elements at once.
322    ///
323    /// Returns an error if an element does not exist, or if the same key was passed more than
324    /// once.
325    pub fn get_disjoint_mut<const N: usize>(
326        &mut self,
327        indices: [K; N],
328    ) -> Result<[&mut V; N], slice::GetDisjointMutError> {
329        self.elems.get_disjoint_mut(indices.map(|k| k.index()))
330    }
331
332    /// Performs a binary search on the values with a key extraction function.
333    ///
334    /// Assumes that the values are sorted by the key extracted by the function.
335    ///
336    /// If the value is found then `Ok(K)` is returned, containing the entity key
337    /// of the matching value.
338    ///
339    /// If there are multiple matches, then any one of the matches could be returned.
340    ///
341    /// If the value is not found then Err(K) is returned, containing the entity key
342    /// where a matching element could be inserted while maintaining sorted order.
343    pub fn binary_search_values_by_key<'a, B, F>(&'a self, b: &B, f: F) -> Result<K, K>
344    where
345        F: FnMut(&'a V) -> B,
346        B: Ord,
347    {
348        self.elems
349            .binary_search_by_key(b, f)
350            .map(|i| K::new(i))
351            .map_err(|i| K::new(i))
352    }
353
354    /// Analog of `get_raw` except that a raw pointer is returned rather than a
355    /// mutable reference.
356    ///
357    /// The default accessors of items in [`PrimaryMap`] will invalidate all
358    /// previous borrows obtained from the map according to miri. This function
359    /// can be used to acquire a pointer and then subsequently acquire a second
360    /// pointer later on without invalidating the first one. In other words
361    /// this is only here to help borrow two elements simultaneously with miri.
362    pub fn get_raw_mut(&mut self, k: K) -> Option<*mut V> {
363        if k.index() < self.elems.len() {
364            // SAFETY: the `add` function requires that the index is in-bounds
365            // with respect to the allocation which is satisfied here due to
366            // the bounds-check above.
367            unsafe { Some(self.elems.as_mut_ptr().add(k.index())) }
368        } else {
369            None
370        }
371    }
372}
373
374impl<K, V> Default for PrimaryMap<K, V>
375where
376    K: EntityRef,
377{
378    fn default() -> PrimaryMap<K, V> {
379        PrimaryMap::new()
380    }
381}
382
383/// Immutable indexing into an `PrimaryMap`.
384/// The indexed value must be in the map.
385impl<K, V> Index<K> for PrimaryMap<K, V>
386where
387    K: EntityRef,
388{
389    type Output = V;
390
391    fn index(&self, k: K) -> &V {
392        &self.elems[k.index()]
393    }
394}
395
396/// Mutable indexing into an `PrimaryMap`.
397impl<K, V> IndexMut<K> for PrimaryMap<K, V>
398where
399    K: EntityRef,
400{
401    fn index_mut(&mut self, k: K) -> &mut V {
402        &mut self.elems[k.index()]
403    }
404}
405
406impl<K, V> IntoIterator for PrimaryMap<K, V>
407where
408    K: EntityRef,
409{
410    type Item = (K, V);
411    type IntoIter = IntoIter<K, V>;
412
413    fn into_iter(self) -> Self::IntoIter {
414        IntoIter::new(self.elems.into_iter())
415    }
416}
417
418impl<'a, K, V> IntoIterator for &'a PrimaryMap<K, V>
419where
420    K: EntityRef,
421{
422    type Item = (K, &'a V);
423    type IntoIter = Iter<'a, K, V>;
424
425    fn into_iter(self) -> Self::IntoIter {
426        Iter::new(self.elems.iter())
427    }
428}
429
430impl<'a, K, V> IntoIterator for &'a mut PrimaryMap<K, V>
431where
432    K: EntityRef,
433{
434    type Item = (K, &'a mut V);
435    type IntoIter = IterMut<'a, K, V>;
436
437    fn into_iter(self) -> Self::IntoIter {
438        IterMut::new(self.elems.iter_mut())
439    }
440}
441
442impl<K, V> FromIterator<V> for PrimaryMap<K, V>
443where
444    K: EntityRef,
445{
446    fn from_iter<T>(iter: T) -> Self
447    where
448        T: IntoIterator<Item = V>,
449    {
450        Self {
451            elems: Vec::from_iter(iter),
452            unused: PhantomData,
453        }
454    }
455}
456
457impl<K, V> Extend<V> for PrimaryMap<K, V>
458where
459    K: EntityRef,
460{
461    fn extend<T>(&mut self, iter: T)
462    where
463        T: IntoIterator<Item = V>,
464    {
465        self.elems.extend(iter);
466    }
467}
468
469impl<K, V> From<Vec<V>> for PrimaryMap<K, V>
470where
471    K: EntityRef,
472{
473    fn from(elems: Vec<V>) -> Self {
474        Self {
475            elems,
476            unused: PhantomData,
477        }
478    }
479}
480
481impl<K, V> From<PrimaryMap<K, V>> for Vec<V>
482where
483    K: EntityRef,
484{
485    fn from(map: PrimaryMap<K, V>) -> Self {
486        map.elems
487    }
488}
489
490impl<K: EntityRef + fmt::Debug, V: fmt::Debug> fmt::Debug for PrimaryMap<K, V> {
491    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
492        let mut struct_ = f.debug_struct("PrimaryMap");
493        for (k, v) in self {
494            struct_.field(&alloc::format!("{k:?}"), v);
495        }
496        struct_.finish()
497    }
498}
499
500/// A double-ended iterator over entity references.
501///
502/// When `core::iter::Step` is stabilized, `Keys` could be implemented as a wrapper around
503/// `core::ops::Range`, but for now, we implement it manually.
504
505/// Iterate over all keys in order.
506pub struct Keys<K: EntityRef> {
507    pos: usize,
508    rev_pos: usize,
509    unused: PhantomData<K>,
510}
511
512impl<K: EntityRef> Keys<K> {
513    /// Create a `Keys` iterator that visits `len` entities starting from 0.
514    pub fn with_len(len: usize) -> Self {
515        Self {
516            pos: 0,
517            rev_pos: len,
518            unused: PhantomData,
519        }
520    }
521}
522
523impl<K: EntityRef> Iterator for Keys<K> {
524    type Item = K;
525
526    fn next(&mut self) -> Option<Self::Item> {
527        if self.pos < self.rev_pos {
528            let k = K::new(self.pos);
529            self.pos += 1;
530            Some(k)
531        } else {
532            None
533        }
534    }
535
536    fn size_hint(&self) -> (usize, Option<usize>) {
537        let size = self.rev_pos - self.pos;
538        (size, Some(size))
539    }
540}
541
542impl<K: EntityRef> DoubleEndedIterator for Keys<K> {
543    fn next_back(&mut self) -> Option<Self::Item> {
544        if self.rev_pos > self.pos {
545            let k = K::new(self.rev_pos - 1);
546            self.rev_pos -= 1;
547            Some(k)
548        } else {
549            None
550        }
551    }
552}
553
554impl<K: EntityRef> ExactSizeIterator for Keys<K> {}
555
556/// Iterate over all keys in order.
557pub struct Iter<'a, K: EntityRef, V>
558where
559    V: 'a,
560{
561    enumerate: Enumerate<slice::Iter<'a, V>>,
562    unused: PhantomData<K>,
563}
564
565impl<'a, K: EntityRef, V> Iter<'a, K, V> {
566    /// Create an `Iter` iterator that visits the `PrimaryMap` keys and values
567    /// of `iter`.
568    pub fn new(iter: slice::Iter<'a, V>) -> Self {
569        Self {
570            enumerate: iter.enumerate(),
571            unused: PhantomData,
572        }
573    }
574}
575
576impl<'a, K: EntityRef, V> Iterator for Iter<'a, K, V> {
577    type Item = (K, &'a V);
578
579    fn next(&mut self) -> Option<Self::Item> {
580        self.enumerate.next().map(|(i, v)| (K::new(i), v))
581    }
582
583    fn size_hint(&self) -> (usize, Option<usize>) {
584        self.enumerate.size_hint()
585    }
586}
587
588impl<'a, K: EntityRef, V> DoubleEndedIterator for Iter<'a, K, V> {
589    fn next_back(&mut self) -> Option<Self::Item> {
590        self.enumerate.next_back().map(|(i, v)| (K::new(i), v))
591    }
592}
593
594impl<'a, K: EntityRef, V> ExactSizeIterator for Iter<'a, K, V> {}
595
596/// Iterate over all keys in order.
597pub struct IterMut<'a, K: EntityRef, V>
598where
599    V: 'a,
600{
601    enumerate: Enumerate<slice::IterMut<'a, V>>,
602    unused: PhantomData<K>,
603}
604
605impl<'a, K: EntityRef, V> IterMut<'a, K, V> {
606    /// Create an `IterMut` iterator that visits the `PrimaryMap` keys and values
607    /// of `iter`.
608    pub fn new(iter: slice::IterMut<'a, V>) -> Self {
609        Self {
610            enumerate: iter.enumerate(),
611            unused: PhantomData,
612        }
613    }
614}
615
616impl<'a, K: EntityRef, V> Iterator for IterMut<'a, K, V> {
617    type Item = (K, &'a mut V);
618
619    fn next(&mut self) -> Option<Self::Item> {
620        self.enumerate.next().map(|(i, v)| (K::new(i), v))
621    }
622
623    fn size_hint(&self) -> (usize, Option<usize>) {
624        self.enumerate.size_hint()
625    }
626}
627
628impl<'a, K: EntityRef, V> DoubleEndedIterator for IterMut<'a, K, V> {
629    fn next_back(&mut self) -> Option<Self::Item> {
630        self.enumerate.next_back().map(|(i, v)| (K::new(i), v))
631    }
632}
633
634impl<'a, K: EntityRef, V> ExactSizeIterator for IterMut<'a, K, V> {}
635
636/// Iterate over all keys in order.
637pub struct IntoIter<K: EntityRef, V> {
638    enumerate: Enumerate<vec::IntoIter<V>>,
639    unused: PhantomData<K>,
640}
641
642impl<K: EntityRef, V> IntoIter<K, V> {
643    /// Create an `IntoIter` iterator that visits the `PrimaryMap` keys and values
644    /// of `iter`.
645    pub fn new(iter: vec::IntoIter<V>) -> Self {
646        Self {
647            enumerate: iter.enumerate(),
648            unused: PhantomData,
649        }
650    }
651}
652
653impl<K: EntityRef, V> Iterator for IntoIter<K, V> {
654    type Item = (K, V);
655
656    fn next(&mut self) -> Option<Self::Item> {
657        self.enumerate.next().map(|(i, v)| (K::new(i), v))
658    }
659
660    fn size_hint(&self) -> (usize, Option<usize>) {
661        self.enumerate.size_hint()
662    }
663}
664
665impl<K: EntityRef, V> DoubleEndedIterator for IntoIter<K, V> {
666    fn next_back(&mut self) -> Option<Self::Item> {
667        self.enumerate.next_back().map(|(i, v)| (K::new(i), v))
668    }
669}
670
671impl<K: EntityRef, V> ExactSizeIterator for IntoIter<K, V> {}