1 files changed, 13 insertions, 1 deletions
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