Add break statements

5 files changed, 61 insertions, 21 deletions

diff --git a/AST.hs b/AST.hs
index 3e80830..5217d46 100644
--- a/AST.hs
+++ b/AST.hs
@@ -27,6 +27,7 @@ data Statement
     | SAs AsExpression Expression
     | SIf Expression Block Block
     | SWhile Expression Block
+    | SBreak Int
     | SReturn (Maybe Expression)
     | SExpr Expression
   deriving (Show, Eq)
@@ -111,6 +112,10 @@ instance Pretty Statement where
         "if " ++ prettyI i c ++ " " ++ prettyI i b1 ++ " else " ++ prettyI i b2
     prettyI i (SWhile c b) =
         "while " ++ prettyI i c ++ " " ++ prettyI i b
+    prettyI _ (SBreak 0) =
+        "break;"
+    prettyI _ (SBreak n) =
+        "break " ++ show n ++ ";"
     prettyI _ (SReturn Nothing) =
         "return;"
     prettyI i (SReturn (Just e)) =
diff --git a/BuildIR.hs b/BuildIR.hs
index 665dd33..a4be797 100644
--- a/BuildIR.hs
+++ b/BuildIR.hs
@@ -21,6 +21,7 @@ type Scope = Map.Map Name (Ref, Type)
 data BuildState = BuildState
     { nextId :: Id,
       scopeStack :: [Scope],
+      loopStack :: [Id],
       currentBlock :: Id,
       blockMap :: Map.Map Id BB }
 
@@ -28,6 +29,7 @@ initBuildState :: BuildState
 initBuildState = BuildState
     { nextId = 0,
       scopeStack = [],
+      loopStack = [],
       currentBlock = undefined,
       blockMap = Map.empty }
 
@@ -81,6 +83,13 @@ getAllBlocks = liftM Map.elems (gets blockMap)
 switchBlock :: Id -> BuildM ()
 switchBlock bid = modify $ \s -> s {currentBlock = bid}
 
+withLoop :: Id -> BuildM a -> BuildM a
+withLoop i act = do
+    modify $ \s -> s {loopStack = i : loopStack s}
+    res <- act
+    modify $ \s -> s {loopStack = tail (loopStack s)}
+    return res
+
 withScope :: BuildM a -> BuildM a
 withScope act = do
     modify $ \s -> s {scopeStack = Map.empty : scopeStack s}
@@ -172,9 +181,12 @@ convertStatement (SWhile c b) nextnext = do
     switchBlock cend
     setTerm $ IJcc CNeq cref (Constant (refSize cref) 0) body nextnext
     switchBlock body
-    convertBlock b bodyend
+    withLoop nextnext $ convertBlock b bodyend
     switchBlock bodyend
     setTerm $ IJmp cond
+convertStatement (SBreak n) _ = do
+    ls <- gets loopStack
+    setTerm $ IJmp (ls !! n)
 convertStatement (SReturn Nothing) _ = do
     setTerm IRet
 convertStatement (SReturn (Just e)) _ = do
diff --git a/ProgramParser.hs b/ProgramParser.hs
index 2cacaf5..38da21f 100644
--- a/ProgramParser.hs
+++ b/ProgramParser.hs
@@ -2,6 +2,7 @@ module ProgramParser(parseProgram) where
 
 import Control.Monad
 import Data.Char
+import Data.Maybe
 import Text.Parsec
 import qualified Text.Parsec.Expr as E
 
@@ -70,7 +71,7 @@ pBlock = do
     return $ Block body
 
 pStatement :: Parser Statement
-pStatement = pSIf <|> pSWhile <|> pSReturn <|> pSDecl <|> pSAs <|> pSExpr
+pStatement = pSIf <|> pSWhile <|> pSReturn <|> pSBreak <|> pSDecl <|> pSAs <|> pSExpr
 
 pSDecl :: Parser Statement
 pSDecl = do
@@ -112,8 +113,16 @@ pSWhile = do
 pSReturn :: Parser Statement
 pSReturn = do
     symbol "return"
-    SReturn <$> ((symbol ";" >> return Nothing) <|>
-                 ((Just <$> pExpression) <* symbol ";"))
+    m <- optionMaybe pExpression
+    symbol ";"
+    return $ SReturn m
+
+pSBreak :: Parser Statement
+pSBreak = do
+    symbol "break"
+    m <- optionMaybe pIntegerInt
+    symbol ";"
+    return $ SBreak (fromMaybe 0 m)
 
 pSExpr :: Parser Statement
 pSExpr = do
@@ -231,6 +240,14 @@ pName = do
 pInteger :: Parser Integer
 pInteger = read <$> many1 (satisfy isDigit) <* pWhiteComment
 
+pIntegerInt :: Parser Int
+pIntegerInt = do
+    i <- pInteger
+    when (i > (fromIntegral (maxBound :: Int) :: Integer) ||
+          i < (fromIntegral (minBound :: Int) :: Integer)) $
+        unexpected $ "Integer literal " ++ show i ++ " does not fit in an Int"
+    return $ fromIntegral i
+
 
 symbol :: String -> Parser ()
 symbol "" = error "symbol \"\""
