Merge pull request #36141 from Chaosus/shader_struct_member_arrays

Added support for arrays as shader struct members

3 files changed, 270 insertions, 16 deletions

diff --git a/servers/visual/rasterizer_rd/shader_compiler_rd.cpp b/servers/visual/rasterizer_rd/shader_compiler_rd.cpp
index 851deb367f..b2cbac8a09 100644
--- a/servers/visual/rasterizer_rd/shader_compiler_rd.cpp
+++ b/servers/visual/rasterizer_rd/shader_compiler_rd.cpp
@@ -390,6 +390,11 @@ String ShaderCompilerRD::_dump_node_code(const SL::Node *p_node, int p_level, Ge
 					}
 					struct_code += " ";
 					struct_code += m->name;
+					if (m->array_size > 0) {
+						struct_code += "[";
+						struct_code += itos(m->array_size);
+						struct_code += "]";
+					}
 					struct_code += ";\n";
 				}
 				struct_code += "}";
@@ -701,6 +706,26 @@ String ShaderCompilerRD::_dump_node_code(const SL::Node *p_node, int p_level, Ge
 			}
 
 		} break;
+		case SL::Node::TYPE_ARRAY_CONSTRUCT: {
+			SL::ArrayConstructNode *acnode = (SL::ArrayConstructNode *)p_node;
+			int sz = acnode->initializer.size();
+			if (acnode->datatype == SL::TYPE_STRUCT) {
+				code += _mkid(acnode->struct_name);
+			} else {
+				code += _typestr(acnode->datatype);
+			}
+			code += "[";
+			code += itos(acnode->initializer.size());
+			code += "]";
+			code += "(";
+			for (int i = 0; i < sz; i++) {
+				code += _dump_node_code(acnode->initializer[i], p_level, r_gen_code, p_actions, p_default_actions, p_assigning);
+				if (i != sz - 1) {
+					code += ", ";
+				}
+			}
+			code += ")";
+		} break;
 		case SL::Node::TYPE_ARRAY_DECLARATION: {
 
 			SL::ArrayDeclarationNode *adnode = (SL::ArrayDeclarationNode *)p_node;
@@ -1000,6 +1025,11 @@ String ShaderCompilerRD::_dump_node_code(const SL::Node *p_node, int p_level, Ge
 		case SL::Node::TYPE_MEMBER: {
 			SL::MemberNode *mnode = (SL::MemberNode *)p_node;
 			code = _dump_node_code(mnode->owner, p_level, r_gen_code, p_actions, p_default_actions, p_assigning) + "." + mnode->name;
+			if (mnode->index_expression != NULL) {
+				code += "[";
+				code += _dump_node_code(mnode->index_expression, p_level, r_gen_code, p_actions, p_default_actions, p_assigning);
+				code += "]";
+			}
 
 		} break;
 	}
diff --git a/servers/visual/shader_language.cpp b/servers/visual/shader_language.cpp
index c2bdc6d7e4..2a0492709c 100644
--- a/servers/visual/shader_language.cpp
+++ b/servers/visual/shader_language.cpp
@@ -3200,17 +3200,159 @@ ShaderLanguage::Node *ShaderLanguage::_parse_expression(BlockNode *p_block, cons
 					func->arguments.push_back(funcname);
 
 					for (int i = 0; i < pstruct->members.size(); i++) {
-						Node *nexpr = _parse_and_reduce_expression(p_block, p_builtin_types);
-						if (!nexpr) {
-							return NULL;
-						}
-						Node *node = pstruct->members[i];
+						Node *nexpr;
 
-						if (!_compare_datatypes_in_nodes(pstruct->members[i], nexpr)) {
-							String type_name = nexpr->get_datatype() == TYPE_STRUCT ? nexpr->get_datatype_name() : get_datatype_name(nexpr->get_datatype());
-							String type_name2 = node->get_datatype() == TYPE_STRUCT ? node->get_datatype_name() : get_datatype_name(node->get_datatype());
-							_set_error("Invalid assignment of '" + type_name + "' to '" + type_name2 + "'");
-							return NULL;
+						if (pstruct->members[i]->array_size != 0) {
+
+							DataType type = pstruct->members[i]->get_datatype();
+							String struct_name = pstruct->members[i]->struct_name;
+							int array_size = pstruct->members[i]->array_size;
+
+							DataType type2;
+							String struct_name2 = "";
+							int array_size2 = 0;
+
+							bool auto_size = false;
+
+							tk = _get_token();
+
+							if (tk.type == TK_CURLY_BRACKET_OPEN) {
+								auto_size = true;
+							} else {
+
+								if (shader->structs.has(tk.text)) {
+									type2 = TYPE_STRUCT;
+									struct_name2 = tk.text;
+								} else {
+									if (!is_token_variable_datatype(tk.type)) {
+										_set_error("Invalid data type for array");
+										return NULL;
+									}
+									type2 = get_token_datatype(tk.type);
+								}
+
+								tk = _get_token();
+								if (tk.type == TK_BRACKET_OPEN) {
+									TkPos pos2 = _get_tkpos();
+									tk = _get_token();
+									if (tk.type == TK_BRACKET_CLOSE) {
+										array_size2 = array_size;
+										tk = _get_token();
+									} else {
+										_set_tkpos(pos2);
+
+										Node *n = _parse_and_reduce_expression(p_block, p_builtin_types);
+										if (!n || n->type != Node::TYPE_CONSTANT || n->get_datatype() != TYPE_INT) {
+											_set_error("Expected single integer constant > 0");
+											return NULL;
+										}
+
+										ConstantNode *cnode = (ConstantNode *)n;
+										if (cnode->values.size() == 1) {
+											array_size2 = cnode->values[0].sint;
+											if (array_size2 <= 0) {
+												_set_error("Expected single integer constant > 0");
+												return NULL;
+											}
+										} else {
+											_set_error("Expected single integer constant > 0");
+											return NULL;
+										}
+
+										tk = _get_token();
+										if (tk.type != TK_BRACKET_CLOSE) {
+											_set_error("Expected ']'");
+											return NULL;
+										} else {
+											tk = _get_token();
+										}
+									}
+								} else {
+									_set_error("Expected '['");
+									return NULL;
+								}
+
+								if (type != type2 || struct_name != struct_name2 || array_size != array_size2) {
+									String error_str = "Cannot convert from '";
+									if (type2 == TYPE_STRUCT) {
+										error_str += struct_name2;
+									} else {
+										error_str += get_datatype_name(type2);
+									}
+									error_str += "[";
+									error_str += itos(array_size2);
+									error_str += "]'";
+									error_str += " to '";
+									if (type == TYPE_STRUCT) {
+										error_str += struct_name;
+									} else {
+										error_str += get_datatype_name(type);
+									}
+									error_str += "[";
+									error_str += itos(array_size);
+									error_str += "]'";
+									_set_error(error_str);
+									return NULL;
+								}
+							}
+
+							ArrayConstructNode *an = alloc_node<ArrayConstructNode>();
+							an->datatype = type;
+							an->struct_name = struct_name;
+
+							if (tk.type == TK_PARENTHESIS_OPEN || auto_size) { // initialization
+								while (true) {
+
+									Node *n = _parse_and_reduce_expression(p_block, p_builtin_types);
+									if (!n) {
+										return NULL;
+									}
+
+									if (type != n->get_datatype() || struct_name != n->get_datatype_name()) {
+										_set_error("Invalid assignment of '" + (n->get_datatype() == TYPE_STRUCT ? n->get_datatype_name() : get_datatype_name(n->get_datatype())) + "' to '" + (type == TYPE_STRUCT ? struct_name : get_datatype_name(type)) + "'");
+										return NULL;
+									}
+
+									tk = _get_token();
+									if (tk.type == TK_COMMA) {
+										an->initializer.push_back(n);
+										continue;
+									} else if (!auto_size && tk.type == TK_PARENTHESIS_CLOSE) {
+										an->initializer.push_back(n);
+										break;
+									} else if (auto_size && tk.type == TK_CURLY_BRACKET_CLOSE) {
+										an->initializer.push_back(n);
+										break;
+									} else {
+										if (auto_size)
+											_set_error("Expected '}' or ','");
+										else
+											_set_error("Expected ')' or ','");
+										return NULL;
+									}
+								}
+								if (an->initializer.size() != array_size) {
+									_set_error("Array size mismatch");
+									return NULL;
+								}
+							} else {
+								_set_error("Expected array initialization!");
+								return NULL;
+							}
+
+							nexpr = an;
+						} else {
+							nexpr = _parse_and_reduce_expression(p_block, p_builtin_types);
+							if (!nexpr) {
+								return NULL;
+							}
+							Node *node = pstruct->members[i];
+							if (!_compare_datatypes_in_nodes(pstruct->members[i], nexpr)) {
+								String type_name = nexpr->get_datatype() == TYPE_STRUCT ? nexpr->get_datatype_name() : get_datatype_name(nexpr->get_datatype());
+								String type_name2 = node->get_datatype() == TYPE_STRUCT ? node->get_datatype_name() : get_datatype_name(node->get_datatype());
+								_set_error("Invalid assignment of '" + type_name + "' to '" + type_name2 + "'");
+								return NULL;
+							}
 						}
 
 						if (i + 1 < pstruct->members.size()) {
@@ -3463,7 +3605,6 @@ ShaderLanguage::Node *ShaderLanguage::_parse_expression(BlockNode *p_block, cons
 					expr = varname;
 				}
 			}
-
 		} else if (tk.type == TK_OP_ADD) {
 			continue; //this one does nothing
 		} else if (tk.type == TK_OP_SUB || tk.type == TK_OP_NOT || tk.type == TK_OP_BIT_INVERT || tk.type == TK_OP_INCREMENT || tk.type == TK_OP_DECREMENT) {
@@ -3482,7 +3623,6 @@ ShaderLanguage::Node *ShaderLanguage::_parse_expression(BlockNode *p_block, cons
 
 			expression.push_back(e);
 			continue;
-
 		} else {
 			_set_error("Expected expression, found: " + get_token_text(tk));
 			return NULL;
@@ -3526,6 +3666,7 @@ ShaderLanguage::Node *ShaderLanguage::_parse_expression(BlockNode *p_block, cons
 				bool ok = true;
 				DataType member_type = TYPE_VOID;
 				StringName member_struct_name = "";
+				int array_size = 0;
 				switch (dt) {
 					case TYPE_STRUCT: {
 						ok = false;
@@ -3535,6 +3676,7 @@ ShaderLanguage::Node *ShaderLanguage::_parse_expression(BlockNode *p_block, cons
 							for (List<MemberNode *>::Element *E = n->members.front(); E; E = E->next()) {
 								if (String(E->get()->name) == member_name) {
 									member_type = E->get()->datatype;
+									array_size = E->get()->array_size;
 									if (member_type == TYPE_STRUCT) {
 										member_struct_name = E->get()->struct_name;
 									}
@@ -3673,8 +3815,52 @@ ShaderLanguage::Node *ShaderLanguage::_parse_expression(BlockNode *p_block, cons
 				mn->datatype = member_type;
 				mn->base_struct_name = st;
 				mn->struct_name = member_struct_name;
+				mn->array_size = array_size;
 				mn->name = ident;
 				mn->owner = expr;
+				if (array_size > 0) {
+
+					tk = _get_token();
+					if (tk.type == TK_PERIOD) {
+						_set_error("Nested array length() is not yet implemented");
+						return NULL;
+					} else if (tk.type == TK_BRACKET_OPEN) {
+
+						Node *index_expression = _parse_and_reduce_expression(p_block, p_builtin_types);
+						if (!index_expression)
+							return NULL;
+
+						if (index_expression->get_datatype() != TYPE_INT && index_expression->get_datatype() != TYPE_UINT) {
+							_set_error("Only integer expressions are allowed for indexing");
+							return NULL;
+						}
+
+						if (index_expression->type == Node::TYPE_CONSTANT) {
+							ConstantNode *cnode = (ConstantNode *)index_expression;
+							if (cnode) {
+								if (!cnode->values.empty()) {
+									int value = cnode->values[0].sint;
+									if (value < 0 || value >= array_size) {
+										_set_error(vformat("Index [%s] out of range [%s..%s]", value, 0, array_size - 1));
+										return NULL;
+									}
+								}
+							}
+						}
+
+						tk = _get_token();
+						if (tk.type != TK_BRACKET_CLOSE) {
+							_set_error("Expected ']'");
+							return NULL;
+						}
+						mn->index_expression = index_expression;
+
+					} else {
+						_set_error("Expected '[' or '.'");
+						return NULL;
+					}
+				}
+
 				expr = mn;
 
 				//todo
@@ -3769,7 +3955,7 @@ ShaderLanguage::Node *ShaderLanguage::_parse_expression(BlockNode *p_block, cons
 						}
 						break;
 					default: {
-						_set_error("Object of type '" + get_datatype_name(expr->get_datatype()) + "' can't be indexed");
+						_set_error("Object of type '" + (expr->get_datatype() == TYPE_STRUCT ? expr->get_datatype_name() : get_datatype_name(expr->get_datatype())) + "' can't be indexed");
 						return NULL;
 					}
 				}
@@ -5353,11 +5539,33 @@ Error ShaderLanguage::_parse_shader(const Map<StringName, FunctionInfo> &p_funct
 						member->datatype = type;
 						member->struct_name = struct_name;
 						member->name = tk.text;
-						st_node->members.push_back(member);
 
 						tk = _get_token();
+						if (tk.type == TK_BRACKET_OPEN) {
+							tk = _get_token();
+							if (tk.type == TK_INT_CONSTANT && tk.constant > 0) {
+								member->array_size = (int)tk.constant;
+
+								tk = _get_token();
+								if (tk.type == TK_BRACKET_CLOSE) {
+									tk = _get_token();
+									if (tk.type != TK_SEMICOLON) {
+										_set_error("Expected ';'");
+										return ERR_PARSE_ERROR;
+									}
+								} else {
+									_set_error("Expected ']'");
+									return ERR_PARSE_ERROR;
+								}
+							} else {
+								_set_error("Expected single integer constant > 0");
+								return ERR_PARSE_ERROR;
+							}
+						}
+						st_node->members.push_back(member);
+
 						if (tk.type != TK_SEMICOLON) {
-							_set_error("Expected ';'");
+							_set_error("Expected ']' or ';'");
 							return ERR_PARSE_ERROR;
 						}
 						member_count++;
diff --git a/servers/visual/shader_language.h b/servers/visual/shader_language.h
index 4d474086a6..aac5795e85 100644
--- a/servers/visual/shader_language.h
+++ b/servers/visual/shader_language.h
@@ -332,6 +332,7 @@ public:
 			TYPE_MEMBER,
 			TYPE_ARRAY,
 			TYPE_ARRAY_DECLARATION,
+			TYPE_ARRAY_CONSTRUCT,
 			TYPE_STRUCT,
 		};
 
@@ -427,6 +428,17 @@ public:
 				is_const(false) {}
 	};
 
+	struct ArrayConstructNode : public Node {
+		DataType datatype;
+		String struct_name;
+		Vector<Node *> initializer;
+
+		ArrayConstructNode() :
+				Node(TYPE_ARRAY_CONSTRUCT),
+				datatype(TYPE_VOID) {
+		}
+	};
+
 	struct ArrayDeclarationNode : public Node {
 		DataPrecision precision;
 		DataType datatype;
@@ -520,9 +532,11 @@ public:
 		StringName base_struct_name;
 		DataPrecision precision;
 		DataType datatype;
+		int array_size;
 		StringName struct_name;
 		StringName name;
 		Node *owner;
+		Node *index_expression;
 
 		virtual DataType get_datatype() const { return datatype; }
 		virtual String get_datatype_name() const { return String(struct_name); }
@@ -531,7 +545,9 @@ public:
 				Node(TYPE_MEMBER),
 				basetype(TYPE_VOID),
 				datatype(TYPE_VOID),
-				owner(NULL) {}
+				array_size(0),
+				owner(NULL),
+				index_expression(NULL) {}
 	};
 
 	struct StructNode : public Node {