-- | Infer which functions can trigger GC.

module Gibbon.Passes.InferFunAllocs
  ( inferFunAllocs ) where

import Data.Map as M
import Gibbon.Common
import Gibbon.L2.Syntax

--------------------------------------------------------------------------------

-- | Chatter level for this module:
lvl :: Int
lvl :: Int
lvl = Int
5

type FunEnv = TyEnv FunMeta

inferFunAllocs :: Prog2 -> PassM Prog2
inferFunAllocs :: Prog2 -> PassM Prog2
inferFunAllocs prg :: Prog2
prg@Prog{FunDefs Exp2
fundefs :: FunDefs Exp2
fundefs :: forall ex. Prog ex -> FunDefs ex
fundefs} = do
  let finalMetas :: FunEnv
finalMetas = Int -> FunDefs Exp2 -> FunEnv -> FunEnv
fixpoint Int
1 FunDefs Exp2
fundefs ((FunDef Exp2 -> FunMeta) -> FunDefs Exp2 -> FunEnv
forall a b k. (a -> b) -> Map k a -> Map k b
M.map FunDef Exp2 -> FunMeta
forall ex. FunDef ex -> FunMeta
funMeta FunDefs Exp2
fundefs)
      funs :: FunDefs Exp2
funs = (FunDef Exp2 -> FunDef Exp2) -> FunDefs Exp2 -> FunDefs Exp2
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (\fn :: FunDef Exp2
fn@FunDef{Var
funName :: Var
funName :: forall ex. FunDef ex -> Var
funName} ->
                       FunDef Exp2
fn { funMeta :: FunMeta
funMeta = FunEnv
finalMetas FunEnv -> Var -> FunMeta
forall k a. Ord k => Map k a -> k -> a
! Var
funName })
             FunDefs Exp2
fundefs
  Prog2 -> PassM Prog2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Prog2 -> PassM Prog2) -> Prog2 -> PassM Prog2
forall a b. (a -> b) -> a -> b
$ Prog2
prg { fundefs :: FunDefs Exp2
fundefs = FunDefs Exp2
funs }
  where
    fixpoint :: Int -> FunDefs2 -> FunEnv -> FunEnv
    fixpoint :: Int -> FunDefs Exp2 -> FunEnv -> FunEnv
fixpoint Int
iter FunDefs Exp2
funs FunEnv
fenv =
       let metas :: FunEnv
metas = (FunDef Exp2 -> FunMeta) -> FunDefs Exp2 -> FunEnv
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (FunEnv -> FunDef Exp2 -> FunMeta
inferFunDef FunEnv
fenv) FunDefs Exp2
funs
       in if FunEnv
fenv FunEnv -> FunEnv -> Bool
forall a. Eq a => a -> a -> Bool
== FunEnv
metas
          then Int -> [Char] -> FunEnv -> FunEnv
