Begin implementing for loop initializer propagation.

.clang-format

@@ -162,3 +162,6 @@ TabWidth: 4
 
 # The way to use tab characters in the resulting file. Possible values: Never, ForIndentation, Always.
 UseTab: ForIndentation
+
+# Do not reflow comments
+ReflowComments: false

spirv_cfg.hpp

@@ -56,6 +56,24 @@ public:
 
 	uint32_t find_common_dominator(uint32_t a, uint32_t b) const;
 
+	const std::vector<uint32_t> &get_preceding_edges(uint32_t block) const
+	{
+		return preceding_edges[block];
+	}
+
+	const std::vector<uint32_t> &get_succeeding_edges(uint32_t block) const
+	{
+		return succeeding_edges[block];
+	}
+
+	template <typename Op>
+	void walk_from(uint32_t block, const Op &op) const
+	{
+		op(block);
+		for (auto b : succeeding_edges[block])
+			walk_from(b, op);
+	}
+
 private:
 	Compiler &compiler;
 	const SPIRFunction &func;

spirv_common.hpp

@@ -445,6 +445,11 @@ struct SPIRBlock : IVariant
 	// All access to these variables are dominated by this block,
 	// so before branching anywhere we need to make sure that we declare these variables.
 	std::vector<uint32_t> dominated_variables;
+
+	// These are variables which should be declared in a for loop header, if we
+	// fail to use a classic for-loop,
+	// we remove these variables, and fall back to regular variables outside the loop.
+	std::vector<uint32_t> loop_variables;
 };
 
 struct SPIRFunction : IVariant
@@ -553,6 +558,15 @@ struct SPIRVariable : IVariant
 	bool remapped_variable = false;
 	uint32_t remapped_components = 0;
 
+	// The block which dominates all access to this variable.
+	uint32_t dominator = 0;
+	// If true, this variable is a loop variable, when accessing the variable
+	// outside a loop,
+	// we should statically forward it.
+	bool loop_variable = false;
+	// Set to true while we're inside the for loop.
+	bool loop_variable_enable = false;
+
 	SPIRFunction::Parameter *parameter = nullptr;
 };
 

spirv_cross.cpp

@@ -2909,6 +2909,8 @@ void Compiler::analyze_variable_scope(SPIRFunction &entry)
 	// Compute the control flow graph for this function.
 	CFG cfg(*this, entry);
 
+	unordered_map<uint32_t, uint32_t> potential_loop_variables;
+
 	// For each variable which is statically accessed.
 	for (auto &var : handler.accessed_variables_to_block)
 	{
@@ -2917,7 +2919,22 @@ void Compiler::analyze_variable_scope(SPIRFunction &entry)
 
 		// Figure out which block is dominating all accesses of those variables.
 		for (auto &block : blocks)
+		{
+			// If we're accessing a variable inside a continue block, this variable
+			// might be a loop variable.
+			if (is_continue(block))
+			{
+				// The variable is used in multiple continue blocks, this is not a loop
+				// candidate, signal that by setting block to -1u.
+				auto &potential = potential_loop_variables[var.first];
+
+				if (potential == 0)
+					potential = block;
+				else
+					potential = -1u;
+			}
 			builder.add_block(block);
+		}
 
 		builder.lift_continue_block_dominator();
 
@@ -2929,6 +2946,84 @@ void Compiler::analyze_variable_scope(SPIRFunction &entry)
 		{
 			auto &block = this->get<SPIRBlock>(dominating_block);
 			block.dominated_variables.push_back(var.first);
+			get<SPIRVariable>(var.first).dominator = dominating_block;
 		}
 	}
+
+	// Now, try to analyze whether or not these variables are actually loop variables.
+	for (auto &loop_variable : potential_loop_variables)
+	{
+		auto &var = get<SPIRVariable>(loop_variable.first);
+		auto dominator = var.dominator;
+		auto block = loop_variable.second;
+
+		// The variable was accessed in multiple continue blocks, ignore.
+		if (block == -1u || block == 0)
+			continue;
+
+		// Dead code.
+		if (dominator == 0)
+			continue;
+
+		uint32_t header = 0;
+
+		// Find the loop header for this block.
+		for (auto b : loop_blocks)
+		{
+			auto &potential_header = get<SPIRBlock>(b);
+			if (potential_header.continue_block == block)
+			{
+				header = b;
+				break;
+			}
+		}
+
+		assert(header);
+		auto &header_block = get<SPIRBlock>(header);
+
+		// Now, there are two conditions we need to meet for the variable to be a loop variable.
+		// 1. The dominating block must have a branch-free path to the loop header,
+		// this way we statically know which expression should be part of the loop variable initializer.
+
+		// Walk from the dominator, if there is one straight edge connecting
+		// dominator and loop header, we statically know the loop initializer.
+		bool static_loop_init = true;
+		while (dominator != header)
+		{
+			auto &succ = cfg.get_succeeding_edges(dominator);
+			if (succ.size() != 1)
+			{
+				static_loop_init = false;
+				break;
+			}
+
+			auto &pred = cfg.get_preceding_edges(succ.front());
+			if (pred.size() != 1 || pred.front() != dominator)
+			{
+				static_loop_init = false;
+				break;
+			}
+
+			dominator = succ.front();
+		}
+
+		if (!static_loop_init)
+			continue;
+
+		// The second condition we need to meet is that no access after the loop
+		// merge can occur. Walk the CFG to see if we find anything.
+		auto &blocks = handler.accessed_variables_to_block[loop_variable.first];
+		cfg.walk_from(header_block.merge_block, [&](uint32_t walk_block) {
+			// We found a block which accesses the variable outside the loop.
+			if (blocks.find(walk_block) != end(blocks))
+				static_loop_init = false;
+		});
+
+		if (!static_loop_init)
+			continue;
+
+		// We have a loop variable.
+		header_block.loop_variables.push_back(loop_variable.first);
+		get<SPIRVariable>(loop_variable.first).loop_variable = true;
+	}
 }

