Lab 6: Control Flow (Milestone 4)

Background

You will add if/elseif/else, while, the numeric for, break, and return, plus the short-circuiting and/or. The Flow enum arrives, and with it the instruction budget — four labs before Section 5 needs it.

By the end, FizzBuzz runs.

Why This Lab Matters

  • Flow is the mechanism a tree walker needs and a VM does not. Building it, and then watching it evaporate into a JUMP in Lab 10, is the clearest single argument for bytecode in the curriculum.
  • The numeric for has an overflow trap that Lua 5.4 fixes with a precomputed iteration count. The obvious implementation loops forever on for i = math.maxinteger - 1, math.maxinteger do.
  • The budget check goes on the loop back-edge now. Retrofitting a budget into control flow later is how one construct ends up able to loop forever.

Prerequisites

Predict First

  1. if 0 then return "a" end return "b" — which?
  2. return nil or 5 and return 0 and "x" — what values, and what types?
  3. while true do break end — but with the break inside do ... end inside if. Does a naive Rust break work?
  4. for i = 1, 3 do end return i — what is i after the loop?
  5. for i = 1, 10, 0 do end — what happens?
  6. for i = 1, 3.5 do end — how many iterations, and what is i's subtype?
  7. for i = math.maxinteger - 1, math.maxinteger do end — how many iterations? Now think about what i = i + 1 does at the top of the range.

Step 1: if / elseif / else

#![allow(unused)]
fn main() {
fn if_stat(&mut self) -> Result<Stmt> {
    let start = self.expect(TokenKind::If)?.span;
    let mut arms = Vec::new();
    let cond = self.expr()?;
    self.expect(TokenKind::Then)?;
    arms.push((cond, self.block()?));
    while self.eat(TokenKind::Elseif) {
        let c = self.expr()?;
        self.expect(TokenKind::Then)?;
        arms.push((c, self.block()?));               // FLAT, not nested — see the AST chapter
    }
    let else_ = if self.eat(TokenKind::Else) { Some(self.block()?) } else { None };
    self.expect(TokenKind::End)?;
    Ok(Stmt::If { arms, else_, span: start.merge(self.prev_span()) })
}
}

The part that matters. arms is a flat Vec, so a 500-arm elseif chain is a 500-element vector rather than a 500-deep tree. That keeps it under MAX_PARSE_DEPTH and out of the recursive- Drop hazard, and in Lab 10 it becomes a flat list of jump patches instead of nested ones.

Note: Ember has no dangling-else problem, because end terminates the construct explicitly. C's famous ambiguity comes from optional braces. Lua chose end in 1993 partly for this reason. The cost is four extra keystrokes; the benefit is that the grammar is unambiguous here.

Step 2: while, break, and Flow

Add the Flow enum and change every eval_stmt/eval_block signature to Result<Flow>, per the concept chapter.

#![allow(unused)]
fn main() {
Stmt::While { cond, body, .. } => {
    loop {
        if !self.eval_expr(cond)?.is_truthy() { break; }
        self.budget.tick_loop()?;                      // ← the back-edge check. See Step 5.
        match self.eval_block(body)? {
            Flow::Normal => {}
            Flow::Break  => break,                     // ABSORBED
            r @ Flow::Return(_) => return Ok(r),       // FORWARDED
        }
    }
    Flow::Normal
}
}

break outside a loop must be a parse error. Track loop depth in the parser:

#![allow(unused)]
fn main() {
fn break_stat(&mut self) -> Result<Stmt> {
    let t = self.expect(TokenKind::Break)?;
    if self.loop_depth == 0 {
        return Err(EmberError { kind: ErrorKind::Parse,
            message: "'break' outside a loop".into(),
            span: Some(t.span), traceback: vec![] });
    }
    Ok(Stmt::Break { span: t.span })
}
}

Checkpoint question. loop_depth must be reset — not just decremented — when the parser enters a function body, because break cannot cross a function boundary. Where does that reset go, and what program exposes the bug if you forget it?

Step 3: The Numeric for

This is the construct with the trap. Read Lua 5.4 §3.3.5 before writing it.

