5. Kaleidoscope: Control Flow

This chapter focuses on adding the support necessary to implement the if-then-else and for loop control flow support in the Kaleidoscope language. Without some sort of control flow the Kaleidoscope language is not particularly useful. So, this chapter completes the core language support to make it a usable language.

if-then-else

It is worth re-visiting the discussion of the intended syntax and semantics for conditional flow in Chapter 2. This will help in understanding the language functionality to implement.

The ultimate goal of the changes to support code generation for control flow constructs is to transform Kaleidoscope code such as:

extern foo();
extern bar();
def baz(x) if x then foo() else bar();

and generate LLVM like this (unoptimized):

declare double @foo()

declare double @bar()

define double @baz(double %x) {
entry:
  %ifcond = fcmp one double %x, 0.000000e+00
  br i1 %ifcond, label %then, label %else

then:       ; preds = %entry
  %calltmp = call double @foo()
  br label %ifcont

else:       ; preds = %entry
  %calltmp1 = call double @bar()
  br label %ifcont

ifcont:     ; preds = %else, %then
  %iftmp = phi double [ %calltmp, %then ], [ %calltmp1, %else ]
  ret double %iftmp
}

The entry code will convert the input x into an llvm i1 value to use as the condition for a branch. This is done by comparing the input value of x to 0.0 to get the condition boolean value. Then the condition is used to branch to either the 'then' block or the 'else' block. The two target blocks contain the generated code for the expressions for each part of the conditional and a final branch to a continuation block.

Since the code branch could flow into the continuation block from either the 'else' or 'end' blocks a phi instruction is placed at the beginning of the continuation block with appropriate values for the result from each of the two predecessor blocks. The resulting value is then provided as the return of the function. It is important to note that using the phi node in this fashion does not require generating all of the code in SSA form. In fact, doing that in the front end is strongly discouraged. Generally speaking there are only two reasons where a phi node may crop up:

Mutable variables like x = 1; x = x + 1;
Values that are part of the structure of the language (usually for control flow)

Chapter 7 Covers the mutable variables case in detail and the techniques for generating the code without using a phi node. For cases like this one where it is straight forward and easy to insert the phi node directly then there's no reason not to. Though, the solution provided in Chapter 7 can, and does, eliminate the need to manually insert the phi node here as well.

Code Generation

Generating the code for the condition expression follows the pattern shown above with the following high level steps:

Generate the code for the condition value expression
Emit conversion of the result of the condition to an LLVM i1 by comparing to 0.0
Create a block for the then expression
Create a block for the else expression
Create a block for the if continuation
Emit conditional branch to the then, else blocks
Switch to the then expression block
Emit code for the then expression
Capture the insertion block location as generating the then expression may add new blocks
Emit a branch to the if continuation block
Switch to the else block
Emit code for the else expression
Emit a branch to the if continuation block
Capture the insertion block location as generating the else expression may add new blocks
Switch to the if continuation block
Emit phi node with the results of the insertion blocks and result values captured after generating each of the sub expressions
Use the result...

That's a bit more complex than the other language constructs seen so far, but is still pretty straight forward once you get the general gist of how LLVM IR works. There's one extra trick repeated in steps 9 and again in 14, where after generating the IR for the sub expressions, the current block insertion point is captured. This is needed as the generation for the sub expression may include another conditional expression, which may contain a conditional sub expression, ... Thus, the 'current block' may well have changed from the starting block. The phi node needs the immediate predecessor block and the value it produced, so the current block is captured after generation, before switching the block to the next one for generation to ensure that the correct block is used with the value.

The actual code follows the description pretty closely and should now be fairly easy to follow:

