Compiler: Fix for/for each variable index remapping

Variable index of for each was not remapped, causing for and for-each in precompiled modules to cause AstAlreadyUsedIndex and AstInvalidIndex errors when importing such modules
NazaraEngine · Oct 14, 2023 · 19a610b · 19a610b
1 parent 02148a1
commit 19a610b
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 0 deletions.
diff --git a/include/NZSL/Ast/IndexRemapperVisitor.hpp b/include/NZSL/Ast/IndexRemapperVisitor.hpp
@@ -51,6 +51,8 @@ namespace nzsl::Ast
 			StatementPtr Clone(DeclareOptionStatement& node) override;
 			StatementPtr Clone(DeclareStructStatement& node) override;
 			StatementPtr Clone(DeclareVariableStatement& node) override;
+			StatementPtr Clone(ForStatement& node) override;
+			StatementPtr Clone(ForEachStatement& node) override;
 
 			ExpressionPtr Clone(AliasValueExpression& node) override;
 			ExpressionPtr Clone(ConstantExpression& node) override;

diff --git a/src/NZSL/Ast/Cloner.cpp b/src/NZSL/Ast/Cloner.cpp
@@ -262,6 +262,7 @@ namespace nzsl::Ast
 		clone->toExpr = CloneExpression(node.toExpr);
 		clone->statement = CloneStatement(node.statement);
 		clone->unroll = Clone(node.unroll);
+		clone->varIndex = node.varIndex;
 		clone->varName = node.varName;
 
 		clone->sourceLocation = node.sourceLocation;
@@ -275,6 +276,7 @@ namespace nzsl::Ast
 		clone->expression = CloneExpression(node.expression);
 		clone->statement = CloneStatement(node.statement);
 		clone->unroll = Clone(node.unroll);
+		clone->varIndex = node.varIndex;
 		clone->varName = node.varName;
 
 		clone->sourceLocation = node.sourceLocation;

diff --git a/src/NZSL/Ast/IndexRemapperVisitor.cpp b/src/NZSL/Ast/IndexRemapperVisitor.cpp
@@ -202,6 +202,44 @@ namespace nzsl::Ast
 		return clone;
 	}
 
+	StatementPtr IndexRemapperVisitor::Clone(ForStatement& node)
+	{
+		// We have to handle the for each var index before its content
+		std::optional<std::size_t> varIndex = node.varIndex;
+		if (varIndex)
+		{
+			std::size_t newVarIndex = m_context->options->varIndexGenerator(*varIndex);
+			UniqueInsert(m_context->newVarIndices, *varIndex, newVarIndex);
+			varIndex = newVarIndex;
+		}
+		else if (m_context->options->forceIndexGeneration)
+			varIndex = m_context->options->varIndexGenerator(std::numeric_limits<std::size_t>::max());
+
+		ForStatementPtr clone = Nz::StaticUniquePointerCast<ForStatement>(Cloner::Clone(node));
+		clone->varIndex = varIndex;
+
+		return clone;
+	}
+
+	StatementPtr IndexRemapperVisitor::Clone(ForEachStatement& node)
+	{
+		// We have to handle the for each var index before its content
+		std::optional<std::size_t> varIndex = node.varIndex;
+		if (varIndex)
+		{
+			std::size_t newVarIndex = m_context->options->varIndexGenerator(*varIndex);
+			UniqueInsert(m_context->newVarIndices, *varIndex, newVarIndex);
+			varIndex = newVarIndex;
+		}
+		else if (m_context->options->forceIndexGeneration)
+			varIndex = m_context->options->varIndexGenerator(std::numeric_limits<std::size_t>::max());
+
+		ForEachStatementPtr clone = Nz::StaticUniquePointerCast<ForEachStatement>(Cloner::Clone(node));
+		clone->varIndex = varIndex;
+
+		return clone;
+	}
+
 	ExpressionPtr IndexRemapperVisitor::Clone(AliasValueExpression& node)
 	{
 		auto it = m_context->newAliasIndices.find(node.aliasId);