[naga] Improve algorithm for Module compaction.

Identify reachable function expressions, constant expressions, and
types using a single pass over each arena, taking advantage of the
fact that expressions and types may only refer to other entries that
precede them within their arena. Only walking the statement tree still
requires a worklist/recursion.

In addition to presumably being faster, this change slightly reduces
the number of non-comment lines of code in `src/compact`.

naga/src/compact/expressions.rs

@@ -1,8 +1,7 @@
 use super::{HandleMap, HandleSet, ModuleMap};
-use crate::arena::{Arena, Handle, UniqueArena};
+use crate::arena::{Arena, Handle};
 
 pub struct ExpressionTracer<'tracer> {
-    pub types: &'tracer UniqueArena<crate::Type>,
     pub constants: &'tracer Arena<crate::Constant>,
 
     /// The arena in which we are currently tracing expressions.
@@ -21,34 +20,51 @@ pub struct ExpressionTracer<'tracer> {
     /// the module's constant expression arena.
     pub expressions_used: &'tracer mut HandleSet<crate::Expression>,
 
-    /// The constant expression arena and its used map, if we haven't
-    /// switched to tracing constant expressions yet.
-    pub const_expressions: Option<(
-        &'tracer Arena<crate::Expression>,
-        &'tracer mut HandleSet<crate::Expression>,
-    )>,
+    /// The used set for the module's `const_expressions` arena.
+    ///
+    /// If `None`, we are already tracing the constant expressions,
+    /// and `expressions_used` already refers to their handle set.
+    pub const_expressions_used: Option<&'tracer mut HandleSet<crate::Expression>>,
 }
 
 impl<'tracer> ExpressionTracer<'tracer> {
-    pub fn trace_expression(&mut self, expr: Handle<crate::Expression>) {
+    /// Propagate usage through `self.expressions`, starting with `self.expressions_used`.
+    ///
+    /// Treat `self.expressions_used` as the initial set of "known
+    /// live" expressions, and follow through to identify all
+    /// transitively used expressions.
+    ///
+    /// Mark types, constants, and constant expressions used directly
+    /// by `self.expressions` as used. Items used indirectly are not
+    /// marked.
+    ///
+    /// [fe]: crate::Function::expressions
+    /// [ce]: crate::Module::const_expressions
+    pub fn trace_expressions(&mut self) {
         log::trace!(
             "entering trace_expression of {}",
-            if self.const_expressions.is_some() {
+            if self.const_expressions_used.is_some() {
                 "function expressions"
             } else {
                 "const expressions"
             }
         );
-        let mut work_list = vec![expr];
-        while let Some(expr) = work_list.pop() {
-            // If we've already seen this expression, no need to trace further.
-            if !self.expressions_used.insert(expr) {
+
+        // We don't need recursion or a work list. Because an
+        // expression may only refer to other expressions that precede
+        // it in the arena, it suffices to make a single pass over the
+        // arena from back to front, marking the referents of used
+        // expressions as used themselves.
+        for (handle, expr) in self.expressions.iter().rev() {
+            // If this expression isn't used, it doesn't matter what it uses.
+            if !self.expressions_used.contains(handle) {
                 continue;
             }
+
             log::trace!("tracing new expression {:?}", expr);
 
             use crate::Expression as Ex;
-            match self.expressions[expr] {
+            match *expr {
                 // Expressions that do not contain handles that need to be traced.
                 Ex::Literal(_)
                 | Ex::FunctionArgument(_)
@@ -59,24 +75,34 @@ impl<'tracer> ExpressionTracer<'tracer> {
 
                 Ex::Constant(handle) => {
                     self.constants_used.insert(handle);
-                    let constant = &self.constants[handle];
-                    self.trace_type(constant.ty);
-                    self.trace_const_expression(constant.init);
+                    // Constants and expressions are mutually recursive, which
+                    // complicates our nice one-pass algorithm. However, since
+                    // constants don't refer to each other, we can get around
+                    // this by looking *through* each constant and marking its
+                    // initializer as used. Since `expr` refers to the constant,
+                    // and the constant refers to the initializer, it must
+                    // precede `expr` in the arena.
+                    let init = self.constants[handle].init;
+                    match self.const_expressions_used {
+                        Some(ref mut used) => used.insert(init),
+                        None => self.expressions_used.insert(init),
+                    }
                 }
-                Ex::ZeroValue(ty) => self.trace_type(ty),
+                Ex::ZeroValue(ty) => self.types_used.insert(ty),
                 Ex::Compose { ty, ref components } => {
-                    self.trace_type(ty);
-                    work_list.extend(components);
+                    self.types_used.insert(ty);
+                    self.expressions_used
+                        .insert_iter(components.iter().cloned());
                 }
-                Ex::Access { base, index } => work_list.extend([base, index]),
-                Ex::AccessIndex { base, index: _ } => work_list.push(base),
-                Ex::Splat { size: _, value } => work_list.push(value),
+                Ex::Access { base, index } => self.expressions_used.insert_iter([base, index]),
+                Ex::AccessIndex { base, index: _ } => self.expressions_used.insert(base),
+                Ex::Splat { size: _, value } => self.expressions_used.insert(value),
                 Ex::Swizzle {
                     size: _,
                     vector,
                     pattern: _,
-                } => work_list.push(vector),
-                Ex::Load { pointer } => work_list.push(pointer),
+                } => self.expressions_used.insert(vector),
+                Ex::Load { pointer } => self.expressions_used.insert(pointer),
                 Ex::ImageSample {
                     image,
                     sampler,
@@ -87,20 +113,20 @@ impl<'tracer> ExpressionTracer<'tracer> {
                     ref level,
                     depth_ref,
                 } => {
-                    work_list.push(image);
-                    work_list.push(sampler);
-                    work_list.push(coordinate);
-                    work_list.extend(array_index);
-                    if let Some(offset) = offset {
-                        self.trace_const_expression(offset);
+                    self.expressions_used
+                        .insert_iter([image, sampler, coordinate]);
+                    self.expressions_used.insert_iter(array_index);
+                    match self.const_expressions_used {
+                        Some(ref mut used) => used.insert_iter(offset),
+                        None => self.expressions_used.insert_iter(offset),
                     }
                     use crate::SampleLevel as Sl;
                     match *level {
                         Sl::Auto | Sl::Zero => {}
-                        Sl::Exact(expr) | Sl::Bias(expr) => work_list.push(expr),
-                        Sl::Gradient { x, y } => work_list.extend([x, y]),
+                        Sl::Exact(expr) | Sl::Bias(expr) => self.expressions_used.insert(expr),
+                        Sl::Gradient { x, y } => self.expressions_used.insert_iter([x, y]),
                     }
-                    work_list.extend(depth_ref);
+                    self.expressions_used.insert_iter(depth_ref);
                 }
                 Ex::ImageLoad {
                     image,
@@ -109,95 +135,64 @@ impl<'tracer> ExpressionTracer<'tracer> {
                     sample,
                     level,
                 } => {
-                    work_list.push(image);
-                    work_list.push(coordinate);
-                    work_list.extend(array_index);
-                    work_list.extend(sample);
-                    work_list.extend(level);
+                    self.expressions_used.insert(image);
+                    self.expressions_used.insert(coordinate);
+                    self.expressions_used.insert_iter(array_index);
+                    self.expressions_used.insert_iter(sample);
+                    self.expressions_used.insert_iter(level);
                 }
                 Ex::ImageQuery { image, ref query } => {
-                    work_list.push(image);
+                    self.expressions_used.insert(image);
                     use crate::ImageQuery as Iq;
                     match *query {
-                        Iq::Size { level } => work_list.extend(level),
+                        Iq::Size { level } => self.expressions_used.insert_iter(level),
                         Iq::NumLevels | Iq::NumLayers | Iq::NumSamples => {}
                     }
                 }
-                Ex::Unary { op: _, expr } => work_list.push(expr),
-                Ex::Binary { op: _, left, right } => work_list.extend([left, right]),
+                Ex::Unary { op: _, expr } => self.expressions_used.insert(expr),
+                Ex::Binary { op: _, left, right } => {
+                    self.expressions_used.insert_iter([left, right]);
+                }
                 Ex::Select {
                     condition,
                     accept,
                     reject,
-                } => work_list.extend([condition, accept, reject]),
+                } => self
+                    .expressions_used
+                    .insert_iter([condition, accept, reject]),
                 Ex::Derivative {
                     axis: _,
                     ctrl: _,
                     expr,
-                } => work_list.push(expr),
-                Ex::Relational { fun: _, argument } => work_list.push(argument),
+                } => self.expressions_used.insert(expr),
+                Ex::Relational { fun: _, argument } => self.expressions_used.insert(argument),
                 Ex::Math {
                     fun: _,
                     arg,
                     arg1,
                     arg2,
                     arg3,
                 } => {
-                    work_list.push(arg);
-                    work_list.extend(arg1);
-                    work_list.extend(arg2);
-                    work_list.extend(arg3);
+                    self.expressions_used.insert(arg);
+                    self.expressions_used.insert_iter(arg1);
+                    self.expressions_used.insert_iter(arg2);
+                    self.expressions_used.insert_iter(arg3);
                 }
                 Ex::As {
                     expr,
                     kind: _,
                     convert: _,
-                } => work_list.push(expr),
-                Ex::AtomicResult { ty, comparison: _ } => self.trace_type(ty),
-                Ex::WorkGroupUniformLoadResult { ty } => self.trace_type(ty),
-                Ex::ArrayLength(expr) => work_list.push(expr),
+                } => self.expressions_used.insert(expr),
+                Ex::AtomicResult { ty, comparison: _ } => self.types_used.insert(ty),
+                Ex::WorkGroupUniformLoadResult { ty } => self.types_used.insert(ty),
+                Ex::ArrayLength(expr) => self.expressions_used.insert(expr),
                 Ex::RayQueryGetIntersection {
                     query,
                     committed: _,
-                } => work_list.push(query),
+                } => self.expressions_used.insert(query),
             }
         }
     }
-
-    fn trace_type(&mut self, ty: Handle<crate::Type>) {
-        let mut types_used = super::types::TypeTracer {
-            types: self.types,
-            types_used: self.types_used,
-        };
-        types_used.trace_type(ty);
-    }
-
-    pub fn as_const_expression(&mut self) -> ExpressionTracer {
-        match self.const_expressions {
-            Some((ref mut exprs, ref mut exprs_used)) => ExpressionTracer {
-                expressions: exprs,
-                expressions_used: exprs_used,
-                types: self.types,
-                constants: self.constants,
-                types_used: self.types_used,
-                constants_used: self.constants_used,
-                const_expressions: None,
-            },
-            None => ExpressionTracer {
-                types: self.types,
-                constants: self.constants,
-                expressions: self.expressions,
-                types_used: self.types_used,
-                constants_used: self.constants_used,
-                expressions_used: self.expressions_used,
-                const_expressions: None,
-            },
-        }
-    }
-
-    fn trace_const_expression(&mut self, const_expr: Handle<crate::Expression>) {
-        self.as_const_expression().trace_expression(const_expr);
-    }
 }
 
 impl ModuleMap {

naga/src/compact/functions.rs

@@ -1,10 +1,9 @@
 use super::handle_set_map::HandleSet;
 use super::{FunctionMap, ModuleMap};
-use crate::arena::Handle;
 
 pub struct FunctionTracer<'a> {
-    pub module: &'a crate::Module,
     pub function: &'a crate::Function,
+    pub constants: &'a crate::Arena<crate::Constant>,
 
     pub types_used: &'a mut HandleSet<crate::Type>,
     pub constants_used: &'a mut HandleSet<crate::Constant>,
@@ -17,57 +16,43 @@ pub struct FunctionTracer<'a> {
 impl<'a> FunctionTracer<'a> {
     pub fn trace(&mut self) {
         for argument in self.function.arguments.iter() {
-            self.trace_type(argument.ty);
+            self.types_used.insert(argument.ty);
         }
 
         if let Some(ref result) = self.function.result {
-            self.trace_type(result.ty);
+            self.types_used.insert(result.ty);
         }
 
         for (_, local) in self.function.local_variables.iter() {
-            self.trace_type(local.ty);
+            self.types_used.insert(local.ty);
             if let Some(init) = local.init {
-                self.trace_expression(init);
+                self.expressions_used.insert(init);
             }
         }
 
         // Treat named expressions as alive, for the sake of our test suite,
         // which uses `let blah = expr;` to exercise lots of things.
-        for (value, _name) in &self.function.named_expressions {
-            self.trace_expression(*value);
+        for (&value, _name) in &self.function.named_expressions {
+            self.expressions_used.insert(value);
         }
 
         self.trace_block(&self.function.body);
-    }
-
-    pub fn trace_type(&mut self, ty: Handle<crate::Type>) {
-        self.as_type().trace_type(ty)
-    }
 
-    pub fn trace_expression(&mut self, expr: Handle<crate::Expression>) {
-        self.as_expression().trace_expression(expr);
-    }
-
-    fn as_type(&mut self) -> super::types::TypeTracer {
-        super::types::TypeTracer {
-            types: &self.module.types,
-            types_used: self.types_used,
-        }
+        // Given that `trace_block` has marked the expressions used
+        // directly by statements, walk the arena to find all
+        // expressions used, directly or indirectly.
+        self.as_expression().trace_expressions();
     }
 
     fn as_expression(&mut self) -> super::expressions::ExpressionTracer {
         super::expressions::ExpressionTracer {
-            types: &self.module.types,
-            constants: &self.module.constants,
+            constants: self.constants,
             expressions: &self.function.expressions,
 
             types_used: self.types_used,
             constants_used: self.constants_used,
             expressions_used: &mut self.expressions_used,
-            const_expressions: Some((
-                &self.module.const_expressions,
-                &mut self.const_expressions_used,
-            )),
+            const_expressions_used: Some(&mut self.const_expressions_used),
         }
     }
 }

naga/src/compact/handle_set_map.rs

@@ -26,13 +26,23 @@ impl<T> HandleSet<T> {
     }
 
     /// Add `handle` to the set.
-    ///
-    /// Return `true` if the handle was not already in the set. In
-    /// other words, return true if it was newly inserted.
-    pub fn insert(&mut self, handle: Handle<T>) -> bool {
+    pub fn insert(&mut self, handle: Handle<T>) {
+        // Note that, oddly, `Handle::index` does not return a 1-based
+        // `Index`, but rather a zero-based `usize`.
+        self.members.insert(handle.index());
+    }
+
+    /// Add handles from `iter` to the set.
+    pub fn insert_iter(&mut self, iter: impl IntoIterator<Item = Handle<T>>) {
+        for handle in iter {
+            self.insert(handle);
+        }
+    }
+
+    pub fn contains(&self, handle: Handle<T>) -> bool {
         // Note that, oddly, `Handle::index` does not return a 1-based
         // `Index`, but rather a zero-based `usize`.
-        self.members.insert(handle.index())
+        self.members.contains(handle.index())
     }
 }
 
@@ -148,6 +158,8 @@ impl<T: 'static> HandleMap<T> {
                 // Build a zero-based end-exclusive range, given one-based handle indices.
                 compacted = first1.get() - 1..last1.get();
             } else {
+                // The range contains only a single live handle, which
+                // we identified with the first `find_map` call.
                 compacted = first1.get() - 1..first1.get();
             }
         } else {