[wgsl-in] Remove unused expressions from Module::const_expressions.

CHANGELOG.md

@@ -58,6 +58,11 @@ For naga changelogs at or before v0.14.0. See [naga's changelog](naga/CHANGELOG.
 
 - Fix issue where local variables were sometimes using variable names from previous functions.
 
+### Changes
+#### Naga
+
+- [wgsl-in] Omit unused constant expressions from the Module's `const_expressions` arena, in preparation for supporting abstract types. By @jimblandy in [#4648](https://github.com/gfx-rs/wgpu/pull/4648).
+
 ## v0.18.0 (2023-10-25)
 
 ### Desktop OpenGL 3.3+ Support on Windows

naga/src/arena.rs

@@ -424,7 +424,7 @@ impl<T> Arena<T> {
         Ok(())
     }
 
-    #[cfg(feature = "compact")]
+    #[cfg(any(feature = "compact", feature = "wgsl-in"))]
     pub(crate) fn retain_mut<P>(&mut self, mut predicate: P)
     where
         P: FnMut(Handle<T>, &mut T) -> bool,

naga/src/front/wgsl/lower/compact.rs

@@ -0,0 +1,160 @@
+use crate::Handle;
+use bit_set::BitSet;
+
+/// Remove all unused expressions from `module.const_expressions`.
+///
+/// Abstract values make extensive use of 64-bit literals, which the
+/// validator should never see.
+pub fn compact(module: &mut crate::Module) {
+    // Trace which expressions in `module.const_expressions` are actually used.
+    let used = trace(module);
+
+    // Assuming unused expressions are squeezed out of the arena,
+    // compute a map from pre-squeeze indices to post-squeeze
+    // expression handles.
+    let mut next = 0;
+    let compacted: Vec<Option<Handle<crate::Expression>>> = module
+        .const_expressions
+        .iter()
+        .map(|(handle, _)| {
+            if used.contains(handle.index()) {
+                let index = next;
+                next += 1;
+                std::num::NonZeroU32::new(index + 1).map(Handle::new)
+            } else {
+                None
+            }
+        })
+        .collect();
+
+    adjust(module, &compacted);
+}
+
+fn trace(module: &crate::Module) -> BitSet {
+    let mut used = BitSet::new();
+
+    // Note uses by global constants.
+    for (_, constant) in module.constants.iter() {
+        used.insert(constant.init.index());
+    }
+
+    // Note uses by global variable initializers.
+    for (_, global) in module.global_variables.iter() {
+        if let Some(init) = global.init {
+            used.insert(init.index());
+        }
+    }
+
+    // Note uses by functions' expressions.
+    for (_, fun) in module.functions.iter() {
+        trace_function(fun, &mut used);
+    }
+
+    // Note uses by entry points' expressions.
+    for entry_point in &module.entry_points {
+        trace_function(&entry_point.function, &mut used);
+    }
+
+    // Note transitive uses by other used constant expressions.
+    //
+    // Since we know that expressions only refer to other expressions
+    // appearing before them in the arena, we can do this without a
+    // stack by making a single pass from back to front.
+    for (handle, expr) in module.const_expressions.iter().rev() {
+        if used.contains(handle.index()) {
+            match *expr {
+                crate::Expression::Compose { ref components, .. } => {
+                    for component in components {
+                        used.insert(component.index());
+                    }
+                }
+                crate::Expression::Splat { value, .. } => {
+                    used.insert(value.index());
+                }
+                _ => {}
+            }
+        }
+    }
+
+    used
+}
+
+fn trace_function(function: &crate::Function, used: &mut BitSet) {
+    for (_, expr) in function.expressions.iter() {
+        match *expr {
+            crate::Expression::ImageSample {
+                offset: Some(offset),
+                ..
+            } => {
+                used.insert(offset.index());
+            }
+            _ => {}
+        }
+    }
+}
+
+fn adjust(module: &mut crate::Module, compacted: &[Option<Handle<crate::Expression>>]) {
+    // Remove unused expressions from the constant arena,
+    // and adjust the handles in retained expressions.
+    module.const_expressions.retain_mut(|handle, expr| {
+        match compacted[handle.index()] {
+            Some(_) => {
+                // This expression is used, and thus its handles are worth adjusting.
+                match *expr {
+                    crate::Expression::Compose {
+                        ref mut components, ..
+                    } => {
+                        for component in components {
+                            *component = compacted[component.index()].unwrap();
+                        }
+                    }
+                    crate::Expression::Splat { ref mut value, .. } => {
+                        *value = compacted[value.index()].unwrap();
+                    }
+                    _ => {}
+                }
+                true
+            }
+            None => false,
+        }
+    });
+
+    // Adjust uses by global constants.
+    for (_, constant) in module.constants.iter_mut() {
+        constant.init = compacted[constant.init.index()].unwrap();
+    }
+
+    // Adjust uses by global variable initializers.
+    for (_, global) in module.global_variables.iter_mut() {
+        if let Some(ref mut init) = global.init {
+            *init = compacted[init.index()].unwrap();
+        }
+    }
+
+    // Adjust uses by functions' expressions.
+    for (_, fun) in module.functions.iter_mut() {
+        adjust_function(fun, compacted);
+    }
+
+    // Adjust uses by entry points' expressions.
+    for entry_point in &mut module.entry_points {
+        adjust_function(&mut entry_point.function, compacted);
+    }
+}
+
+fn adjust_function(
+    function: &mut crate::Function,
+    compacted: &[Option<Handle<crate::Expression>>],
+) {
+    for (_, expr) in function.expressions.iter_mut() {
+        match *expr {
+            crate::Expression::ImageSample {
+                offset: Some(ref mut offset),
+                ..
+            } => {
+                *offset = compacted[offset.index()].unwrap();
+            }
+            _ => {}
+        }
+    }
+}

naga/src/front/wgsl/lower/mod.rs

@@ -11,6 +11,7 @@ use crate::proc::{
 };
 use crate::{Arena, FastHashMap, FastIndexMap, Handle, Span};
 
+mod compact;
 mod construction;
 
 /// Resolves the inner type of a given expression.
@@ -971,6 +972,8 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
             }
         }
 
+        compact::compact(&mut module);
+
         Ok(module)
     }
 

naga/tests/out/ir/access.ron

@@ -412,32 +412,6 @@
                 6,
             ],
         ),
-        Literal(I32(8)),
-        Literal(I32(2)),
-        Literal(I32(10)),
-        Literal(I32(2)),
-        Literal(I32(0)),
-        Literal(I32(0)),
-        Literal(I32(0)),
-        Literal(I32(1)),
-        Literal(I32(0)),
-        Literal(I32(2)),
-        Literal(I32(2)),
-        Literal(I32(0)),
-        Literal(I32(3)),
-        Literal(I32(2)),
-        Literal(I32(2)),
-        Literal(I32(10)),
-        Literal(I32(5)),
-        Literal(I32(5)),
-        Literal(I32(10)),
-        Literal(I32(5)),
-        Literal(I32(0)),
-        Literal(I32(2)),
-        Literal(I32(2)),
-        Literal(I32(2)),
-        Literal(I32(2)),
-        Literal(I32(1)),
     ],
     functions: [
         (

naga/tests/out/ir/collatz.ron

@@ -58,11 +58,7 @@
             init: None,
         ),
     ],
-    const_expressions: [
-        Literal(I32(0)),
-        Literal(I32(0)),
-        Literal(I32(1)),
-    ],
+    const_expressions: [],
     functions: [
         (
             name: Some("collatz_iterations"),