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}