From 96b37aff1f0377f76eb9ddae55ed46121aac2bb1 Mon Sep 17 00:00:00 2001
From: Gordon-F <ishaposhnik@icloud.com>
Date: Wed, 16 Jun 2021 10:04:40 +0300
Subject: [PATCH] [hlsl-out] Reduce heap allocation

---
 src/back/hlsl/writer.rs | 55 ++++++++++++++++++++++++++++++++---------
 tests/out/quad.hlsl     |  2 +-
 tests/snapshots.rs      |  5 +++-
 3 files changed, 49 insertions(+), 13 deletions(-)

diff --git a/src/back/hlsl/writer.rs b/src/back/hlsl/writer.rs
index c2260cf187..019260339e 100644
--- a/src/back/hlsl/writer.rs
+++ b/src/back/hlsl/writer.rs
@@ -100,6 +100,19 @@ impl<'a, W: Write> Writer<'a, W> {
         self.reset(module);
 
         // Write all constants
+        // For example, input wgsl shader:
+        // ```wgsl
+        // let c_scale: f32 = 1.2;
+        // return VertexOutput(uv, vec4<f32>(c_scale * pos, 0.0, 1.0));
+        // ```
+        //
+        // Output shader:
+        // ```hlsl
+        // static const float c_scale = 1.2;
+        // const VertexOutput vertexoutput1 = { vertexinput.uv3, float4((c_scale * vertexinput.pos1), 0.0, 1.0) };
+        // ```
+        //
+        // If we remove `write_global_constant` `c_scale` will be inlined.
         for (handle, constant) in module.constants.iter() {
             if constant.name.is_some() {
                 self.write_global_constant(module, &constant.inner, handle)?;
@@ -304,18 +317,29 @@ impl<'a, W: Write> Writer<'a, W> {
                 width: _,
                 ref value,
             } => {
+                write!(self.out, "static const ")?;
+                // Write type
+                match *value {
+                    crate::ScalarValue::Sint(_) => write!(self.out, "int")?,
+                    crate::ScalarValue::Uint(_) => write!(self.out, "uint")?,
+                    crate::ScalarValue::Float(_) => write!(self.out, "float")?,
+                    crate::ScalarValue::Bool(_) => write!(self.out, "bool")?,
+                };
                 let name = &self.names[&NameKey::Constant(handle)];
-                let (ty, value) = match *value {
-                    crate::ScalarValue::Sint(value) => ("int", format!("{}", value)),
-                    crate::ScalarValue::Uint(value) => ("uint", format!("{}", value)),
+                write!(self.out, " {} = ", name)?;
+
+                // Second match required to avoid heap allocation by `format!()`
+                match *value {
+                    crate::ScalarValue::Sint(value) => write!(self.out, "{}", value)?,
+                    crate::ScalarValue::Uint(value) => write!(self.out, "{}", value)?,
                     crate::ScalarValue::Float(value) => {
                         // Floats are written using `Debug` instead of `Display` because it always appends the
                         // decimal part even it's zero
-                        ("float", format!("{:?}", value))
+                        write!(self.out, "{:?}", value)?
                     }
-                    crate::ScalarValue::Bool(value) => ("bool", format!("{}", value)),
+                    crate::ScalarValue::Bool(value) => write!(self.out, "{}", value)?,
                 };
-                writeln!(self.out, "static const {} {} = {};", ty, name, value)?;
+                writeln!(self.out, ";")?;
             }
             ConstantInner::Composite { .. } => {
                 return Err(Error::Unimplemented(format!(
@@ -773,19 +797,24 @@ impl<'a, W: Write> Writer<'a, W> {
                 }
             }
             Expression::FunctionArgument(pos) => {
-                let name = match func_ctx.ty {
+                match func_ctx.ty {
                     FunctionType::Function(handle) => {
-                        self.names[&NameKey::FunctionArgument(handle, pos)].clone()
+                        let name = &self.names[&NameKey::FunctionArgument(handle, pos)];
+                        write!(self.out, "{}", name)?;
                     }
                     FunctionType::EntryPoint(index) => {
                         // EntryPoint arguments wrapped into structure
                         // We can safery unwrap here, because if we write function arguments it means, that ep_input struct already exists
                         let ep_input = self.ep_inputs[index as usize].as_ref().unwrap();
                         let member_name = &ep_input.members[pos as usize].name;
-                        format!("{}.{}", &ep_input.name.to_lowercase(), member_name)
+                        write!(
+                            self.out,
+                            "{}.{}",
+                            &ep_input.name.to_lowercase(),
+                            member_name
+                        )?
                     }
                 };
-                write!(self.out, "{}", name)?;
             }
             Expression::ImageSample {
                 image,
@@ -826,7 +855,11 @@ impl<'a, W: Write> Writer<'a, W> {
                 width: _,
                 ref value,
             } => {
-                self.write_scalar_value(*value)?;
+                if constant.name.is_some() {
+                    write!(self.out, "{}", &self.names[&NameKey::Constant(handle)])?;
+                } else {
+                    self.write_scalar_value(*value)?;
+                }
             }
             _ => {
                 return Err(Error::Unimplemented(format!(
diff --git a/tests/out/quad.hlsl b/tests/out/quad.hlsl
index c678e3e8f9..c9685d2346 100644
--- a/tests/out/quad.hlsl
+++ b/tests/out/quad.hlsl
@@ -19,7 +19,7 @@ struct FragmentInput {
 
 VertexOutput vert_main(VertexInput vertexinput)
 {
-    const VertexOutput vertexoutput1 = { vertexinput.uv3, float4((1.2 * vertexinput.pos1), 0.0, 1.0) };
+    const VertexOutput vertexoutput1 = { vertexinput.uv3, float4((c_scale * vertexinput.pos1), 0.0, 1.0) };
     return vertexoutput1;
 }
 
diff --git a/tests/snapshots.rs b/tests/snapshots.rs
index d3076391b1..06c48e218d 100644
--- a/tests/snapshots.rs
+++ b/tests/snapshots.rs
@@ -238,7 +238,10 @@ fn check_output_hlsl(module: &naga::Module, info: &naga::valid::ModuleInfo, dest
             }
             naga::ShaderStage::Compute => ("compute", "cs_5_0", &options.compute_entry_point_name),
         };
-        config_str = format!("{}{}={}\n{}_name={}\n", config_str, stage_str, profile, stage_str, ep_name);
+        config_str = format!(
+            "{}{}={}\n{}_name={}\n",
+            config_str, stage_str, profile, stage_str, ep_name
+        );
     }
     fs::write(destination.with_extension("hlsl.config"), config_str).unwrap();
 }