Merge pull request #71564 from dalexeev/gds-optimize-for-range

GDScript: Optimize non-constant `for`-`range`

Author: Repiteo

modules/gdscript/gdscript_analyzer.cpp

@@ -2195,101 +2195,39 @@ void GDScriptAnalyzer::resolve_if(GDScriptParser::IfNode *p_if) {
 }
 
 void GDScriptAnalyzer::resolve_for(GDScriptParser::ForNode *p_for) {
-	bool list_resolved = false;
-
-	// Optimize constant range() call to not allocate an array.
-	// Use int, Vector2i, Vector3i instead, which also can be used as range iterators.
-	if (p_for->list && p_for->list->type == GDScriptParser::Node::CALL) {
-		GDScriptParser::CallNode *call = static_cast<GDScriptParser::CallNode *>(p_for->list);
-		GDScriptParser::Node::Type callee_type = call->get_callee_type();
-		if (callee_type == GDScriptParser::Node::IDENTIFIER) {
-			GDScriptParser::IdentifierNode *callee = static_cast<GDScriptParser::IdentifierNode *>(call->callee);
-			if (callee->name == "range") {
-				list_resolved = true;
-				if (call->arguments.is_empty()) {
-					push_error(R"*(Invalid call for "range()" function. Expected at least 1 argument, none given.)*", call->callee);
-				} else if (call->arguments.size() > 3) {
-					push_error(vformat(R"*(Invalid call for "range()" function. Expected at most 3 arguments, %d given.)*", call->arguments.size()), call->callee);
-				} else {
-					// Now we can optimize it.
-					bool can_reduce = true;
-					Vector<Variant> args;
-					args.resize(call->arguments.size());
-					for (int i = 0; i < call->arguments.size(); i++) {
-						GDScriptParser::ExpressionNode *argument = call->arguments[i];
-						reduce_expression(argument);
-
-						if (argument->is_constant) {
-							if (argument->reduced_value.get_type() != Variant::INT && argument->reduced_value.get_type() != Variant::FLOAT) {
-								can_reduce = false;
-								push_error(vformat(R"*(Invalid argument for "range()" call. Argument %d should be int or float but "%s" was given.)*", i + 1, Variant::get_type_name(argument->reduced_value.get_type())), argument);
-							}
-							if (can_reduce) {
-								args.write[i] = argument->reduced_value;
-							}
-						} else {
-							can_reduce = false;
-							GDScriptParser::DataType argument_type = argument->get_datatype();
-							if (argument_type.is_variant() || !argument_type.is_hard_type()) {
-								mark_node_unsafe(argument);
-							}
-							if (!argument_type.is_variant() && (argument_type.builtin_type != Variant::INT && argument_type.builtin_type != Variant::FLOAT)) {
-								if (!argument_type.is_hard_type()) {
-									downgrade_node_type_source(argument);
-								} else {
-									push_error(vformat(R"*(Invalid argument for "range()" call. Argument %d should be int or float but "%s" was given.)*", i + 1, argument_type.to_string()), argument);
-								}
-							}
-						}
-					}
+	GDScriptParser::DataType variable_type;
+	GDScriptParser::DataType list_type;
 
-					Variant reduced;
+	if (p_for->list) {
+		resolve_node(p_for->list, false);
 
-					if (can_reduce) {
-						switch (args.size()) {
-							case 1:
-								reduced = (int32_t)args[0];
-								break;
-							case 2:
-								reduced = Vector2i(args[0], args[1]);
-								break;
-							case 3:
-								reduced = Vector3i(args[0], args[1], args[2]);
-								break;
-						}
-						p_for->list->is_constant = true;
-						p_for->list->reduced_value = reduced;
+		bool is_range = false;
+		if (p_for->list->type == GDScriptParser::Node::CALL) {
+			GDScriptParser::CallNode *call = static_cast<GDScriptParser::CallNode *>(p_for->list);
+			if (call->get_callee_type() == GDScriptParser::Node::IDENTIFIER) {
+				if (static_cast<GDScriptParser::IdentifierNode *>(call->callee)->name == "range") {
+					if (call->arguments.is_empty()) {
+						push_error(R"*(Invalid call for "range()" function. Expected at least 1 argument, none given.)*", call);
+					} else if (call->arguments.size() > 3) {
+						push_error(vformat(R"*(Invalid call for "range()" function. Expected at most 3 arguments, %d given.)*", call->arguments.size()), call);
 					}
-				}
-
-				if (p_for->list->is_constant) {
-					p_for->list->set_datatype(type_from_variant(p_for->list->reduced_value, p_for->list));
-				} else {
-					GDScriptParser::DataType list_type;
-					list_type.type_source = GDScriptParser::DataType::ANNOTATED_EXPLICIT;
-					list_type.kind = GDScriptParser::DataType::BUILTIN;
-					list_type.builtin_type = Variant::ARRAY;
-					p_for->list->set_datatype(list_type);
+					is_range = true;
+					variable_type.type_source = GDScriptParser::DataType::ANNOTATED_INFERRED;
+					variable_type.kind = GDScriptParser::DataType::BUILTIN;
+					variable_type.builtin_type = Variant::INT;
 				}
 			}
 		}
-	}
 
-	GDScriptParser::DataType variable_type;
-	String list_visible_type = "<unresolved type>";
-	if (list_resolved) {
-		variable_type.type_source = GDScriptParser::DataType::ANNOTATED_INFERRED;
-		variable_type.kind = GDScriptParser::DataType::BUILTIN;
-		variable_type.builtin_type = Variant::INT;
-		list_visible_type = "Array[int]"; // NOTE: `range()` has `Array` return type.
-	} else if (p_for->list) {
-		resolve_node(p_for->list, false);
-		GDScriptParser::DataType list_type = p_for->list->get_datatype();
-		list_visible_type = list_type.to_string();
+		list_type = p_for->list->get_datatype();
+
 		if (!list_type.is_hard_type()) {
 			mark_node_unsafe(p_for->list);
 		}
-		if (list_type.is_variant()) {
+
+		if (is_range) {
+			// Already solved.
+		} else if (list_type.is_variant()) {
 			variable_type.kind = GDScriptParser::DataType::VARIANT;
 			mark_node_unsafe(p_for->list);
 		} else if (list_type.has_container_element_type(0)) {
@@ -2342,7 +2280,7 @@ void GDScriptAnalyzer::resolve_for(GDScriptParser::ForNode *p_for) {
 						mark_node_unsafe(p_for->variable);
 						p_for->use_conversion_assign = true;
 					} else {
-						push_error(vformat(R"(Unable to iterate on value of type "%s" with variable of type "%s".)", list_visible_type, specified_type.to_string()), p_for->datatype_specifier);
+						push_error(vformat(R"(Unable to iterate on value of type "%s" with variable of type "%s".)", list_type.to_string(), specified_type.to_string()), p_for->datatype_specifier);
 					}
 				} else if (!is_type_compatible(specified_type, variable_type)) {
 					p_for->use_conversion_assign = true;

modules/gdscript/gdscript_byte_codegen.cpp

@@ -1542,16 +1542,35 @@ void GDScriptByteCodeGenerator::write_end_jump_if_shared() {
 	if_jmp_addrs.pop_back();
 }
 
-void GDScriptByteCodeGenerator::start_for(const GDScriptDataType &p_iterator_type, const GDScriptDataType &p_list_type) {
+void GDScriptByteCodeGenerator::start_for(const GDScriptDataType &p_iterator_type, const GDScriptDataType &p_list_type, bool p_is_range) {
 	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);
 
 	// Store state.
 	for_counter_variables.push_back(counter);
-	for_container_variables.push_back(container);
+
+	if (p_is_range) {
+		GDScriptDataType int_type;
+		int_type.has_type = true;
+		int_type.kind = GDScriptDataType::BUILTIN;
+		int_type.builtin_type = Variant::INT;
+
+		Address range_from(Address::LOCAL_VARIABLE, add_local("@range_from", int_type), int_type);
+		Address range_to(Address::LOCAL_VARIABLE, add_local("@range_to", int_type), int_type);
+		Address range_step(Address::LOCAL_VARIABLE, add_local("@range_step", int_type), int_type);
+
+		// Store state.
+		for_range_from_variables.push_back(range_from);
+		for_range_to_variables.push_back(range_to);
+		for_range_step_variables.push_back(range_step);
+	} else {
+		Address container(Address::LOCAL_VARIABLE, add_local("@container_pos", p_list_type), p_list_type);
+
+		// Store state.
+		for_container_variables.push_back(container);
+	}
 }
 
-void GDScriptByteCodeGenerator::write_for_assignment(const Address &p_list) {
+void GDScriptByteCodeGenerator::write_for_list_assignment(const Address &p_list) {
 	const Address &container = for_container_variables.back()->get();
 
 	// Assign container.
@@ -1560,16 +1579,33 @@ void GDScriptByteCodeGenerator::write_for_assignment(const Address &p_list) {
 	append(p_list);
 }
 
-void GDScriptByteCodeGenerator::write_for(const Address &p_variable, bool p_use_conversion) {
+void GDScriptByteCodeGenerator::write_for_range_assignment(const Address &p_from, const Address &p_to, const Address &p_step) {
+	const Address &range_from = for_range_from_variables.back()->get();
+	const Address &range_to = for_range_to_variables.back()->get();
+	const Address &range_step = for_range_step_variables.back()->get();
+
+	// Assign range args.
+	write_assign(range_from, p_from);
+	write_assign(range_to, p_to);
+	write_assign(range_step, p_step);
+}
+
+void GDScriptByteCodeGenerator::write_for(const Address &p_variable, bool p_use_conversion, bool p_is_range) {
 	const Address &counter = for_counter_variables.back()->get();
-	const Address &container = for_container_variables.back()->get();
+	const Address &container = p_is_range ? Address() : for_container_variables.back()->get();
+	const Address &range_from = p_is_range ? for_range_from_variables.back()->get() : Address();
+	const Address &range_to = p_is_range ? for_range_to_variables.back()->get() : Address();
+	const Address &range_step = p_is_range ? for_range_step_variables.back()->get() : Address();
 
 	current_breaks_to_patch.push_back(List<int>());
 
 	GDScriptFunction::Opcode begin_opcode = GDScriptFunction::OPCODE_ITERATE_BEGIN;
 	GDScriptFunction::Opcode iterate_opcode = GDScriptFunction::OPCODE_ITERATE;
 
-	if (container.type.has_type) {
+	if (p_is_range) {
+		begin_opcode = GDScriptFunction::OPCODE_ITERATE_BEGIN_RANGE;
+		iterate_opcode = GDScriptFunction::OPCODE_ITERATE_RANGE;
+	} else if (container.type.has_type) {
 		if (container.type.kind == GDScriptDataType::BUILTIN) {
 			switch (container.type.builtin_type) {
 				case Variant::INT:
@@ -1665,19 +1701,30 @@ void GDScriptByteCodeGenerator::write_for(const Address &p_variable, bool p_use_
 	// Begin loop.
 	append_opcode(begin_opcode);
 	append(counter);
-	append(container);
+	if (p_is_range) {
+		append(range_from);
+		append(range_to);
+		append(range_step);
+	} else {
+		append(container);
+	}
 	append(p_use_conversion ? temp : p_variable);
 	for_jmp_addrs.push_back(opcodes.size());
 	append(0); // End of loop address, will be patched.
 	append_opcode(GDScriptFunction::OPCODE_JUMP);
-	append(opcodes.size() + 6); // Skip over 'continue' code.
+	append(opcodes.size() + (p_is_range ? 7 : 6)); // Skip over 'continue' code.
 
 	// Next iteration.
 	int continue_addr = opcodes.size();
 	continue_addrs.push_back(continue_addr);
 	append_opcode(iterate_opcode);
 	append(counter);
-	append(container);
+	if (p_is_range) {
+		append(range_to);
+		append(range_step);
+	} else {
+		append(container);
+	}
 	append(p_use_conversion ? temp : p_variable);
 	for_jmp_addrs.push_back(opcodes.size());
 	append(0); // Jump destination, will be patched.
@@ -1690,7 +1737,7 @@ void GDScriptByteCodeGenerator::write_for(const Address &p_variable, bool p_use_
 	}
 }
 
-void GDScriptByteCodeGenerator::write_endfor() {
+void GDScriptByteCodeGenerator::write_endfor(bool p_is_range) {
 	// Jump back to loop check.
 	append_opcode(GDScriptFunction::OPCODE_JUMP);
 	append(continue_addrs.back()->get());
@@ -1710,7 +1757,13 @@ void GDScriptByteCodeGenerator::write_endfor() {
 
 	// Pop state.
 	for_counter_variables.pop_back();
-	for_container_variables.pop_back();
+	if (p_is_range) {
+		for_range_from_variables.pop_back();
+		for_range_to_variables.pop_back();
+		for_range_step_variables.pop_back();
+	} else {
+		for_container_variables.pop_back();
+	}
 }
 
 void GDScriptByteCodeGenerator::start_while_condition() {

modules/gdscript/gdscript_byte_codegen.h

@@ -144,6 +144,9 @@ class GDScriptByteCodeGenerator : public GDScriptCodeGenerator {
 	List<int> for_jmp_addrs;
 	List<Address> for_counter_variables;
 	List<Address> for_container_variables;
+	List<Address> for_range_from_variables;
+	List<Address> for_range_to_variables;
+	List<Address> for_range_step_variables;
 	List<int> while_jmp_addrs;
 	List<int> continue_addrs;
 
@@ -535,10 +538,11 @@ class GDScriptByteCodeGenerator : public GDScriptCodeGenerator {
 	virtual void write_endif() override;
 	virtual void write_jump_if_shared(const Address &p_value) override;
 	virtual void write_end_jump_if_shared() override;
-	virtual void start_for(const GDScriptDataType &p_iterator_type, const GDScriptDataType &p_list_type) override;
-	virtual void write_for_assignment(const Address &p_list) override;
-	virtual void write_for(const Address &p_variable, bool p_use_conversion) override;
-	virtual void write_endfor() override;
+	virtual void start_for(const GDScriptDataType &p_iterator_type, const GDScriptDataType &p_list_type, bool p_is_range) override;
+	virtual void write_for_list_assignment(const Address &p_list) override;
+	virtual void write_for_range_assignment(const Address &p_from, const Address &p_to, const Address &p_step) override;
+	virtual void write_for(const Address &p_variable, bool p_use_conversion, bool p_is_range) override;
+	virtual void write_endfor(bool p_is_range) override;
 	virtual void start_while_condition() override;
 	virtual void write_while(const Address &p_condition) override;
 	virtual void write_endwhile() override;

modules/gdscript/gdscript_codegen.h

@@ -148,10 +148,11 @@ class GDScriptCodeGenerator {
 	virtual void write_endif() = 0;
 	virtual void write_jump_if_shared(const Address &p_value) = 0;
 	virtual void write_end_jump_if_shared() = 0;
-	virtual void start_for(const GDScriptDataType &p_iterator_type, const GDScriptDataType &p_list_type) = 0;
-	virtual void write_for_assignment(const Address &p_list) = 0;
-	virtual void write_for(const Address &p_variable, bool p_use_conversion) = 0;
-	virtual void write_endfor() = 0;
+	virtual void start_for(const GDScriptDataType &p_iterator_type, const GDScriptDataType &p_list_type, bool p_is_range) = 0;
+	virtual void write_for_list_assignment(const Address &p_list) = 0;
+	virtual void write_for_range_assignment(const Address &p_from, const Address &p_to, const Address &p_step) = 0;
+	virtual void write_for(const Address &p_variable, bool p_use_conversion, bool p_is_range) = 0;
+	virtual void write_endfor(bool p_is_range) = 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;
 	virtual void write_endwhile() = 0;

modules/gdscript/gdscript_compiler.cpp

@@ -2045,20 +2045,64 @@ Error GDScriptCompiler::_parse_block(CodeGen &codegen, const GDScriptParser::Sui
 
 				GDScriptCodeGenerator::Address iterator = codegen.add_local(for_n->variable->name, _gdtype_from_datatype(for_n->variable->get_datatype(), codegen.script));
 
-				gen->start_for(iterator.type, _gdtype_from_datatype(for_n->list->get_datatype(), codegen.script));
-
-				GDScriptCodeGenerator::Address list = _parse_expression(codegen, err, for_n->list);
-				if (err) {
-					return err;
+				// Optimize `range()` call to not allocate an array.
+				GDScriptParser::CallNode *range_call = nullptr;
+				if (for_n->list && for_n->list->type == GDScriptParser::Node::CALL) {
+					GDScriptParser::CallNode *call = static_cast<GDScriptParser::CallNode *>(for_n->list);
+					if (call->get_callee_type() == GDScriptParser::Node::IDENTIFIER) {
+						if (static_cast<GDScriptParser::IdentifierNode *>(call->callee)->name == "range") {
+							range_call = call;
+						}
+					}
 				}
 
-				gen->write_for_assignment(list);
+				gen->start_for(iterator.type, _gdtype_from_datatype(for_n->list->get_datatype(), codegen.script), range_call != nullptr);
 
-				if (list.mode == GDScriptCodeGenerator::Address::TEMPORARY) {
-					codegen.generator->pop_temporary();
+				if (range_call != nullptr) {
+					Vector<GDScriptCodeGenerator::Address> args;
+					args.resize(range_call->arguments.size());
+
+					for (int j = 0; j < args.size(); j++) {
+						args.write[j] = _parse_expression(codegen, err, range_call->arguments[j]);
+						if (err) {
+							return err;
+						}
+					}
+
+					switch (args.size()) {
+						case 1:
+							gen->write_for_range_assignment(codegen.add_constant(0), args[0], codegen.add_constant(1));
+							break;
+						case 2:
+							gen->write_for_range_assignment(args[0], args[1], codegen.add_constant(1));
+							break;
+						case 3:
+							gen->write_for_range_assignment(args[0], args[1], args[2]);
+							break;
+						default:
+							_set_error(R"*(Analyzer bug: Wrong "range()" argument count.)*", range_call);
+							return ERR_BUG;
+					}
+
+					for (int j = 0; j < args.size(); j++) {
+						if (args[j].mode == GDScriptCodeGenerator::Address::TEMPORARY) {
+							codegen.generator->pop_temporary();
+						}
+					}
+				} else {
+					GDScriptCodeGenerator::Address list = _parse_expression(codegen, err, for_n->list);
+					if (err) {
+						return err;
+					}
+
+					gen->write_for_list_assignment(list);
+
+					if (list.mode == GDScriptCodeGenerator::Address::TEMPORARY) {
+						codegen.generator->pop_temporary();
+					}
 				}
 
-				gen->write_for(iterator, for_n->use_conversion_assign);
+				gen->write_for(iterator, for_n->use_conversion_assign, range_call != nullptr);
 
 				// Loop variables must be cleared even when `break`/`continue` is used.
 				List<GDScriptCodeGenerator::Address> loop_locals = _add_block_locals(codegen, for_n->loop);
@@ -2070,7 +2114,7 @@ Error GDScriptCompiler::_parse_block(CodeGen &codegen, const GDScriptParser::Sui
 					return err;
 				}
 
-				gen->write_endfor();
+				gen->write_endfor(range_call != nullptr);
 
 				_clear_block_locals(codegen, loop_locals); // Outside loop, after block - for `break` and normal exit.