Parse For Stmts

4 files changed, 281 insertions, 4 deletions

diff --git a/pkg/lang/ast/stmt.go b/pkg/lang/ast/stmt.go
index 6395f57..65b2532 100644
--- a/pkg/lang/ast/stmt.go
+++ b/pkg/lang/ast/stmt.go
@@ -11,6 +11,8 @@ const (
 	StmtKindObjectDecl
 	StmtKindVarDecl
 	StmtKindIf
+	StmtKindForCond
+	StmtKindForIn
 	StmtKindTry
 	StmtKindReturn
 	StmtKindContinue
@@ -47,7 +49,18 @@ type StmtVarDecl struct {
 }
 
 type StmtIf struct {
-	Conds  []CondNode
+	Conds []CondNode
+}
+
+type StmtForCond struct {
+	Cond Expr
+	Do   BlockNode
+}
+
+type StmtForIn struct {
+	Name       IdentNode
+	Collection Expr
+	Do         BlockNode
 }
 
 type StmtTry struct {
diff --git a/pkg/lang/parser/exprs.go b/pkg/lang/parser/exprs.go
index 8c2b475..56904a1 100644
--- a/pkg/lang/parser/exprs.go
+++ b/pkg/lang/parser/exprs.go
@@ -261,6 +261,10 @@ func (p *Parser) parseArrayLitExpr() (ast.Expr, error) {
 		}
 	}
 
+	if _, err := p.expect(token.RBracket); err != nil {
+		return ast.Expr{}, err
+	}
+
 	return ast.Expr{
 		At:   bracketTok.At,
 		Kind: ast.ExprKindArrayLit,
diff --git a/pkg/lang/parser/parser_test.go b/pkg/lang/parser/parser_test.go
index b756b1b..e02ab5a 100644
--- a/pkg/lang/parser/parser_test.go
+++ b/pkg/lang/parser/parser_test.go
@@ -106,7 +106,7 @@ func TestLeftAssocBinary4Expr(t *testing.T) {
 
 	require.Equal(t, 1, len(program.Stmts))
 	require.Equal(t, ast.Stmt{
-		Kind: 11,
+		Kind: ast.StmtKindExpr,
 		Value: ast.StmtExpr{
 			Value: ast.Expr{
 				Kind: ast.ExprKindBinary,
@@ -420,6 +420,204 @@ func TestIfStmt(t *testing.T) {
 	}, program.Stmts[0])
 }
 
+func TestForCondStmt(t *testing.T) {
+	src := sourceify(
+		`for x <= 5 {`,
+		`	x = x + 1`,
+		`}`,
+	)
+
+	p := cheatWithScanner(t, src)
+	program, err := p.Parse()
+	require.NoError(t, err)
+
+	require.Equal(t, 1, len(program.Stmts))
+	require.Equal(t, ast.Stmt{
+		At:   source.NewLoc(0, 0),
+		Kind: ast.StmtKindForCond,
+		Value: ast.StmtForCond{
+			Cond: ast.Expr{
+				At:   source.NewLoc(0, 4),
+				Kind: ast.ExprKindBinary,
+				Value: ast.ExprBinary{
+					Left: ast.Expr{
+						At:   source.NewLoc(0, 4),
+						Kind: ast.ExprKindIdent,
+						Value: ast.ExprIdent{
+							Value: ast.IdentNode{
+								At:    source.NewLoc(0, 4),
+								Value: "x",
+							},
+						},
+					},
+					Op: ast.BinOpLte,
+					Right: ast.Expr{
+						At:   source.NewLoc(0, 9),
+						Kind: ast.ExprKindIntLit,
+						Value: ast.ExprIntLit{
+							Value: 5,
+						},
+					},
+				},
+			},
+			Do: ast.BlockNode{
+				At: source.NewLoc(0, 11),
+				Stmts: []ast.Stmt{
+					{
+						At:    source.NewLoc(0, 12),
+						Kind:  ast.StmtKindEmpty,
+						Value: ast.StmtEmpty{},
+					},
+					{
+						At:   source.NewLoc(1, 1),
+						Kind: ast.StmtKindExpr,
+						Value: ast.StmtExpr{
+							Value: ast.Expr{
+								At:   source.NewLoc(1, 1),
+								Kind: ast.ExprKindBinary,
+								Value: ast.ExprBinary{
+									Left: ast.Expr{
+										At:   source.NewLoc(1, 1),
+										Kind: ast.ExprKindIdent,
+										Value: ast.ExprIdent{
+											Value: ast.IdentNode{
+												At:    source.NewLoc(1, 1),
+												Value: "x",
+											},
+										},
+									},
+									Op: ast.BinOpAssign,
+									Right: ast.Expr{
+										At:   source.NewLoc(1, 5),
+										Kind: ast.ExprKindBinary,
+										Value: ast.ExprBinary{
+											Left: ast.Expr{
+												At:   source.NewLoc(1, 5),
+												Kind: ast.ExprKindIdent,
+												Value: ast.ExprIdent{
+													Value: ast.IdentNode{
+														At:    source.NewLoc(1, 5),
+														Value: "x",
+													},
+												},
+											},
+											Op: ast.BinOpPlus,
+											Right: ast.Expr{
+												At:   source.NewLoc(1, 9),
+												Kind: ast.ExprKindIntLit,
+												Value: ast.ExprIntLit{
+													Value: 1,
+												},
+											},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+	}, program.Stmts[0])
+}
+
+func TestForInStmt(t *testing.T) {
+	src := sourceify(
+		`for x in [1, 2, 3] {`,
+		`	say(x)`,
+		`}`,
+	)
+
+	p := cheatWithScanner(t, src)
+	program, err := p.Parse()
+	require.NoError(t, err)
+
+	require.Equal(t, 1, len(program.Stmts))
+	require.Equal(t, ast.Stmt{
+		At:   source.NewLoc(0, 0),
+		Kind: ast.StmtKindForIn,
+		Value: ast.StmtForIn{
+			Name: ast.IdentNode{
+				At:    source.NewLoc(0, 4),
+				Value: "x",
+			},
+			Collection: ast.Expr{
+				At:   source.NewLoc(0, 9),
+				Kind: ast.ExprKindArrayLit,
+				Value: ast.ExprArrayLit{
+					Values: []ast.Expr{
+						{
+							At:   source.NewLoc(0, 10),
+							Kind: ast.ExprKindIntLit,
+							Value: ast.ExprIntLit{
+								Value: 1,
+							},
+						},
+						{
+							At:   source.NewLoc(0, 13),
+							Kind: ast.ExprKindIntLit,
+							Value: ast.ExprIntLit{
+								Value: 2,
+							},
+						},
+						{
+							At:   source.NewLoc(0, 16),
+							Kind: ast.ExprKindIntLit,
+							Value: ast.ExprIntLit{
+								Value: 3,
+							},
+						},
+					},
+				},
+			},
+			Do: ast.BlockNode{
+				At: source.NewLoc(0, 19),
+				Stmts: []ast.Stmt{
+					{
+						At:    source.NewLoc(0, 20),
+						Kind:  ast.StmtKindEmpty,
+						Value: ast.StmtEmpty{},
+					},
+					{
+						At:   source.NewLoc(1, 1),
+						Kind: ast.StmtKindExpr,
+						Value: ast.StmtExpr{
+							Value: ast.Expr{
+								At:   source.NewLoc(1, 1),
+								Kind: ast.ExprKindCall,
+								Value: ast.ExprCall{
+									Callee: ast.Expr{
+										At:   source.NewLoc(1, 1),
+										Kind: ast.ExprKindIdent,
+										Value: ast.ExprIdent{
+											Value: ast.IdentNode{
+												At:    source.NewLoc(1, 1),
+												Value: "say",
+											},
+										},
+									},
+									Args: []ast.Expr{
+										{
+											At:   source.NewLoc(1, 5),
+											Kind: ast.ExprKindIdent,
+											Value: ast.ExprIdent{
+												Value: ast.IdentNode{
+													At:    source.NewLoc(1, 5),
+													Value: "x",
+												},
+											},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+	}, program.Stmts[0])
+}
+
 func sourceify(lines ...string) string {
 	return strings.Join(lines, "\n")
 }
diff --git a/pkg/lang/parser/stmts.go b/pkg/lang/parser/stmts.go
index 4b1077e..2a32906 100644
--- a/pkg/lang/parser/stmts.go
+++ b/pkg/lang/parser/stmts.go
@@ -18,6 +18,8 @@ func (p *Parser) parseStmt() (ast.Stmt, error) {
 		return p.parseVarDeclStmt()
 	case token.KwIf:
 		return p.parseIfStmt()
+	case token.KwFor:
+		return p.parseForStmt()
 	case token.KwTry:
 		return p.parseTryStmt()
 	case token.KwReturn:
@@ -192,7 +194,6 @@ func (p *Parser) parseIfStmt() (ast.Stmt, error) {
 		})
 	}
 
-
 	if p.peek().Kind == token.KwElse {
 		elseTok, err := p.expect(token.KwElse)
 		if err != nil {
@@ -205,7 +206,7 @@ func (p *Parser) parseIfStmt() (ast.Stmt, error) {
 		}
 
 		conds = append(conds, ast.CondNode{
-			At: elseTok.At,
+			At:   elseTok.At,
 			Cond: ast.Expr{},
 			Then: elseThen,
 		})
@@ -220,6 +221,67 @@ func (p *Parser) parseIfStmt() (ast.Stmt, error) {
 	}, nil
 }
 
+func (p *Parser) parseForStmt() (ast.Stmt, error) {
+	// ForStmt = ForCondStmt | ForInStmt
+	// ForCondStmt = "for" Expr? Block
+	// ForInStmt = "for" Ident "in" Expr Block
+
+	forTok, err := p.expect(token.KwFor)
+	if err != nil {
+		return ast.Stmt{}, err
+	}
+
+	expr, err := p.parseExpr()
+	if err != nil {
+		return ast.Stmt{}, err
+	}
+
+	// ForInStmt
+	if expr.Kind == ast.ExprKindIdent && p.peek().Kind == token.KwIn {
+		_, err := p.expect(token.KwIn)
+		if err != nil {
+			return ast.Stmt{}, err
+		}
+
+		collection, err := p.parseExpr()
+		if err != nil {
+			return ast.Stmt{}, err
+		}
+
+		do, err := p.parseBlock()
+		if err != nil {
+			return ast.Stmt{}, err
+		}
+
+		return ast.Stmt{
+			At:   forTok.At,
+			Kind: ast.StmtKindForIn,
+			Value: ast.StmtForIn{
+				Name:       expr.Value.(ast.ExprIdent).Value,
+				Collection: collection,
+				Do:         do,
+			},
+		}, nil
+	}
+
+	// ForCondStmt
+	cond := expr
+
+	do, err := p.parseBlock()
+	if err != nil {
+		return ast.Stmt{}, err
+	}
+
+	return ast.Stmt{
+		At:   forTok.At,
+		Kind: ast.StmtKindForCond,
+		Value: ast.StmtForCond{
+			Cond: cond,
+			Do:   do,
+		},
+	}, nil
+}
+
 func (p *Parser) parseTryStmt() (ast.Stmt, error) {
 	// TryStmt = "try" Block "catch" Ident ("finally" Block)?