diff --git a/TypeCheck.hs b/TypeCheck.hs
index 36d98c8..13d33c9 100644
--- a/TypeCheck.hs
+++ b/TypeCheck.hs
@@ -70,24 +70,27 @@ annotateDVar db (DVar toptype name expr) = do
         else Left $ "Cannot assign a value of type " ++ pretty typ ++
                     " to a variable of type " ++ pretty toptype
 
+data State = State {stDfunc :: DFunc, stLoopDepth :: Int}
+  deriving Show
+
 annotateDFunc :: TypeDB -> DFunc -> Error DFunc
 annotateDFunc db dfunc@(DFunc rettype name arglist block) = do
     when (name == "main" && rettype /= Just TInt) $
         Left $ "Function 'main' should return an int"
     let db' = foldl registerArg db arglist
-    block' <- annotateBlock dfunc db' block
+    block' <- annotateBlock (State dfunc 0) db' block
     return $ DFunc rettype name arglist block'
   where
     registerArg :: TypeDB -> (Type, Name) -> TypeDB
     registerArg db' (t, n) = dbSet db' n (DBVar t)
 
-annotateBlock :: DFunc -> TypeDB -> Block -> Error Block
-annotateBlock dfunc db (Block sts) =
+annotateBlock :: State -> TypeDB -> Block -> Error Block
+annotateBlock state db (Block sts) =
     Block . snd <$> foldM (\(db', l) st ->
-                               (\(db'', st') -> (db'', l ++ [st'])) <$> annotateStatement dfunc db' st)
+                               (\(db'', st') -> (db'', l ++ [st'])) <$> annotateStatement state db' st)
                           (db, []) sts
 
-annotateStatement :: DFunc -> TypeDB -> Statement -> Error (TypeDB, Statement)
+annotateStatement :: State -> TypeDB -> Statement -> Error (TypeDB, Statement)
 annotateStatement _ db (SDecl toptype name expr) = do
     expr' <- annotateExpr db expr
     when (isNothing (typeof expr')) $
@@ -112,28 +115,32 @@ annotateStatement _ db (SAs ae expr) = do
         Left $ "Cannot assign a value of type " ++ pretty typ ++
                " to a location of type " ++ pretty aetyp
     return (db, SAs ae' expr')
-annotateStatement dfunc db (SIf expr bl1 bl2) = do
+annotateStatement st db (SIf expr bl1 bl2) = do
     expr' <- annotateExpr db expr
     when (isNothing (typeof expr')) $
         Left $ "Cannot use void value in 'if' condition"
     when (not $ canCoerce (fromJust (typeof expr')) TInt) $
         Left $ "Cannot use type " ++ pretty (fromJust (typeof expr')) ++ " in 'if' condition"
-    bl1' <- withScope db $ flip (annotateBlock dfunc) bl1
-    bl2' <- withScope db $ flip (annotateBlock dfunc) bl2
+    bl1' <- withScope db $ flip (annotateBlock st) bl1
+    bl2' <- withScope db $ flip (annotateBlock st) bl2
     return (db, SIf expr' bl1' bl2')
-annotateStatement dfunc db (SWhile expr bl) = do
+annotateStatement (State dfunc ld) db (SWhile expr bl) = do
     expr' <- annotateExpr db expr
     when (isNothing (typeof expr')) $
         Left $ "Cannot use void value in 'while' condition"
     when (not $ canCoerce (fromJust (typeof expr')) TInt) $
         Left $ "Cannot use type " ++ pretty (fromJust (typeof expr')) ++ " in 'while' condition"
-    bl' <- withScope db $ flip (annotateBlock dfunc) bl
+    bl' <- withScope db $ flip (annotateBlock (State dfunc (ld+1))) bl
     return (db, SWhile expr' bl')
-annotateStatement (DFunc Nothing _ _ _) db (SReturn Nothing) =
+annotateStatement (State {stLoopDepth = ld}) db (SBreak n) =
+    if n > ld
+        then Left $ "'break " ++ show n ++ "' while in only " ++ show ld ++ " loops"
+        else return (db, SBreak n)
+annotateStatement (State {stDfunc = DFunc Nothing _ _ _}) db (SReturn Nothing) =
     return (db, SReturn Nothing)
-annotateStatement (DFunc (Just _) _ _ _) _ (SReturn Nothing) =
+annotateStatement (State {stDfunc = DFunc (Just _) _ _ _}) _ (SReturn Nothing) =
     Left "Cannot return void value from non-void function"
-annotateStatement (DFunc mrt _ _ _) db (SReturn (Just expr)) = do
+annotateStatement (State {stDfunc = DFunc mrt _ _ _}) db (SReturn (Just expr)) = do
     expr' <- annotateExpr db expr
     case mrt of
         Nothing -> Left $ "Cannot return non-void value from void function"
diff --git a/bf.lang b/bf.lang
index 71a5d8a..e2f6788 100644
--- a/bf.lang
+++ b/bf.lang
@@ -85,17 +85,16 @@ func int main() {
 	char[] source := new char[4088];
 
 	int sourcelen := 0;
-	int done := 0;
-	while (done != 1) {
+	while (1) {
 		int c := getc();
 		// putc(char(c));
 		if (c < 0) {
-			done = 1;
+			break;
 		} else {
 			source[sourcelen] = char(c);
 			sourcelen = sourcelen + 1;
 			if (sourcelen >= bufsize - 1) {
-				done = 1;
+				break;
 			}
 		}
 	}