naga/compact/handle_set_map.rs
1use alloc::vec::Vec;
2
3use crate::arena::{Arena, Handle, HandleSet, Range};
4
5type Index = crate::non_max_u32::NonMaxU32;
6
7/// A map keyed by handles.
8///
9/// In most cases, this is used to map from old handle indices to new,
10/// compressed handle indices.
11#[derive(Debug)]
12pub struct HandleMap<T, U = Index> {
13 /// The indices assigned to handles in the compacted module.
14 ///
15 /// If `new_index[i]` is `Some(n)`, then `n` is the `Index` of the
16 /// compacted `Handle` corresponding to the pre-compacted `Handle`
17 /// whose index is `i`.
18 new_index: Vec<Option<U>>,
19
20 /// This type is indexed by values of type `T`.
21 as_keys: core::marker::PhantomData<T>,
22}
23
24impl<T, U> HandleMap<T, U> {
25 pub fn with_capacity(capacity: usize) -> Self {
26 Self {
27 new_index: Vec::with_capacity(capacity),
28 as_keys: core::marker::PhantomData,
29 }
30 }
31
32 pub fn get(&self, handle: Handle<T>) -> Option<&U> {
33 self.new_index.get(handle.index()).unwrap_or(&None).as_ref()
34 }
35
36 pub fn insert(&mut self, handle: Handle<T>, value: U) -> Option<U> {
37 if self.new_index.len() <= handle.index() {
38 self.new_index.resize_with(handle.index() + 1, || None);
39 }
40 self.new_index[handle.index()].replace(value)
41 }
42}
43
44impl<T: 'static> HandleMap<T> {
45 pub fn from_set(set: HandleSet<T>) -> Self {
46 let mut next_index = Index::new(0).unwrap();
47 Self {
48 new_index: set
49 .all_possible()
50 .map(|handle| {
51 if set.contains(handle) {
52 // This handle will be retained in the compacted version,
53 // so assign it a new index.
54 let this = next_index;
55 next_index = next_index.checked_add(1).unwrap();
56 Some(this)
57 } else {
58 // This handle will be omitted in the compacted version.
59 None
60 }
61 })
62 .collect(),
63 as_keys: core::marker::PhantomData,
64 }
65 }
66
67 /// Return true if `old` is used in the compacted module.
68 pub fn used(&self, old: Handle<T>) -> bool {
69 self.new_index[old.index()].is_some()
70 }
71
72 /// Return the counterpart to `old` in the compacted module.
73 ///
74 /// If we thought `old` wouldn't be used in the compacted module, return
75 /// `None`.
76 pub fn try_adjust(&self, old: Handle<T>) -> Option<Handle<T>> {
77 log::trace!(
78 "adjusting {} handle [{}] -> [{:?}]",
79 core::any::type_name::<T>(),
80 old.index(),
81 self.new_index[old.index()]
82 );
83 self.new_index[old.index()].map(Handle::new)
84 }
85
86 /// Return the counterpart to `old` in the compacted module.
87 ///
88 /// If we thought `old` wouldn't be used in the compacted module, panic.
89 pub fn adjust(&self, handle: &mut Handle<T>) {
90 *handle = self.try_adjust(*handle).unwrap();
91 }
92
93 /// Like `adjust`, but for optional handles.
94 pub fn adjust_option(&self, handle: &mut Option<Handle<T>>) {
95 if let Some(ref mut handle) = *handle {
96 self.adjust(handle);
97 }
98 }
99
100 /// Shrink `range` to include only used handles.
101 ///
102 /// Fortunately, compaction doesn't arbitrarily scramble the expressions
103 /// in the arena, but instead preserves the order of the elements while
104 /// squeezing out unused ones. That means that a contiguous range in the
105 /// pre-compacted arena always maps to a contiguous range in the
106 /// post-compacted arena. So we just need to adjust the endpoints.
107 ///
108 /// Compaction may have eliminated the endpoints themselves.
109 ///
110 /// Use `compacted_arena` to bounds-check the result.
111 pub fn adjust_range(&self, range: &mut Range<T>, compacted_arena: &Arena<T>) {
112 let mut index_range = range.index_range();
113 let compacted;
114 if let Some(first) = index_range.find_map(|i| self.new_index[i as usize]) {
115 // The first call to `find_map` mutated `index_range` to hold the
116 // remainder of original range, which is exactly the range we need
117 // to search for the new last handle.
118 if let Some(last) = index_range.rev().find_map(|i| self.new_index[i as usize]) {
119 // Build an end-exclusive range, given the two included indices
120 // `first` and `last`.
121 compacted = first.get()..last.get() + 1;
122 } else {
123 // The range contains only a single live handle, which
124 // we identified with the first `find_map` call.
125 compacted = first.get()..first.get() + 1;
126 }
127 } else {
128 compacted = 0..0;
129 };
130 *range = Range::from_index_range(compacted, compacted_arena);
131 }
132}