public override Value? Visit( ConditionalExpression conditionalExpression )
{
    conditionalExpression.ValidateNotNull( nameof( conditionalExpression ) );
    var condition = conditionalExpression.Condition.Accept( this );
    if( condition == null )
    {
        return null;
    }

    var condBool = InstructionBuilder.Compare( RealPredicate.OrderedAndNotEqual, condition, Context.CreateConstant( 0.0 ) )
                                     .RegisterName( "ifcond" );

    var function = InstructionBuilder.InsertFunction;
    if( function is null )
    {
        throw new InternalCodeGeneratorException( "ICE: expected block that is attached to a function at this point" );
    }

    var thenBlock = function.AppendBasicBlock( "then" );
    var elseBlock = function.AppendBasicBlock( "else" );
    var continueBlock = function.AppendBasicBlock( "ifcont" );
    InstructionBuilder.Branch( condBool, thenBlock, elseBlock );

    // generate then block instructions
    InstructionBuilder.PositionAtEnd( thenBlock );

    // InstructionBuilder.InserBlock after this point is !null
    Debug.Assert( InstructionBuilder.InsertBlock != null, "expected non-null InsertBlock" );
    var thenValue = conditionalExpression.ThenExpression.Accept( this );
    if( thenValue == null )
    {
        return null;
    }

    InstructionBuilder.Branch( continueBlock );

    // capture the insert in case generating else adds new blocks
    var thenResultBlock = InstructionBuilder.InsertBlock;

    // generate else block
    InstructionBuilder.PositionAtEnd( elseBlock );
    var elseValue = conditionalExpression.ElseExpression.Accept( this );
    if( elseValue == null )
    {
        return null;
    }

    InstructionBuilder.Branch( continueBlock );
    var elseResultBlock = InstructionBuilder.InsertBlock;

    // generate continue block
    InstructionBuilder.PositionAtEnd( continueBlock );
    var phiNode = InstructionBuilder.PhiNode( function.Context.DoubleType )
                                    .RegisterName( "ifresult" );

    phiNode.AddIncoming( (thenValue, thenResultBlock), (elseValue, elseResultBlock) );
    return phiNode;
}

For Loop

Now that the basics of control flow are available it is possible to leverage the same concepts to implement the for loop constructs for the language.

The general idea is to transform the loops in Kaleidoscope such as this:

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

# print 100 '*' characters
printstar(100);

In LLVM IR (unoptimized) that should look like this:

declare double @putchard(double)

define double @printstar(double %n) {
entry:
  ; initial value = 1.0 (inlined into phi)
  br label %loop

loop:       ; preds = %loop, %entry
  %i = phi double [ 1.000000e+00, %entry ], [ %nextvar, %loop ]
  ; body
  %calltmp = call double @putchard(double 4.200000e+01)
  ; increment
  %nextvar = fadd double %i, 1.000000e+00

  ; termination test
  %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:      ; preds = %loop
  ; loop always returns 0.0
  ret double 0.000000e+00
}

Thus, the basic pattern to generate the for loop code consists of the following steps:

Create block for loop header
Switch to the loop header block
Emit code to Initialize start value with starting value from initialization expression
Create block for the loop body
Create block for the loop end
Emit unconditional branch to the loop body
Switch to the loop body block
Emit phi node for the loop value with the loop header block and initial value as first predecessor
Push a new scope for named values as the loop body represents a new scope
Add the variable for the loop to the current scope
Emit the body expression, which may create new blocks
Emit the code to compute the next value (e.g. next = current + step )
Emit code for the end condition
Emit code to convert the result of the condition to an LLVM i1 for a conditional branch
Capture loop end block for PHI node
Create after loop block
Emit conditional branch to the loop body block or after loop block depending on the result of the end condition
Add an incoming predecessor to the phi node at the beginning of the loop body for the next loop value and the loop end block it comes from
Switch to after block
Create constant value of 0.0 as the result expression of the for loop

That's a few more steps than even the if-then-else but the basic concepts of blocks, conditional branches and direct phi-nodes remains the same.

The code to generate a for loop follows this pattern pretty closely.

