GDScript: Allow enum values to be set to constant expressions

Also allow them to access previous values wihout referencing the enum.

4 files changed, 96 insertions, 16 deletions

diff --git a/modules/gdscript/gdscript_analyzer.cpp b/modules/gdscript/gdscript_analyzer.cpp
index 3a04777bd5..2082339ba7 100644
--- a/modules/gdscript/gdscript_analyzer.cpp
+++ b/modules/gdscript/gdscript_analyzer.cpp
@@ -602,17 +602,67 @@ void GDScriptAnalyzer::resolve_class_interface(GDScriptParser::ClassNode *p_clas
 				enum_type.is_meta_type = true;
 				enum_type.is_constant = true;
 
+				// Enums can't be nested, so we can safely override this.
+				current_enum = member.m_enum;
+
 				for (int j = 0; j < member.m_enum->values.size(); j++) {
-					enum_type.enum_values[member.m_enum->values[j].identifier->name] = member.m_enum->values[j].value;
+					GDScriptParser::EnumNode::Value &element = member.m_enum->values.write[j];
+
+					if (element.custom_value) {
+						reduce_expression(element.custom_value);
+						if (!element.custom_value->is_constant) {
+							push_error(R"(Enum values must be constant.)", element.custom_value);
+						} else if (element.custom_value->reduced_value.get_type() != Variant::INT) {
+							push_error(R"(Enum values must be integers.)", element.custom_value);
+						} else {
+							element.value = element.custom_value->reduced_value;
+							element.resolved = true;
+						}
+					} else {
+						if (element.index > 0) {
+							element.value = element.parent_enum->values[element.index - 1].value + 1;
+						} else {
+							element.value = 0;
+						}
+						element.resolved = true;
+					}
+
+					enum_type.enum_values[element.identifier->name] = element.value;
 				}
 
+				current_enum = nullptr;
+
 				member.m_enum->set_datatype(enum_type);
 			} break;
 			case GDScriptParser::ClassNode::Member::FUNCTION:
 				resolve_function_signature(member.function);
 				break;
-			case GDScriptParser::ClassNode::Member::ENUM_VALUE:
-				break; // Nothing to do, type and value set in parser.
+			case GDScriptParser::ClassNode::Member::ENUM_VALUE: {
+				if (member.enum_value.custom_value) {
+					current_enum = member.enum_value.parent_enum;
+					reduce_expression(member.enum_value.custom_value);
+					current_enum = nullptr;
+
+					if (!member.enum_value.custom_value->is_constant) {
+						push_error(R"(Enum values must be constant.)", member.enum_value.custom_value);
+					} else if (member.enum_value.custom_value->reduced_value.get_type() != Variant::INT) {
+						push_error(R"(Enum values must be integers.)", member.enum_value.custom_value);
+					} else {
+						member.enum_value.value = member.enum_value.custom_value->reduced_value;
+						member.enum_value.resolved = true;
+					}
+				} else {
+					if (member.enum_value.index > 0) {
+						member.enum_value.value = member.enum_value.parent_enum->values[member.enum_value.index - 1].value + 1;
+					} else {
+						member.enum_value.value = 0;
+					}
+					member.enum_value.resolved = true;
+				}
+				// Also update the original references.
+				member.enum_value.parent_enum->values.write[member.enum_value.index] = member.enum_value;
+				p_class->members.write[i].enum_value = member.enum_value;
+			} break;
 			case GDScriptParser::ClassNode::Member::CLASS:
 				break; // Done later.
 			case GDScriptParser::ClassNode::Member::UNDEFINED:
