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

naga/compact/
handle_set_map.rs
