Type/Expression Tandem

CHANGELOG.md

@@ -324,6 +324,7 @@ By @ErichDonGubler in [#6456](https://github.com/gfx-rs/wgpu/pull/6456), [#6148]
 
 - Fix crash when a texture argument is missing. By @aedm in [#6486](https://github.com/gfx-rs/wgpu/pull/6486)
 - Emit an error in constant evaluation, rather than crash, in certain cases where `vecN` constructors have less than N arguments. By @ErichDonGubler in [#6508](https://github.com/gfx-rs/wgpu/pull/6508).
+- Fix a leak by ensuring that types that depend on expressions are correctly compacted. By @KentSlaney in [#6806](https://github.com/gfx-rs/wgpu/pull/6806).
 - Fix an error in template list matching `>=` in `a<b>=c`. By @KentSlaney in [#6898](https://github.com/gfx-rs/wgpu/pull/6898).
 - Correctly validate handles in override-sized array types. By @jimblandy in [#6882](https://github.com/gfx-rs/wgpu/pull/6882).
 - Clean up validation of `Statement::ImageStore`. By @jimblandy in [#6729](https://github.com/gfx-rs/wgpu-pull/6729).

naga/src/compact/expressions.rs

@@ -1,7 +1,14 @@
 use super::{HandleMap, HandleSet, ModuleMap};
-use crate::arena::{Arena, Handle};
+use crate::arena::{Arena, Handle, UniqueArena};
+use crate::compact::types::TypeTracer;
 
 pub struct ExpressionTracer<'tracer> {
+    /// The type arena used by `expressions`.
+    ///
+    /// Having this set as None will result in a shallow types_used set.
+    /// `Some` is used to ensure propagation through interleaved type
+    /// and expression dependencies.
+    pub types: Option<&'tracer UniqueArena<crate::Type>>,
     pub constants: &'tracer Arena<crate::Constant>,
 
     /// The arena in which we are currently tracing expressions.
@@ -28,6 +35,28 @@ pub struct ExpressionTracer<'tracer> {
 }
 
 impl ExpressionTracer<'_> {
+    fn types_used_insert(&mut self, used: Handle<crate::Type>) {
+        if self.types.is_some() {
+            self.trace_type(used);
+        }
+        self.types_used.insert(used);
+    }
+
+    fn trace_type(&mut self, tracing: Handle<crate::Type>) {
+        fn type_used(caller: &mut TypeTracer, used: Handle<crate::Type>) {
+            if caller.types_used.insert(used) {
+                caller.trace_type(&caller.types[used], type_used);
+            }
+        }
+
+        TypeTracer {
+            types: self.types.unwrap(),
+            types_used: self.types_used,
+            expressions_used: self.expressions_used,
+        }
+        .trace_type(&self.types.unwrap()[tracing], type_used)
+    }
+
     /// Propagate usage through `self.expressions`, starting with `self.expressions_used`.
     ///
     /// Treat `self.expressions_used` as the initial set of "known
@@ -62,168 +91,171 @@ impl ExpressionTracer<'_> {
             }
 
             log::trace!("tracing new expression {:?}", expr);
+            self.trace_expression(expr);
+        }
+    }
 
-            use crate::Expression as Ex;
-            match *expr {
-                // Expressions that do not contain handles that need to be traced.
-                Ex::Literal(_)
-                | Ex::FunctionArgument(_)
-                | Ex::GlobalVariable(_)
-                | Ex::LocalVariable(_)
-                | Ex::CallResult(_)
-                | Ex::SubgroupBallotResult
-                | Ex::RayQueryProceedResult => {}
+    pub fn trace_expression(&mut self, expr: &crate::Expression) {
+        use crate::Expression as Ex;
+        match *expr {
+            // Expressions that do not contain handles that need to be traced.
+            Ex::Literal(_)
+            | Ex::FunctionArgument(_)
+            | Ex::GlobalVariable(_)
+            | Ex::LocalVariable(_)
+            | Ex::CallResult(_)
+            | Ex::SubgroupBallotResult
+            | Ex::RayQueryProceedResult => {}
 
-                Ex::Constant(handle) => {
-                    self.constants_used.insert(handle);
-                    // 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.global_expressions_used {
-                        Some(ref mut used) => used.insert(init),
-                        None => self.expressions_used.insert(init),
-                    };
-                }
-                Ex::Override(_) => {
-                    // All overrides are considered used by definition. We mark
-                    // their types and initialization expressions as used in
-                    // `compact::compact`, so we have no more work to do here.
-                }
-                Ex::ZeroValue(ty) => {
-                    self.types_used.insert(ty);
-                }
-                Ex::Compose { ty, ref components } => {
-                    self.types_used.insert(ty);
-                    self.expressions_used
-                        .insert_iter(components.iter().cloned());
-                }
-                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: _,
-                } => {
-                    self.expressions_used.insert(vector);
-                }
-                Ex::Load { pointer } => {
-                    self.expressions_used.insert(pointer);
-                }
-                Ex::ImageSample {
-                    image,
-                    sampler,
-                    gather: _,
-                    coordinate,
-                    array_index,
-                    offset,
-                    ref level,
-                    depth_ref,
-                } => {
-                    self.expressions_used
-                        .insert_iter([image, sampler, coordinate]);
-                    self.expressions_used.insert_iter(array_index);
-                    match self.global_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) => {
-                            self.expressions_used.insert(expr);
-                        }
-                        Sl::Gradient { x, y } => self.expressions_used.insert_iter([x, y]),
-                    }
-                    self.expressions_used.insert_iter(depth_ref);
-                }
-                Ex::ImageLoad {
-                    image,
-                    coordinate,
-                    array_index,
-                    sample,
-                    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::Constant(handle) => {
+                self.constants_used.insert(handle);
+                // 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.global_expressions_used {
+                    Some(ref mut used) => used.insert(init),
+                    None => self.expressions_used.insert(init),
+                };
+            }
+            Ex::Override(_) => {
+                // All overrides are considered used by definition. We mark
+                // their types and initialization expressions as used in
+                // `compact::compact`, so we have no more work to do here.
+            }
+            Ex::ZeroValue(ty) => {
+                self.types_used_insert(ty);
+            }
+            Ex::Compose { ty, ref components } => {
+                self.types_used_insert(ty);
+                self.expressions_used
+                    .insert_iter(components.iter().cloned());
+            }
+            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: _,
+            } => {
+                self.expressions_used.insert(vector);
+            }
+            Ex::Load { pointer } => {
+                self.expressions_used.insert(pointer);
+            }
+            Ex::ImageSample {
+                image,
+                sampler,
+                gather: _,
+                coordinate,
+                array_index,
+                offset,
+                ref level,
+                depth_ref,
+            } => {
+                self.expressions_used
+                    .insert_iter([image, sampler, coordinate]);
+                self.expressions_used.insert_iter(array_index);
+                match self.global_expressions_used {
+                    Some(ref mut used) => used.insert_iter(offset),
+                    None => self.expressions_used.insert_iter(offset),
                 }
-                Ex::ImageQuery { image, ref query } => {
-                    self.expressions_used.insert(image);
-                    use crate::ImageQuery as Iq;
-                    match *query {
-                        Iq::Size { level } => self.expressions_used.insert_iter(level),
-                        Iq::NumLevels | Iq::NumLayers | Iq::NumSamples => {}
+                use crate::SampleLevel as Sl;
+                match *level {
+                    Sl::Auto | Sl::Zero => {}
+                    Sl::Exact(expr) | Sl::Bias(expr) => {
+                        self.expressions_used.insert(expr);
                     }
+                    Sl::Gradient { x, y } => self.expressions_used.insert_iter([x, y]),
                 }
-                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,
-                } => self
-                    .expressions_used
-                    .insert_iter([condition, accept, reject]),
-                Ex::Derivative {
-                    axis: _,
-                    ctrl: _,
-                    expr,
-                } => {
-                    self.expressions_used.insert(expr);
-                }
-                Ex::Relational { fun: _, argument } => {
-                    self.expressions_used.insert(argument);
-                }
-                Ex::Math {
-                    fun: _,
-                    arg,
-                    arg1,
-                    arg2,
-                    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: _,
-                } => {
-                    self.expressions_used.insert(expr);
-                }
-                Ex::ArrayLength(expr) => {
-                    self.expressions_used.insert(expr);
-                }
-                Ex::AtomicResult { ty, comparison: _ }
-                | Ex::WorkGroupUniformLoadResult { ty }
-                | Ex::SubgroupOperationResult { ty } => {
-                    self.types_used.insert(ty);
-                }
-                Ex::RayQueryGetIntersection {
-                    query,
-                    committed: _,
-                } => {
-                    self.expressions_used.insert(query);
+                self.expressions_used.insert_iter(depth_ref);
+            }
+            Ex::ImageLoad {
+                image,
+                coordinate,
+                array_index,
+                sample,
+                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 } => {
+                self.expressions_used.insert(image);
+                use crate::ImageQuery as Iq;
+                match *query {
+                    Iq::Size { level } => self.expressions_used.insert_iter(level),
+                    Iq::NumLevels | Iq::NumLayers | Iq::NumSamples => {}
                 }
             }
+            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,
+            } => self
+                .expressions_used
+                .insert_iter([condition, accept, reject]),
+            Ex::Derivative {
+                axis: _,
+                ctrl: _,
+                expr,
+            } => {
+                self.expressions_used.insert(expr);
+            }
+            Ex::Relational { fun: _, argument } => {
+                self.expressions_used.insert(argument);
+            }
+            Ex::Math {
+                fun: _,
+                arg,
+                arg1,
+                arg2,
+                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: _,
+            } => {
+                self.expressions_used.insert(expr);
+            }
+            Ex::ArrayLength(expr) => {
+                self.expressions_used.insert(expr);
+            }
+            Ex::AtomicResult { ty, comparison: _ }
+            | Ex::WorkGroupUniformLoadResult { ty }
+            | Ex::SubgroupOperationResult { ty } => {
+                self.types_used_insert(ty);
+            }
+            Ex::RayQueryGetIntersection {
+                query,
+                committed: _,
+            } => {
+                self.expressions_used.insert(query);
+            }
         }
     }
 }

naga/src/compact/functions.rs

@@ -46,6 +46,7 @@ impl FunctionTracer<'_> {
 
     fn as_expression(&mut self) -> super::expressions::ExpressionTracer {
         super::expressions::ExpressionTracer {
+            types: None,
             constants: self.constants,
             expressions: &self.function.expressions,