Add while loop support to parser and runtime

Author: mgomes

ROADMAP.md

@@ -209,7 +209,7 @@ Goal: improve language ergonomics for complex script logic and recovery behavior
 
 ### Control Flow
 
-- [ ] Add `while` loops.
+- [x] Add `while` loops.
 - [ ] Add `until` loops.
 - [ ] Add loop control keywords: `break` and `next`.
 - [ ] Add `case/when` expression support (if approved).

vibes/ast.go

@@ -122,6 +122,15 @@ type ForStmt struct {
 func (s *ForStmt) stmtNode()     {}
 func (s *ForStmt) Pos() Position { return s.position }
 
+type WhileStmt struct {
+	Condition Expression
+	Body      []Statement
+	position  Position
+}
+
+func (s *WhileStmt) stmtNode()     {}
+func (s *WhileStmt) Pos() Position { return s.position }
+
 type Identifier struct {
 	Name     string
 	position Position

vibes/execution.go

@@ -336,6 +336,8 @@ func (exec *Execution) evalStatement(stmt Statement, env *Env) (Value, bool, err
 		return NewNil(), false, nil
 	case *ForStmt:
 		return exec.evalForStatement(s, env)
+	case *WhileStmt:
+		return exec.evalWhileStatement(s, env)
 	default:
 		return NewNil(), false, exec.errorAt(stmt.Pos(), "unsupported statement")
 	}
@@ -1124,6 +1126,33 @@ func (exec *Execution) evalForStatement(stmt *ForStmt, env *Env) (Value, bool, e
 	return last, false, nil
 }
 
+func (exec *Execution) evalWhileStatement(stmt *WhileStmt, env *Env) (Value, bool, error) {
+	last := NewNil()
+	for {
+		if err := exec.step(); err != nil {
+			return NewNil(), false, exec.wrapError(err, stmt.Pos())
+		}
+		condition, err := exec.evalExpression(stmt.Condition, env)
+		if err != nil {
+			return NewNil(), false, err
+		}
+		if err := exec.checkMemoryWith(condition); err != nil {
+			return NewNil(), false, err
+		}
+		if !condition.Truthy() {
+			return last, false, nil
+		}
+		val, returned, err := exec.evalStatements(stmt.Body, env)
+		if err != nil {
+			return NewNil(), false, err
+		}
+		if returned {
+			return val, true, nil
+		}
+		last = val
+	}
+}
+
 func (exec *Execution) getMember(obj Value, property string, pos Position) (Value, error) {
 	switch obj.Kind() {
 	case KindHash, KindObject:

vibes/lexer.go

@@ -415,6 +415,8 @@ func lookupIdent(ident string) TokenType {
 		return tokenDo
 	case "for":
 		return tokenFor
+	case "while":
+		return tokenWhile
 	case "in":
 		return tokenIn
 	case "if":

vibes/parser.go

@@ -126,6 +126,8 @@ func (p *parser) parseStatement() Statement {
 		return p.parseIfStatement()
 	case tokenFor:
 		return p.parseForStatement()
+	case tokenWhile:
+		return p.parseWhileStatement()
 	case tokenIdent:
 		if p.curToken.Literal == "assert" {
 			return p.parseAssertStatement()
@@ -386,6 +388,21 @@ func (p *parser) parseForStatement() Statement {
 	return &ForStmt{Iterator: iterator, Iterable: iterable, Body: body, position: pos}
 }
 
+func (p *parser) parseWhileStatement() Statement {
+	pos := p.curToken.Pos
+	p.nextToken()
+	condition := p.parseExpression(lowestPrec)
+
+	p.nextToken()
+	body := p.parseBlock(tokenEnd)
+
+	if p.curToken.Type != tokenEnd {
+		p.errorExpected(p.curToken, "end")
+	}
+
+	return &WhileStmt{Condition: condition, Body: body, position: pos}
+}
+
 func (p *parser) parseBlock(stop ...TokenType) []Statement {
 	stmts := []Statement{}
 	stopSet := make(map[TokenType]struct{}, len(stop))

vibes/runtime_test.go

@@ -753,6 +753,60 @@ func TestIntTimes(t *testing.T) {
 	}
 }
 
+func TestWhileLoops(t *testing.T) {
+	script := compileScript(t, `
+    def countdown(n)
+      out = []
+      while n > 0
+        out = out + [n]
+        n = n - 1
+      end
+      out
+    end
+
+    def first_positive(n)
+      while n > 0
+        return n
+      end
+      0
+    end
+
+    def skip_false()
+      while false
+        1
+      end
+    end
+    `)
+
+	countdown := callFunc(t, script, "countdown", []Value{NewInt(3)})
+	compareArrays(t, countdown, []Value{NewInt(3), NewInt(2), NewInt(1)})
+
+	if got := callFunc(t, script, "first_positive", []Value{NewInt(4)}); !got.Equal(NewInt(4)) {
+		t.Fatalf("first_positive mismatch for positive input: %v", got)
+	}
+	if got := callFunc(t, script, "first_positive", []Value{NewInt(0)}); !got.Equal(NewInt(0)) {
+		t.Fatalf("first_positive mismatch for zero input: %v", got)
+	}
+	if got := callFunc(t, script, "skip_false", nil); !got.Equal(NewNil()) {
+		t.Fatalf("skip_false expected nil, got %v", got)
+	}
+
+	engine := MustNewEngine(Config{StepQuota: 40})
+	spinScript, err := engine.Compile(`
+    def spin()
+      while true
+      end
+    end
+    `)
+	if err != nil {
+		t.Fatalf("compile error: %v", err)
+	}
+	_, err = spinScript.Call(context.Background(), "spin", nil, CallOptions{})
+	if err == nil || !strings.Contains(err.Error(), "step quota exceeded") {
+		t.Fatalf("expected step quota error for infinite while loop, got %v", err)
+	}
+}
+
 func TestDurationMethods(t *testing.T) {
 	script := compileScript(t, `
     def duration_helpers()

vibes/token.go

@@ -56,6 +56,7 @@ const (
 	tokenYield    TokenType = "YIELD"
 	tokenDo       TokenType = "DO"
 	tokenFor      TokenType = "FOR"
+	tokenWhile    TokenType = "WHILE"
 	tokenIn       TokenType = "IN"
 	tokenIf       TokenType = "IF"
 	tokenElsif    TokenType = "ELSIF"