From c1fc5d6424c3188b6cbf39edaa86d8febdbf4082 Mon Sep 17 00:00:00 2001
From: Dzmitry Malyshau <kvark@fastmail.com>
Date: Fri, 12 Feb 2021 01:39:08 -0500
Subject: [PATCH] [msl] re-use expression results based on the analysis

---
 examples/convert.rs       |   5 +-
 src/back/glsl/mod.rs      |   2 +-
 src/back/msl/mod.rs       |  10 +-
 src/back/msl/writer.rs    | 197 ++++++++++++++++++++++++++++++++------
 src/front/global_usage.rs |   4 +-
 src/lib.rs                |   2 +
 src/proc/analyzer.rs      |  26 ++++-
 src/proc/interface.rs     |  10 +-
 tests/out/shadow.msl.snap |   7 +-
 tests/out/skybox.msl.snap |   5 +-
 tests/snapshots.rs        |  14 ++-
 11 files changed, 230 insertions(+), 52 deletions(-)

diff --git a/examples/convert.rs b/examples/convert.rs
index 82c03a7349..4cb419fd5c 100644
--- a/examples/convert.rs
+++ b/examples/convert.rs
@@ -147,7 +147,8 @@ fn main() {
     };
 
     // validate the IR
-    naga::proc::Validator::new()
+    #[cfg_attr(not(feature = "msl-out"), allow(unused_variables))]
+    let analysis = naga::proc::Validator::new()
         .validate(&module)
         .unwrap_pretty();
 
@@ -190,7 +191,7 @@ fn main() {
                 spirv_cross_compatibility: false,
                 binding_map,
             };
-            let (msl, _) = msl::write_string(&module, &options).unwrap();
+            let (msl, _) = msl::write_string(&module, &analysis, &options).unwrap();
             fs::write(&args[2], msl).unwrap();
         }
         #[cfg(feature = "spv-out")]
