1 files changed, 14 insertions, 4 deletions
diff --git a/app/Compiler.hs b/app/Compiler.hs
index afa7e71..f77a969 100644
--- a/app/Compiler.hs
+++ b/app/Compiler.hs
@@ -53,15 +53,25 @@ squashVars x =
         [(idx, VarI idx' 0) | ((idx, n), idx') <- zip occurs [1 ..], n > 1]
    in renumVars (m' M.!?) x
 
-compileGoals :: Int -> Int -> PrlgInt -> [Code]
-compileGoals andop cut = go'
+squashChoices :: [Code] -> Code
+squashChoices = out . concatMap go
+  where
+    go [Choices cs] = cs
+    go x = [x]
+    out [] = error "choice compilation"
+    out [x] = x
+    out xs = [Choices xs]
+
+compileGoals :: Int -> Int -> Int -> PrlgInt -> Code
+compileGoals andop orop cut = (++ [Done]) . go'
   where
     go' = go . struct2goal
     go p@(CallI x args@[_, _])
       | x == andop = concatMap go' args
+      | x == orop = squashChoices $ map go' args
     go p@(CallI x [])
-      | x == cut = [[Cut]]
-    go x = [compileGoal x]
+      | x == cut = [Cut]
+    go x = compileGoal x
 
 compileGoal :: PrlgInt -> Code
 compileGoal = compileArg . struct2goal