Compile If Stmts

2 files changed, 44 insertions, 32 deletions

diff --git a/pkg/lang/parser/parser_test.go b/pkg/lang/parser/parser_test.go
index f48b6a5..b756b1b 100644
--- a/pkg/lang/parser/parser_test.go
+++ b/pkg/lang/parser/parser_test.go
@@ -384,34 +384,38 @@ func TestIfStmt(t *testing.T) {
 	require.Equal(t, ast.Stmt{
 		Kind: ast.StmtKindIf,
 		Value: ast.StmtIf{
-			Cond: ast.Expr{
-				At:    source.NewLoc(0, 3),
-				Kind:  ast.ExprKindBoolLit,
-				Value: ast.ExprBoolLit{Value: false},
-			},
-			Then: ast.BlockNode{
-				At: source.NewLoc(0, 9),
-				Stmts: []ast.Stmt{
-					{
-						At:    source.NewLoc(0, 10),
-						Kind:  ast.StmtKindEmpty,
-						Value: ast.StmtEmpty{},
+			Conds: []ast.CondNode{
+				{
+					At: source.NewLoc(0, 0),
+					Cond: ast.Expr{
+						At:    source.NewLoc(0, 3),
+						Kind:  ast.ExprKindBoolLit,
+						Value: ast.ExprBoolLit{Value: false},
 					},
-					{
-						At:   source.NewLoc(1, 1),
-						Kind: ast.StmtKindVarDecl,
-						Value: ast.StmtVarDecl{
-							Name: ast.IdentNode{At: source.NewLoc(1, 5), Value: "x"},
-							Value: ast.Expr{
-								At:    source.NewLoc(1, 9),
-								Kind:  ast.ExprKindIntLit,
-								Value: ast.ExprIntLit{Value: 2},
+					Then: ast.BlockNode{
+						At: source.NewLoc(0, 9),
+						Stmts: []ast.Stmt{
+							{
+								At:    source.NewLoc(0, 10),
+								Kind:  ast.StmtKindEmpty,
+								Value: ast.StmtEmpty{},
+							},
+							{
+								At:   source.NewLoc(1, 1),
+								Kind: ast.StmtKindVarDecl,
+								Value: ast.StmtVarDecl{
+									Name: ast.IdentNode{At: source.NewLoc(1, 5), Value: "x"},
+									Value: ast.Expr{
+										At:    source.NewLoc(1, 9),
+										Kind:  ast.ExprKindIntLit,
+										Value: ast.ExprIntLit{Value: 2},
+									},
+								},
 							},
 						},
 					},
 				},
 			},
-			Elifs: []ast.CondNode{},
 		},
 	}, program.Stmts[0])
 }
diff --git a/pkg/lang/parser/stmts.go b/pkg/lang/parser/stmts.go
index 52e590c..4b1077e 100644
--- a/pkg/lang/parser/stmts.go
+++ b/pkg/lang/parser/stmts.go
@@ -151,7 +151,9 @@ func (p *Parser) parseIfStmt() (ast.Stmt, error) {
 		return ast.Stmt{}, err
 	}
 
-	cond, err := p.parseExpr()
+	conds := []ast.CondNode{}
+
+	ifCond, err := p.parseExpr()
 	if err != nil {
 		return ast.Stmt{}, err
 	}
@@ -161,7 +163,11 @@ func (p *Parser) parseIfStmt() (ast.Stmt, error) {
 		return ast.Stmt{}, err
 	}
 
-	elifs := []ast.CondNode{}
+	conds = append(conds, ast.CondNode{
+		At:   ifTok.At,
+		Cond: ifCond,
+		Then: then,
+	})
 
 	for p.peek().Kind == token.KwElif {
 		elifTok, err := p.expect(token.KwElif)
@@ -179,35 +185,37 @@ func (p *Parser) parseIfStmt() (ast.Stmt, error) {
 			return ast.Stmt{}, err
 		}
 
-		elifs = append(elifs, ast.CondNode{
+		conds = append(conds, ast.CondNode{
 			At:   elifTok.At,
 			Cond: elifCond,
 			Then: elifThen,
 		})
 	}
 
-	elseThen := ast.BlockNode{}
 
 	if p.peek().Kind == token.KwElse {
-		_, err := p.expect(token.KwElse)
+		elseTok, err := p.expect(token.KwElse)
 		if err != nil {
 			return ast.Stmt{}, err
 		}
 
-		elseThen, err = p.parseBlock()
+		elseThen, err := p.parseBlock()
 		if err != nil {
 			return ast.Stmt{}, err
 		}
+
+		conds = append(conds, ast.CondNode{
+			At: elseTok.At,
+			Cond: ast.Expr{},
+			Then: elseThen,
+		})
 	}
 
 	return ast.Stmt{
 		At:   ifTok.At,
 		Kind: ast.StmtKindIf,
 		Value: ast.StmtIf{
-			Cond:  cond,
-			Then:  then,
-			Elifs: elifs,
-			Else:  elseThen,
+			Conds: conds,
 		},
 	}, nil
 }