spirv_glsl.cpp

@@ -3490,6 +3490,8 @@ void CompilerGLSL::emit_instruction(const Instruction &instruction)
 
 		if (var && var->statically_assigned)
 			var->static_expression = ops[1];
+		else if (var && var->loop_variable && !var->loop_variable_enable)
+			var->static_expression = ops[1];
 		else
 		{
 			auto lhs = to_expression(ops[0]);
@@ -4920,7 +4922,9 @@ string CompilerGLSL::variable_decl(const SPIRVariable &variable)
 	// Ignore the pointer type since GLSL doesn't have pointers.
 	auto &type = get<SPIRType>(variable.basetype);
 	auto res = join(to_qualifiers_glsl(variable.self), variable_decl(type, to_name(variable.self)));
-	if (variable.initializer)
+	if (variable.loop_variable)
+		res += join(" = ", to_expression(variable.static_expression));
+	else if (variable.initializer)
 		res += join(" = ", to_expression(variable.initializer));
 	return res;
 }
@@ -5345,6 +5349,16 @@ void CompilerGLSL::emit_function(SPIRFunction &func, uint64_t return_flags)
 	begin_scope();
 
 	current_function = &func;
+	auto &entry_block = get<SPIRBlock>(func.entry_block);
+
+	if (!func.analyzed_variable_scope)
+	{
+		if (options.cfg_analysis)
+			analyze_variable_scope(func);
+		else
+			entry_block.dominated_variables = func.local_variables;
+		func.analyzed_variable_scope = true;
+	}
 
 	for (auto &v : func.local_variables)
 	{
@@ -5365,24 +5379,19 @@ void CompilerGLSL::emit_function(SPIRFunction &func, uint64_t return_flags)
 		}
 		else
 		{
-			// HACK: SPIRV likes to use samplers and images as local variables, but GLSL does not allow
-			// this. For these types (non-lvalue), we enforce forwarding through a shadowed variable.
+			// HACK: SPIRV likes to use samplers and images as local variables, but GLSL does not allow this.
+			// For these types (non-lvalue), we enforce forwarding through a shadowed variable.
 			// This means that when we OpStore to these variables, we just write in the expression ID directly.
 			// This breaks any kind of branching, since the variable must be statically assigned.
 			// Branching on samplers and images would be pretty much impossible to fake in GLSL.
 			var.statically_assigned = true;
 		}
-	}
 
-	auto &entry_block = get<SPIRBlock>(func.entry_block);
+		var.loop_variable_enable = false;
 
-	if (!func.analyzed_variable_scope)
-	{
-		if (options.cfg_analysis)
-			analyze_variable_scope(func);
-		else
-			entry_block.dominated_variables = func.local_variables;
-		func.analyzed_variable_scope = true;
+		// Loop variables are never declared outside their for-loop, so block any implicit declaration.
+		if (var.loop_variable)
+			var.deferred_declaration = false;
 	}
 
 	entry_block.loop_dominator = SPIRBlock::NoDominator;
@@ -5612,6 +5621,36 @@ string CompilerGLSL::emit_continue_block(uint32_t continue_block)
 	return merge(statements);
 }
 
