Fix using post-init shader array constructors

3 files changed, 179 insertions, 144 deletions

diff --git a/servers/rendering/renderer_rd/shader_compiler_rd.cpp b/servers/rendering/renderer_rd/shader_compiler_rd.cpp
index c6ea906ea9..ae392b7586 100644
--- a/servers/rendering/renderer_rd/shader_compiler_rd.cpp
+++ b/servers/rendering/renderer_rd/shader_compiler_rd.cpp
@@ -920,7 +920,7 @@ String ShaderCompilerRD::_dump_node_code(const SL::Node *p_node, int p_level, Ge
 			if (adnode->datatype == SL::TYPE_STRUCT) {
 				declaration += _mkid(adnode->struct_name);
 			} else {
-				declaration = _prestr(adnode->precision) + _typestr(adnode->datatype);
+				declaration += _prestr(adnode->precision) + _typestr(adnode->datatype);
 			}
 			for (int i = 0; i < adnode->declarations.size(); i++) {
 				if (i > 0) {
@@ -990,12 +990,13 @@ String ShaderCompilerRD::_dump_node_code(const SL::Node *p_node, int p_level, Ge
 			if (anode->call_expression != nullptr) {
 				code += ".";
 				code += _dump_node_code(anode->call_expression, p_level, r_gen_code, p_actions, p_default_actions, p_assigning, false);
-			}
-
-			if (anode->index_expression != nullptr) {
+			} else if (anode->index_expression != nullptr) {
 				code += "[";
 				code += _dump_node_code(anode->index_expression, p_level, r_gen_code, p_actions, p_default_actions, p_assigning);
 				code += "]";
+			} else if (anode->assign_expression != nullptr) {
+				code += "=";
+				code += _dump_node_code(anode->assign_expression, p_level, r_gen_code, p_actions, p_default_actions, true, false);
 			}
 
 			if (anode->name == time_name) {
@@ -1233,8 +1234,10 @@ String ShaderCompilerRD::_dump_node_code(const SL::Node *p_node, int p_level, Ge
 				code += "[";
 				code += _dump_node_code(mnode->index_expression, p_level, r_gen_code, p_actions, p_default_actions, p_assigning);
 				code += "]";
+			} else if (mnode->assign_expression != nullptr) {
+				code += "=";
+				code += _dump_node_code(mnode->assign_expression, p_level, r_gen_code, p_actions, p_default_actions, true, false);
 			}
-
 		} break;
 	}
 
diff --git a/servers/rendering/shader_language.cpp b/servers/rendering/shader_language.cpp
index 3b73f6a8d6..c6eef48dfa 100644
--- a/servers/rendering/shader_language.cpp
+++ b/servers/rendering/shader_language.cpp
@@ -913,6 +913,7 @@ void ShaderLanguage::clear() {
 	char_idx = 0;
 	error_set = false;
 	error_str = "";
+	last_const = false;
 	while (nodes) {
 		Node *n = nodes;
 		nodes = nodes->next;
@@ -3252,6 +3253,137 @@ bool ShaderLanguage::_propagate_function_call_sampler_builtin_reference(StringNa
 	ERR_FAIL_V(false); //bug? function not found
 }
 
+ShaderLanguage::Node *ShaderLanguage::_parse_array_constructor(BlockNode *p_block, const FunctionInfo &p_function_info, DataType p_type, const StringName &p_struct_name, int p_array_size) {
+	DataType type = TYPE_VOID;
+	String struct_name = "";
+	int array_size = 0;
+	bool auto_size = false;
+	Token tk = _get_token();
+
+	if (tk.type == TK_CURLY_BRACKET_OPEN) {
+		auto_size = true;
+	} else {
+		if (shader->structs.has(tk.text)) {
+			type = TYPE_STRUCT;
+			struct_name = tk.text;
+		} else {
+			if (!is_token_variable_datatype(tk.type)) {
+				_set_error("Invalid data type for array");
+				return nullptr;
+			}
+			type = get_token_datatype(tk.type);
+		}
+		tk = _get_token();
+		if (tk.type == TK_BRACKET_OPEN) {
+			TkPos pos = _get_tkpos();
+			tk = _get_token();
+			if (tk.type == TK_BRACKET_CLOSE) {
+				array_size = p_array_size;
+				tk = _get_token();
+			} else {
+				_set_tkpos(pos);
+
+				Node *n = _parse_and_reduce_expression(p_block, p_function_info);
+				if (!n || n->type != Node::TYPE_CONSTANT || n->get_datatype() != TYPE_INT) {
+					_set_error("Expected single integer constant > 0");
+					return nullptr;
+				}
+
+				ConstantNode *cnode = (ConstantNode *)n;
+				if (cnode->values.size() == 1) {
+					array_size = cnode->values[0].sint;
+					if (array_size <= 0) {
+						_set_error("Expected single integer constant > 0");
+						return nullptr;
+					}
+				} else {
+					_set_error("Expected single integer constant > 0");
+					return nullptr;
+				}
+
+				tk = _get_token();
+				if (tk.type != TK_BRACKET_CLOSE) {
+					_set_error("Expected ']'");
+					return nullptr;
+				} else {
+					tk = _get_token();
+				}
+			}
+		} else {
+			_set_error("Expected '['");
+			return nullptr;
+		}
+
+		if (type != p_type || struct_name != p_struct_name || array_size != p_array_size) {
+			String error_str = "Cannot convert from '";
+			if (type == TYPE_STRUCT) {
+				error_str += struct_name;
+			} else {
+				error_str += get_datatype_name(type);
+			}
+			error_str += "[";
+			error_str += itos(array_size);
+			error_str += "]'";
+			error_str += " to '";
+			if (type == TYPE_STRUCT) {
+				error_str += p_struct_name;
+			} else {
+				error_str += get_datatype_name(p_type);
+			}
+			error_str += "[";
+			error_str += itos(p_array_size);
+			error_str += "]'";
+			_set_error(error_str);
+			return nullptr;
+		}
+	}
+
+	ArrayConstructNode *an = alloc_node<ArrayConstructNode>();
+	an->datatype = p_type;
+	an->struct_name = p_struct_name;
+
+	if (tk.type == TK_PARENTHESIS_OPEN || auto_size) { // initialization
+		while (true) {
+			Node *n = _parse_and_reduce_expression(p_block, p_function_info);
+			if (!n) {
+				return nullptr;
+			}
+
+			if (p_type != n->get_datatype() || p_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 nullptr;
+			}
+
+			tk = _get_token();
+			if (tk.type == TK_COMMA) {
+				an->initializer.push_back(n);
+			} 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 nullptr;
+			}
+		}
+		if (an->initializer.size() != p_array_size) {
+			_set_error("Array size mismatch");
+			return nullptr;
+		}
+	} else {
+		_set_error("Expected array initialization!");
+		return nullptr;
+	}
+
+	return an;
+}
+
 ShaderLanguage::Node *ShaderLanguage::_parse_expression(BlockNode *p_block, const FunctionInfo &p_function_info) {
 	Vector<Expression> expression;
 
@@ -3395,142 +3527,10 @@ ShaderLanguage::Node *ShaderLanguage::_parse_expression(BlockNode *p_block, cons
 						Node *nexpr;
 
 						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 nullptr;
-									}
-									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_function_info);
-										if (!n || n->type != Node::TYPE_CONSTANT || n->get_datatype() != TYPE_INT) {
-											_set_error("Expected single integer constant > 0");
-											return nullptr;
-										}
-
-										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 nullptr;
-											}
-										} else {
-											_set_error("Expected single integer constant > 0");
-											return nullptr;
-										}
-
-										tk = _get_token();
-										if (tk.type != TK_BRACKET_CLOSE) {
-											_set_error("Expected ']'");
-											return nullptr;
-										} else {
-											tk = _get_token();
-										}
-									}
-								} else {
-									_set_error("Expected '['");
-									return nullptr;
-								}
-
-								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 nullptr;
-								}
-							}
-
-							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_function_info);
-									if (!n) {
-										return nullptr;
-									}
-
-									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 nullptr;
-									}
-
-									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 nullptr;
-									}
-								}
-								if (an->initializer.size() != array_size) {
-									_set_error("Array size mismatch");
-									return nullptr;
-								}
-							} else {
-								_set_error("Expected array initialization!");
+							nexpr = _parse_array_constructor(p_block, p_function_info, pstruct->members[i]->get_datatype(), pstruct->members[i]->struct_name, pstruct->members[i]->array_size);
+							if (!nexpr) {
 								return nullptr;
 							}
-
-							nexpr = an;
 						} else {
 							nexpr = _parse_and_reduce_expression(p_block, p_function_info);
 							if (!nexpr) {
@@ -3733,6 +3733,7 @@ ShaderLanguage::Node *ShaderLanguage::_parse_expression(BlockNode *p_block, cons
 			} else {
 				//an identifier
 
+				last_const = false;
 				_set_tkpos(pos);
 
 				DataType data_type;
@@ -3760,6 +3761,7 @@ ShaderLanguage::Node *ShaderLanguage::_parse_expression(BlockNode *p_block, cons
 						_set_error("Unknown identifier in expression: " + String(identifier));
 						return nullptr;
 					}
+					last_const = is_const;
 
 					if (ident_type == IDENTIFIER_FUNCTION) {
 						_set_error("Can't use function as identifier: " + String(identifier));
@@ -3769,16 +3771,30 @@ ShaderLanguage::Node *ShaderLanguage::_parse_expression(BlockNode *p_block, cons
 
 				Node *index_expression = nullptr;
 				Node *call_expression = nullptr;
+				Node *assign_expression = nullptr;
 
 				if (array_size > 0) {
 					tk = _get_token();
 
-					if (tk.type != TK_BRACKET_OPEN && tk.type != TK_PERIOD) {
-						_set_error("Expected '[' or '.'");
+					if (tk.type != TK_BRACKET_OPEN && tk.type != TK_PERIOD && tk.type != TK_OP_ASSIGN) {
+						_set_error("Expected '[','.' or '='");
 						return nullptr;
 					}
 
-					if (tk.type == TK_PERIOD) {
+					if (tk.type == TK_OP_ASSIGN) {
+						if (is_const) {
+							_set_error("Constants cannot be modified.");
+							return nullptr;
+						}
+						if (shader->varyings.has(identifier) && current_function != String("vertex")) {
+							_set_error("Varyings can only be assigned in vertex function.");
+							return nullptr;
+						}
+						assign_expression = _parse_array_constructor(p_block, p_function_info, data_type, struct_name, array_size);
+						if (!assign_expression) {
+							return nullptr;
+						}
+					} else if (tk.type == TK_PERIOD) {
 						completion_class = TAG_ARRAY;
 						p_block->block_tag = SubClassTag::TAG_ARRAY;
 						call_expression = _parse_and_reduce_expression(p_block, p_function_info);
@@ -3825,6 +3841,7 @@ ShaderLanguage::Node *ShaderLanguage::_parse_expression(BlockNode *p_block, cons
 					arrname->struct_name = struct_name;
 					arrname->index_expression = index_expression;
 					arrname->call_expression = call_expression;
+					arrname->assign_expression = assign_expression;
 					arrname->is_const = is_const;
 					expr = arrname;
 
@@ -4165,7 +4182,18 @@ ShaderLanguage::Node *ShaderLanguage::_parse_expression(BlockNode *p_block, cons
 
 				if (array_size > 0) {
 					tk = _get_token();
-					if (tk.type == TK_PERIOD) {
+					if (tk.type == TK_OP_ASSIGN) {
+						if (last_const) {
+							last_const = false;
+							_set_error("Constants cannot be modified.");
+							return nullptr;
+						}
+						Node *assign_expression = _parse_array_constructor(p_block, p_function_info, member_type, member_struct_name, array_size);
+						if (!assign_expression) {
+							return nullptr;
+						}
+						mn->assign_expression = assign_expression;
+					} else if (tk.type == TK_PERIOD) {
 						_set_error("Nested array length() is not yet implemented");
 						return nullptr;
 					} else if (tk.type == TK_BRACKET_OPEN) {
@@ -4200,7 +4228,7 @@ ShaderLanguage::Node *ShaderLanguage::_parse_expression(BlockNode *p_block, cons
 						mn->index_expression = index_expression;
 
 					} else {
-						_set_error("Expected '[' or '.'");
+						_set_error("Expected '[','.' or '='");
 						return nullptr;
 					}
 				}
diff --git a/servers/rendering/shader_language.h b/servers/rendering/shader_language.h
index 61bd83bc72..d8fc316f47 100644
--- a/servers/rendering/shader_language.h
+++ b/servers/rendering/shader_language.h
@@ -414,6 +414,7 @@ public:
 		StringName name;
 		Node *index_expression = nullptr;
 		Node *call_expression = nullptr;
+		Node *assign_expression = nullptr;
 		bool is_const = false;
 
 		virtual DataType get_datatype() const { return datatype_cache; }
@@ -528,6 +529,7 @@ public:
 		StringName name;
 		Node *owner = nullptr;
 		Node *index_expression = nullptr;
+		Node *assign_expression = nullptr;
 		bool has_swizzling_duplicates = false;
 
 		virtual DataType get_datatype() const { return datatype; }
@@ -776,6 +778,7 @@ private:
 	int tk_line;
 
 	StringName current_function;
+	bool last_const = false;
 
 	struct TkPos {
 		int char_idx;
@@ -863,6 +866,7 @@ private:
 	bool _propagate_function_call_sampler_builtin_reference(StringName p_name, int p_argument, const StringName &p_builtin);
 
 	Node *_parse_expression(BlockNode *p_block, const FunctionInfo &p_function_info);
+	Node *_parse_array_constructor(BlockNode *p_block, const FunctionInfo &p_function_info, DataType p_type, const StringName &p_struct_name, int p_array_size);
 	ShaderLanguage::Node *_reduce_expression(BlockNode *p_block, ShaderLanguage::Node *p_node);
 
 	Node *_parse_and_reduce_expression(BlockNode *p_block, const FunctionInfo &p_function_info);