compact: Don't assume Emitted expressions are live.

src/arena.rs

@@ -199,7 +199,7 @@ impl<T> Range<T> {
         }
     }
 
-    /// Return the first and last handles included in `self`.
+    /// return the first and last handles included in `self`.
     ///
     /// If `self` is an empty range, there are no handles included, so
     /// return `None`.
@@ -216,6 +216,24 @@ impl<T> Range<T> {
             None
         }
     }
+
+    /// Return the zero-based index range covered by `self`.
+    pub fn zero_based_index_range(&self) -> ops::Range<u32> {
+        self.inner.clone()
+    }
+
+    /// Construct a `Range` that covers the zero-based indices in `inner`.
+    pub fn from_zero_based_index_range(inner: ops::Range<u32>, arena: &Arena<T>) -> Self {
+        // Since `inner` is a `Range<u32>`, we only need to check that
+        // the start and end are well-ordered, and that the end fits
+        // within `arena`.
+        assert!(inner.start <= inner.end);
+        assert!(inner.end as usize <= arena.len());
+        Self {
+            inner,
+            marker: Default::default(),
+        }
+    }
 }
 
 /// An arena holding some kind of component (e.g., type, constant,
@@ -387,18 +405,26 @@ impl<T> Arena<T> {
 
     /// Assert that `range` is valid for this arena.
     pub fn check_contains_range(&self, range: &Range<T>) -> Result<(), BadRangeError> {
-        // Since `range.inner` is a `Range<u32>`, we only need to
-        // check that the start precedes the end, and that the end is
-        // in range.
-        if range.inner.start > range.inner.end
-            || self
-                .check_contains_handle(Handle::new(range.inner.end.try_into().unwrap()))
-                .is_err()
-        {
-            Err(BadRangeError::new(range.clone()))
-        } else {
-            Ok(())
+        // Since `range.inner` is a `Range<u32>`, we only need to check that the
+        // start precedes the end, and that the end is in range.
+        if range.inner.start > range.inner.end {
+            return Err(BadRangeError::new(range.clone()));
+        }
+
+        // Empty ranges are tolerated: they can be produced by compaction.
+        if range.inner.start == range.inner.end {
+            return Ok(());
+        }
+
+        // `range.inner` is zero-based, but end-exclusive, so `range.inner.end`
+        // is actually the right one-based index for the last handle within the
+        // range.
+        let last_handle = Handle::new(range.inner.end.try_into().unwrap());
+        if self.check_contains_handle(last_handle).is_err() {
+            return Err(BadRangeError::new(range.clone()));
         }
+
+        Ok(())
     }
 
     #[cfg(feature = "compact")]

src/compact/expressions.rs

@@ -271,7 +271,6 @@ impl ModuleMap {
                 adjust(sampler);
                 adjust(coordinate);
                 operand_map.adjust_option(array_index);
-                // TEST: try adjusting this with plain operand_map
                 if let Some(ref mut offset) = *offset {
                     self.const_expressions.adjust(offset);
                 }

src/compact/functions.rs

@@ -2,16 +2,16 @@ use super::handle_set_map::HandleSet;
 use super::{FunctionMap, ModuleMap};
 use crate::arena::Handle;
 
-pub(super) struct FunctionTracer<'a> {
-    pub(super) module: &'a crate::Module,
-    pub(super) function: &'a crate::Function,
+pub struct FunctionTracer<'a> {
+    pub module: &'a crate::Module,
+    pub function: &'a crate::Function,
 
-    pub(super) types_used: &'a mut HandleSet<crate::Type>,
-    pub(super) constants_used: &'a mut HandleSet<crate::Constant>,
-    pub(super) const_expressions_used: &'a mut HandleSet<crate::Expression>,
+    pub types_used: &'a mut HandleSet<crate::Type>,
+    pub constants_used: &'a mut HandleSet<crate::Constant>,
+    pub const_expressions_used: &'a mut HandleSet<crate::Expression>,
 
     /// Function-local expressions used.
-    pub(super) expressions_used: HandleSet<crate::Expression>,
+    pub expressions_used: HandleSet<crate::Expression>,
 }
 
 impl<'a> FunctionTracer<'a> {
@@ -27,7 +27,6 @@ impl<'a> FunctionTracer<'a> {
         for (_, local) in self.function.local_variables.iter() {
             self.trace_type(local.ty);
             if let Some(init) = local.init {
-                // TEST: try changing this to trace_expression
                 self.trace_const_expression(init);
             }
         }
@@ -55,12 +54,6 @@ impl<'a> FunctionTracer<'a> {
             .trace_expression(expr);
     }
 
-    /*
-        pub fn trace_const_expression(&mut self, const_expr: Handle<crate::Expression>) {
-            self.as_expression().as_const_expression().trace_expression(const_expr);
-    }
-        */
-
     fn as_type(&mut self) -> super::types::TypeTracer {
         super::types::TypeTracer {
             types: &self.module.types,
@@ -129,6 +122,6 @@ impl FunctionMap {
         assert!(reuse.is_empty());
 
         // Adjust statements.
-        self.adjust_block(&mut function.body);
+        self.adjust_body(function);
     }
 }

src/compact/handle_set_map.rs

@@ -1,4 +1,4 @@
-use crate::arena::{Arena, Handle, UniqueArena};
+use crate::arena::{Arena, Handle, Range, UniqueArena};
 
 type Index = std::num::NonZeroU32;
 
@@ -122,4 +122,37 @@ impl<T: 'static> HandleMap<T> {
             self.adjust(handle);
         }
     }
+
+    /// Shrink `range` to include only used handles.
+    ///
+    /// Fortunately, compaction doesn't arbitrarily scramble the expressions
+    /// in the arena, but instead preserves the order of the elements while
+    /// squeezing out unused ones. That means that a contiguous range in the
+    /// pre-compacted arena always maps to a contiguous range in the
+    /// post-compacted arena. So we just need to adjust the endpoints.
+    ///
+    /// Compaction may have eliminated the endpoints themselves.
+    ///
+    /// Use `compacted_arena` to bounds-check the result.
+    pub fn adjust_range(&self, range: &mut Range<T>, compacted_arena: &Arena<T>) {
+        let mut index_range = range.zero_based_index_range();
+        let compacted;
+        // Remember that the indices we retrieve from `new_index` are 1-based
+        // compacted indices, but the index range we're computing is zero-based
+        // compacted indices.
+        if let Some(first1) = index_range.find_map(|i| self.new_index[i as usize]) {
+            // The first call to `find_map` mutated `index_range` to hold the
+            // remainder of original range, which is exactly the range we need
+            // to search for the new last handle.
+            if let Some(last1) = index_range.rev().find_map(|i| self.new_index[i as usize]) {
+                // Build a zero-based end-exclusive range, given one-based handle indices.
+                compacted = first1.get() - 1..last1.get();
+            } else {
+                compacted = first1.get() - 1..first1.get();
+            }
+        } else {
+            compacted = 0..0;
+        };
+        *range = Range::from_zero_based_index_range(compacted, compacted_arena);
+    }
 }

src/compact/statements.rs

@@ -9,11 +9,12 @@ impl FunctionTracer<'_> {
             for stmt in last {
                 use crate::Statement as St;
                 match *stmt {
-                    St::Emit(ref range) => {
-                        for expr in range.clone() {
-                            log::trace!("trace emitted expression {:?}", expr);
-                            self.trace_expression(expr);
-                        }
+                    St::Emit(ref _range) => {
+                        // If we come across a statement that actually uses an
+                        // expression in this range, it'll get traced from
+                        // there. But since evaluating expressions has no
+                        // effect, we don't need to assume that everything
+                        // emitted is live.
                     }
                     St::Block(ref block) => worklist.push(block),
                     St::If {
@@ -140,7 +141,8 @@ impl FunctionTracer<'_> {
 }
 
 impl FunctionMap {
-    pub fn adjust_block(&self, block: &mut [crate::Statement]) {
+    pub fn adjust_body(&self, function: &mut crate::Function) {
+        let block = &mut function.body;
         let mut worklist: Vec<&mut [crate::Statement]> = vec![block];
         let adjust = |handle: &mut Handle<crate::Expression>| {
             self.expressions.adjust(handle);
@@ -150,17 +152,7 @@ impl FunctionMap {
                 use crate::Statement as St;
                 match *stmt {
                     St::Emit(ref mut range) => {
-                        // Fortunately, compaction doesn't arbitrarily scramble
-                        // the expressions in the arena, but instead preserves
-                        // the order of the elements while squeezing out unused
-                        // ones. That means that a contiguous range in the
-                        // pre-compacted arena always maps to a contiguous range
-                        // in the post-compacted arena. So we just need to
-                        // adjust the endpoints.
-                        let (mut first, mut last) = range.first_and_last().unwrap();
-                        self.expressions.adjust(&mut first);
-                        self.expressions.adjust(&mut last);
-                        *range = crate::arena::Range::new_from_bounds(first, last);
+                        self.expressions.adjust_range(range, &function.expressions);
                     }
                     St::Block(ref mut block) => worklist.push(block),
                     St::If {