From: Luke Lau <luke_lau@icloud.com>
Date: Fri, 8 Nov 2019 15:07:55 +0000 (+0000)
Subject: Add control flow
X-Git-Url: http://git.lukelau.me/?a=commitdiff_plain;h=705b0b8458673edb7273e6e19914390a55e51d23;hp=7b8ef6725099f09b81f8c39ca6f00dec14213bed;p=kaleidoscope-hs.git

Add control flow

So far our language can parse and evaluate floating point expressions.
However each expression and it's subexpressions always gets evaluated. To
do more complex computations, we want to be able to control what gets
evaluated and what doesn't.

This is known as control flow, and the two most famous imperative
constructs are probably the if statement and for loop. We will add them
to our language in this chapter, and finally put our boolean expressions
(42 < i) to good use.
---

diff --git a/AST.hs b/AST.hs
index 9ff555a..b019a2a 100644
--- a/AST.hs
+++ b/AST.hs
@@ -2,12 +2,14 @@ module AST where
 
 import Data.Char
 import Text.Read 
-import Text.ParserCombinators.ReadP hiding ((+++), choice)
+import Text.ParserCombinators.ReadP hiding ((+++), (<++), choice)
 
 data Expr = Num Double
           | Var String
           | BinOp BinOp Expr Expr
           | Call String [Expr]
+          | If Expr Expr Expr
+          | For String Expr Expr (Maybe Expr) Expr
   deriving Show
 
 data BinOp = Add | Sub | Mul | Cmp Ordering
@@ -17,6 +19,8 @@ instance Read Expr where
   readPrec = parens $ choice [ parseNum
                              , parseVar
                              , parseCall
+                             , parseIf
+                             , parseFor
                              , parseBinOp "<" 10 (Cmp LT)
                              , parseBinOp ">" 10 (Cmp GT)
                              , parseBinOp "==" 10 (Cmp EQ)
@@ -28,10 +32,7 @@ instance Read Expr where
           parseVar = Var <$> lift (munch1 isAlpha)
           parseBinOp s prc op = prec prc $ do
             a <- step readPrec
-            lift $ do
-              skipSpaces
-              string s
-              skipSpaces
+            spaced $ string s
             b <- readPrec
             return (BinOp op a b)
           parseCall = do
@@ -40,6 +41,27 @@ instance Read Expr where
                         sepBy (readS_to_P reads)
                               (skipSpaces >> char ',' >> skipSpaces)
             return (Call func params)
+          parseIf = do
+            spaced $ string "if" 
+            cond <- readPrec
+            spaced $ string "then"
+            thenE <- readPrec
+            spaced $ string "else"
+            elseE <- readPrec
+            return (If cond thenE elseE)
+          parseFor = do
+            spaced $ string "for"
+            identifier <- lift (munch1 isAlpha)
+            spaced $ char '='
+            start <- readPrec
+            spaced $ char ','
+            cond <- readPrec
+            stp <- (spaced (char ',') >> Just <$> step readPrec)
+                     <++ pure Nothing
+            spaced $ string "in"
+            body <- readPrec
+            return (For identifier start cond stp body)
+          spaced f = lift $ skipSpaces >> f >> skipSpaces
             
 data Prototype = Prototype String [String]
   deriving Show
diff --git a/Main.hs b/Main.hs
index 468573d..97407fe 100644
--- a/Main.hs
+++ b/Main.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE RecursiveDo #-}
 
 import AST as K -- K for Kaleidoscope
 import Utils
@@ -141,3 +142,60 @@ buildExpr (Call callee params) = do
       ptrTyp = Type.PointerType typ (AddrSpace 0)
       ref = GlobalReference ptrTyp nam
   call (ConstantOperand ref) (zip paramOps (repeat []))
+
+buildExpr (If cond thenE elseE) = mdo
+  _ifB <- block `named` "if"
+
+  -- since everything is a double, false == 0
+  let zero = ConstantOperand (Float (Double 0))
+  condV <- buildExpr cond
+  cmp <- fcmp ONE zero condV `named` "cmp"
+
+  condBr cmp thenB elseB
+
+  thenB <- block `named` "then"
+  thenOp <- buildExpr thenE
+  br mergeB
+
+  elseB <- block `named` "else"
+  elseOp <- buildExpr elseE
+  br mergeB
+
+  mergeB <- block `named` "ifcont"
+  phi [(thenOp, thenB), (elseOp, elseB)]
+
+buildExpr (For name init cond mStep body) = mdo
+  preheaderB <- block `named` "preheader"
+
+  initV <- buildExpr init `named` "init"
+  
+  -- build the condition expression with 'i' in the bindings
+  initCondV <- withReaderT (Map.insert name initV) $
+                (buildExpr cond >>= fcmp ONE zero) `named` "initcond"
+
+  -- skip the loop if we don't meet the condition with the init
+  condBr initCondV loopB afterB
+
+  loopB <- block `named` "loop"
+  i <- phi [(initV, preheaderB), (nextVar, loopB)] `named` "i"
+
+  -- build the body expression with 'i' in the bindings
+  withReaderT (Map.insert name i) $ buildExpr body `named` "body"
+
+  -- default to 1 if there's no step defined
+  stepV <- case mStep of
+    Just step -> buildExpr step
+    Nothing -> return $ ConstantOperand (Float (Double 1))
+
+  nextVar <- fadd i stepV `named` "nextvar"
+
+  let zero = ConstantOperand (Float (Double 0))
+  -- again we need 'i' in the bindings
+  condV <- withReaderT (Map.insert name i) $
+            (buildExpr cond >>= fcmp ONE zero) `named` "cond"
+  condBr condV loopB afterB
+
+  afterB <- block `named` "after"
+  -- since a for loop doesn't really have a value, return 0
+  return $ ConstantOperand (Float (Double 0))
+