@@ -2125,6 +2175,33 @@ void GDScriptAnalyzer::reduce_identifier_from_base(GDScriptParser::IdentifierNod
 
 void GDScriptAnalyzer::reduce_identifier(GDScriptParser::IdentifierNode *p_identifier, bool can_be_builtin) {
 	// TODO: This is opportunity to further infer types.
+
+	// Check if we are inside and enum. This allows enum values to access other elements of the same enum.
+	if (current_enum) {
+		for (int i = 0; i < current_enum->values.size(); i++) {
+			const GDScriptParser::EnumNode::Value &element = current_enum->values[i];
+			if (element.identifier->name == p_identifier->name) {
+				GDScriptParser::DataType type;
+				type.type_source = GDScriptParser::DataType::ANNOTATED_EXPLICIT;
+				type.kind = element.parent_enum->identifier ? GDScriptParser::DataType::ENUM_VALUE : GDScriptParser::DataType::BUILTIN;
+				type.builtin_type = Variant::INT;
+				type.is_constant = true;
+				if (element.parent_enum->identifier) {
+					type.enum_type = element.parent_enum->identifier->name;
+				}
+				p_identifier->set_datatype(type);
+
+				if (element.resolved) {
+					p_identifier->is_constant = true;
+					p_identifier->reduced_value = element.value;
+				} else {
+					push_error(R"(Cannot use another enum element before it was declared.)", p_identifier);
+				}
+				return; // Found anyway.
+			}
+		}
+	}
+
 	// Check if identifier is local.
 	// If that's the case, the declaration already was solved before.
 	switch (p_identifier->source) {
diff --git a/modules/gdscript/gdscript_analyzer.h b/modules/gdscript/gdscript_analyzer.h
index 06d3530cb6..da767522ad 100644
--- a/modules/gdscript/gdscript_analyzer.h
+++ b/modules/gdscript/gdscript_analyzer.h
@@ -41,6 +41,8 @@ class GDScriptAnalyzer {
 	GDScriptParser *parser = nullptr;
 	HashMap<String, Ref<GDScriptParserRef>> depended_parsers;
 
+	const GDScriptParser::EnumNode *current_enum = nullptr;
+
 	Error resolve_inheritance(GDScriptParser::ClassNode *p_class, bool p_recursive = true);
 	GDScriptParser::DataType resolve_datatype(GDScriptParser::TypeNode *p_type);
 
diff --git a/modules/gdscript/gdscript_parser.cpp b/modules/gdscript/gdscript_parser.cpp
index a0f37fcf23..e53a1090fe 100644
--- a/modules/gdscript/gdscript_parser.cpp
+++ b/modules/gdscript/gdscript_parser.cpp
@@ -1022,8 +1022,6 @@ GDScriptParser::EnumNode *GDScriptParser::parse_enum() {
 	push_multiline(true);
 	consume(GDScriptTokenizer::Token::BRACE_OPEN, vformat(R"(Expected "{" after %s.)", named ? "enum name" : R"("enum")"));
 
-	int current_value = 0;
-
 	do {
 		if (check(GDScriptTokenizer::Token::BRACE_CLOSE)) {
 			break; // Allow trailing comma.
@@ -1031,6 +1029,9 @@ GDScriptParser::EnumNode *GDScriptParser::parse_enum() {
 		if (consume(GDScriptTokenizer::Token::IDENTIFIER, R"(Expected identifer for enum key.)")) {
 			EnumNode::Value item;
 			item.identifier = parse_identifier();
+			item.parent_enum = enum_node;
+			item.line = previous.start_line;
+			item.leftmost_column = previous.leftmost_column;
 
 			if (!named) {
 				// TODO: Abstract this recursive member check.
@@ -1045,18 +1046,15 @@ GDScriptParser::EnumNode *GDScriptParser::parse_enum() {
 			}
 
 			if (match(GDScriptTokenizer::Token::EQUAL)) {
-				if (consume(GDScriptTokenizer::Token::LITERAL, R"(Expected integer value after "=".)")) {
-					item.custom_value = parse_literal();
-
-					if (item.custom_value->value.get_type() != Variant::INT) {
-						push_error(R"(Expected integer value after "=".)");
-						item.custom_value = nullptr;
-					} else {
-						current_value = item.custom_value->value;
-					}
+				ExpressionNode *value = parse_expression(false);
+				if (value == nullptr) {
+					push_error(R"(Expected expression value after "=".)");
 				}
+				item.custom_value = value;
 			}
-			item.value = current_value++;
+			item.rightmost_column = previous.rightmost_column;
+
+			item.index = enum_node->values.size();
 			enum_node->values.push_back(item);
 			if (!named) {
 				// Add as member of current class.
diff --git a/modules/gdscript/gdscript_parser.h b/modules/gdscript/gdscript_parser.h
index edfe330c0c..7d8ae7fc55 100644
--- a/modules/gdscript/gdscript_parser.h
+++ b/modules/gdscript/gdscript_parser.h
@@ -405,7 +405,10 @@ public:
 	struct EnumNode : public Node {
 		struct Value {
 			IdentifierNode *identifier = nullptr;
-			LiteralNode *custom_value = nullptr;
+			ExpressionNode *custom_value = nullptr;
+			EnumNode *parent_enum = nullptr;
+			int index = -1;
+			bool resolved = false;
 			int value = 0;
 			int line = 0;
 			int leftmost_column = 0;