#![allow(unused)]
fn main() {
Stmt::NumericFor { var, start, stop, step, body, span } => {
    // All three control expressions are evaluated ONCE, before the loop.
    let start_v = self.eval_expr(start)?;
    let stop_v  = self.eval_expr(stop)?;
    let step_v  = match step { Some(e) => self.eval_expr(e)?, None => Value::Integer(1) };

    // Lua 5.4: each must be a number, and the step must not be zero.
    let (mut i, limit, st) = self.prepare_for(start_v, stop_v, step_v, *span)?;

    // INTEGER LOOPS USE A PRECOMPUTED COUNT. See the warning below.
    let mut remaining = for_iteration_count(i, limit, st);

    while remaining > 0 {
        remaining -= 1;
        self.budget.tick_loop()?;
        // A FRESH BINDING PER ITERATION. Not a mutated one. This matters in
        // Lab 14, where closures created in the loop must capture different
        // variables. It is a SCOPE decision, made here.
        let _scope = self.env.scope_guard();
        self.env.declare(var.clone(), i);

        match self.eval_block(body)? {
            Flow::Normal => {}
            Flow::Break  => break,
            r @ Flow::Return(_) => return Ok(r),
        }
        i = advance(i, st);
    }
    Flow::Normal
}
}

Warning — the overflow trap. The obvious loop is while i <= limit { body; i += step }. Now run it on for i = math.maxinteger - 1, math.maxinteger do end. On the last iteration i + 1 wraps to i64::MIN, which is <= limit, and the loop runs forever — quietly, at full speed, inside your host process.

Lua 5.4 fixed this by computing the iteration count up front as an unsigned value and counting down. for_iteration_count is ((limit - i) / step) + 1 computed in a wider or unsigned domain, with the empty-loop case handled first. Work out the four sign cases yourself, then compare with forprep in Lua's lvm.c (rg -n 'forprep|forlimit|OP_FORPREP' lvm.c).

Verify: lua -e 'local n=0 for i = math.maxinteger-1, math.maxinteger do n=n+1 end print(n)' prints 2. Yours must too, and it must terminate.

The mixed integer/float rule, also from §3.3.5:

Start, stop, stepLoop variableNotes
all integersintegerprecomputed count; no overflow
any floatfloatall three converted to float first
a non-number—error: 'for' initial value must be a number
step is 0—error: 'for' step is zero
lua -e 'for i = 1, 3.5 do io.write(i, " ", math.type(i), "  ") end'   # 1.0 float 2.0 float 3.0 float
lua -e 'for i = 1, 3 do io.write(math.type(i), " ") end'              # integer integer integer
lua -e 'for i = 1, 10, 0 do end'                                       # error

Step 4: and / or, and return

Add the short-circuit arms above the generic Binary arm in eval_expr (why), and:

#![allow(unused)]
fn main() {
Stmt::Return { exprs, .. } => Flow::Return(self.eval_exprlist_all(exprs)?),
}

The top-level runner absorbs Flow::Return:

#![allow(unused)]
fn main() {
pub fn run(&mut self, chunk: &Block) -> Result<Vec<Value>> {
    match self.eval_block(chunk)? {
        Flow::Return(vs) => Ok(vs),
        Flow::Normal     => Ok(vec![]),
        Flow::Break      => Err(internal("break escaped the chunk")),  // parser prevents this
    }
}
}

Step 5: The Instruction Budget

Concept. A script must not be able to run forever.

#![allow(unused)]
fn main() {
// src/limits.rs
pub struct Budget { pub remaining: u64, pub enabled: bool }

impl Budget {
    /// Called on every loop back-edge and every function call. NOT on every
    /// statement: this is the minimal set of places that can create an
    /// unbounded execution, and checking only there keeps the check cheap.
    #[inline]
    pub fn tick_loop(&mut self) -> Result<()> {
        if !self.enabled { return Ok(()); }
        if self.remaining == 0 {
            return Err(EmberError { kind: ErrorKind::Limit,
                message: "instruction budget exhausted".into(), span: None, traceback: vec![] });
        }
        self.remaining -= 1;
        Ok(())
    }
}
}

Why now and not in Section 5. There are exactly three ways an Ember program can run forever: a while back-edge, a for back-edge, and recursion. Two of them are created in this lab. Adding the check when the construct is created means the set is complete by construction; adding it later means auditing every construct and missing one. Section 5 replaces the counter with a real budget and adds the memory and depth limits, but the call sites are all here.

Note: In the VM this check moves to the fetch position in the dispatch loop and covers everything at once — see the mental model. The tree walker cannot do that because it has no fetch. A budget is another thing bytecode gives you for free, and noting that now makes the Section 3 payoff clearer.

The Trace

$ cat > /tmp/fizz.ember <<'EOF'
for i = 1, 15 do
  if i % 15 == 0 then      print("FizzBuzz")
  elseif i % 3 == 0 then   print("Fizz")
  elseif i % 5 == 0 then   print("Buzz")
  else                     print(i)
  end
end
EOF

print arrives in Lab 21, so until then use ember run --print-returns and a table of results, or add a temporary print native. Either way the control flow is what is being traced here:

$ ember ast /tmp/fizz.ember | head -14

Block                                  @0..167
└── NumericFor                         @0..166
    ├── var: "i"
    ├── start: Int(1)                  @8..9
    ├── stop:  Int(15)                 @11..13
    ├── step:  (none → 1)
    └── Block                          @19..161
        └── If                         @19..159
            ├── arm 0: Binary(Eq) @22..33  →  Block
            ├── arm 1: Binary(Eq) @54..64  →  Block      ← FLAT, not nested
            ├── arm 2: Binary(Eq) @89..99  →  Block
            └── else:                     Block

Now the break trace, which is the one to study:

$ ember run --trace-flow -e '
local n = 0
while true do
  do
    if n > 2 then break end
  end
  n = n + 1
end
return n'

While        enter
  Block(1)   enter                    [while body]
    Block(2) enter                    [do..end]
      If     cond=false → Normal
    Block(2) ← Normal
    Assign n = 1
  Block(1)   ← Normal
  ... n=1, n=2 ...
  Block(1)   enter
    Block(2) enter
      If     cond=TRUE → Block → Flow::Break
    Block(2) ← Flow::Break            ← forwarded through 2 frames
  Block(1)   ← Flow::Break            ← forwarded
While        Flow::Break ABSORBED → loop exits
3

The two ← Flow::Break lines are the whole lab. A Rust break written inside eval_stmt would have broken for stmt in &b.stmts — the wrong loop, three frames too early. Keep this trace; in Lab 10 you will put the disassembly of the same program next to it, and the entire mechanism will be JUMP 0016.

Expected Output

$ ember run -e 'if 0 then return "zero is truthy" end return "unreachable"'
zero is truthy
$ ember run -e 'return nil or 5'
5
$ ember run --types -e 'return 0 and "x"'
x (string)
$ ember run -e 'local n=0 for i=1,3 do for j=1,3 do if j==2 then break end n=n+1 end end return n'
3
$ ember run -e 'for i = 1, 10, 0 do end'
<argv>:1:1: error: 'for' step is zero

   1 │ for i = 1, 10, 0 do end
     │ ^^^
$ ember run -e 'local n=0 for i = 9223372036854775806, 9223372036854775807 do n=n+1 end return n'
2
$ ember run --max-instructions 1000 -e 'while true do end'
error: instruction budget exhausted
$ echo $?
1
$ ember run -e 'break'
<argv>:1:1: error: 'break' outside a loop

The nested-break case returning 3 (not 9, not 0) is the one that catches a broken implementation. The maxinteger case returning 2 is the one that catches the overflow trap.

Debugging Steps

break exits everything, or nothing

Flow::Break is being absorbed by eval_block instead of forwarded, or forwarded by While instead of absorbed. Check the absorbs/forwards table.

break in a nested block does nothing

eval_block is discarding the Flow from eval_stmt (self.eval_stmt(s)?; with the result dropped). It must match and return non-Normal upward.

nil or 5 returns true

and/or are reaching the generic Binary arm. Move their arms above it.

false and boom() calls boom

Same cause. Add a test with a side effect on the right-hand side; equality tests alone will not catch it.

for i = maxinteger-1, maxinteger hangs

The overflow trap. Precompute the iteration count.

The loop variable is visible after the loop

You declared it in the enclosing scope instead of a per-iteration scope.

Closures made in a loop all see the last value (you will hit this in Lab 14)

One binding mutated instead of a fresh one per iteration. Fix it here, in the scope handling, not in Lab 14.

while true do end never stops even with --max-instructions

tick_loop is on the for back-edge but not the while one, or it is inside the body block rather than around it — so a loop with an empty body never ticks.

Experiment

CLAIM. The Flow mechanism costs something measurable, and it is a cost bytecode does not have.

METHOD. Benchmark a tight loop — local n=0 for i=1,1000000 do n=n+1 end — with criterion. Then, in a scratch branch, change eval_stmt to return Result<()> and implement break with a Cell<bool> flag checked by the loop (a real alternative design). Benchmark again.

PREDICTION. Which is faster, and by how much? Is the difference bigger than the noise? Which is easier to get right?

RESULT. Record both numbers and your judgment. Then throw the branch away. The point is that "faster" and "correct-by-construction" are separable, and you now have data on how much you paid.

Test

