add array variable assignment (#867)

internal/engine/integration/procedure_test.go

@@ -344,6 +344,16 @@ func Test_Procedures(t *testing.T) {
 			}`,
 			outputs: [][]any{{false}},
 		},
+		{
+			name: "array",
+			procedure: `procedure assign_array() public view returns (ints int[]) {
+				$arr int[] := [1, 2, 3];
+
+				$arr[2] := 4;
+				return $arr;
+			}`,
+			outputs: [][]any{{[]any{int64(1), int64(4), int64(3)}}},
+		},
 	}
 
 	for _, test := range tests {

parse/analyze.go

@@ -2234,12 +2234,18 @@ func (p *procedureAnalyzer) VisitProcedureStmtAssignment(p0 *ProcedureStmtAssign
 		return zeroProcedureReturn()
 	}
 
-	_, ok = p.variables[p0.Variable.String()]
-	if !ok {
-		// if it does not exist, we can declare it here.
-		p.variables[p0.Variable.String()] = dt
-		p.markDeclared(p0.Variable, p0.Variable.String(), dt)
-		return zeroProcedureReturn()
+	// the variable can be either an ExpressionVariable or an ExpressionArrayAccess
+	// If it is an ExpressionVariable, we need to declare it
+
+	exprVar, ok := p0.Variable.(*ExpressionVariable)
+	if ok {
+		_, ok = p.variables[exprVar.String()]
+		if !ok {
+			// if it does not exist, we can declare it here.
+			p.variables[exprVar.String()] = dt
+			p.markDeclared(p0.Variable, exprVar.String(), dt)
+			return zeroProcedureReturn()
+		}
 	}
 
 	// the type can be inferred from the value.

parse/antlr.go

@@ -2018,10 +2018,17 @@ func (s *schemaVisitor) VisitVariable_or_underscore(ctx *gen.Variable_or_undersc
 }
 
 func (s *schemaVisitor) VisitStmt_variable_assignment(ctx *gen.Stmt_variable_assignmentContext) any {
-	stmt := &ProcedureStmtAssign{
-		Variable: varFromTerminalNode(ctx.VARIABLE()),
-		Value:    ctx.Procedure_expr().Accept(s).(Expression),
+	stmt := &ProcedureStmtAssign{}
+
+	stmt.Variable = ctx.Procedure_expr(0).Accept(s).(Expression)
+	// this can either be a variable or an array access
+	switch v := stmt.Variable.(type) {
+	case *ExpressionVariable, *ExpressionArrayAccess:
+		// ok
+	default:
+		s.errs.RuleErr(ctx.Procedure_expr(0), ErrSyntax, "cannot assign to %T", v)
 	}
+	stmt.Value = ctx.Procedure_expr(1).Accept(s).(Expression)
 
 	if ctx.Type_() != nil {
 		stmt.Type = ctx.Type_().Accept(s).(*types.DataType)

parse/ast.go

@@ -836,7 +836,7 @@ func (p *ProcedureStmtDeclaration) Accept(v Visitor) any {
 type ProcedureStmtAssign struct {
 	baseProcedureStmt
 	// Variable is the variable that is being assigned.
-	Variable *ExpressionVariable
+	Variable Expression
 	// Type is the type of the variable.
 	// It can be nil if the variable is not being assigned,
 	// or if the type should be inferred.

parse/grammar/KuneiformParser.g4

@@ -373,7 +373,7 @@ proc_statement:
     VARIABLE type SCOL                                                                                  # stmt_variable_declaration
     // stmt_procedure_call must go above stmt_variable_assignment due to lexer ambiguity
     | ((variable_or_underscore) (COMMA (variable_or_underscore))* ASSIGN)? procedure_function_call SCOL # stmt_procedure_call
-    | VARIABLE type? ASSIGN procedure_expr SCOL                                                         # stmt_variable_assignment
+    | procedure_expr type? ASSIGN procedure_expr SCOL                                                         # stmt_variable_assignment
     | FOR receiver=VARIABLE IN (range|target_variable=variable|sql_statement) LBRACE proc_statement* RBRACE  # stmt_for_loop
     | IF if_then_block (ELSEIF if_then_block)* (ELSE LBRACE proc_statement* RBRACE)?                         # stmt_if
     | sql_statement SCOL                                                                                # stmt_sql