+string CompilerGLSL::emit_for_loop_initializers(const SPIRBlock &block)
+{
+	if (block.loop_variables.empty())
+		return "";
+
+	if (block.loop_variables.size() == 1)
+	{
+		return variable_decl(get<SPIRVariable>(block.loop_variables.front()));
+	}
+	else
+	{
+		auto &var = get<SPIRVariable>(block.loop_variables.front());
+		auto &type = get<SPIRType>(var.basetype);
+
+		// Don't remap the type here as we have multiple names,
+		// doesn't make sense to remap types for loop variables anyways.
+		// It is assumed here that all relevant qualifiers are equal for all loop variables.
+		string expr = join(to_qualifiers_glsl(var.self), type_to_glsl(type), " ");
+
+		for (auto &loop_var : block.loop_variables)
+		{
+			auto &v = get<SPIRVariable>(loop_var);
+			expr += join(to_name(loop_var), " = ", to_expression(v.static_expression));
+			if (&loop_var != &block.loop_variables.back())
+				expr += ", ";
+		}
+		return expr;
+	}
+}
+
 bool CompilerGLSL::attempt_emit_loop_header(SPIRBlock &block, SPIRBlock::Method method)
 {
 	SPIRBlock::ContinueBlockType continue_type = continue_block_type(get<SPIRBlock>(block.continue_block));
@@ -5633,8 +5672,12 @@ bool CompilerGLSL::attempt_emit_loop_header(SPIRBlock &block, SPIRBlock::Method
 			switch (continue_type)
 			{
 			case SPIRBlock::ForLoop:
-				statement("for (; ", to_expression(block.condition), "; ", emit_continue_block(block.continue_block),
-				          ")");
+				// If we have loop variables, stop masking out access to the variable now.
+				for (auto var : block.loop_variables)
+					get<SPIRVariable>(var).loop_variable_enable = true;
+
+				statement("for (", emit_for_loop_initializers(block), "; ", to_expression(block.condition), "; ",
+				          emit_continue_block(block.continue_block), ")");
 				break;
 
 			case SPIRBlock::WhileLoop:
@@ -5680,8 +5723,12 @@ bool CompilerGLSL::attempt_emit_loop_header(SPIRBlock &block, SPIRBlock::Method
 			switch (continue_type)
 			{
 			case SPIRBlock::ForLoop:
-				statement("for (; ", to_expression(child.condition), "; ", emit_continue_block(block.continue_block),
-				          ")");
+				// If we have loop variables, stop masking out access to the variable now.
+				for (auto var : block.loop_variables)
+					get<SPIRVariable>(var).loop_variable_enable = true;
+
+				statement("for (", emit_for_loop_initializers(block), "; ", to_expression(child.condition), "; ",
+				          emit_continue_block(block.continue_block), ")");
 				break;
 
 			case SPIRBlock::WhileLoop:
@@ -5725,6 +5772,7 @@ void CompilerGLSL::emit_block_chain(SPIRBlock &block)
 
 	bool select_branch_to_true_block = false;
 	bool skip_direct_branch = false;
+	bool emitted_for_loop_header = false;
 
 	// If we need to force temporaries for certain IDs due to continue blocks, do it before starting loop header.
 	for (auto &tmp : block.declare_temporary)
@@ -5744,9 +5792,9 @@ void CompilerGLSL::emit_block_chain(SPIRBlock &block)
 		flush_undeclared_variables(block);
 		if (attempt_emit_loop_header(block, SPIRBlock::MergeToSelectForLoop))
 		{
-			// The body of while, is actually just the true block, so always branch there
-			// unconditionally.
+			// The body of while, is actually just the true block, so always branch there unconditionally.
 			select_branch_to_true_block = true;
+			emitted_for_loop_header = true;
 		}
 	}
 	// This is the newer loop behavior in glslang which branches from Loop header directly to
@@ -5755,7 +5803,10 @@ void CompilerGLSL::emit_block_chain(SPIRBlock &block)
 	{
 		flush_undeclared_variables(block);
 		if (attempt_emit_loop_header(block, SPIRBlock::MergeToDirectForLoop))
+		{
 			skip_direct_branch = true;
+			emitted_for_loop_header = true;
+		}
 	}
 	else if (continue_type == SPIRBlock::DoWhileLoop)
 	{
@@ -5783,6 +5834,16 @@ void CompilerGLSL::emit_block_chain(SPIRBlock &block)
 			emit_instruction(op);
 	}
 
+	// If we didn't successfully emit a loop header and we had loop variable candidates, we have a problem
+	// as writes to said loop variables might have been masked out, we need a recompile.
+	if (!emitted_for_loop_header && !block.loop_variables.empty())
+	{
+		force_recompile = true;
+		for (auto var : block.loop_variables)
+			get<SPIRVariable>(var).loop_variable = false;
+		block.loop_variables.clear();
+	}
+
 	flush_undeclared_variables(block);
 	bool emit_next_block = true;
 

spirv_glsl.hpp

@@ -388,6 +388,8 @@ protected:
 	void check_function_call_constraints(const uint32_t *args, uint32_t length);
 	void handle_invalid_expression(uint32_t id);
 	void find_static_extensions();
+
+	std::string emit_for_loop_initializers(const SPIRBlock &block);
 };
 }