-- | Aggressive dead code elimination.  No regard for termination,
-- effects, or partiality.

-- [Populated with duplicated code from InlinePacked]

module Gibbon.Passes.ShakeTree
    (shakeTree) where

import Prelude hiding (exp)
import qualified Data.Map as M
import qualified Data.Set as S

import Gibbon.Common (PassM, dbgTrace)
import Gibbon.L1.Syntax
import Gibbon.L3.Syntax


-- | Drop all unreferenced let-bindings.
shakeTree :: Prog3 -> PassM Prog3
shakeTree :: Prog3 -> PassM Prog3
shakeTree prg :: Prog3
prg@Prog{FunDefs Exp3
fundefs :: FunDefs Exp3
fundefs :: forall ex. Prog ex -> FunDefs ex
fundefs,Maybe (Exp3, TyOf Exp3)
mainExp :: Maybe (Exp3, TyOf Exp3)
mainExp :: forall ex. Prog ex -> Maybe (ex, TyOf ex)
mainExp} = Prog3 -> PassM Prog3
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Prog3 -> PassM Prog3) -> Prog3 -> PassM Prog3
forall a b. (a -> b) -> a -> b
$
  Prog3
prg { fundefs :: FunDefs Exp3
fundefs = (FunDef Exp3 -> FunDef Exp3) -> FunDefs Exp3 -> FunDefs Exp3
forall a b k. (a -> b) -> Map k a -> Map k b
M.map FunDef Exp3 -> FunDef Exp3
fd FunDefs Exp3
fundefs
      , mainExp :: Maybe (Exp3, TyOf Exp3)
mainExp = case Maybe (Exp3, TyOf Exp3)
mainExp of
                    Maybe (Exp3, TyOf Exp3)
Nothing      -> Maybe (Exp3, TyOf Exp3)
Maybe (Exp3, UrTy ())
forall a. Maybe a
Nothing
                    (Just (Exp3
e,TyOf Exp3
t)) -> (Exp3, UrTy ()) -> Maybe (Exp3, UrTy ())
forall a. a -> Maybe a
Just (Exp3 -> Exp3
shakeTreeExp Exp3
e, TyOf Exp3
UrTy ()
t)
      }
 where
   fd :: FunDef Exp3 -> FunDef Exp3
fd f :: FunDef Exp3
f@FunDef{Exp3
funBody :: Exp3
funBody :: forall ex. FunDef ex -> ex
funBody} = FunDef Exp3
f { funBody :: Exp3
funBody = Exp3 -> Exp3
shakeTreeExp Exp3
funBody }

shakeTreeExp :: Exp3 -> Exp3
shakeTreeExp :: Exp3 -> Exp3
shakeTreeExp = Exp3 -> Exp3
go
  where

  go :: Exp3 -> Exp3
  go :: Exp3 -> Exp3
go Exp3
e0 =
   -- dbgTrace 5 ("Inline, processing with env:\n "++sdoc env++"\n exp: "++sdoc e0) $
   case Exp3
e0 of

    (LetE (Var
v,[()]
locs,UrTy ()
t,Exp3
rhs) Exp3
bod) ->
        let bod' :: Exp3
bod' = Exp3 -> Exp3
go Exp3
bod
            fv :: Set Var
fv   = Exp3 -> Set Var
forall a. FreeVars a => a -> Set Var
gFreeVars Exp3
bod'
        in
        if Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Var
v Set Var
fv Bool -> Bool -> Bool
|| Exp3 -> Bool
hasEffect Exp3
rhs
        then (Var, [()], UrTy (), Exp3) -> Exp3 -> Exp3
