Add for loops
[kaleidoscope-hs-old.git] / AST.hs
diff --git a/AST.hs b/AST.hs
index 99b6d42a23b9af254c8767ac969accc8248f03bb..e5ea72ad3e059e33b71cb387b82a100b4c207a3a 100644 (file)
--- a/AST.hs
+++ b/AST.hs
@@ -2,7 +2,7 @@ module AST where
 
 import Data.Char
 import Text.Read
-import Text.ParserCombinators.ReadP hiding ((+++), choice)
+import Text.ParserCombinators.ReadP hiding ((+++), (<++), choice)
 
 newtype Program = Program [AST]
   deriving Show
@@ -26,6 +26,7 @@ data Expr = Num Double
           | Var String
           | Call String [Expr]
           | If Expr Expr Expr
+          | For String Expr Expr (Maybe Expr) Expr
   deriving Show
 data BinOpType = Add | Sub | Mul | Cmp Ordering
   deriving Show
@@ -52,6 +53,7 @@ instance Read Expr where
                     , parseVar
                     , parseCall
                     , parseIf
+                    , parseFor
                     , parseBinOp "+" Add
                     , parseBinOp "-" Sub
                     , parseBinOp "*" Mul
@@ -92,3 +94,18 @@ instance Read Expr where
               skipSpaces
             elseE <- step readPrec
             return (If cond thenE elseE)
+          parseFor = do
+            lift $ do
+              string "for"
+              skipSpaces
+            identifier <- lift (munch1 isAlpha)
+            lift $ skipSpaces >> char '=' >> skipSpaces
+            start <- step readPrec
+            lift $ skipSpaces >> char ',' >> skipSpaces
+            cond <- step readPrec
+            step' <- (do
+              lift $ skipSpaces >> char ',' >> skipSpaces
+              Just <$> step readPrec) <++ pure Nothing
+            lift $ skipSpaces >> string "in" >> skipSpaces
+            body <- step readPrec
+            return (For identifier start cond step' body)