Repository: gfx-rs/wgpu (language: rust)

Issue:
[naga] compaction may miss used constants' types
I think part of the compaction code is either wrong or redundant.

When `compact::compact` marks a name constants as used by definition, it also marks its initialization expression as used. But then later in that function, it marks the types of all used constants as used. I'm not sure why these two cases, type and initialization expression, should be handled differently. It seems to me that either the type marking can be moved up (removing that entire loop altogether) or the initializer marking must be moved down, because new constants are getting marked and we'd better include their initialization expressions.

Whatever the case, overrides deserve similar scrutiny.

(I'm not sure this is actually a bug, but there's an asymmetry to `naga::compact::compact` that seems wrong, and I'm afraid I'll forget about this doubt if I don't write it down.)


Discussion / hints:
We have this comment:

```rust
    // Constants' initializers are taken care of already, because
    // expression tracing sees through constants. But we still need to
    // note type usage.
```

I'm not sure it's right.

Fix patch:
diff --git a/naga/src/compact/expressions.rs b/naga/src/compact/expressions.rs
index df4f173ce1..7ddb46619e 100644
--- a/naga/src/compact/expressions.rs
+++ b/naga/src/compact/expressions.rs
@@ -94,15 +94,18 @@ impl ExpressionTracer<'_> {
             // it.
             Ex::Constant(handle) => {
                 self.constants_used.insert(handle);
-                let init = self.constants[handle].init;
+                let constant = &self.constants[handle];
+                self.types_used.insert(constant.ty);
                 match self.global_expressions_used {
-                    Some(ref mut used) => used.insert(init),
-                    None => self.expressions_used.insert(init),
+                    Some(ref mut used) => used.insert(constant.init),
+                    None => self.expressions_used.insert(constant.init),
                 };
             }
             Ex::Override(handle) => {
                 self.overrides_used.insert(handle);
-                if let Some(init) = self.overrides[handle].init {
+                let r#override = &self.overrides[handle];
+                self.types_used.insert(r#override.ty);
+                if let Some(init) = r#override.init {
                     match self.global_expressions_used {
                         Some(ref mut used) => used.insert(init),
                         None => self.expressions_used.insert(init),
diff --git a/naga/src/compact/mod.rs b/naga/src/compact/mod.rs
index 930f7a39c9..c9a7c93977 100644
--- a/naga/src/compact/mod.rs
+++ b/naga/src/compact/mod.rs
@@ -10,12 +10,20 @@ use handle_set_map::HandleMap;
 
 /// Remove unused types, expressions, and constants from `module`.
 ///
-/// Assuming that all globals, named constants, special types,
-/// functions and entry points in `module` are used, determine which
-/// types, constants, and expressions (both function-local and global
+/// Assume that the following are used by definition:
+/// - global variables
+/// - named constants and overrides
+/// - special types
+/// - functions and entry points, and within those:
+///     - arguments
+///     - local variables
+///     - named expressions
+///
+/// Given those assumptions, determine which types, constants,
+/// overrides, and expressions (both function-local and global
 /// constant expressions) are actually used, and remove the rest,
-/// adjusting all handles as necessary. The result should be a module
-/// functionally identical to the original.
+/// adjusting all handles as necessary. The result should always be a
+/// module functionally identical to the original.
 ///
 /// This may be useful to apply to modules generated in the snapshot
 /// tests. Our backends often generate temporary names based on handle
@@ -28,8 +36,62 @@ use handle_set_map::HandleMap;
 ///
 /// # Panics
 ///
-/// If `module` has not passed validation, this may panic.
+/// If `module` would not pass validation, this may panic.
 pub fn compact(module: &mut crate::Module) {
+    // The trickiest part of compaction is determining what is used and what is
+    // not. Once we have computed that correctly, it's easy enough to call
+    // `retain_mut` on each arena, drop unused elements, and fix up the handles
+    // in what's left.
+    //
+    // Some arenas' contents are considered used by definition, like
+    // `Module::global_variables` and `Module::functions`, so those are never
+    // compacted.
+    //
+    // But for every compactable arena in a `Module`, whether global to the
+    // `Module` or local to a function or entry point, the `ModuleTracer` type
+    // holds a bitmap indicating which elements of that arena are used. Our task
+    // is to populate those bitmaps correctly.
+    //
+    // First, we mark everything that is considered used by definition, as
+    // described in this function's documentation.
+    //
+    // Since functions and entry points are considered used by definition, we
+    // traverse their statement trees, and mark the referents of all handles
+    // appearing in those statements as used.
+    //
+    // Once we've marked which elements of an arena are referred to directly by
+    // handles elsewhere (for example, which of a function's expressions are
+    // referred to by handles in its body statements), we can mark all the other
+    // arena elements that are used indirectly in a single pass, traversing the
+    // arena from back to front. Since Naga allows arena elements to refer only
+    // to prior elements, we know that by the time we reach an element, all
+    // other elements that could possibly refer to it have already been visited.
+    // Thus, if the present element has not been marked as used, then it is
+    // definitely unused, and compaction can remove it. Otherwise, the element
+    // is used and must be retained, so we must mark everything it refers to.
+    //
+    // The final step is to mark the global expressions and types, which must be
+    // traversed simultaneously; see `ModuleTracer::type_expression_tandem`'s
+    // documentation for details.
+    //
+    // # A definition and a rule of thumb
+    //
+    // In this module, to "trace" something is to mark everything else it refers
+    // to as used, on the assumption that the thing itself is used. For example,
+    // to trace an `Expression` is to mark its subexpressions as used, as well
+    // as any types, constants, overrides, etc. that it refers to. This is what
+    // `ExpressionTracer::trace_expression` does.
+    //
+    // Given that we we want to visit each thing only once (to keep compaction
+    // linear in the size of the module), this definition of "trace" implies
+    // that things that are not "used by definition" must be marked as used
+    // *before* we trace them.
+    //
+    // Thus, whenever you are marking something as used, it's a good idea to ask
+    // yourself how you know that thing will be traced in the future. If you're
+    // not sure, then you could be marking it too late to be noticed. The thing
+    // itself will be retained by compaction, but since it will not be traced,
+    // anything it refers to could be compacted away.
     let mut module_tracer = ModuleTracer::new(module);
 
     // We treat all globals as used by definition.
@@ -54,6 +116,7 @@ pub fn compact(module: &mut crate::Module) {
         if constant.name.is_some() {
             log::trace!("tracing constant {:?}", constant.name.as_ref().unwrap());
             module_tracer.constants_used.insert(handle);
+            module_tracer.types_used.insert(constant.ty);
             module_tracer.global_expressions_used.insert(constant.init);
         }
     }
@@ -64,6 +127,7 @@ pub fn compact(module: &mut crate::Module) {
         if r#override.name.is_some() {
             log::trace!("tracing override {:?}", r#override.name.as_ref().unwrap());
             module_tracer.overrides_used.insert(handle);
+            module_tracer.types_used.insert(r#override.ty);
             if let Some(init) = r#override.init {
                 module_tracer.global_expressions_used.insert(init);
             }
@@ -108,24 +172,6 @@ pub fn compact(module: &mut crate::Module) {
         })
         .collect();
 
-    // Constants' initializers are taken care of already, because
-    // expression tracing sees through constants. But we still need to
-    // note type usage.
-    for (handle, constant) in module.constants.iter() {
-        if module_tracer.constants_used.contains(handle) {
-            module_tracer.types_used.insert(constant.ty);
-        }
-    }
-
-    // Overrides' initializers are taken care of already, because
-    // expression tracing sees through overrides. But we still need to
-    // note type usage.
-    for (handle, r#override) in module.overrides.iter() {
-        if module_tracer.overrides_used.contains(handle) {
-            module_tracer.types_used.insert(r#override.ty);
-        }
-    }
-
     // Treat all named types as used.
     for (handle, ty) in module.types.iter() {
         log::trace!("tracing type {:?}, name {:?}", handle, ty.name);
@@ -880,3 +926,147 @@ fn local_expression_override() {
     compact(&mut module);
     assert!(validator.validate(&module).is_ok());
 }
+
+#[test]
+fn unnamed_constant_type() {
+    let mut module = crate::Module::default();
+    let nowhere = crate::Span::default();
+
+    // This type is used only by the unnamed constant.
+    let ty_u32 = module.types.insert(
+        crate::Type {
+            name: None,
+            inner: crate::TypeInner::Scalar(crate::Scalar::U32),
+        },
+        nowhere,
+    );
+
+    // This type is used by the named constant.
+    let ty_vec_u32 = module.types.insert(
+        crate::Type {
+            name: None,
+            inner: crate::TypeInner::Vector {
+                size: crate::VectorSize::Bi,
+                scalar: crate::Scalar::U32,
+            },
+        },
+        nowhere,
+    );
+
+    let unnamed_init = module
+        .global_expressions
+        .append(crate::Expression::Literal(crate::Literal::U32(0)), nowhere);
+
+    let unnamed_constant = module.constants.append(
+        crate::Constant {
+            name: None,
+            ty: ty_u32,
+            init: unnamed_init,
+        },
+        nowhere,
+    );
+
+    // The named constant is initialized using a Splat expression, to
+    // give the named constant a type distinct from the unnamed
+    // constant's.
+    let unnamed_constant_expr = module
+        .global_expressions
+        .append(crate::Expression::Constant(unnamed_constant), nowhere);
+    let named_init = module.global_expressions.append(
+        crate::Expression::Splat {
+            size: crate::VectorSize::Bi,
+            value: unnamed_constant_expr,
+        },
+        nowhere,
+    );
+
+    let _named_constant = module.constants.append(
+        crate::Constant {
+            name: Some("totally_named".to_string()),
+            ty: ty_vec_u32,
+            init: named_init,
+        },
+        nowhere,
+    );
+
+    let mut validator = super::valid::Validator::new(
+        super::valid::ValidationFlags::all(),
+        super::valid::Capabilities::all(),
+    );
+
+    assert!(validator.validate(&module).is_ok());
+    compact(&mut module);
+    assert!(validator.validate(&module).is_ok());
+}
+
+#[test]
+fn unnamed_override_type() {
+    let mut module = crate::Module::default();
+    let nowhere = crate::Span::default();
+
+    // This type is used only by the unnamed override.
+    let ty_u32 = module.types.insert(
+        crate::Type {
+            name: None,
+            inner: crate::TypeInner::Scalar(crate::Scalar::U32),
+        },
+        nowhere,
+    );
+
+    // This type is used by the named override.
+    let ty_i32 = module.types.insert(
+        crate::Type {
+            name: None,
+            inner: crate::TypeInner::Scalar(crate::Scalar::I32),
+        },
+        nowhere,
+    );
+
+    let unnamed_init = module
+        .global_expressions
+        .append(crate::Expression::Literal(crate::Literal::U32(0)), nowhere);
+
+    let unnamed_override = module.overrides.append(
+        crate::Override {
+            name: None,
+            id: Some(42),
+            ty: ty_u32,
+            init: Some(unnamed_init),
+        },
+        nowhere,
+    );
+
+    // The named override is initialized using a Splat expression, to
+    // give the named override a type distinct from the unnamed
+    // override's.
+    let unnamed_override_expr = module
+        .global_expressions
+        .append(crate::Expression::Override(unnamed_override), nowhere);
+    let named_init = module.global_expressions.append(
+        crate::Expression::As {
+            expr: unnamed_override_expr,
+            kind: crate::ScalarKind::Sint,
+            convert: None,
+        },
+        nowhere,
+    );
+
+    let _named_override = module.overrides.append(
+        crate::Override {
+            name: Some("totally_named".to_string()),
+            id: None,
+            ty: ty_i32,
+            init: Some(named_init),
+        },
+        nowhere,
+    );
+
+    let mut validator = super::valid::Validator::new(
+        super::valid::ValidationFlags::all(),
+        super::valid::Capabilities::all(),
+    );
+
+    assert!(validator.validate(&module).is_ok());
+    compact(&mut module);
+    assert!(validator.validate(&module).is_ok());
+}


Test patch: