‘for’ Loop Expressions

Now that we know how to add basic control flow constructs to the language, we have the tools to add more powerful things. Let’s add something more aggressive, a ‘for’ expression:

extern putchard(char);
def printstar(n)
  for i = 1, i < n, 1.0 in
    putchard(42);  # ascii 42 = '*'
# print 100 '*' characters
printstar(100);

This expression defines a new variable (i in this case) which iterates from a starting value, while the condition (i < n in this case) is true, incrementing by an optional step value (1.0 in this case). While the loop is true, it executes its body expression. Because we don’t have anything better to return, we’ll just define the loop as always returning 0.0. In the future when we have mutable variables, it will get more useful.

To get started, we again extend our lexer with new reserved names “for” and “in”.

lexer :: Tok.TokenParser ()
lexer = Tok.makeTokenParser style
  where
    ops = ["+","*","-","/",";",",","<"]
    names = ["def","extern","if","then","else","in","for"]
    style = emptyDef {
               Tok.commentLine = "#"
             , Tok.reservedOpNames = ops
             , Tok.reservedNames = names
             }

As before, let’s talk about the changes that we need to Kaleidoscope to support this. The AST node is just as simple. It basically boils down to capturing the variable name and the constituent expressions in the node.

data Expr
  ...
  | For Name Expr Expr Expr Expr
  deriving (Eq, Ord, Show)

The parser code captures a named value for the iterator variable and the four expressions objects for the parameters of the loop parameters.

for :: Parser Expr
for = do
  reserved "for"
  var <- identifier
  reservedOp "="
  start <- expr
  reservedOp ","
  cond <- expr
  reservedOp ","
  step <- expr
  reserved "in"
  body <- expr
  return $ For var start cond step body

Now we get to the good part: the LLVM IR we want to generate for this thing. With the simple example above, we get this LLVM IR (note that this dump is generated with optimizations disabled for clarity):

declare double @putchard(double)
define double @printstar(double %n) {
entry:
   br label %loop
loop:
   %i = phi double [ 1.000000e+00, %entry ], [ %nextvar, %loop ]
   %calltmp = call double @putchard(double 4.200000e+01)
   %nextvar = fadd double %i, 1.000000e+00
   %cmptmp = fcmp ult double %i, %n
   %booltmp = uitofp i1 %cmptmp to double
   %loopcond = fcmp one double %booltmp, 0.000000e+00
   br i1 %loopcond, label %loop, label %afterloop
afterloop:
   ret double 0.000000e+00 
}

The code to generate this is only slightly more complicated than the above “if” statement.

cgen (S.For ivar start cond step body) = do
  forloop <- addBlock "for.loop"
  forexit <- addBlock "for.exit"
  -- %entry
  ------------------
  i <- alloca double
  istart <- cgen start           -- Generate loop variable initial value
  stepval <- cgen step           -- Generate loop variable step
  store i istart                 -- Store the loop variable initial value
  assign ivar i                  -- Assign loop variable to the variable name
  br forloop                     -- Branch to the loop body block
  -- for.loop
  ------------------
  setBlock forloop
  cgen body                      -- Generate the loop body
  ival <- load i                 -- Load the current loop iteration
  inext <- fadd ival stepval     -- Increment loop variable
  store i inext
  cond <- cgen cond              -- Generate the loop condition
  test <- fcmp FP.ONE false cond -- Test if the loop condition is True ( 1.0 )
  cbr test forloop forexit       -- Generate the loop condition

The first step is to set up the LLVM basic block for the start of the loop body. In the case above, the whole loop body is one block, but remember that the generating code for the body of the loop could consist of multiple blocks (e.g. if it contains an if/then/else or a for/in expression).

  forloop <- addBlock "for.loop"
  forexit <- addBlock "for.exit"

Next we allocate the iteration variable and generate the code for the constant initial value and step.

  i <- alloca double
  istart <- cgen start           -- Generate loop variable initial value
  stepval <- cgen step           -- Generate loop variable step

Now the code starts to get more interesting. Our ‘for’ loop introduces a new variable to the symbol table. This means that our symbol table can now contain either function arguments or loop variables. Once the loop variable is set into the symbol table, the code recursively codegen’s the body. This allows the body to use the loop variable: any references to it will naturally find it in the symbol table.

  store i istart                 -- Store the loop variable initial value
  assign ivar i                  -- Assign loop variable to the variable name
  br forloop                     -- Branch to the loop body block

Now that the “preheader” for the loop is set up, we switch to emitting code for the loop body.

  setBlock forloop
  cgen body                      -- Generate the loop body

The body will contain the iteration variable scoped with its code generation. After loading its current state we increment it by the step value and store the value.

  ival <- load i                 -- Load the current loop iteration
  inext <- fadd ival stepval     -- Increment loop variable
  store i inext

Finally, we evaluate the exit test of the loop, and conditionally either branch back to the same block or exit the loop.

  cond <- cgen cond              -- Generate the loop condition
  test <- fcmp FP.ONE false cond -- Test if the loop condition is True ( 1.0 )
  cbr test forloop forexit       -- Generate the loop condition

Finally, code generation of the for loop always returns 0.0. Also note that the loop variable remains in scope even after the function exits.

  setBlock forexit
  return zero

We can now generate the assembly for our printstar function, for example the body of our function will generate code like the following on x86.

printstar:                              # @printstar
    .cfi_startproc
# BB#0:                                 # %entry
    subq    $24, %rsp
.Ltmp1:
    .cfi_def_cfa_offset 32
    vmovsd  %xmm0, 8(%rsp)          # 8-byte Spill
    vmovsd  .LCPI0_0(%rip), %xmm0
    vmovapd %xmm0, %xmm1
    .align  16, 0x90
.LBB0_1:                                # %loop
                                        # =>This Inner Loop Header: Depth=1
    vmovsd  %xmm1, 16(%rsp)         # 8-byte Spill
    vmovsd  .LCPI0_1(%rip), %xmm0
    callq   putchard
    vmovsd  16(%rsp), %xmm1         # 8-byte Reload
    vucomisd    8(%rsp), %xmm1  # 8-byte Folded Reload
    sbbl    %eax, %eax
    andl    $1, %eax
    vcvtsi2sd   %eax, %xmm0, %xmm0
    vaddsd  .LCPI0_0(%rip), %xmm1, %xmm1
    vucomisd    .LCPI0_2, %xmm0
    jne .LBB0_1
# BB#2:                                 # %afterloop
    vxorpd  %xmm0, %xmm0, %xmm0
    addq    $24, %rsp
    ret