Merge pull request #43775 from vnen/gdscript-fix-stack

GDScript: Fix mishandling of stack pointers

4 files changed, 104 insertions, 35 deletions

diff --git a/modules/gdscript/gdscript_byte_codegen.cpp b/modules/gdscript/gdscript_byte_codegen.cpp
index 5f1c738207..a81b3298ef 100644
--- a/modules/gdscript/gdscript_byte_codegen.cpp
+++ b/modules/gdscript/gdscript_byte_codegen.cpp
@@ -77,11 +77,22 @@ uint32_t GDScriptByteCodeGenerator::add_or_get_name(const StringName &p_name) {
 
 uint32_t GDScriptByteCodeGenerator::add_temporary() {
 	current_temporaries++;
-	return increase_stack();
+	int idx = increase_stack();
+#ifdef DEBUG_ENABLED
+	temp_stack.push_back(idx);
+#endif
+	return idx;
 }
 
 void GDScriptByteCodeGenerator::pop_temporary() {
+	ERR_FAIL_COND(current_temporaries == 0);
 	current_stack_size--;
+#ifdef DEBUG_ENABLED
+	if (temp_stack.back()->get() != current_stack_size) {
+		ERR_PRINT("Mismatched popping of temporary value");
+	}
+	temp_stack.pop_back();
+#endif
 	current_temporaries--;
 }
 
@@ -111,6 +122,11 @@ void GDScriptByteCodeGenerator::write_start(GDScript *p_script, const StringName
 }
 
 GDScriptFunction *GDScriptByteCodeGenerator::write_end() {
+#ifdef DEBUG_ENABLED
+	if (current_temporaries != 0) {
+		ERR_PRINT("Non-zero temporary variables at end of function: " + itos(current_temporaries));
+	}
+#endif
 	append(GDScriptFunction::OPCODE_END, 0);
 
 	if (constant_map.size()) {
@@ -905,23 +921,39 @@ void GDScriptByteCodeGenerator::write_endif() {
 	if_jmp_addrs.pop_back();
 }
 
-void GDScriptByteCodeGenerator::write_for(const Address &p_variable, const Address &p_list) {
-	int counter_pos = add_temporary() | (GDScriptFunction::ADDR_TYPE_STACK << GDScriptFunction::ADDR_BITS);
-	int container_pos = add_temporary() | (GDScriptFunction::ADDR_TYPE_STACK << GDScriptFunction::ADDR_BITS);
+void GDScriptByteCodeGenerator::start_for(const GDScriptDataType &p_iterator_type, const GDScriptDataType &p_list_type) {
+	Address counter(Address::LOCAL_VARIABLE, add_local("@counter_pos", p_iterator_type), p_iterator_type);
+	Address container(Address::LOCAL_VARIABLE, add_local("@container_pos", p_list_type), p_list_type);
 
-	current_breaks_to_patch.push_back(List<int>());
+	// Store state.
+	for_counter_variables.push_back(counter);
+	for_container_variables.push_back(container);
+}
+
+void GDScriptByteCodeGenerator::write_for_assignment(const Address &p_variable, const Address &p_list) {
+	const Address &container = for_container_variables.back()->get();
 
 	// Assign container.
 	append(GDScriptFunction::OPCODE_ASSIGN, 2);
-	append(container_pos);
+	append(container);
 	append(p_list);
 
+	for_iterator_variables.push_back(p_variable);
+}
+
+void GDScriptByteCodeGenerator::write_for() {
+	const Address &iterator = for_iterator_variables.back()->get();
+	const Address &counter = for_counter_variables.back()->get();
+	const Address &container = for_container_variables.back()->get();
+
+	current_breaks_to_patch.push_back(List<int>());
+
 	GDScriptFunction::Opcode begin_opcode = GDScriptFunction::OPCODE_ITERATE_BEGIN;
 	GDScriptFunction::Opcode iterate_opcode = GDScriptFunction::OPCODE_ITERATE;
 
-	if (p_list.type.has_type) {
-		if (p_list.type.kind == GDScriptDataType::BUILTIN) {
-			switch (p_list.type.builtin_type) {
+	if (container.type.has_type) {
+		if (container.type.kind == GDScriptDataType::BUILTIN) {
+			switch (container.type.builtin_type) {
 				case Variant::INT:
 					begin_opcode = GDScriptFunction::OPCODE_ITERATE_BEGIN_INT;
 					iterate_opcode = GDScriptFunction::OPCODE_ITERATE_INT;
@@ -1005,9 +1037,9 @@ void GDScriptByteCodeGenerator::write_for(const Address &p_variable, const Addre
 
 	// Begin loop.
 	append(begin_opcode, 3);
-	append(counter_pos);
-	append(container_pos);
-	append(p_variable);
+	append(counter);
+	append(container);
+	append(iterator);
 	for_jmp_addrs.push_back(opcodes.size());
 	append(0); // End of loop address, will be patched.
 	append(GDScriptFunction::OPCODE_JUMP, 0);
@@ -1017,9 +1049,9 @@ void GDScriptByteCodeGenerator::write_for(const Address &p_variable, const Addre
 	int continue_addr = opcodes.size();
 	continue_addrs.push_back(continue_addr);
 	append(iterate_opcode, 3);
-	append(counter_pos);
-	append(container_pos);
-	append(p_variable);
+	append(counter);
+	append(container);
+	append(iterator);
 	for_jmp_addrs.push_back(opcodes.size());
 	append(0); // Jump destination, will be patched.
 }
@@ -1042,9 +1074,10 @@ void GDScriptByteCodeGenerator::write_endfor() {
 	}
 	current_breaks_to_patch.pop_back();
 
-	// Remove loop temporaries.
-	pop_temporary();
-	pop_temporary();
+	// Pop state.
+	for_iterator_variables.pop_back();
+	for_counter_variables.pop_back();
+	for_container_variables.pop_back();
 }
 
 void GDScriptByteCodeGenerator::start_while_condition() {
diff --git a/modules/gdscript/gdscript_byte_codegen.h b/modules/gdscript/gdscript_byte_codegen.h
index a546920fb2..0173b7f820 100644
--- a/modules/gdscript/gdscript_byte_codegen.h
+++ b/modules/gdscript/gdscript_byte_codegen.h
@@ -43,6 +43,7 @@ class GDScriptByteCodeGenerator : public GDScriptCodeGenerator {
 	Vector<int> opcodes;
 	List<Map<StringName, int>> stack_id_stack;
 	Map<StringName, int> stack_identifiers;
+	List<int> stack_identifiers_counts;
 	Map<StringName, int> local_constants;
 
 	List<GDScriptFunction::StackDebug> stack_debug;
@@ -51,11 +52,16 @@ class GDScriptByteCodeGenerator : public GDScriptCodeGenerator {
 
 	int current_stack_size = 0;
 	int current_temporaries = 0;
+	int current_locals = 0;
 	int current_line = 0;
 	int stack_max = 0;
 	int instr_args_max = 0;
 	int ptrcall_max = 0;
 
+#ifdef DEBUG_ENABLED
+	List<int> temp_stack;
+#endif
+
 	HashMap<Variant, int, VariantHasher, VariantComparator> constant_map;
 	Map<StringName, int> name_map;
 #ifdef TOOLS_ENABLED
@@ -72,8 +78,12 @@ class GDScriptByteCodeGenerator : public GDScriptCodeGenerator {
 	Map<Variant::ValidatedConstructor, int> constructors_map;
 	Map<MethodBind *, int> method_bind_map;
 
-	List<int> if_jmp_addrs; // List since this can be nested.
+	// Lists since these can be nested.
+	List<int> if_jmp_addrs;
 	List<int> for_jmp_addrs;
+	List<Address> for_iterator_variables;
+	List<Address> for_counter_variables;
+	List<Address> for_container_variables;
 	List<int> while_jmp_addrs;
 	List<int> continue_addrs;
 
@@ -89,6 +99,7 @@ class GDScriptByteCodeGenerator : public GDScriptCodeGenerator {
 	List<List<int>> match_continues_to_patch;
 
 	void add_stack_identifier(const StringName &p_id, int p_stackpos) {
+		current_locals++;
 		stack_identifiers[p_id] = p_stackpos;
 		if (debug_stack) {
 			block_identifiers[p_id] = p_stackpos;
@@ -102,17 +113,26 @@ class GDScriptByteCodeGenerator : public GDScriptCodeGenerator {
 	}
 
 	void push_stack_identifiers() {
+		stack_identifiers_counts.push_back(current_locals);
 		stack_id_stack.push_back(stack_identifiers);
 		if (debug_stack) {
-			block_identifier_stack.push_back(block_identifiers);
+			Map<StringName, int> block_ids(block_identifiers);
+			block_identifier_stack.push_back(block_ids);
 			block_identifiers.clear();
 		}
 	}
 
 	void pop_stack_identifiers() {
+		current_locals = stack_identifiers_counts.back()->get();
+		stack_identifiers_counts.pop_back();
 		stack_identifiers = stack_id_stack.back()->get();
-		current_stack_size = stack_identifiers.size() + current_temporaries;
 		stack_id_stack.pop_back();
+#ifdef DEBUG_ENABLED
+		if (current_temporaries != 0) {
+			ERR_PRINT("Leaving block with non-zero temporary variables: " + itos(current_temporaries));
+		}
+#endif
+		current_stack_size = current_locals;
 
 		if (debug_stack) {
 			for (Map<StringName, int>::Element *E = block_identifiers.front(); E; E = E->next()) {
@@ -397,7 +417,9 @@ public:
 	virtual void write_if(const Address &p_condition) override;
 	virtual void write_else() override;
 	virtual void write_endif() override;
-	virtual void write_for(const Address &p_variable, const Address &p_list) override;
+	virtual void start_for(const GDScriptDataType &p_iterator_type, const GDScriptDataType &p_list_type) override;
+	virtual void write_for_assignment(const Address &p_variable, const Address &p_list) override;
+	virtual void write_for() override;
 	virtual void write_endfor() override;
 	virtual void start_while_condition() override;
 	virtual void write_while(const Address &p_condition) override;
diff --git a/modules/gdscript/gdscript_codegen.h b/modules/gdscript/gdscript_codegen.h
index 2616a34719..decab3df0b 100644
--- a/modules/gdscript/gdscript_codegen.h
+++ b/modules/gdscript/gdscript_codegen.h
@@ -139,7 +139,9 @@ public:
 	// virtual void write_elseif(const Address &p_condition) = 0; This kind of makes things more difficult for no real benefit.
 	virtual void write_else() = 0;
 	virtual void write_endif() = 0;
-	virtual void write_for(const Address &p_variable, const Address &p_list) = 0;
+	virtual void start_for(const GDScriptDataType &p_iterator_type, const GDScriptDataType &p_list_type) = 0;
+	virtual void write_for_assignment(const Address &p_variable, const Address &p_list) = 0;
+	virtual void write_for() = 0;
 	virtual void write_endfor() = 0;
 	virtual void start_while_condition() = 0; // Used to allow a jump to the expression evaluation.
 	virtual void write_while(const Address &p_condition) = 0;
diff --git a/modules/gdscript/gdscript_compiler.cpp b/modules/gdscript/gdscript_compiler.cpp
index 69fe6d27cc..d47c9accc9 100644
--- a/modules/gdscript/gdscript_compiler.cpp
+++ b/modules/gdscript/gdscript_compiler.cpp
@@ -1474,13 +1474,27 @@ Error GDScriptCompiler::_parse_block(CodeGen &codegen, const GDScriptParser::Sui
 				codegen.start_block();
 
 				// Evaluate the match expression.
+				GDScriptCodeGenerator::Address value_local = codegen.add_local("@match_value", _gdtype_from_datatype(match->test->get_datatype()));
 				GDScriptCodeGenerator::Address value = _parse_expression(codegen, error, match->test);
 				if (error) {
 					return error;
 				}
 
+				// Assign to local.
+				// TODO: This can be improved by passing the target to parse_expression().
+				gen->write_assign(value_local, value);
+
+				if (value.mode == GDScriptCodeGenerator::Address::TEMPORARY) {
+					codegen.generator->pop_temporary();
+				}
+
 				// Then, let's save the type of the value in the stack too, so we can reuse for later comparisons.
-				GDScriptCodeGenerator::Address type = codegen.add_temporary();
+				GDScriptDataType typeof_type;
+				typeof_type.has_type = true;
+				typeof_type.kind = GDScriptDataType::BUILTIN;
+				typeof_type.builtin_type = Variant::INT;
+				GDScriptCodeGenerator::Address type = codegen.add_local("@match_type", typeof_type);
+
 				Vector<GDScriptCodeGenerator::Address> typeof_args;
 				typeof_args.push_back(value);
 				gen->write_call_builtin(type, GDScriptFunctions::TYPE_OF, typeof_args);
@@ -1534,12 +1548,6 @@ Error GDScriptCompiler::_parse_block(CodeGen &codegen, const GDScriptParser::Sui
 					gen->write_endif();
 				}
 
-				gen->pop_temporary();
-
-				if (value.mode == GDScriptCodeGenerator::Address::TEMPORARY) {
-					codegen.generator->pop_temporary();
-				}
-
 				gen->end_match();
 			} break;
 			case GDScriptParser::Node::IF: {
@@ -1577,12 +1585,20 @@ Error GDScriptCompiler::_parse_block(CodeGen &codegen, const GDScriptParser::Sui
 				codegen.start_block();
 				GDScriptCodeGenerator::Address iterator = codegen.add_local(for_n->variable->name, _gdtype_from_datatype(for_n->variable->get_datatype()));
 
+				gen->start_for(iterator.type, _gdtype_from_datatype(for_n->list->get_datatype()));
+
 				GDScriptCodeGenerator::Address list = _parse_expression(codegen, error, for_n->list);
 				if (error) {
 					return error;
 				}
 
-				gen->write_for(iterator, list);
+				gen->write_for_assignment(iterator, list);
+
+				if (list.mode == GDScriptCodeGenerator::Address::TEMPORARY) {
+					codegen.generator->pop_temporary();
+				}
+
+				gen->write_for();
 
 				error = _parse_block(codegen, for_n->loop);
 				if (error) {
@@ -1591,10 +1607,6 @@ Error GDScriptCompiler::_parse_block(CodeGen &codegen, const GDScriptParser::Sui
 
 				gen->write_endfor();
 
-				if (list.mode == GDScriptCodeGenerator::Address::TEMPORARY) {
-					codegen.generator->pop_temporary();
-				}
-
 				codegen.end_block();
 			} break;
 			case GDScriptParser::Node::WHILE: {