diff --git a/src/back/glsl/mod.rs b/src/back/glsl/mod.rs
index 4762b79759..528deed9f9 100644
--- a/src/back/glsl/mod.rs
+++ b/src/back/glsl/mod.rs
@@ -1908,7 +1908,7 @@ struct TextureMappingVisitor<'a> {
 }
 
 impl<'a> Visitor for TextureMappingVisitor<'a> {
-    fn visit_expr(&mut self, expr: &crate::Expression) {
+    fn visit_expr(&mut self, _: Handle<crate::Expression>, expr: &crate::Expression) {
         // We only care about `ImageSample` and `ImageLoad`
         //
         // Both `image` and `sampler` are `Expression::GlobalVariable` otherwise the module is
diff --git a/src/back/msl/mod.rs b/src/back/msl/mod.rs
index 53c8ddf38c..14dbaaed80 100644
--- a/src/back/msl/mod.rs
+++ b/src/back/msl/mod.rs
@@ -14,7 +14,11 @@ the output struct. If there is a structure in the outputs, and it contains any b
 we move them up to the root output structure that we define ourselves.
 !*/
 
-use crate::{arena::Handle, proc::ResolveError, FastHashMap};
+use crate::{
+    arena::Handle,
+    proc::{analyzer::Analysis, ResolveError},
+    FastHashMap,
+};
 use std::{
     io::{Error as IoError, Write},
     string::FromUtf8Error,
@@ -66,6 +70,7 @@ pub enum Error {
     UnexpectedSampleLevel(crate::SampleLevel),
     UnsupportedCall(String),
     UnsupportedDynamicArrayLength,
+    FeatureNotImplemented(String),
     /// The source IR is not valid.
     Validation,
 }
@@ -222,10 +227,11 @@ pub struct TranslationInfo {
 
 pub fn write_string(
     module: &crate::Module,
+    analysis: &Analysis,
     options: &Options,
 ) -> Result<(String, TranslationInfo), Error> {
     let mut w = writer::Writer::new(Vec::new());
-    let info = w.write(module, options)?;
+    let info = w.write(module, analysis, options)?;
     let string = String::from_utf8(w.finish())?;
     Ok((string, info))
 }
diff --git a/src/back/msl/writer.rs b/src/back/msl/writer.rs
index 8fb22119fc..c88f71b024 100644
--- a/src/back/msl/writer.rs
+++ b/src/back/msl/writer.rs
@@ -1,18 +1,23 @@
 use super::{keywords::RESERVED, Error, LocationMode, Options, TranslationInfo};
 use crate::{
     arena::Handle,
-    proc::{EntryPointIndex, NameKey, Namer, ResolveContext, Typifier},
+    proc::{
+        analyzer::{Analysis, FunctionInfo},
+        EntryPointIndex, Interface, NameKey, Namer, ResolveContext, Typifier, Visitor,
+    },
     FastHashMap,
 };
+use bit_set::BitSet;
 use std::{
     fmt::{Display, Error as FmtError, Formatter},
     io::Write,
-    iter,
+    iter, mem,
 };
 
 const NAMESPACE: &str = "metal";
 const INDENT: &str = "    ";
 
+#[derive(Clone)]
 struct Level(usize);
 impl Level {
     fn next(&self) -> Self {
@@ -64,8 +69,10 @@ impl<'a> TypedGlobalVariable<'a> {
 pub struct Writer<W> {
     out: W,
     names: FastHashMap<NameKey, String>,
+    named_expressions: BitSet,
     typifier: Typifier,
     namer: Namer,
+    temp_bake_handles: Vec<Handle<crate::Expression>>,
 }
 
 fn scalar_kind_string(kind: crate::ScalarKind) -> &'static str {
@@ -113,6 +120,39 @@ impl crate::StorageClass {
     }
 }
 
+struct BakeExpressionVisitor<'a> {
+    named_expressions: &'a mut BitSet,
+    bake_handles: &'a mut Vec<Handle<crate::Expression>>,
+    fun_info: &'a FunctionInfo,
+    exclude: Option<Handle<crate::Expression>>,
+}
+impl Visitor for BakeExpressionVisitor<'_> {
+    fn visit_expr(&mut self, handle: Handle<crate::Expression>, expr: &crate::Expression) {
+        use crate::Expression as E;
+        // filter out the expressions that don't need to bake
+        let min_ref_count = match *expr {
+            // The following expressions can be inlined nicely.
+            E::AccessIndex { .. }
+            | E::Constant(_)
+            | E::FunctionArgument(_)
+            | E::GlobalVariable(_)
+            | E::LocalVariable(_) => !0,
+            // Image sampling and function calling are nice to isolate
+            // into separate statements even when done only once.
+            E::ImageSample { .. } | E::ImageLoad { .. } | E::Call { .. } => 1,
+            // Bake only expressions referenced more than once.
+            _ => 2,
+        };
+
+        let modifier = if self.exclude == Some(handle) { 1 } else { 0 };
+        if self.fun_info[handle].ref_count - modifier >= min_ref_count
+            && self.named_expressions.insert(handle.index())
+        {
+            self.bake_handles.push(handle);
+        }
+    }
+}
+
 enum FunctionOrigin {
     Handle(Handle<crate::Function>),
     EntryPoint(EntryPointIndex),
@@ -124,14 +164,22 @@ struct ExpressionContext<'a> {
     module: &'a crate::Module,
 }
 
+struct StatementContext<'a> {
+    expression: ExpressionContext<'a>,
+    fun_info: &'a FunctionInfo,
+    return_value: Option<&'a str>,
+}
+
 impl<W: Write> Writer<W> {
     /// Creates a new `Writer` instance.
     pub fn new(out: W) -> Self {
         Writer {
             out,
             names: FastHashMap::default(),
+            named_expressions: BitSet::new(),
             typifier: Typifier::new(),
             namer: Namer::default(),
+            temp_bake_handles: Vec::new(),
         }
     }
 
@@ -216,6 +264,11 @@ impl<W: Write> Writer<W> {
         expr_handle: Handle<crate::Expression>,
         context: &ExpressionContext,
     ) -> Result<(), Error> {
+        if self.named_expressions.contains(expr_handle.index()) {
+            write!(self.out, "expr{}", expr_handle.index())?;
+            return Ok(());
+        }
+
         let expression = &context.function.expressions[expr_handle];
         log::trace!("expression {:?} = {:?}", expr_handle, expression);
         match *expression {
@@ -684,12 +737,76 @@ impl<W: Write> Writer<W> {
         Ok(())
     }
 
+    // Write down any required intermediate results
+    fn prepare_expression(
+        &mut self,
+        level: Level,
+        root_handle: Handle<crate::Expression>,
+        context: &StatementContext,
+        exclude_root: bool,
+    ) -> Result<(), Error> {
+        // set up the search
+        let mut interface = Interface {
+            expressions: &context.expression.function.expressions,
+            local_variables: &context.expression.function.local_variables,
+            visitor: BakeExpressionVisitor {
+                named_expressions: &mut self.named_expressions,
+                bake_handles: &mut self.temp_bake_handles,
+                fun_info: context.fun_info,
+                exclude: if exclude_root {
+                    Some(root_handle)
+                } else {
+                    None
+                },
+            },
+        };
+        // populate the bake handles
+        interface.traverse_expr(root_handle);
+        // bake
+        let mut temp_bake_handles = mem::replace(&mut self.temp_bake_handles, Vec::new());
+        for handle in temp_bake_handles.drain(..).rev() {
+            write!(self.out, "{}", level)?;
+            match self.typifier.get_handle(handle) {
+                Ok(ty_handle) => {
+                    let ty_name = &self.names[&NameKey::Type(ty_handle)];
+                    write!(self.out, "{}", ty_name)?;
+                }
+                Err(&crate::TypeInner::Scalar { kind, .. }) => {
+                    write!(self.out, "{}", scalar_kind_string(kind))?;
+                }
+                Err(&crate::TypeInner::Vector { size, kind, .. }) => {
+                    write!(
+                        self.out,
+                        "{}::{}{}",
+                        NAMESPACE,
+                        scalar_kind_string(kind),
+                        vector_size_string(size)
+                    )?;
+                }
+                Err(other) => {
+                    log::error!("Type {:?} isn't a known local", other);
+                    return Err(Error::FeatureNotImplemented("weird local type".to_string()));
+                }
+            }
+
+            //TODO: figure out the naming scheme that wouldn't collide with user names.
+            write!(self.out, " expr{} = ", handle.index())?;
+            // Make sure to temporarily unblock the expression before writing it down.
+            self.named_expressions.remove(handle.index());
+            self.put_expression(handle, &context.expression)?;
+            self.named_expressions.insert(handle.index());
+            writeln!(self.out, ";")?;
+        }
+
+        self.temp_bake_handles = temp_bake_handles;
+        Ok(())
+    }
+
     fn put_block(
         &mut self,
         level: Level,
         statements: &[crate::Statement],
-        context: &ExpressionContext,
-        return_value: Option<&str>,
+        context: &StatementContext,
     ) -> Result<(), Error> {
         for statement in statements {
             log::trace!("statement[{}] {:?}", level.0, statement);
@@ -697,7 +814,7 @@ impl<W: Write> Writer<W> {
                 crate::Statement::Block(ref block) => {
                     if !block.is_empty() {
                         writeln!(self.out, "{}{{", level)?;
-                        self.put_block(level.next(), block, context, return_value)?;
+                        self.put_block(level.next(), block, context)?;
                         writeln!(self.out, "{}}}", level)?;
                     }
                 }
@@ -706,13 +823,14 @@ impl<W: Write> Writer<W> {
                     ref accept,
                     ref reject,
                 } => {
+                    self.prepare_expression(level.clone(), condition, context, false)?;
                     write!(self.out, "{}if (", level)?;
-                    self.put_expression(condition, context)?;
+                    self.put_expression(condition, &context.expression)?;
                     writeln!(self.out, ") {{")?;
-                    self.put_block(level.next(), accept, context, return_value)?;
+                    self.put_block(level.next(), accept, context)?;
                     if !reject.is_empty() {
                         writeln!(self.out, "{}}} else {{", level)?;
-                        self.put_block(level.next(), reject, context, return_value)?;
+                        self.put_block(level.next(), reject, context)?;
                     }
                     writeln!(self.out, "{}}}", level)?;
                 }
@@ -721,20 +839,21 @@ impl<W: Write> Writer<W> {
                     ref cases,
                     ref default,
                 } => {
+                    self.prepare_expression(level.clone(), selector, context, false)?;
                     write!(self.out, "{}switch(", level)?;
-                    self.put_expression(selector, context)?;
+                    self.put_expression(selector, &context.expression)?;
                     writeln!(self.out, ") {{")?;
                     let lcase = level.next();
                     for case in cases.iter() {
                         writeln!(self.out, "{}case {}: {{", lcase, case.value)?;
-                        self.put_block(lcase.next(), &case.body, context, return_value)?;
+                        self.put_block(lcase.next(), &case.body, context)?;
                         if case.fall_through {
                             writeln!(self.out, "{}break;", lcase.next())?;
                         }
                         writeln!(self.out, "{}}}", lcase)?;
                     }
                     writeln!(self.out, "{}default: {{", lcase)?;
-                    self.put_block(lcase.next(), default, context, return_value)?;
+                    self.put_block(lcase.next(), default, context)?;
                     writeln!(self.out, "{}}}", lcase)?;
                     writeln!(self.out, "{}}}", level)?;
                 }
@@ -748,13 +867,13 @@ impl<W: Write> Writer<W> {
                         writeln!(self.out, "{}while(true) {{", level)?;
                         let lif = level.next();
                         writeln!(self.out, "{}if (!{}) {{", lif, gate_name)?;
-                        self.put_block(lif.next(), continuing, context, return_value)?;
+                        self.put_block(lif.next(), continuing, context)?;
                         writeln!(self.out, "{}}}", lif)?;
                         writeln!(self.out, "{}{} = false;", lif, gate_name)?;
                     } else {
                         writeln!(self.out, "{}while(true) {{", level)?;
                     }
-                    self.put_block(level.next(), body, context, return_value)?;
+                    self.put_block(level.next(), body, context)?;
                     writeln!(self.out, "{}}}", level)?;
                 }
                 crate::Statement::Break => {
@@ -766,8 +885,9 @@ impl<W: Write> Writer<W> {
                 crate::Statement::Return {
                     value: Some(expr_handle),
                 } => {
+                    self.prepare_expression(level.clone(), expr_handle, context, true)?;
                     write!(self.out, "{}return ", level)?;
-                    self.put_expression(expr_handle, context)?;
+                    self.put_expression(expr_handle, &context.expression)?;
                     writeln!(self.out, ";")?;
                 }
                 crate::Statement::Return { value: None } => {
@@ -775,25 +895,28 @@ impl<W: Write> Writer<W> {
                         self.out,
                         "{}return {};",
                         level,
-                        return_value.unwrap_or_default(),
+                        context.return_value.unwrap_or_default(),
                     )?;
                 }
                 crate::Statement::Kill => {
                     writeln!(self.out, "{}discard_fragment();", level)?;
                 }
                 crate::Statement::Store { pointer, value } => {
-                    //write!(self.out, "{}*", INDENT)?;
+                    self.prepare_expression(level.clone(), value, context, true)?;
                     write!(self.out, "{}", level)?;
-                    self.put_expression(pointer, context)?;
+                    self.put_expression(pointer, &context.expression)?;
                     write!(self.out, " = ")?;
-                    self.put_expression(value, context)?;
+                    self.put_expression(value, &context.expression)?;
                     writeln!(self.out, ";")?;
                 }
                 crate::Statement::Call {
                     function,
                     ref arguments,
                 } => {
-                    self.put_local_call(function, arguments, context)?;
+                    for &arg in arguments {
+                        self.prepare_expression(level.clone(), arg, context, false)?;
+                    }
+                    self.put_local_call(function, arguments, &context.expression)?;
                     writeln!(self.out, ";")?;
                 }
             }
@@ -804,6 +927,7 @@ impl<W: Write> Writer<W> {
     pub fn write(
         &mut self,
         module: &crate::Module,
+        analysis: &Analysis,
         options: &Options,
     ) -> Result<TranslationInfo, Error> {
         self.names.clear();
@@ -814,7 +938,7 @@ impl<W: Write> Writer<W> {
         writeln!(self.out)?;
 
         self.write_type_defs(module)?;
-        self.write_functions(module, options)
+        self.write_functions(module, analysis, options)
     }
 
     fn write_type_defs(&mut self, module: &crate::Module) -> Result<(), Error> {
@@ -956,6 +1080,7 @@ impl<W: Write> Writer<W> {
     fn write_functions(
         &mut self,
         module: &crate::Module,
+        analysis: &Analysis,
         options: &Options,
     ) -> Result<TranslationInfo, Error> {
         let mut pass_through_globals = Vec::new();
@@ -1022,12 +1147,17 @@ impl<W: Write> Writer<W> {
                 writeln!(self.out, ";")?;
             }
 
-            let context = ExpressionContext {
-                function: fun,
-                origin: FunctionOrigin::Handle(fun_handle),
-                module,
+            let context = StatementContext {
+                expression: ExpressionContext {
+                    function: fun,
+                    origin: FunctionOrigin::Handle(fun_handle),
+                    module,
+                },
+                fun_info: &analysis[fun_handle],
+                return_value: None,
             };
-            self.put_block(Level(1), &fun.body, &context, None)?;
+            self.named_expressions.clear();
+            self.put_block(Level(1), &fun.body, &context)?;
             writeln!(self.out, "}}")?;
             writeln!(self.out)?;
         }
@@ -1035,7 +1165,7 @@ impl<W: Write> Writer<W> {
         let mut info = TranslationInfo {
             entry_point_names: Vec::with_capacity(module.entry_points.len()),
         };
-        for (ep_index, (&(stage, _), ep)) in module.entry_points.iter().enumerate() {
+        for (ep_index, (&(stage, ref ep_name), ep)) in module.entry_points.iter().enumerate() {
             let fun = &ep.function;
             self.typifier.resolve_all(
                 &fun.expressions,
@@ -1222,12 +1352,17 @@ impl<W: Write> Writer<W> {
                 writeln!(self.out, ";")?;
             }
 
-            let context = ExpressionContext {
-                function: fun,
-                origin: FunctionOrigin::EntryPoint(ep_index as _),
-                module,
+            let context = StatementContext {
+                expression: ExpressionContext {
+                    function: fun,
+                    origin: FunctionOrigin::EntryPoint(ep_index as _),
+                    module,
+                },
+                fun_info: analysis.get_entry_point(stage, ep_name),
+                return_value,
             };
-            self.put_block(Level(1), &fun.body, &context, return_value)?;
+            self.named_expressions.clear();
+            self.put_block(Level(1), &fun.body, &context)?;
             writeln!(self.out, "}}")?;
             let is_last = ep_index == module.entry_points.len() - 1;
             if !is_last {
diff --git a/src/front/global_usage.rs b/src/front/global_usage.rs
index 0cc5fa8666..e5a43509c0 100644
--- a/src/front/global_usage.rs
+++ b/src/front/global_usage.rs
@@ -9,13 +9,13 @@ struct GlobalUseVisitor<'a> {
 }
 
 impl Visitor for GlobalUseVisitor<'_> {
-    fn visit_expr(&mut self, expr: &crate::Expression) {
+    fn visit_expr(&mut self, _: Handle<crate::Expression>, expr: &crate::Expression) {
         if let crate::Expression::GlobalVariable(handle) = expr {
             self.usage[handle.index()] |= crate::GlobalUse::READ;
         }
     }
 
-    fn visit_lhs_expr(&mut self, expr: &crate::Expression) {
+    fn visit_lhs_expr(&mut self, _: Handle<crate::Expression>, expr: &crate::Expression) {
         if let crate::Expression::GlobalVariable(handle) = expr {
             self.usage[handle.index()] |= crate::GlobalUse::WRITE;
         }
diff --git a/src/lib.rs b/src/lib.rs
index d569af58ab..c87af3c533 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -640,6 +640,8 @@ pub enum ImageQuery {
 }
 
 /// An expression that can be evaluated to obtain a value.
+///
+/// This is a Single Static Assignment (SSA) scheme similar to SPIR-V.
 #[derive(Clone, Debug)]
 #[cfg_attr(test, derive(PartialEq))]
 #[cfg_attr(feature = "serialize", derive(Serialize))]
diff --git a/src/proc/analyzer.rs b/src/proc/analyzer.rs
index 6ff4c70fcc..9dcfbf2191 100644
--- a/src/proc/analyzer.rs
+++ b/src/proc/analyzer.rs
@@ -7,6 +7,7 @@ Figures out the following properties:
 !*/
 
 use crate::arena::{Arena, Handle};
+use std::ops;
 
 bitflags::bitflags! {
     #[derive(Default)]
@@ -35,7 +36,14 @@ pub struct ExpressionInfo {
 pub struct FunctionInfo {
     pub control_flags: ControlFlags,
     pub sampling_set: crate::FastHashSet<SamplingKey>,
-    pub expressions: Box<[ExpressionInfo]>,
+    expressions: Box<[ExpressionInfo]>,
+}
+
+impl ops::Index<Handle<crate::Expression>> for FunctionInfo {
+    type Output = ExpressionInfo;
+    fn index(&self, handle: Handle<crate::Expression>) -> &ExpressionInfo {
+        &self.expressions[handle.index()]
+    }
 }
 
 #[derive(Clone, Debug, thiserror::Error)]
@@ -355,6 +363,22 @@ impl Analysis {
 
         Ok(this)
     }
+
+    pub fn get_entry_point(&self, stage: crate::ShaderStage, name: &str) -> &FunctionInfo {
+        let (_, info) = self
+            .entry_points
+            .iter()
+            .find(|(key, _)| key.0 == stage && key.1 == name)
+            .unwrap();
+        info
+    }
+}
+
+impl ops::Index<Handle<crate::Function>> for Analysis {
+    type Output = FunctionInfo;
+    fn index(&self, handle: Handle<crate::Function>) -> &FunctionInfo {
+        &self.functions[handle.index()]
+    }
 }
 
 #[test]
diff --git a/src/proc/interface.rs b/src/proc/interface.rs
index 13f6f1f714..ca6804388c 100644
--- a/src/proc/interface.rs
+++ b/src/proc/interface.rs
@@ -7,8 +7,8 @@ pub struct Interface<'a, T> {
 }
 
 pub trait Visitor {
-    fn visit_expr(&mut self, _: &crate::Expression) {}
-    fn visit_lhs_expr(&mut self, _: &crate::Expression) {}
+    fn visit_expr(&mut self, _: Handle<crate::Expression>, _: &crate::Expression) {}
+    fn visit_lhs_expr(&mut self, _: Handle<crate::Expression>, _: &crate::Expression) {}
     fn visit_fun(&mut self, _: Handle<crate::Function>) {}
 }
 
@@ -16,12 +16,12 @@ impl<'a, T> Interface<'a, T>
 where
     T: Visitor,
 {
-    fn traverse_expr(&mut self, handle: Handle<crate::Expression>) {
+    pub fn traverse_expr(&mut self, handle: Handle<crate::Expression>) {
         use crate::Expression as E;
 
         let expr = &self.expressions[handle];
 
-        self.visitor.visit_expr(expr);
+        self.visitor.visit_expr(handle, expr);
 
         match *expr {
             E::Access { base, index } => {
@@ -200,7 +200,7 @@ where
                             _ => break,
                         }
                     }
-                    self.visitor.visit_lhs_expr(&self.expressions[left]);
+                    self.visitor.visit_lhs_expr(left, &self.expressions[left]);
                     self.traverse_expr(value);
                 }
                 S::Call {
diff --git a/tests/out/shadow.msl.snap b/tests/out/shadow.msl.snap
index 2a8f90f464..ccb2a01ba8 100644
--- a/tests/out/shadow.msl.snap
+++ b/tests/out/shadow.msl.snap
@@ -52,7 +52,8 @@ type7 fetch_shadow(
     if ((homogeneous_coords.w <= 0.0)) {
         return 1.0;
     }
-    return t_shadow.sample_compare(sampler_shadow, (((metal::float2(homogeneous_coords.x, homogeneous_coords.y) * metal::float2(0.5, -0.5)) * (1.0 / homogeneous_coords.w)) + metal::float2(0.5, 0.5)), static_cast<int>(light_id), (homogeneous_coords.z * (1.0 / homogeneous_coords.w)));
+    float expr15 = (1.0 / homogeneous_coords.w);
+    return t_shadow.sample_compare(sampler_shadow, (((metal::float2(homogeneous_coords.x, homogeneous_coords.y) * metal::float2(0.5, -0.5)) * expr15) + metal::float2(0.5, 0.5)), static_cast<int>(light_id), (homogeneous_coords.z * expr15));
 }
 
 struct fs_mainInput {
@@ -83,7 +84,9 @@ fragment fs_mainOutput fs_main(
         if ((i >= metal::min(u_globals.num_lights.x, 10u))) {
             break;
         }
-        color1 = (color1 + ((fetch_shadow(i, (s_lights.data[i].proj * input.in_position_fs), t_shadow, sampler_shadow) * metal::max(0.0, metal::dot(metal::normalize(input.in_normal_fs), metal::normalize((metal::float3(s_lights.data[i].pos.x, s_lights.data[i].pos.y, s_lights.data[i].pos.z) - metal::float3(input.in_position_fs.x, input.in_position_fs.y, input.in_position_fs.z)))))) * metal::float3(s_lights.data[i].color.x, s_lights.data[i].color.y, s_lights.data[i].color.z)));
+        Light expr18 = s_lights.data[i];
+        type7 expr21 = fetch_shadow(i, (expr18.proj * input.in_position_fs), t_shadow, sampler_shadow);
+        color1 = (color1 + ((expr21 * metal::max(0.0, metal::dot(metal::normalize(input.in_normal_fs), metal::normalize((metal::float3(expr18.pos.x, expr18.pos.y, expr18.pos.z) - metal::float3(input.in_position_fs.x, input.in_position_fs.y, input.in_position_fs.z)))))) * metal::float3(expr18.color.x, expr18.color.y, expr18.color.z)));
     }
     output.out_color_fs = metal::float4(color1, 1.0);
     return output;
diff --git a/tests/out/skybox.msl.snap b/tests/out/skybox.msl.snap
index 1faf661cf6..c05cba30ce 100644
--- a/tests/out/skybox.msl.snap
+++ b/tests/out/skybox.msl.snap
@@ -47,9 +47,10 @@ vertex vs_mainOutput vs_main(
     type unprojected;
     tmp1_ = (static_cast<int>(in_vertex_index) / 2);
     tmp2_ = (static_cast<int>(in_vertex_index) & 1);
-    unprojected = (r_data.proj_inv * metal::float4(((static_cast<float>(tmp1_) * 4.0) - 1.0), ((static_cast<float>(tmp2_) * 4.0) - 1.0), 0.0, 1.0));
+    type expr24 = metal::float4(((static_cast<float>(tmp1_) * 4.0) - 1.0), ((static_cast<float>(tmp2_) * 4.0) - 1.0), 0.0, 1.0);
+    unprojected = (r_data.proj_inv * expr24);
     output.out_uv = (metal::transpose(metal::float3x3(metal::float3(r_data.view[0].x, r_data.view[0].y, r_data.view[0].z), metal::float3(r_data.view[1].x, r_data.view[1].y, r_data.view[1].z), metal::float3(r_data.view[2].x, r_data.view[2].y, r_data.view[2].z))) * metal::float3(unprojected.x, unprojected.y, unprojected.z));
-    output.out_position = metal::float4(((static_cast<float>(tmp1_) * 4.0) - 1.0), ((static_cast<float>(tmp2_) * 4.0) - 1.0), 0.0, 1.0);
+    output.out_position = expr24;
     return output;
 }
 
diff --git a/tests/snapshots.rs b/tests/snapshots.rs
index 7523bb0b9e..8fae779a66 100644
--- a/tests/snapshots.rs
+++ b/tests/snapshots.rs
@@ -75,7 +75,12 @@ fn check_output_spv(module: &naga::Module, name: &str, params: &Parameters) {
 }
 
 #[cfg(feature = "msl-out")]
-fn check_output_msl(module: &naga::Module, name: &str, params: &Parameters) {
+fn check_output_msl(
+    module: &naga::Module,
+    analysis: &naga::proc::analyzer::Analysis,
+    name: &str,
+    params: &Parameters,
+) {
     use naga::back::msl;
 
     let mut binding_map = msl::BindingMap::default();
@@ -104,7 +109,7 @@ fn check_output_msl(module: &naga::Module, name: &str, params: &Parameters) {
         binding_map,
     };
 
-    let (msl, _) = msl::write_string(&module, &options).unwrap();
+    let (msl, _) = msl::write_string(module, analysis, &options).unwrap();
 
     with_snapshot_settings(|| {
         insta::assert_snapshot!(format!("{}.msl", name), msl);
@@ -143,7 +148,8 @@ fn convert_wgsl(name: &str, language: Language) {
             .expect("Couldn't find wgsl file"),
     )
     .unwrap();
-    naga::proc::Validator::new().validate(&module).unwrap();
+    #[cfg_attr(not(feature = "msl-out"), allow(unused_variables))]
+    let analysis = naga::proc::Validator::new().validate(&module).unwrap();
 
     #[cfg(feature = "spv-out")]
     {
@@ -154,7 +160,7 @@ fn convert_wgsl(name: &str, language: Language) {
     #[cfg(feature = "msl-out")]
     {
         if language.contains(Language::METAL) {
-            check_output_msl(&module, name, &params);
+            check_output_msl(&module, &analysis, name, &params);
         }
     }
     #[cfg(feature = "glsl-out")]