{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ConstraintKinds       #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE FlexibleInstances     #-}

-- | Compiler pass to inline trivials.
module Gibbon.Passes.InlineTriv (inlineTriv, inlineTrivExp) where

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

import           Gibbon.Common
import           Gibbon.Language

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

-- | Inline trivial let bindings (binding a var to a var or int), mainly to clean up
--   the output of `flatten`.
inlineTriv :: (HasSimplifiable e l d)
           => Prog (PreExp e l d) -> PassM (Prog (PreExp e l d))
inlineTriv :: forall (e :: * -> * -> *) l d.
HasSimplifiable e l d =>
Prog (PreExp e l d) -> PassM (Prog (PreExp e l d))
inlineTriv (Prog DDefs (TyOf (PreExp e l d))
ddefs FunDefs (PreExp e l d)
funs Maybe (PreExp e l d, TyOf (PreExp e l d))
main) =
    Prog (PreExp e l d) -> PassM (Prog (PreExp e l d))
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (DDefs (TyOf (PreExp e l d))
-> FunDefs (PreExp e l d)
-> Maybe (PreExp e l d, TyOf (PreExp e l d))
-> Prog (PreExp e l d)
forall ex.
DDefs (TyOf ex) -> FunDefs ex -> Maybe (ex, TyOf ex) -> Prog ex
Prog DDefs (TyOf (PreExp e l d))
ddefs ((FunDef (PreExp e l d) -> FunDef (PreExp e l d))
-> FunDefs (PreExp e l d) -> FunDefs (PreExp e l d)
forall a b. (a -> b) -> Map Var a -> Map Var b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (FunDef (PreExp e l d) -> FunDef (PreExp e l d)
forall {e :: * -> * -> *} {l} {d}.
(Expression (e l d), SimplifiableExt (PreExp e l d) (e l d),
 Show l, Show d, Out l, Out d) =>
FunDef (PreExp e l d) -> FunDef (PreExp e l d)
inlineTrivFun (FunDef (PreExp e l d) -> FunDef (PreExp e l d))
-> (FunDef (PreExp e l d) -> FunDef (PreExp e l d))
-> FunDef (PreExp e l d)
-> FunDef (PreExp e l d)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FunDef (PreExp e l d) -> FunDef (PreExp e l d)
forall {e :: * -> * -> *} {l} {d}.
(Expression (e l d), SimplifiableExt (PreExp e l d) (e l d),
 Show l, Show d, Out l, Out d) =>
FunDef (PreExp e l d) -> FunDef (PreExp e l d)
inlineTrivFun) FunDefs (PreExp e l d)
funs) Maybe (PreExp e l d, d)
Maybe (PreExp e l d, TyOf (PreExp e l d))
main')
  where
    inlineTrivFun :: FunDef (PreExp e l d) -> FunDef (PreExp e l d)
inlineTrivFun (FunDef Var
nam [Var]
narg ArrowTy (TyOf (PreExp e l d))
ty PreExp e l d
bod FunMeta
meta) =
      Var
-> [Var]
-> ArrowTy (TyOf (PreExp e l d))
-> PreExp e l d
-> FunMeta
-> FunDef (PreExp e l d)
forall ex.
Var -> [Var] -> ArrowTy (TyOf ex) -> ex -> FunMeta -> FunDef ex
FunDef Var
nam [Var]
narg ArrowTy (TyOf (PreExp e l d))
ty (ExpEnv e l d -> PreExp e l d -> PreExp e l d
forall (e :: * -> * -> *) l d.
HasSimplifiable e l d =>
ExpEnv e l d -> PreExp e l d -> PreExp e l d
inlineTrivExp ExpEnv e l d
forall k a. Map k a
M.empty PreExp e l d
bod) FunMeta
meta

    main' :: Maybe (PreExp e l d, d)
main' = case Maybe (PreExp e l d, TyOf (PreExp e l d))
main of
              Maybe (PreExp e l d, TyOf (PreExp e l d))
Nothing -> Maybe (PreExp e l d, d)
forall a. Maybe a
Nothing
              Just (PreExp e l d
m,TyOf (PreExp e l d)
ty) -> (PreExp e l d, d) -> Maybe (PreExp e l d, d)
forall a. a -> Maybe a
Just (ExpEnv e l d -> PreExp e l d -> PreExp e l d
forall (e :: * -> * -> *) l d.
HasSimplifiable e l d =>
ExpEnv e l d -> PreExp e l d -> PreExp e l d
inlineTrivExp ExpEnv e l d
forall k a. Map k a
M.empty PreExp e l d
m, d
TyOf (PreExp e l d)
ty)

type ExpEnv e l d = M.Map Var (PreExp e l d)

inlineTrivExp :: forall e l d. HasSimplifiable e l d
              => ExpEnv e l d -> (PreExp e l d) -> (PreExp e l d)
inlineTrivExp :: forall (e :: * -> * -> *) l d.
HasSimplifiable e l d =>
ExpEnv e l d -> PreExp e l d -> PreExp e l d
inlineTrivExp = ExpEnv e l d -> PreExp e l d -> PreExp e l d
go
  where
  go :: ExpEnv e l d -> (PreExp e l d) -> (PreExp e l d)
  go :: ExpEnv e l d -> PreExp e l d -> PreExp e l d
go ExpEnv e l d
env PreExp e l d
e0 =
    case PreExp e l d
e0 of
      VarE Var
v    -> case Var -> ExpEnv e l d -> Maybe (PreExp e l d)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Var
v ExpEnv e l d
env of
                     Maybe (PreExp e l d)
Nothing -> Var -> PreExp e l d
forall (ext :: * -> * -> *) loc dec. Var -> PreExp ext loc dec
VarE Var
v
                     Just PreExp e l d
e  -> PreExp e l d
e
      Ext e l d
ext   -> e l d -> PreExp e l d
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (e l d -> PreExp e l d) -> e l d -> PreExp e l d
forall a b. (a -> b) -> a -> b
$ ExpEnv e l d -> e l d -> e l d
forall e ext. SimplifiableExt e ext => Map Var e -> ext -> ext
gInlineTrivExt ExpEnv e l d
env e l d
ext
      LitE{}    -> PreExp e l d
e0
      CharE{}   -> PreExp e l d
e0
      FloatE{}  -> PreExp e l d
e0
      LitSymE{} -> PreExp e l d
e0

      AppE Var
v [l]
lvs [PreExp e l d]
es -> Var -> [l] -> [PreExp e l d] -> PreExp e l d
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE Var
v [l]
lvs ([PreExp e l d] -> PreExp e l d) -> [PreExp e l d] -> PreExp e l d
forall a b. (a -> b) -> a -> b
$ (PreExp e l d -> PreExp e l d) -> [PreExp e l d] -> [PreExp e l d]
forall a b. (a -> b) -> [a] -> [b]
map (ExpEnv e l d -> PreExp e l d -> PreExp e l d
go ExpEnv e l d
env) [PreExp e l d]
es
      PrimAppE Prim d
p [PreExp e l d]
es -> Prim d -> [PreExp e l d] -> PreExp e l d
forall (ext :: * -> * -> *) loc dec.
Prim dec -> [PreExp ext loc dec] -> PreExp ext loc dec
PrimAppE Prim d
p ([PreExp e l d] -> PreExp e l d) -> [PreExp e l d] -> PreExp e l d
forall a b. (a -> b) -> a -> b
$ (PreExp e l d -> PreExp e l d) -> [PreExp e l d] -> [PreExp e l d]
forall a b. (a -> b) -> [a] -> [b]
map (ExpEnv e l d -> PreExp e l d -> PreExp e l d
go ExpEnv e l d
env) [PreExp e l d]
es

      LetE (Var
v,[l]
lvs,d
t,PreExp e l d
e') PreExp e l d
e ->
       case PreExp e l d
e' of
         (VarE Var
v') ->
           case Var -> ExpEnv e l d -> Maybe (PreExp e l d)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Var
v' ExpEnv e l d
env of
             Maybe (PreExp e l d)
Nothing -> ExpEnv e l d -> PreExp e l d -> PreExp e l d
go (Var -> PreExp e l d -> ExpEnv e l d -> ExpEnv e l d
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
v PreExp e l d
e' ExpEnv e l d
env) PreExp e l d
e
             Just PreExp e l d
pr -> ExpEnv e l d -> PreExp e l d -> PreExp e l d
go (Var -> PreExp e l d -> ExpEnv e l d -> ExpEnv e l d
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
v PreExp e l d
pr ExpEnv e l d
env) PreExp e l d
e
         PreExp e l d
et | PreExp e l d -> Bool
forall e. Expression e => e -> Bool
isTrivial PreExp e l d
et ->
                -- Apply existing renames:
                let et' :: PreExp e l d
et' = ExpEnv e l d -> PreExp e l d -> PreExp e l d
go ExpEnv e l d
env PreExp e l d
et in
                ExpEnv e l d -> PreExp e l d -> PreExp e l d
go (Var -> PreExp e l d -> ExpEnv e l d -> ExpEnv e l d
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
v PreExp e l d
et' ExpEnv e l d
env) PreExp e l d
e
         PreExp e l d
_ -> (Var, [l], d, PreExp e l d) -> PreExp e l d -> PreExp e l d
forall (ext :: * -> * -> *) loc dec.
(Var, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (Var
v,[l]
lvs,d
t,ExpEnv e l d -> PreExp e l d -> PreExp e l d
go ExpEnv e l d
env PreExp e l d
e') (ExpEnv e l d -> PreExp e l d -> PreExp e l d
go ExpEnv e l d
env PreExp e l d
e)

      IfE PreExp e l d
e1 PreExp e l d
e2 PreExp e l d
e3 -> PreExp e l d -> PreExp e l d -> PreExp e l d -> PreExp e l d
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> PreExp ext loc dec -> PreExp ext loc dec -> PreExp ext loc dec
IfE (ExpEnv e l d -> PreExp e l d -> PreExp e l d
go ExpEnv e l d
env PreExp e l d
e1) (ExpEnv e l d -> PreExp e l d -> PreExp e l d
go ExpEnv e l d
env PreExp e l d
e2) (ExpEnv e l d -> PreExp e l d -> PreExp e l d
go ExpEnv e l d
env PreExp e l d
e3)

      -- TODO: Type check here:
      ProjE Int
i PreExp e l d
e -> Int -> PreExp e l d -> PreExp e l d
forall (e :: * -> * -> *) l d. Int -> PreExp e l d -> PreExp e l d
mkProj Int
i (PreExp e l d -> PreExp e l d) -> PreExp e l d -> PreExp e l d
forall a b. (a -> b) -> a -> b
$ ExpEnv e l d -> PreExp e l d -> PreExp e l d
go ExpEnv e l d
env PreExp e l d
e

      MkProdE [PreExp e l d]
es -> [PreExp e l d] -> PreExp e l d
forall (ext :: * -> * -> *) loc dec.
[PreExp ext loc dec] -> PreExp ext loc dec
MkProdE ([PreExp e l d] -> PreExp e l d) -> [PreExp e l d] -> PreExp e l d
forall a b. (a -> b) -> a -> b
$ (PreExp e l d -> PreExp e l d) -> [PreExp e l d] -> [PreExp e l d]
forall a b. (a -> b) -> [a] -> [b]
map (ExpEnv e l d -> PreExp e l d -> PreExp e l d
go ExpEnv e l d
env) [PreExp e l d]
es
      CaseE PreExp e l d
e [(DataCon, [(Var, l)], PreExp e l d)]
mp ->
       let e' :: PreExp e l d
e' = ExpEnv e l d -> PreExp e l d -> PreExp e l d
go ExpEnv e l d
env PreExp e l d
e
           mp' :: [(DataCon, [(Var, l)], PreExp e l d)]
mp' = ((DataCon, [(Var, l)], PreExp e l d)
 -> (DataCon, [(Var, l)], PreExp e l d))
-> [(DataCon, [(Var, l)], PreExp e l d)]
-> [(DataCon, [(Var, l)], PreExp e l d)]
forall a b. (a -> b) -> [a] -> [b]
map (\(DataCon
c,[(Var, l)]
args,PreExp e l d
ae) -> (DataCon
c,[(Var, l)]
args,ExpEnv e l d -> PreExp e l d -> PreExp e l d
go ExpEnv e l d
env PreExp e l d
ae)) [(DataCon, [(Var, l)], PreExp e l d)]
mp
       in PreExp e l d
-> [(DataCon, [(Var, l)], PreExp e l d)] -> PreExp e l d
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> [(DataCon, [(Var, loc)], PreExp ext loc dec)]
-> PreExp ext loc dec
CaseE PreExp e l d
e' [(DataCon, [(Var, l)], PreExp e l d)]
mp'

      DataConE l
loc DataCon
c [PreExp e l d]
es -> l -> DataCon -> [PreExp e l d] -> PreExp e l d
forall (ext :: * -> * -> *) loc dec.
loc -> DataCon -> [PreExp ext loc dec] -> PreExp ext loc dec
DataConE l
loc DataCon
c ([PreExp e l d] -> PreExp e l d) -> [PreExp e l d] -> PreExp e l d
forall a b. (a -> b) -> a -> b
$ (PreExp e l d -> PreExp e l d) -> [PreExp e l d] -> [PreExp e l d]
forall a b. (a -> b) -> [a] -> [b]
map (ExpEnv e l d -> PreExp e l d -> PreExp e l d
go ExpEnv e l d
env) [PreExp e l d]
es
      TimeIt PreExp e l d
e d
t Bool
b -> PreExp e l d -> d -> Bool -> PreExp e l d
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec -> dec -> Bool -> PreExp ext loc dec
TimeIt (ExpEnv e l d -> PreExp e l d -> PreExp e l d
go ExpEnv e l d
env PreExp e l d
e) d
t Bool
b
      SpawnE Var
fn [l]
locs [PreExp e l d]
args -> Var -> [l] -> [PreExp e l d] -> PreExp e l d
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
SpawnE Var
fn [l]
locs ([PreExp e l d] -> PreExp e l d) -> [PreExp e l d] -> PreExp e l d
forall a b. (a -> b) -> a -> b
$ (PreExp e l d -> PreExp e l d) -> [PreExp e l d] -> [PreExp e l d]
forall a b. (a -> b) -> [a] -> [b]
map (ExpEnv e l d -> PreExp e l d -> PreExp e l d
go ExpEnv e l d
env) [PreExp e l d]
args
      PreExp e l d
SyncE               -> PreExp e l d
forall (ext :: * -> * -> *) loc dec. PreExp ext loc dec
SyncE
      WithArenaE Var
v PreExp e l d
e -> Var -> PreExp e l d -> PreExp e l d
forall (ext :: * -> * -> *) loc dec.
Var -> PreExp ext loc dec -> PreExp ext loc dec
WithArenaE Var
v (ExpEnv e l d -> PreExp e l d -> PreExp e l d
go ExpEnv e l d
env PreExp e l d
e)
      MapE (Var
v,d
t,PreExp e l d
e') PreExp e l d
e -> (Var, d, PreExp e l d) -> PreExp e l d -> PreExp e l d
forall (ext :: * -> * -> *) loc dec.
(Var, dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
MapE (Var
v,d
t,ExpEnv e l d -> PreExp e l d -> PreExp e l d
go ExpEnv e l d
env PreExp e l d
e') (ExpEnv e l d -> PreExp e l d -> PreExp e l d
go ExpEnv e l d
env PreExp e l d
e)
      FoldE (Var
v1,d
t1,PreExp e l d
e1) (Var
v2,d
t2,PreExp e l d
e2) PreExp e l d
e3 ->
       (Var, d, PreExp e l d)
-> (Var, d, PreExp e l d) -> PreExp e l d -> PreExp e l d
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,d
t1,ExpEnv e l d -> PreExp e l d -> PreExp e l d
go ExpEnv e l d
env PreExp e l d
e1) (Var
v2,d
t2,ExpEnv e l d -> PreExp e l d -> PreExp e l d
go ExpEnv e l d
env PreExp e l d
e2) (ExpEnv e l d -> PreExp e l d -> PreExp e l d
go ExpEnv e l d
env PreExp e l d
e3)

instance HasSimplifiable e l d => Simplifiable (PreExp e l d) where
  gInlineTrivExp :: Map Var (PreExp e l d) -> PreExp e l d -> PreExp e l d
gInlineTrivExp = Map Var (PreExp e l d) -> PreExp e l d -> PreExp e l d
forall (e :: * -> * -> *) l d.
HasSimplifiable e l d =>
ExpEnv e l d -> PreExp e l d -> PreExp e l d
inlineTrivExp