Ensure that transitive dependencies of non-inlined functions get comp…

…iled.

examples/eval-tests.dx

@@ -822,3 +822,25 @@ def fromLeftFloat (x:(Float | Int)) : Float =
 
 :p fromLeftFloat $ Left 1.2
 > 1.2
+
+@noinline
+def f1 (x:Int) : Int = x + 1
+
+@noinline
+def f2 (x:Int) : Int = f1 $ f1 $ f1 $ f1 $ f1 $ f1 $ f1 $ f1 $ f1 $ f1 $ x
+
+@noinline
+def f3 (x:Int) : Int = f2 $ f2 $ f2 $ f2 $ f2 $ f2 $ f2 $ f2 $ f2 $ f2 $ x
+
+@noinline
+def f4 (x:Int) : Int = f3 $ f3 $ f3 $ f3 $ f3 $ f3 $ f3 $ f3 $ f3 $ f3 $ x
+
+@noinline
+def f5 (x:Int) : Int = f4 $ f4 $ f4 $ f4 $ f4 $ f4 $ f4 $ f4 $ f4 $ f4 $ x
+
+@noinline
+def f6 (x:Int) : Int = f5 $ f5 $ f5 $ f5 $ f5 $ f5 $ f5 $ f5 $ f5 $ f5 $ x
+
+-- This will compile extremely slowly if non-inlining is broken
+:p f6 0
+> 100000

src/lib/Imp.hs

@@ -80,9 +80,8 @@ toImpModule :: TopEnv -> Backend -> CallingConvention -> Name
             -> (ImpFunction, ImpModule, AtomRecon)
 toImpModule env backend cc entryName argBinders maybeDest block = do
   let standaloneFunctions =
-        for (envPairs $ envIntersect (freeVars block) env) $
-           \(v, (_, LetBound _ (Atom f))) ->
-              runImpM initCtx inVarScope $ toImpStandalone v f
+        for (requiredFunctions env block) $ \(v, f) ->
+          runImpM initCtx inVarScope $ toImpStandalone v f
   runImpM initCtx inVarScope $ do
     (reconAtom, impBlock) <- scopedBlock $ translateTopLevel (maybeDest, block)
     otherFunctions <- toList <$> looks envFunctions
@@ -97,6 +96,18 @@ toImpModule env backend cc entryName argBinders maybeDest block = do
     destScope = fromMaybe mempty $ fmap freeVars maybeDest
     initCtx = ImpCtx backend CPU TopLevel
 
+requiredFunctions :: HasVars a => Scope -> a -> [(Name, Atom)]
+requiredFunctions scope expr =
+  for (transitiveClosure getParents immediateParents) $ \fname -> do
+    let (_, LetBound _ (Atom f)) = scope ! fname
+    (fname, f)
+  where
+    getParents :: Name -> [Name]
+    getParents fname = envNames $ freeVars $ scope ! fname
+
+    immediateParents :: [Name]
+    immediateParents = envNames $ freeVars expr `envIntersect` scope
+
 -- We don't emit any results when a destination is provided, since they are already
 -- going to be available through the dest.
 translateTopLevel :: WithDest Block -> ImpM (AtomRecon, [IExpr])

src/lib/Util.hs

@@ -5,10 +5,11 @@
 -- https://developers.google.com/open-source/licenses/bsd
 
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 
 module Util (IsBool (..), group, ungroup, pad, padLeft, delIdx, replaceIdx,
              insertIdx, mvIdx, mapFst, mapSnd, splitOn, scan,
-             scanM, composeN, mapMaybe, uncons, repeated,
+             scanM, composeN, mapMaybe, uncons, repeated, transitiveClosure,
              showErr, listDiff, splitMap, enumerate, restructure,
              onSnd, onFst, highlightRegion, findReplace, swapAt, uncurry3,
              bindM2, foldMapM, lookupWithIdx, (...), zipWithT, for) where
@@ -21,6 +22,8 @@ import qualified Data.Set as Set
 import qualified Data.Map.Strict as M
 import Control.Monad.State.Strict
 
+import Cat
+
 class IsBool a where
   toBool :: a -> Bool
 
@@ -218,3 +221,14 @@ zipWithT f trav args = flip evalStateT (toList args) $ flip traverse trav $ \e -
 
 for :: Functor f => f a -> (a -> b) -> f b
 for = flip fmap
+
+transitiveClosure :: forall a. Ord a => (a -> [a]) -> [a] -> [a]
+transitiveClosure getParents seeds =
+  toList $ snd $ runCat (mapM_ go seeds) mempty
+  where
+    go :: a -> Cat (Set.Set a) ()
+    go x = do
+      visited <- look
+      unless (x `Set.member` visited) $ do
+        extend $ Set.singleton x
+        mapM_ go $ getParents x