forall a. Int -> [Char] -> a -> a
dbgTrace Int
lvl ([Char]
"\n<== Fixpoint completed after iteration "[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++Int -> [Char]
forall a. Show a => a -> [Char]
show Int
iter[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++[Char]
" ==>") (FunEnv -> FunEnv) -> FunEnv -> FunEnv
forall a b. (a -> b) -> a -> b
$ FunEnv
fenv
          else Int -> FunDefs Exp2 -> FunEnv -> FunEnv
fixpoint (Int
iterInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) FunDefs Exp2
funs FunEnv
metas

inferFunDef :: FunEnv -> FunDef2 -> FunMeta
inferFunDef :: FunEnv -> FunDef Exp2 -> FunMeta
inferFunDef FunEnv
fenv FunDef{Exp2
funBody :: Exp2
funBody :: forall ex. FunDef ex -> ex
funBody,FunMeta
funMeta :: forall ex. FunDef ex -> FunMeta
funMeta :: FunMeta
funMeta} =
  FunMeta
funMeta { funCanTriggerGC :: Bool
funCanTriggerGC = FunEnv -> Exp2 -> Bool
inferExp FunEnv
fenv Exp2
funBody }

inferExp :: FunEnv -> Exp2 -> Bool
inferExp :: FunEnv -> Exp2 -> Bool
inferExp FunEnv
fenv Exp2
expr =
  case Exp2
expr of
    VarE{}    -> Bool
False
    LitE{}    -> Bool
False
    CharE{}   -> Bool
False
    FloatE{}  -> Bool
False
    LitSymE{} -> Bool
False
    AppE Var
v [Var]
_locs [Exp2]
_e -> FunMeta -> Bool
funCanTriggerGC (FunEnv
fenv FunEnv -> Var -> FunMeta
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
v)
    PrimAppE Prim Ty2
_ [Exp2]
ls   -> (Exp2 -> Bool) -> [Exp2] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Exp2 -> Bool
go [Exp2]
ls
    LetE (Var
_,[Var]
_,Ty2
_,Exp2
rhs) Exp2
bod -> Exp2 -> Bool
go Exp2
rhs Bool -> Bool -> Bool
|| Exp2 -> Bool
go Exp2
bod
    IfE Exp2
tst Exp2
consq Exp2
alt    -> Exp2 -> Bool
go Exp2
tst Bool -> Bool -> Bool
|| Exp2 -> Bool
go Exp2
consq Bool -> Bool -> Bool
|| Exp2 -> Bool
go Exp2
alt
    MkProdE [Exp2]
ls           -> (Exp2 -> Bool) -> [Exp2] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Exp2 -> Bool
go [Exp2]
ls
    SpawnE Var
v [Var]
_locs [Exp2]
_e    -> FunMeta -> Bool
funCanTriggerGC (FunEnv
fenv FunEnv -> Var -> FunMeta
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
v)
    Exp2
SyncE      -> Bool
False
    ProjE Int
_n Exp2
e -> Exp2 -> Bool
go Exp2
e
    CaseE Exp2
e [([Char], [(Var, Var)], Exp2)]
mp -> Exp2 -> Bool
go Exp2
e Bool -> Bool -> Bool
|| (([Char], [(Var, Var)], Exp2) -> Bool)
-> [([Char], [(Var, Var)], Exp2)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\([Char]
_,[(Var, Var)]
_,Exp2
rhs) -> Exp2 -> Bool
go Exp2
rhs) [([Char], [(Var, Var)], Exp2)]
mp
    DataConE Var
_loc [Char]
_dcon [Exp2]
es -> (Exp2 -> Bool) -> [Exp2] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Exp2 -> Bool
go [Exp2]
es
    TimeIt Exp2
e Ty2
_ Bool
_           -> Exp2 -> Bool
go Exp2
e
    WithArenaE Var
_v Exp2
e -> Exp2 -> Bool
go Exp2
e
    MapE{}  -> [Char] -> Bool
forall a. HasCallStack => [Char] -> a
error [Char]
"inferFunAllocs: MapE not handled."
    FoldE{} -> [Char] -> Bool
forall a. HasCallStack => [Char] -> a
error [Char]
"inferFunAllocs: FoldE not handled."
    Ext (LetRegionE{})     -> Bool
True
    Ext (LetParRegionE{})  -> Bool
True
    Ext (LetLocE Var
_ PreLocExp Var
_ Exp2
rhs)  -> Exp2 -> Bool
go Exp2
rhs
    Ext (RetE [Var]
_ Var
_)         -> Bool
False
    Ext (FromEndE Var
_ )      -> Bool
False
    Ext (IndirectionE{})   -> Bool
False
    Ext (BoundsCheck{})    -> Bool
_todo -- (S.empty, Nothing)
    Ext (AddFixed{})       -> Bool
False
    Ext (E2Ext Var Ty2
GetCilkWorkerNum) -> Bool
False
    Ext (LetAvail [Var]
_ Exp2
e)     -> Exp2 -> Bool
go Exp2
e
    Ext (StartOfPkdCursor{})    -> Bool
False
    Ext (TagCursor{})           -> Bool
False
    Ext (AllocateTagHere{})     -> Bool
False
    Ext (AllocateScalarsHere{}) -> Bool
False
    Ext (SSPush{})              -> Bool
False
    Ext (SSPop{})               -> Bool
False
  where
    go :: Exp2 -> Bool
go = FunEnv -> Exp2 -> Bool
inferExp FunEnv
fenv