Allow calling recursive functions at load time

Mark the refs in lams of decls as refs of the nearest app in the call
stack (if any) rather than refs of the current decl, since they just
need to have been defined when the app is evaluated rather than when the
current decl is loaded. This allows (mutually) recursive functions to be
called in apps in decls but still disallows cyclic decls.

src/Language/Dtfpl/Simplify/Reorder.hs

@@ -1,9 +1,11 @@
-{-# LANGUAGE BlockArguments        #-}
-{-# LANGUAGE DataKinds             #-}
-{-# LANGUAGE FlexibleContexts      #-}
-{-# LANGUAGE FlexibleInstances     #-}
-{-# LANGUAGE MultiParamTypeClasses #-}
-{-# LANGUAGE TypeFamilies          #-}
+{-# LANGUAGE BlockArguments             #-}
+{-# LANGUAGE DataKinds                  #-}
+{-# LANGUAGE FlexibleContexts           #-}
+{-# LANGUAGE FlexibleInstances          #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE MultiParamTypeClasses      #-}
+{-# LANGUAGE RecordWildCards            #-}
+{-# LANGUAGE TypeFamilies               #-}
 
 {-# OPTIONS_GHC -fno-warn-orphans #-}
 
@@ -13,9 +15,9 @@ module Language.Dtfpl.Simplify.Reorder () where
 
 import           Control.Monad
 import           Data.Foldable
+import           Data.Functor
 import           Data.Generics.Uniplate.Data
 import qualified Data.List.NonEmpty                  as N
-import           Data.Map.Strict                     (Map)
 import qualified Data.Map.Strict                     as M
 import           Data.Set                            (Set)
 import qualified Data.Set                            as S
@@ -37,77 +39,67 @@ type instance StepEffs 'Reordered = '[Error Err]
 type Preordered = Pred 'Reordered
 type PIdentBind = IdentBind Preordered
 
-data CheckType
-    -- | Check all references in a declaration
-    = CheckLoad
-    -- | Check only identifiers referenced at load time
-    | CheckAll
-    deriving (Eq, Ord, Show)
+newtype AppRefs = AppRefs (Set PIdentBind) deriving (Semigroup, Monoid)
+newtype LamRefs = LamRefs (Set PIdentBind) deriving (Semigroup, Monoid)
+
+data ExprRefs = ExprRefs
+    { loadRefs :: Set PIdentBind
+    , appRefs  :: Set PIdentBind
+    , lamRefs  :: Set PIdentBind }
 
 -- | Perform a (reverse) topological sort on the top-level declarations.
 instance Step (T [] (A TopLevel)) 'Reordered where
     step (T tls) = do
-        let binds = map (fst . splitTLDecl) tls
+        (tls', _) <- runOutputList $
+            execState S.empty $ for_ binds $ visit []
+        T <$> traverse step tls'
+      where binds = map (fst . splitTLDecl) tls
             tlMap = M.fromList $ zip binds tls
-            -- First pass: find identifiers referenced by each declaration
-            visitAddRefs :: Members
-                '[ State (Map PIdentBind (CheckType, Set PIdentBind))
+            visit :: Members
+                '[ Output (A TopLevel Preordered) -- topological ordering
+                 , State (Set PIdentBind)         -- visited nodes
                  , Error Err ] r
-                => CheckType -> [PIdentBind] -> PIdentBind -> Sem r Bool
-            visitAddRefs checkType visiting bind = case M.lookup bind tlMap of
-                Nothing -> pure False
-                Just tl
-                    | bind `elem` visiting -> throw $ SimplifyErr $
-                        CyclicDeclErr (tlToDecl tl) $
+                => [PIdentBind] -> PIdentBind -> Sem r ()
+            visit visiting bind = do
+                visited <- get
+                when (bind `S.notMember` visited)
+                    if bind `elem` visiting
+                        then throw $ SimplifyErr $ CyclicDeclErr (tlToDecl tl) $
                             N.fromList $ map (tlToDecl . (tlMap M.!)) visiting
-                    | otherwise -> do
-                        refMap <- get
-                        case M.lookup bind refMap of
-                            Nothing -> addRefs tl S.empty
-                            Just (ct, old)
-                                | ct < checkType -> addRefs tl old
-                                | otherwise -> pure ()
-                        pure True
-              where tlToDecl (A (TLDecl _ decl) _) = decl
-                    visiting' = bind : visiting
-                    addRefs tl old = do
-                        refs <- case checkType of
-                            CheckLoad ->
-                                let (checkAllRefs, loadExpr) = run $
-                                        runWriterAssocR $ loadEval expr
-                                in  (++) <$> visitRefs CheckAll checkAllRefs
-                                         <*> visitRefs CheckLoad
-                                                (getRefs loadExpr)
-                            CheckAll -> visitRefs CheckAll $ getRefs expr
-                        let refs' = S.union old $ S.fromList refs
-                        modify' $ M.insert bind (checkType, refs')
-                      where (_, A expr _) = splitTLDecl tl
-                            visitRefs ct = filterM $ visitAddRefs ct visiting'
-        refMap <- execState M.empty $
-            for_ binds $ visitAddRefs CheckLoad []
-        -- Second pass: determine topological ordering
-        let visitTopoSort :: Members
-                '[ Output (A TopLevel Preordered)
-                 , State (Set PIdentBind)] r
+                        else do
+                            for_ (refMap M.! bind) $ visit (bind : visiting)
+                            modify' $ S.insert bind
+                            output tl
+              where tl = tlMap M.! bind
+                    tlToDecl (A (TLDecl _ decl) _) = decl
+            exprMap = M.fromList $ map splitTLDecl tls
+            exprRefMap = M.map getExprRefs exprMap
+            refMap = M.map getAllRefs exprRefMap
+            getAllRefs :: ExprRefs -> Set PIdentBind
+            getAllRefs ExprRefs {..} = S.union loadRefs $
+                run $ execState S.empty $ traverse_ getRunRefs appRefs
+            getRunRefs :: Member (State (Set PIdentBind)) r
                 => PIdentBind -> Sem r ()
-            visitTopoSort bind = do
-                visited <- get
-                when (bind `S.notMember` visited) do
-                    let (_, refs) = refMap M.! bind
-                    for_ refs visitTopoSort
+            getRunRefs bind = do
+                refs <- get
+                when (bind `S.notMember` refs) do
                     modify' $ S.insert bind
-                    output $ tlMap M.! bind
-        (tls', _) <- runOutputList $
-            execState S.empty $ for_ binds visitTopoSort
-        T <$> traverse step tls'
-      where getRefs :: Expr Preordered -> [PIdentBind]
-            getRefs expr = [ ib | IdentRef ib _ <- universeBi expr ]
-            loadEval :: Member (Writer [PIdentBind]) r
-                => Expr Preordered -> Sem r (Expr Preordered)
-            loadEval expr = case expr of
-                App _ _ -> do
-                    tell $ getRefs expr
-                    pure dummyExpr
-                LamExpr _ -> pure dummyExpr
-                _ -> descendM loadEval expr
-              where dummyExpr = LitExpr $ genLoc $ NumLit 0
+                    traverse_ getRunRefs $ lamRefs $ exprRefMap M.! bind
+            getExprRefs :: A Expr Preordered -> ExprRefs
+            getExprRefs (A expr _) =
+                let (AppRefs appRefs, (LamRefs lamRefs, loadExpr)) =
+                        run $ runWriter $ runWriter $ go expr
+                    loadRefs = S.union appRefs $ getRefs loadExpr
+                in  ExprRefs {..}
+              where go :: Members '[Writer AppRefs, Writer LamRefs] r
+                        => Expr Preordered -> Sem r (Expr Preordered)
+                    go e = case e of
+                        App _ _   -> tell (AppRefs $ getRefs e) $> dummyExpr
+                        LamExpr _ -> tell (LamRefs $ getRefs e) $> dummyExpr
+                        _         -> descendM go e
+                    getRefs :: Expr Preordered -> Set PIdentBind
+                    getRefs e = S.fromList
+                        [ ib | IdentRef ib _ <- universeBi e
+                             , M.member ib tlMap ]
+                    dummyExpr :: Expr Preordered
+                    dummyExpr = LitExpr $ genLoc $ NumLit 0