#![allow(unused)]
fn main() {
#[test]
fn only_nil_and_false_are_falsy() {
    for truthy in ["0", "0.0", "\"\"", "0/0"] {
        assert_eq!(run(&format!("if {truthy} then return 1 end return 2")), "1");
    }
    for falsy in ["nil", "false"] {
        assert_eq!(run(&format!("if {falsy} then return 1 end return 2")), "2");
    }
}

#[test]
fn and_or_short_circuit_and_return_operands() {
    // Lua 5.4 §3.4.5: they return one of their OPERANDS, not a boolean.
    assert_eq!(run("return nil or 5"), "5");
    assert_eq!(run("return 0 and 'x'"), "x");
    assert_eq!(run("return false or nil"), "nil");
    // The side-effect test is the one that proves short-circuiting.
    assert_eq!(run("local hit = false
                    local function boom() hit = true return 1 end
                    local _ = false and boom()
                    return tostring(hit)"), "false");
}

#[test]
fn break_exits_only_the_innermost_loop_through_nested_blocks() {
    assert_eq!(run("local n=0
                    for i=1,3 do for j=1,3 do if j==2 then break end n=n+1 end end
                    return n"), "3");
    assert_eq!(run("local n=0
                    while true do do do if n>2 then break end end end n=n+1 end
                    return n"), "3");
}

#[test]
fn numeric_for_does_not_overflow() {
    // Lua 5.4 precomputes the iteration count precisely so this terminates.
    // The naive `while i <= limit { i += step }` loops FOREVER here.
    assert_eq!(run("local n=0
                    for i = 9223372036854775806, 9223372036854775807 do n = n + 1 end
                    return n"), "2");
}

#[test]
fn numeric_for_subtype_and_errors() {
    assert_eq!(run("local s='' for i=1,3 do s = s .. math_type(i) end return s"),
               "integerintegerinteger");
    assert_eq!(run("local n=0 for i=1,3.5 do n=n+1 end return n"), "3");
    assert_eq!(err("for i = 1, 10, 0 do end").kind, ErrorKind::Runtime);
    assert_eq!(err("for i = 1, 'x' do end").kind, ErrorKind::Runtime);
}

#[test]
fn break_outside_a_loop_is_a_parse_error() {
    assert_eq!(err("break").kind, ErrorKind::Parse);
    // And it must not cross a function boundary:
    assert_eq!(err("while true do local f = function() break end end").kind, ErrorKind::Parse);
}

#[test]
fn infinite_loops_terminate_under_a_budget() {
    let e = run_with_budget("while true do end", 1000).unwrap_err();
    assert_eq!(e.kind, ErrorKind::Limit);   // NOT Runtime — the host cares
}
}

Challenge Extensions

  1. repeat ... until. Lua's do-while. Its scope rule is unusual: the until condition can see locals declared in the body. Implement it and explain why that requires the scope to be popped after the condition is evaluated.
  2. goto and labels. Lua 5.4 has them, with restrictions (no jumping into a local's scope). Implement the restriction check. This is where you learn why goto is hard for compilers, not for readers.
  3. Loop-invariant budget check. Move tick_loop so that it charges proportionally to work done rather than iterations. Is that better? What does a host actually want to bound?
  4. A continue. Lua does not have one (it uses goto continue). Add it, then write the ADR arguing whether a language this small should.
  5. Constant-condition detection. Detect while true and if false at parse time and note them. In Lab 10 this becomes real dead-code elimination; here, just report it with ember ast --warnings.

Deliverables

  • if/elseif/else, while, numeric for, break, return all work.
  • elseif chains are flat in the AST.
  • Only nil and false are falsy; and/or short-circuit and return operands, tested with a side effect.
  • break exits only the innermost loop, tested through two levels of nested blocks.
  • break outside a loop, and across a function boundary, are parse errors.
  • The numeric for precomputes its iteration count and terminates at math.maxinteger.
  • for mixed integer/float and zero-step rules match Lua, tested.
  • The loop variable is a fresh binding per iteration.
  • --max-instructions terminates while true do end with ErrorKind::Limit.
  • --trace-flow prints the enter/exit/Flow trace shown above.
  • FizzBuzz is in tests/golden/.

Validation / Self-check

  1. Write the absorbs/forwards table for if, do, while, for, and the chunk.
  2. Give the smallest program where a Rust break inside eval_stmt produces the wrong behavior.
  3. Why is elseif flat rather than nested? Give two consequences.
  4. What does for i = math.maxinteger - 1, math.maxinteger do end do naively, and what is Lua's fix?
  5. State the four rules for the numeric for's control expressions.
  6. Why is the loop variable a fresh binding per iteration, and which lab makes that observable?
  7. Why does the budget check go on the back-edge rather than on every statement? Name the three places an Ember program can run forever.
  8. What replaces the entire Flow mechanism in Section 3, and why does the VM not need it?

Next: Lab 7 — Functions and Recursion.