forall (ext :: * -> * -> *) loc dec.
(Var, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (Var
v,[()]
locs,UrTy ()
t, Exp3 -> Exp3
go Exp3
rhs) Exp3
bod'
        else Int -> [Char] -> Exp3 -> Exp3
forall a. Int -> [Char] -> a -> a
dbgTrace Int
4 ([Char]
" [shakeTreeExp] dropping binding: "[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++(Var, UrTy (), Exp3) -> [Char]
forall a. Show a => a -> [Char]
show (Var
v,UrTy ()
t,Exp3
rhs))(Exp3 -> Exp3) -> Exp3 -> Exp3
forall a b. (a -> b) -> a -> b
$ Exp3
bod'

    (VarE Var
v)           -> Var -> Exp3
forall (ext :: * -> * -> *) loc dec. Var -> PreExp ext loc dec
VarE Var
v
    (LitE Int
i)           -> Int -> Exp3
forall (ext :: * -> * -> *) loc dec. Int -> PreExp ext loc dec
LitE Int
i
    (CharE Char
i)           -> Char -> Exp3
forall (ext :: * -> * -> *) loc dec. Char -> PreExp ext loc dec
CharE Char
i
    (FloatE Double
i)         -> Double -> Exp3
forall (ext :: * -> * -> *) loc dec. Double -> PreExp ext loc dec
FloatE Double
i
    (LitSymE Var
v)        -> Var -> Exp3
forall (ext :: * -> * -> *) loc dec. Var -> PreExp ext loc dec
LitSymE Var
v
    (AppE Var
f [()]
locs [Exp3]
es)   -> Var -> [()] -> [Exp3] -> Exp3
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE Var
f [()]
locs ([Exp3] -> Exp3) -> [Exp3] -> Exp3
forall a b. (a -> b) -> a -> b
$ (Exp3 -> Exp3) -> [Exp3] -> [Exp3]
forall a b. (a -> b) -> [a] -> [b]
map Exp3 -> Exp3
go [Exp3]
es
    (PrimAppE Prim (UrTy ())
pr [Exp3]
es)   -> Prim (UrTy ()) -> [Exp3] -> Exp3
forall (ext :: * -> * -> *) loc dec.
Prim dec -> [PreExp ext loc dec] -> PreExp ext loc dec
PrimAppE Prim (UrTy ())
pr ([Exp3] -> Exp3) -> [Exp3] -> Exp3
forall a b. (a -> b) -> a -> b
$ (Exp3 -> Exp3) -> [Exp3] -> [Exp3]
forall a b. (a -> b) -> [a] -> [b]
map Exp3 -> Exp3
go [Exp3]
es
    (IfE Exp3
e1 Exp3
e2 Exp3
e3)     -> Exp3 -> Exp3 -> Exp3 -> Exp3
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> PreExp ext loc dec -> PreExp ext loc dec -> PreExp ext loc dec
IfE (Exp3 -> Exp3
go Exp3
e1) (Exp3 -> Exp3
go Exp3
e2) (Exp3 -> Exp3
go Exp3
e3)

    (ProjE Int
i Exp3
e)  -> Int -> Exp3 -> Exp3
forall (ext :: * -> * -> *) loc dec.
Int -> PreExp ext loc dec -> PreExp ext loc dec
ProjE Int
i (Exp3 -> Exp3) -> Exp3 -> Exp3
forall a b. (a -> b) -> a -> b
$ Exp3 -> Exp3
go Exp3
e
    (MkProdE [Exp3]
es) -> [Exp3] -> Exp3
forall (ext :: * -> * -> *) loc dec.
[PreExp ext loc dec] -> PreExp ext loc dec
MkProdE ([Exp3] -> Exp3) -> [Exp3] -> Exp3
forall a b. (a -> b) -> a -> b
$ (Exp3 -> Exp3) -> [Exp3] -> [Exp3]
forall a b. (a -> b) -> [a] -> [b]
map (Exp3 -> Exp3
go) [Exp3]
es

    -- We don't rename field binders with to/from witness:
    (CaseE Exp3
e [([Char], [(Var, ())], Exp3)]
mp) -> let mp' :: [([Char], [(Var, ())], Exp3)]
mp' = (([Char], [(Var, ())], Exp3) -> ([Char], [(Var, ())], Exp3))
-> [([Char], [(Var, ())], Exp3)] -> [([Char], [(Var, ())], Exp3)]
forall a b. (a -> b) -> [a] -> [b]
map ([Char], [(Var, ())], Exp3) -> ([Char], [(Var, ())], Exp3)
forall {a} {b}. (a, b, Exp3) -> (a, b, Exp3)
dorhs [([Char], [(Var, ())], Exp3)]
mp
                        dorhs :: (a, b, Exp3) -> (a, b, Exp3)
dorhs (a
c,b
args,Exp3
ae) =
                            (a
c,b
args,Exp3 -> Exp3
go Exp3
ae)
                    in Exp3 -> [([Char], [(Var, ())], Exp3)] -> Exp3
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> [([Char], [(Var, loc)], PreExp ext loc dec)]
-> PreExp ext loc dec
CaseE (Exp3 -> Exp3
go Exp3
e) [([Char], [(Var, ())], Exp3)]
mp'

    (DataConE ()
c [Char]
loc [Exp3]
es) -> () -> [Char] -> [Exp3] -> Exp3
forall (ext :: * -> * -> *) loc dec.
loc -> [Char] -> [PreExp ext loc dec] -> PreExp ext loc dec
DataConE ()
c [Char]
loc ([Exp3] -> Exp3) -> [Exp3] -> Exp3
forall a b. (a -> b) -> a -> b
$ (Exp3 -> Exp3) -> [Exp3] -> [Exp3]
forall a b. (a -> b) -> [a] -> [b]
map (Exp3 -> Exp3
go) [Exp3]
es
    (TimeIt Exp3
e UrTy ()
t Bool
b)      -> Exp3 -> UrTy () -> Bool -> Exp3
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec -> dec -> Bool -> PreExp ext loc dec
TimeIt (Exp3 -> Exp3
go Exp3
e) UrTy ()
t Bool
b
    (MapE (Var
v,UrTy ()
t,Exp3
e') Exp3
e)   -> (Var, UrTy (), Exp3) -> Exp3 -> Exp3
forall (ext :: * -> * -> *) loc dec.
(Var, dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
MapE (Var
v,UrTy ()
t,Exp3 -> Exp3
go Exp3
e') (Exp3 -> Exp3
go Exp3
e)
    (FoldE (Var
v1,UrTy ()
t1,Exp3
e1) (Var
v2,UrTy ()
t2,Exp3
e2) Exp3
e3) ->
         (Var, UrTy (), Exp3) -> (Var, UrTy (), Exp3) -> Exp3 -> Exp3
forall (ext :: * -> * -> *) loc dec.
(Var, dec, PreExp ext loc dec)
-> (Var, dec, PreExp ext loc dec)
-> PreExp ext loc dec
-> PreExp ext loc dec
FoldE (Var
v1,UrTy ()
t1,Exp3 -> Exp3
go Exp3
e1) (Var
v2,UrTy ()
t2,Exp3 -> Exp3
go Exp3
e2)
               (Exp3 -> Exp3
go Exp3
e3)

    (WithArenaE{}) -> [Char] -> Exp3
forall a. HasCallStack => [Char] -> a
error [Char]
"shakeTreExp: WithArenaE not handled."
    (SpawnE{}) -> [Char] -> Exp3
forall a. HasCallStack => [Char] -> a
error [Char]
"shakeTreExp: SpawnE not handled."
    (SyncE{}) -> [Char] -> Exp3
forall a. HasCallStack => [Char] -> a
error [Char]
"shakeTreExp: SyncE not handled."

    -- Assume that these are trivial, and always have effects
    Ext E3Ext () (UrTy ())
_ext -> Exp3
e0




-- | On cursors we have affine types rather than linear.  Thus, unfortunately, we don't
-- have the invariant that the cursors returned by WriteTag are actually USED.  Thus we
-- cannot use dataflow alone to determine what must be kept.
--
-- This contains details that are specific to this pass, which is
-- post-cursorize.  It's not really a good general definition of "hasEffect".
hasEffect :: Exp3 -> Bool
hasEffect :: Exp3 -> Bool
hasEffect Exp3
rhs =
    -- Trivials have been inlined, but we're still flat-ish:
    case Exp3
rhs of
      VarE Var
_ -> Bool
False
      LitE Int
_ -> Bool
False
      CharE Char
_ -> Bool
False
      FloatE{}  -> Bool
False
      LitSymE Var
_ -> Bool
False

      -- These might have effects on output cursors, but the output cursors aren't used
      -- again!  We need to tie the knot in dataflow dependencies, making the start (value)
      -- depend on the end (final cursor).
      AppE Var
_ [()]
_ [Exp3]
_ -> Bool
True  -- For now, don't drop.

      PrimAppE Prim (UrTy ())
_ [Exp3]
_ -> Bool
False -- No prims have effects.

      LetE (Var
_,[()]
_,UrTy ()
_,Exp3
e1) Exp3
e2 -> Exp3 -> Bool
hasEffect Exp3
e1 Bool -> Bool -> Bool
|| Exp3 -> Bool
hasEffect Exp3
e2

      ProjE Int
_ Exp3
e    -> Exp3 -> Bool
hasEffect Exp3
e      -- Flattening should make this equivalent to "False"
      MkProdE [Exp3]
ls   -> (Exp3 -> Bool) -> [Exp3] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Exp3 -> Bool
hasEffect [Exp3]
ls -- Flattening should make this equivalent to "False"

      IfE Exp3
a Exp3
b Exp3
c -> Exp3 -> Bool
hasEffect Exp3
a Bool -> Bool -> Bool
|| Exp3 -> Bool
hasEffect Exp3
b Bool -> Bool -> Bool
|| Exp3 -> Bool
hasEffect Exp3
c

      CaseE Exp3
_ [([Char], [(Var, ())], Exp3)]
_ -> Bool
True -- Umm, just don't drop for now. FIXME/ REVISIT THIS!

      DataConE ()
_ [Char]
_ [Exp3]
_ -> Bool
True

      TimeIt{} -> Bool
True -- Yes, has effect of printing!

      MapE (Var, UrTy (), Exp3)
_ Exp3
_ -> [Char] -> Bool
forall a. HasCallStack => [Char] -> a
error [Char]
"hasEffect: FIXME MapE"
      FoldE (Var, UrTy (), Exp3)
_ (Var, UrTy (), Exp3)
_ Exp3
_ -> [Char] -> Bool
forall a. HasCallStack => [Char] -> a
error [Char]
"hasEffect: FIXME FoldE"

      WithArenaE{} -> [Char] -> Bool
forall a. HasCallStack => [Char] -> a
error [Char]
"hasEffect: WithArenaE not handled."
      SpawnE{} -> [Char] -> Bool
forall a. HasCallStack => [Char] -> a
error [Char]
"hasEffect: SpawnE not handled."
      SyncE{} -> [Char] -> Bool
forall a. HasCallStack => [Char] -> a
error [Char]
"hasEffect: SyncE not handled."

      -- always have effects
      Ext E3Ext () (UrTy ())
_ -> Bool
True

      -- oth -> error $" [shakeTrees] unexpected RHS on Let:\n "++sdoc rhs