3 files changed, 42 insertions, 27 deletions

diff --git a/src/LambdaCube/Compiler/DesugaredSource.hs b/src/LambdaCube/Compiler/DesugaredSource.hs
index f0af5437..a2c4e591 100644
--- a/src/LambdaCube/Compiler/DesugaredSource.hs
+++ b/src/LambdaCube/Compiler/DesugaredSource.hs
@@ -470,8 +470,8 @@ data StmtNode = StmtNode
     , snRevChildren :: [StmtNode]
     }
 
-sortDefs :: [Stmt] -> [Stmt]
-sortDefs xs = concatMap (desugarMutual . map snValue) $ scc snId snChildren snRevChildren nodes
+sortDefs :: [Stmt] -> [[Stmt]]
+sortDefs xs = map snValue <$> scc snId snChildren snRevChildren nodes
   where
     nodes = zipWith mkNode [0..] xs
       where
@@ -494,9 +494,6 @@ sortDefs xs = concatMap (desugarMutual . map snValue) $ scc snId snChildren snRe
             Let n _ _ -> [n]
             Data n _ _ cs -> n: map fst cs
 
-desugarMutual [x] = [x]
-desugarMutual xs = xs
-
 -------------------------------------------------------------------------------- module
 
 data Module_ a = Module
diff --git a/src/LambdaCube/Compiler/Parser.hs b/src/LambdaCube/Compiler/Parser.hs
index 8f457e49..50a1632b 100644
--- a/src/LambdaCube/Compiler/Parser.hs
+++ b/src/LambdaCube/Compiler/Parser.hs
@@ -500,18 +500,7 @@ valueDef :: BodyParser [PreStmt]
 valueDef = do
     (dns, p) <- try "pattern" $ longPattern <* reservedOp "="
     checkPattern dns
-    desugarValueDef p =<< setR parseTermLam
-  where
-    desugarValueDef p e = runCheck $ sequence
-        $ pure (FunAlt n [] $ noGuards e)
-        : [ FunAlt x [] . noGuards <$> compileCase (SGlobal n) [(p, noGuards $ SVar x i)]
-          | (i, x) <- zip [0..] dns
-          ]
-      where
-        dns = reverse $ getPVars p
-        n = mangleNames dns
-
-mangleNames xs = SIName (foldMap sourceInfo xs) $ "_" ++ intercalate "_" (sName <$> xs)
+    runCheck . desugarValueDef p =<< setR parseTermLam
 
 -------------------------------------------------------------------------------- modules
 
@@ -587,5 +576,5 @@ runDefParser ds_ dp = do
         f (Uncovered' si x) | not $ null $ filter (not . null . fst) x = Just $ Uncovered si x
         f _ = Nothing
 
-    return (sortDefs defs, catMaybes [f w | Right w <- dns], ds)
+    return (concatMap desugarMutual $ sortDefs defs, catMaybes [f w | Right w <- dns], ds)
 
diff --git a/src/LambdaCube/Compiler/Statements.hs b/src/LambdaCube/Compiler/Statements.hs
index c9dfa1f7..2f1b5859 100644
--- a/src/LambdaCube/Compiler/Statements.hs
+++ b/src/LambdaCube/Compiler/Statements.hs
@@ -17,7 +17,7 @@ import Data.Function
 import qualified Data.Set as Set
 import Control.Monad.Writer
 import Control.Arrow hiding ((<+>))
---import Debug.Trace
+import Debug.Trace
 
 import LambdaCube.Compiler.DeBruijn
 import LambdaCube.Compiler.Pretty hiding (braces, parens)
@@ -47,14 +47,17 @@ instance DeBruijnify SIName PreStmt where
 mkLets :: [Stmt]{-where block-} -> SExp{-main expression-} -> SExp{-big let with lambdas; replaces global names with de bruijn indices-}
 mkLets = mkLets_ SLet
 
-mkLets_ mkLet = mkLets' . sortDefs where
-    mkLets' [] e = e
-    mkLets' (Let n mt x: ds) e
-        = mkLet n (maybe id (flip SAnn) mt x') (deBruijnify [n] $ mkLets' ds e)
-      where
-        x' = if usedS n x then SBuiltin "primFix" `SAppV` SLamV (deBruijnify [n] x) else x
-    mkLets' (PrecDef{}: ds) e = mkLets' ds e
-    mkLets' (x: ds) e = error $ "mkLets: " ++ ppShow x
+mkLets_ mkLet = mkLets' mkLet . concatMap desugarMutual . sortDefs
+
+mkLets' mkLet = f where
+    f [] e = e
+    f (Let n mt x: ds) e = mkLet n (maybe id (flip SAnn) mt (addFix n x)) (deBruijnify [n] $ f ds e)
+    f (PrecDef{}: ds) e = f ds e
+    f (x: ds) e = error $ "mkLets: " ++ ppShow x
+
+addFix n x
+    | usedS n x = SBuiltin "primFix" `SAppV` SLamV (deBruijnify [n] x)
+    | otherwise = x
 
 type DefinedSet = Set.Set SName
 
@@ -118,3 +121,29 @@ funAlt n pats gt = FunAlt n (fst <$> pats) $ compilePatts (map snd pats) gt
 
 funAlt' n ts x gt = FunAlt n ts $ compilePatts x gt
 
+desugarValueDef :: MonadWriter [ParseCheck] m => ParPat -> SExp -> m [PreStmt]
+desugarValueDef p e = sequence
+    $ pure (FunAlt n [] $ noGuards e)
+    : [ FunAlt x [] . noGuards <$> compileCase (SGlobal n) [(p, noGuards $ SVar x i)]
+      | (i, x) <- zip [0..] dns
+      ]
+  where
+    dns = reverse $ getPVars p
+    n = mangleNames dns
+
+getLet (Let x Nothing (SRHS dx)) = Just (x, dx)
+getLet _ = Nothing
+
+desugarMutual :: {-MonadWriter [ParseCheck] m => -} [Stmt] -> [Stmt]
+desugarMutual [x] = [x]
+desugarMutual (traverse getLet -> Just (unzip -> (ns, ds))) = fst{-TODO-} $ runWriter $ do
+    ss <- compileStmt' =<< desugarValueDef (foldr cHCons cHNil $ PVarSimp <$> ns) (SGlobal xy)
+    return $
+        Let xy Nothing (addFix xy $ SRHS $ mkLets' SLet ss $ foldr HCons HNil ds) : ss
+  where
+    xy = mangleNames ns
+desugarMutual xs = error "desugarMutual"
+
+mangleNames xs = SIName (foldMap sourceInfo xs) $ "_" ++ intercalate "_" (sName <$> xs)
+
+