GDScript: Reset local variables on exit from block

8 files changed, 96 insertions, 14 deletions

diff --git a/modules/gdscript/gdscript_compiler.cpp b/modules/gdscript/gdscript_compiler.cpp
index 327e24ef11..8c2751ff7e 100644
--- a/modules/gdscript/gdscript_compiler.cpp
+++ b/modules/gdscript/gdscript_compiler.cpp
@@ -1664,25 +1664,39 @@ GDScriptCodeGenerator::Address GDScriptCompiler::_parse_match_pattern(CodeGen &c
 	ERR_FAIL_V_MSG(p_previous_test, "Reaching the end of pattern compilation without matching a pattern.");
 }
 
-void GDScriptCompiler::_add_locals_in_block(CodeGen &codegen, const GDScriptParser::SuiteNode *p_block) {
+List<GDScriptCodeGenerator::Address> GDScriptCompiler::_add_locals_in_block(CodeGen &codegen, const GDScriptParser::SuiteNode *p_block) {
+	List<GDScriptCodeGenerator::Address> addresses;
 	for (int i = 0; i < p_block->locals.size(); i++) {
 		if (p_block->locals[i].type == GDScriptParser::SuiteNode::Local::PARAMETER || p_block->locals[i].type == GDScriptParser::SuiteNode::Local::FOR_VARIABLE) {
 			// Parameters are added directly from function and loop variables are declared explicitly.
 			continue;
 		}
-		codegen.add_local(p_block->locals[i].name, _gdtype_from_datatype(p_block->locals[i].get_datatype(), codegen.script));
+		addresses.push_back(codegen.add_local(p_block->locals[i].name, _gdtype_from_datatype(p_block->locals[i].get_datatype(), codegen.script)));
 	}
+	return addresses;
 }
 
-Error GDScriptCompiler::_parse_block(CodeGen &codegen, const GDScriptParser::SuiteNode *p_block, bool p_add_locals) {
+// Avoid keeping in the stack long-lived references to objects, which may prevent RefCounted objects from being freed.
+void GDScriptCompiler::_clear_addresses(CodeGen &codegen, const List<GDScriptCodeGenerator::Address> &p_addresses) {
+	for (const List<GDScriptCodeGenerator::Address>::Element *E = p_addresses.front(); E; E = E->next()) {
+		GDScriptDataType type = E->get().type;
+		// If not an object and cannot contain an object, no need to clear.
+		if (type.kind != GDScriptDataType::BUILTIN || type.builtin_type == Variant::ARRAY || type.builtin_type == Variant::DICTIONARY) {
+			codegen.generator->write_assign_false(E->get());
+		}
+	}
+}
+
+Error GDScriptCompiler::_parse_block(CodeGen &codegen, const GDScriptParser::SuiteNode *p_block, bool p_add_locals, bool p_reset_locals) {
 	Error err = OK;
 	GDScriptCodeGenerator *gen = codegen.generator;
+	List<GDScriptCodeGenerator::Address> block_locals;
 
 	gen->clean_temporaries();
 	codegen.start_block();
 
 	if (p_add_locals) {
-		_add_locals_in_block(codegen, p_block);
+		block_locals = _add_locals_in_block(codegen, p_block);
 	}
 
 	for (int i = 0; i < p_block->statements.size(); i++) {
@@ -1738,7 +1752,7 @@ Error GDScriptCompiler::_parse_block(CodeGen &codegen, const GDScriptParser::Sui
 					codegen.start_block(); // Create an extra block around for binds.
 
 					// Add locals in block before patterns, so temporaries don't use the stack address for binds.
-					_add_locals_in_block(codegen, branch->block);
+					List<GDScriptCodeGenerator::Address> branch_locals = _add_locals_in_block(codegen, branch->block);
 
 #ifdef DEBUG_ENABLED
 					// Add a newline before each branch, since the debugger needs those.
@@ -1765,6 +1779,8 @@ Error GDScriptCompiler::_parse_block(CodeGen &codegen, const GDScriptParser::Sui
 						return err;
 					}
 
+					_clear_addresses(codegen, branch_locals);
+
 					codegen.end_block(); // Get out of extra block.
 				}
 
@@ -1949,7 +1965,7 @@ Error GDScriptCompiler::_parse_block(CodeGen &codegen, const GDScriptParser::Sui
 				}
 
 				// Assigns a null for the unassigned variables in loops.
-				if (!initialized && p_block->is_loop) {
+				if (!initialized && p_block->is_in_loop) {
 					codegen.generator->write_construct(local, Variant::NIL, Vector<GDScriptCodeGenerator::Address>());
 				}
 			} break;
@@ -1985,6 +2001,10 @@ Error GDScriptCompiler::_parse_block(CodeGen &codegen, const GDScriptParser::Sui
 		gen->clean_temporaries();
 	}
 
+	if (p_add_locals && p_reset_locals) {
+		_clear_addresses(codegen, block_locals);
+	}
+
 	codegen.end_block();
 	return OK;
 }
@@ -2138,7 +2158,8 @@ GDScriptFunction *GDScriptCompiler::_parse_function(Error &r_error, GDScript *p_
 			codegen.generator->end_parameters();
 		}
 
-		r_error = _parse_block(codegen, p_func->body);
+		// No need to reset locals at the end of the function, the stack will be cleared anyway.
+		r_error = _parse_block(codegen, p_func->body, true, false);
 		if (r_error) {
 			memdelete(codegen.generator);
 			return nullptr;
diff --git a/modules/gdscript/gdscript_compiler.h b/modules/gdscript/gdscript_compiler.h
index 2d15d461fb..c5efbb621e 100644
--- a/modules/gdscript/gdscript_compiler.h
+++ b/modules/gdscript/gdscript_compiler.h
@@ -128,8 +128,9 @@ class GDScriptCompiler {
 	GDScriptCodeGenerator::Address _parse_assign_right_expression(CodeGen &codegen, Error &r_error, const GDScriptParser::AssignmentNode *p_assignmentint, const GDScriptCodeGenerator::Address &p_index_addr = GDScriptCodeGenerator::Address());
 	GDScriptCodeGenerator::Address _parse_expression(CodeGen &codegen, Error &r_error, const GDScriptParser::ExpressionNode *p_expression, bool p_root = false, bool p_initializer = false, const GDScriptCodeGenerator::Address &p_index_addr = GDScriptCodeGenerator::Address());
 	GDScriptCodeGenerator::Address _parse_match_pattern(CodeGen &codegen, Error &r_error, const GDScriptParser::PatternNode *p_pattern, const GDScriptCodeGenerator::Address &p_value_addr, const GDScriptCodeGenerator::Address &p_type_addr, const GDScriptCodeGenerator::Address &p_previous_test, bool p_is_first, bool p_is_nested);
-	void _add_locals_in_block(CodeGen &codegen, const GDScriptParser::SuiteNode *p_block);
-	Error _parse_block(CodeGen &codegen, const GDScriptParser::SuiteNode *p_block, bool p_add_locals = true);
+	List<GDScriptCodeGenerator::Address> _add_locals_in_block(CodeGen &codegen, const GDScriptParser::SuiteNode *p_block);
+	void _clear_addresses(CodeGen &codegen, const List<GDScriptCodeGenerator::Address> &p_addresses);
+	Error _parse_block(CodeGen &codegen, const GDScriptParser::SuiteNode *p_block, bool p_add_locals = true, bool p_reset_locals = true);
 	GDScriptFunction *_parse_function(Error &r_error, GDScript *p_script, const GDScriptParser::ClassNode *p_class, const GDScriptParser::FunctionNode *p_func, bool p_for_ready = false, bool p_for_lambda = false);
 	GDScriptFunction *_make_static_initializer(Error &r_error, GDScript *p_script, const GDScriptParser::ClassNode *p_class);
 	Error _parse_setter_getter(GDScript *p_script, const GDScriptParser::ClassNode *p_class, const GDScriptParser::VariableNode *p_variable, bool p_is_setter);
diff --git a/modules/gdscript/gdscript_parser.cpp b/modules/gdscript/gdscript_parser.cpp
index fdbd505975..03dae967c2 100644
--- a/modules/gdscript/gdscript_parser.cpp
+++ b/modules/gdscript/gdscript_parser.cpp
@@ -1532,6 +1532,11 @@ GDScriptParser::SuiteNode *GDScriptParser::parse_suite(const String &p_context,
 	suite->parent_function = current_function;
 	current_suite = suite;
 
+	if (!p_for_lambda && suite->parent_block != nullptr && suite->parent_block->is_in_loop) {
+		// Do not reset to false if true is set before calling parse_suite().
+		suite->is_in_loop = true;
+	}
+
 	bool multiline = false;
 
 	if (match(GDScriptTokenizer::Token::NEWLINE)) {
@@ -1867,9 +1872,8 @@ GDScriptParser::ForNode *GDScriptParser::parse_for() {
 		}
 		suite->add_local(SuiteNode::Local(n_for->variable, current_function));
 	}
-
+	suite->is_in_loop = true;
 	n_for->loop = parse_suite(R"("for" block)", suite);
-	n_for->loop->is_loop = true;
 	complete_extents(n_for);
 
 	// Reset break/continue state.
@@ -2182,8 +2186,9 @@ GDScriptParser::WhileNode *GDScriptParser::parse_while() {
 	can_break = true;
 	can_continue = true;
 
-	n_while->loop = parse_suite(R"("while" block)");
-	n_while->loop->is_loop = true;
+	SuiteNode *suite = alloc_node<SuiteNode>();
+	suite->is_in_loop = true;
+	n_while->loop = parse_suite(R"("while" block)", suite);
 	complete_extents(n_while);
 
 	// Reset break/continue state.
diff --git a/modules/gdscript/gdscript_parser.h b/modules/gdscript/gdscript_parser.h
index 8f0265510f..19e265a3e9 100644
--- a/modules/gdscript/gdscript_parser.h
+++ b/modules/gdscript/gdscript_parser.h
@@ -1117,7 +1117,7 @@ public:
 		bool has_return = false;
 		bool has_continue = false;
 		bool has_unreachable_code = false; // Just so warnings aren't given more than once per block.
-		bool is_loop = false;
+		bool is_in_loop = false; // The block is nested in a loop (directly or indirectly).
 
 		bool has_local(const StringName &p_name) const;
 		const Local &get_local(const StringName &p_name) const;
diff --git a/modules/gdscript/tests/scripts/runtime/features/reset_local_var_on exit_block.gd b/modules/gdscript/tests/scripts/runtime/features/reset_local_var_on exit_block.gd
new file mode 100644
index 0000000000..c774ebf83c
--- /dev/null
+++ b/modules/gdscript/tests/scripts/runtime/features/reset_local_var_on exit_block.gd
@@ -0,0 +1,10 @@
+# GH-77666
+
+func test():
+	var ref := RefCounted.new()
+	print(ref.get_reference_count())
+
+	if true:
+		var _temp := ref
+
+	print(ref.get_reference_count())
diff --git a/modules/gdscript/tests/scripts/runtime/features/reset_local_var_on exit_block.out b/modules/gdscript/tests/scripts/runtime/features/reset_local_var_on exit_block.out
new file mode 100644
index 0000000000..04b4638adf
--- /dev/null
+++ b/modules/gdscript/tests/scripts/runtime/features/reset_local_var_on exit_block.out
@@ -0,0 +1,3 @@
+GDTEST_OK
+1
+1
diff --git a/modules/gdscript/tests/scripts/runtime/features/reset_unassigned_variables_in_loops.gd b/modules/gdscript/tests/scripts/runtime/features/reset_unassigned_variables_in_loops.gd
new file mode 100644
index 0000000000..c45f8dce48
--- /dev/null
+++ b/modules/gdscript/tests/scripts/runtime/features/reset_unassigned_variables_in_loops.gd
@@ -0,0 +1,28 @@
+# GH-56223, GH-76569
+
+func test():
+	for i in 3:
+		var a
+		if true:
+			var b
+			if true:
+				var c
+				prints("Begin:", i, a, b, c)
+				a = 1
+				b = 1
+				c = 1
+				prints("End:", i, a, b, c)
+	print("===")
+	var j := 0
+	while j < 3:
+		var a
+		if true:
+			var b
+			if true:
+				var c
+				prints("Begin:", j, a, b, c)
+				a = 1
+				b = 1
+				c = 1
+				prints("End:", j, a, b, c)
+		j += 1
diff --git a/modules/gdscript/tests/scripts/runtime/features/reset_unassigned_variables_in_loops.out b/modules/gdscript/tests/scripts/runtime/features/reset_unassigned_variables_in_loops.out
new file mode 100644
index 0000000000..7eddcbf903
--- /dev/null
+++ b/modules/gdscript/tests/scripts/runtime/features/reset_unassigned_variables_in_loops.out
@@ -0,0 +1,14 @@
+GDTEST_OK
+Begin: 0 <null> <null> <null>
+End: 0 1 1 1
+Begin: 1 <null> <null> <null>
+End: 1 1 1 1
+Begin: 2 <null> <null> <null>
+End: 2 1 1 1
+===
+Begin: 0 <null> <null> <null>
+End: 0 1 1 1
+Begin: 1 <null> <null> <null>
+End: 1 1 1 1
+Begin: 2 <null> <null> <null>
+End: 2 1 1 1