public override Value? Visit( ForInExpression forInExpression )
{
    forInExpression.ValidateNotNull( nameof( forInExpression ) );
    var function = InstructionBuilder.InsertFunction;
    if( function is null )
    {
        throw new InternalCodeGeneratorException( "ICE: Expected block attached to a function at this point" );
    }

    string varName = forInExpression.LoopVariable.Name;

    // Emit the start code first, without 'variable' in scope.
    Value? startVal;
    if( forInExpression.LoopVariable.Initializer != null )
    {
        startVal = forInExpression.LoopVariable.Initializer.Accept( this );
        if( startVal is null )
        {
            return null;
        }
    }
    else
    {
        startVal = Context.CreateConstant( 0.0 );
    }

    Debug.Assert( InstructionBuilder.InsertBlock != null, "expected non-null InsertBlock" );

    // Make the new basic block for the loop header, inserting after current
    // block.
    var preHeaderBlock = InstructionBuilder.InsertBlock;
    var loopBlock = function.AppendBasicBlock( "loop" );

    // Insert an explicit fall through from the current block to the loopBlock.
    InstructionBuilder.Branch( loopBlock );

    // Start insertion in loopBlock.
    InstructionBuilder.PositionAtEnd( loopBlock );

    // Start the PHI node with an entry for Start.
    var variable = InstructionBuilder.PhiNode( Context.DoubleType )
                                     .RegisterName( varName );

    variable.AddIncoming( startVal, preHeaderBlock );

    // Within the loop, the variable is defined equal to the PHI node.
    // So, push a new scope for it and any values the body might set
    using( NamedValues.EnterScope( ) )
    {
        NamedValues[ varName ] = variable;

        // Emit the body of the loop.  This, like any other expression, can change the
        // current BB.  Note that we ignore the value computed by the body, but don't
        // allow an error.
        if( forInExpression.Body.Accept( this ) == null )
        {
            return null;
        }

        Value? stepValue = forInExpression.Step.Accept( this );
        if( stepValue == null )
        {
            return null;
        }

        var nextVar = InstructionBuilder.FAdd( variable, stepValue)
                                        .RegisterName( "nextvar" );

        // Compute the end condition.
        Value? endCondition = forInExpression.Condition.Accept( this );
        if( endCondition == null )
        {
            return null;
        }

        // Convert condition to a bool by comparing non-equal to 0.0.
        endCondition = InstructionBuilder.Compare( RealPredicate.OrderedAndNotEqual, endCondition, Context.CreateConstant( 0.0 ) )
                                         .RegisterName( "loopcond" );

        // capture loop end result block for loop variable PHI node
        var loopEndBlock = InstructionBuilder.InsertBlock;

        // Create the "after loop" block and insert it.
        var afterBlock = function.AppendBasicBlock( "afterloop" );

        // Insert the conditional branch into the end of LoopEndBB.
        InstructionBuilder.Branch( endCondition, loopBlock, afterBlock );
        InstructionBuilder.PositionAtEnd( afterBlock );

        // Add a new entry to the PHI node for the back-edge.
        variable.AddIncoming( nextVar, loopEndBlock );

        // for expression always returns 0.0 for consistency, there is no 'void'
        return Context.DoubleType.GetNullValue( );
    }
}

The only new functionality in that is the use of the ScopeStack class to support nested scopes and the named variables within them. ScopeStack is provided in the Ubiquity.NET.Llvm library. It is basically a stack of name to value mapping dictionaries. The EnterScope method will push a new dictionary on to the stack and return an IDisposable that will handle popping it back off. This allows for nested expressions to use variables in the parent scope and to override them with its own value too. That, is the symbols available in a loop include the loop variable and any variables in the parent scope, all the way back to the function parameters. The stack nature allows for deeper scopes to shadow the variable of the same name in the parent, while allowing access to all other variables from other scopes.

Conclusion

Control flow is certainly more complex to generate than any of the other language constructs but it relies on a few basic primitive building block patterns. Thus, it is fairly easy to understand and implement once the basic patterns are understood. With the inclusion of control flow the Kaleidoscope language is now a complete, albeit simplistic, functional language.