From 6f06479b765946bdd3174b2c05e7a1dda24c31c2 Mon Sep 17 00:00:00 2001
From: Mikko Rasa <tdb@tdb.fi>
Date: Sun, 10 Oct 2021 11:54:24 +0300
Subject: [PATCH] Only remove load IDs of variables actually assigned in a loop

Some variables like function arguments are not pointers and can't be
loaded again.
---
 source/glsl/reflect.cpp               | 13 ++++++
 source/glsl/reflect.h                 | 12 +++++
 source/glsl/spirv.cpp                 |  3 +-
 tests/glsl/use_parameter_in_loop.glsl | 67 +++++++++++++++++++++++++++
 4 files changed, 94 insertions(+), 1 deletion(-)
 create mode 100644 tests/glsl/use_parameter_in_loop.glsl

diff --git a/source/glsl/reflect.cpp b/source/glsl/reflect.cpp
index 258244e8..cc2e3a0c 100644
--- a/source/glsl/reflect.cpp
+++ b/source/glsl/reflect.cpp
@@ -380,6 +380,19 @@ void DependencyCollector::visit(FunctionDeclaration &func)
 	}
 }
 
+
+set<VariableDeclaration *> AssignmentCollector::apply(Node &node)
+{
+	node.visit(*this);
+	return assigned_variables;
+}
+
+void AssignmentCollector::visit(Assignment &assign)
+{
+	if(VariableDeclaration *var = dynamic_cast<VariableDeclaration *>(assign.target.declaration))
+		assigned_variables.insert(var);
+}
+
 } // namespace SL
 } // namespace GL
 } // namespace Msp
diff --git a/source/glsl/reflect.h b/source/glsl/reflect.h
index d48fbec1..5b4170a9 100644
--- a/source/glsl/reflect.h
+++ b/source/glsl/reflect.h
@@ -104,6 +104,18 @@ private:
 	virtual void visit(FunctionDeclaration &);
 };
 
+class AssignmentCollector: private TraversingVisitor
+{
+private:
+	std::set<VariableDeclaration *> assigned_variables;
+
+public:
+	std::set<VariableDeclaration *> apply(Node &);
+
+private:
+	virtual void visit(Assignment &);
+};
+
 } // namespace SL
 } // namespace GL
 } // namespace Msp
diff --git a/source/glsl/spirv.cpp b/source/glsl/spirv.cpp
index edbdf635..16168a8c 100644
--- a/source/glsl/spirv.cpp
+++ b/source/glsl/spirv.cpp
@@ -1858,7 +1858,8 @@ void SpirVGenerator::visit(Iteration &iter)
 	if(iter.init_statement)
 		iter.init_statement->visit(*this);
 
-	variable_load_ids.clear();
+	for(VariableDeclaration *v: AssignmentCollector().apply(iter))
+		variable_load_ids.erase(v);
 
 	Id header_id = next_id++;
 	Id continue_id = next_id++;
diff --git a/tests/glsl/use_parameter_in_loop.glsl b/tests/glsl/use_parameter_in_loop.glsl
new file mode 100644
index 00000000..a0bbc5c3
--- /dev/null
+++ b/tests/glsl/use_parameter_in_loop.glsl
@@ -0,0 +1,67 @@
+uniform Lighting
+{
+	vec3 light_dir[3];
+};
+uniform mat4 model_matrix;
+uniform mat4 view_matrix;
+uniform mat4 proj_matrix;
+
+#pragma MSP stage(vertex)
+layout(location=0) in vec4 vertex;
+void main()
+{
+	out vec4 world_vertex = model_matrix*vertex;
+	vec3 eye_pos = inverse(view_matrix)[3].xyz;
+	out vec3 world_look_dir = normalize(world_vertex.xyz-eye_pos);
+	gl_Position = proj_matrix*view_matrix*world_vertex;
+}
+
+#pragma MSP stage(fragment)
+layout(location=0) out vec4 frag_color;
+float calculate_lighting(vec3 look_dir)
+{
+	float intensity = 0.0;
+	for(int i=0; i<3; ++i)
+		intensity += max(dot(look_dir, light_dir[i]), 0.0);
+	return intensity;
+}
+void main()
+{
+	float light_intensity = calculate_lighting(world_look_dir);
+	frag_color = vec4(vec3(light_intensity), 1.0);
+}
+
+/* Expected output: vertex
+layout(location=0) uniform mat4 model_matrix;
+layout(location=4) uniform mat4 view_matrix;
+layout(location=8) uniform mat4 proj_matrix;
+layout(location=0) in vec4 vertex;
+layout(location=0) out vec4 world_vertex;
+layout(location=1) out vec3 world_look_dir;
+void main()
+{
+	world_vertex = model_matrix*vertex;
+	world_look_dir = normalize(world_vertex.xyz-inverse(view_matrix)[3].xyz);
+	gl_Position = proj_matrix*view_matrix*world_vertex;
+}
+*/
+
+/* Expected output: fragment
+layout(binding=5) uniform Lighting
+{
+	vec3 light_dir[3];
+};
+layout(location=0) out vec4 frag_color;
+float calculate_lighting(vec3 look_dir)
+{
+	float intensity = 0.0;
+	for(int i = 0; i<3; ++i)
+		intensity += max(dot(look_dir, light_dir[i]), 0.0);
+	return intensity;
+}
+layout(location=1) in vec3 world_look_dir;
+void main()
+{
+	frag_color = vec4(vec3(calculate_lighting(world_look_dir)), 1.0);
+}
+*/
-- 
2.43.0