-- | Replace calls to copy functions with tagged indirection nodes
module Gibbon.Passes.RemoveCopies where

import qualified Data.Map as M

import Gibbon.Common
import Gibbon.L2.Syntax

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

-- Maps a location to a region
type LocEnv = M.Map LocVar Var

removeCopies :: Prog2 -> PassM Prog2
removeCopies :: Prog2 -> PassM Prog2
removeCopies Prog{DDefs (TyOf Exp2)
ddefs :: DDefs (TyOf Exp2)
ddefs :: forall ex. Prog ex -> DDefs (TyOf ex)
ddefs,FunDefs Exp2
fundefs :: FunDefs Exp2
fundefs :: forall ex. Prog ex -> FunDefs ex
fundefs,Maybe (Exp2, TyOf Exp2)
mainExp :: Maybe (Exp2, TyOf Exp2)
mainExp :: forall ex. Prog ex -> Maybe (ex, TyOf ex)
mainExp} = do
  Map Var (DDef (UrTy Var))
ddefs' <- (DDef (UrTy Var) -> PassM (DDef (UrTy Var)))
-> Map Var (DDef (UrTy Var)) -> PassM (Map Var (DDef (UrTy Var)))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Map Var a -> m (Map Var b)
mapM (\ddf :: DDef (UrTy Var)
ddf@DDef{[(DataCon, [(IsBoxed, UrTy Var)])]
dataCons :: [(DataCon, [(IsBoxed, UrTy Var)])]
dataCons :: forall a. DDef a -> [(DataCon, [(IsBoxed, a)])]
dataCons} -> do
                    DataCon
dcon <- Var -> DataCon
fromVar (Var -> DataCon) -> PassM Var -> PassM DataCon
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Var -> PassM Var
forall (m :: * -> *). MonadState Int m => Var -> m Var
gensym (DataCon -> Var
toVar DataCon
indirectionTag)
                    -- RemoveCopies might run more than once (b/c repairProgram), so
                    -- we ensure that we add the Indirection constructor only once.
                    let datacons :: [(DataCon, [(IsBoxed, UrTy Var)])]
datacons = ((DataCon, [(IsBoxed, UrTy Var)]) -> IsBoxed)
-> [(DataCon, [(IsBoxed, UrTy Var)])]
-> [(DataCon, [(IsBoxed, UrTy Var)])]
forall a. (a -> IsBoxed) -> [a] -> [a]
filter (IsBoxed -> IsBoxed
not (IsBoxed -> IsBoxed)
-> ((DataCon, [(IsBoxed, UrTy Var)]) -> IsBoxed)
-> (DataCon, [(IsBoxed, UrTy Var)])
-> IsBoxed
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataCon -> IsBoxed
isIndirectionTag (DataCon -> IsBoxed)
-> ((DataCon, [(IsBoxed, UrTy Var)]) -> DataCon)
-> (DataCon, [(IsBoxed, UrTy Var)])
-> IsBoxed
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (DataCon, [(IsBoxed, UrTy Var)]) -> DataCon
forall a b. (a, b) -> a
fst) [(DataCon, [(IsBoxed, UrTy Var)])]
dataCons
                    DDef (UrTy Var) -> PassM (DDef (UrTy Var))
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return DDef (UrTy Var)
ddf {dataCons :: [(DataCon, [(IsBoxed, UrTy Var)])]
dataCons = [(DataCon, [(IsBoxed, UrTy Var)])]
datacons [(DataCon, [(IsBoxed, UrTy Var)])]
-> [(DataCon, [(IsBoxed, UrTy Var)])]
-> [(DataCon, [(IsBoxed, UrTy Var)])]
forall a. [a] -> [a] -> [a]
++ [(DataCon
dcon, [(IsBoxed
False, UrTy Var
forall loc. UrTy loc
CursorTy)])]} )
            DDefs (TyOf Exp2)
Map Var (DDef (UrTy Var))
ddefs
  -- Don't process copy* functions
  [FunDef Exp2]
fds' <- (FunDef Exp2 -> PassM (FunDef Exp2))
-> [FunDef Exp2] -> PassM [FunDef Exp2]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (\FunDef Exp2
fn -> if Var -> IsBoxed
isCopyFunName (FunDef Exp2 -> Var
forall ex. FunDef ex -> Var
funName FunDef Exp2
fn)
                       then FunDef Exp2 -> PassM (FunDef Exp2)
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return FunDef Exp2
fn
                       else Map Var (DDef (UrTy Var))
-> FunDefs Exp2 -> FunDef Exp2 -> PassM (FunDef Exp2)
removeCopiesFn Map Var (DDef (UrTy Var))
ddefs' FunDefs Exp2
fundefs FunDef Exp2
fn)
               (FunDefs Exp2 -> [FunDef Exp2]
forall k a. Map k a -> [a]
M.elems FunDefs Exp2
fundefs)
  let fundefs' :: FunDefs Exp2
fundefs' = [(Var, FunDef Exp2)] -> FunDefs Exp2
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Var, FunDef Exp2)] -> FunDefs Exp2)
-> [(Var, FunDef Exp2)] -> FunDefs Exp2
forall a b. (a -> b) -> a -> b
$ (FunDef Exp2 -> (Var, FunDef Exp2))
-> [FunDef Exp2] -> [(Var, FunDef Exp2)]
forall a b. (a -> b) -> [a] -> [b]
map (\FunDef Exp2
f -> (FunDef Exp2 -> Var
forall ex. FunDef ex -> Var
funName FunDef Exp2
f,FunDef Exp2
f)) [FunDef Exp2]
fds'
      env2 :: Env2 (UrTy Var)
env2 = TyEnv (UrTy Var) -> TyEnv (ArrowTy (UrTy Var)) -> Env2 (UrTy Var)
forall a. TyEnv a -> TyEnv (ArrowTy a) -> Env2 a
Env2 TyEnv (UrTy Var)
forall k a. Map k a
M.empty (FunDefs Exp2 -> TyEnv (ArrowTy (TyOf Exp2))
forall a. FunDefs a -> TyEnv (ArrowTy (TyOf a))
initFunEnv FunDefs Exp2
fundefs)
  Maybe (Exp2, UrTy Var)
mainExp' <- case Maybe (Exp2, TyOf Exp2)
mainExp of
                Maybe (Exp2, TyOf Exp2)
Nothing -> Maybe (Exp2, UrTy Var) -> PassM (Maybe (Exp2, UrTy Var))
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Exp2, UrTy Var)
forall a. Maybe a
Nothing
                Just (Exp2
mn, TyOf Exp2
ty) -> (Exp2, UrTy Var) -> Maybe (Exp2, UrTy Var)
forall a. a -> Maybe a
Just ((Exp2, UrTy Var) -> Maybe (Exp2, UrTy Var))
-> (Exp2 -> (Exp2, UrTy Var)) -> Exp2 -> Maybe (Exp2, UrTy Var)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (,TyOf Exp2
UrTy Var
ty) (Exp2 -> Maybe (Exp2, UrTy Var))
-> PassM Exp2 -> PassM (Maybe (Exp2, UrTy Var))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
                  Map Var (DDef (UrTy Var))
-> FunDefs Exp2 -> LocEnv -> Env2 (UrTy Var) -> Exp2 -> PassM Exp2
removeCopiesExp Map Var (DDef (UrTy Var))
ddefs' FunDefs Exp2
fundefs LocEnv
forall k a. Map k a
M.empty Env2 (UrTy Var)
env2 Exp2
mn
  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
$ DDefs (TyOf Exp2)
-> FunDefs Exp2 -> Maybe (Exp2, TyOf Exp2) -> Prog2
forall ex.
DDefs (TyOf ex) -> FunDefs ex -> Maybe (ex, TyOf ex) -> Prog ex
Prog DDefs (TyOf Exp2)
Map Var (DDef (UrTy Var))
ddefs' FunDefs Exp2
fundefs' Maybe (Exp2, TyOf Exp2)
Maybe (Exp2, UrTy Var)
mainExp'

removeCopiesFn :: DDefs Ty2 -> FunDefs2 -> FunDef2 -> PassM FunDef2
removeCopiesFn :: Map Var (DDef (UrTy Var))
-> FunDefs Exp2 -> FunDef Exp2 -> PassM (FunDef Exp2)
removeCopiesFn Map Var (DDef (UrTy Var))
ddefs FunDefs Exp2
fundefs f :: FunDef Exp2
f@FunDef{[Var]
funArgs :: [Var]
funArgs :: forall ex. FunDef ex -> [Var]
funArgs,ArrowTy (TyOf Exp2)
funTy :: ArrowTy (TyOf Exp2)
funTy :: forall ex. FunDef ex -> ArrowTy (TyOf ex)
funTy,Exp2
funBody :: Exp2
funBody :: forall ex. FunDef ex -> ex
funBody} = do
  let initLocEnv :: LocEnv
initLocEnv = [(Var, Var)] -> LocEnv
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Var, Var)] -> LocEnv) -> [(Var, Var)] -> LocEnv
forall a b. (a -> b) -> a -> b
$ (LRM -> (Var, Var)) -> [LRM] -> [(Var, Var)]
forall a b. (a -> b) -> [a] -> [b]
map (\(LRM Var
lc Region
r Modality
_) -> (Var
lc, Region -> Var
regionToVar Region
r)) (ArrowTy2 (UrTy Var) -> [LRM]
forall ty2. ArrowTy2 ty2 -> [LRM]
locVars ArrowTy (TyOf Exp2)
ArrowTy2 (UrTy Var)
funTy)
      initTyEnv :: TyEnv (UrTy Var)
initTyEnv  = [(Var, UrTy Var)] -> TyEnv (UrTy Var)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Var, UrTy Var)] -> TyEnv (UrTy Var))
-> [(Var, UrTy Var)] -> TyEnv (UrTy Var)
forall a b. (a -> b) -> a -> b
$ [Var] -> [UrTy Var] -> [(Var, UrTy Var)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Var]
funArgs (ArrowTy2 (UrTy Var) -> [UrTy Var]
forall ty2. ArrowTy2 ty2 -> [ty2]
arrIns ArrowTy (TyOf Exp2)
ArrowTy2 (UrTy Var)
funTy)
      env2 :: Env2 (UrTy Var)
env2 = TyEnv (UrTy Var) -> TyEnv (ArrowTy (UrTy Var)) -> Env2 (UrTy Var)
forall a. TyEnv a -> TyEnv (ArrowTy a) -> Env2 a
Env2 TyEnv (UrTy Var)
initTyEnv (FunDefs Exp2 -> TyEnv (ArrowTy (TyOf Exp2))
forall a. FunDefs a -> TyEnv (ArrowTy (TyOf a))
initFunEnv FunDefs Exp2
fundefs)
  Exp2
bod' <- Map Var (DDef (UrTy Var))
-> FunDefs Exp2 -> LocEnv -> Env2 (UrTy Var) -> Exp2 -> PassM Exp2
removeCopiesExp Map Var (DDef (UrTy Var))
ddefs FunDefs Exp2
fundefs LocEnv
initLocEnv Env2 (UrTy Var)
env2 Exp2
funBody
  FunDef Exp2 -> PassM (FunDef Exp2)
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (FunDef Exp2 -> PassM (FunDef Exp2))
-> FunDef Exp2 -> PassM (FunDef Exp2)
forall a b. (a -> b) -> a -> b
$ FunDef Exp2
f {funBody :: Exp2
funBody = Exp2
bod'}

-- ASSUMPTION: copy functions would always be called on a single argument.
removeCopiesExp :: DDefs Ty2 -> FunDefs2 -> LocEnv -> Env2 Ty2 -> Exp2 -> PassM Exp2
removeCopiesExp :: Map Var (DDef (UrTy Var))
-> FunDefs Exp2 -> LocEnv -> Env2 (UrTy Var) -> Exp2 -> PassM Exp2
removeCopiesExp Map Var (DDef (UrTy Var))
ddefs FunDefs Exp2
fundefs LocEnv
lenv Env2 (UrTy Var)
env2 Exp2
ex =
  case Exp2
ex of
    -- This AppE copies data from 'lin' to 'lout'. When this becomes an
    -- indirection node, 'lout' is the _pointer_, and 'lin' the _pointee_.
    AppE Var
f [Var
lin,Var
lout] [Exp2
arg] | Var -> IsBoxed
isCopyFunName Var
f -> do
      Var
indirection <- Var -> PassM Var
forall (m :: * -> *). MonadState Int m => Var -> m Var
gensym Var
"indirection"
      let (PackedTy DataCon
tycon Var
_) = DDefs (TyOf Exp2) -> Env2 (TyOf Exp2) -> Exp2 -> TyOf Exp2
forall e.
Typeable e =>
DDefs (TyOf e) -> Env2 (TyOf e) -> e -> TyOf e
gRecoverType DDefs (TyOf Exp2)
Map Var (DDef (UrTy Var))
ddefs Env2 (TyOf Exp2)
Env2 (UrTy Var)
env2 Exp2
ex
          -- the indirection datacon for this type
          indrDcon :: [DataCon]
indrDcon = (DataCon -> IsBoxed) -> [DataCon] -> [DataCon]
forall a. (a -> IsBoxed) -> [a] -> [a]
filter DataCon -> IsBoxed
isIndirectionTag ([DataCon] -> [DataCon]) -> [DataCon] -> [DataCon]
forall a b. (a -> b) -> a -> b
$ Map Var (DDef (UrTy Var)) -> DataCon -> [DataCon]
forall a. Out a => DDefs a -> DataCon -> [DataCon]
getConOrdering Map Var (DDef (UrTy Var))
ddefs DataCon
tycon
      case [DataCon]
indrDcon of
        [] -> DataCon -> PassM Exp2
forall a. HasCallStack => DataCon -> a
error (DataCon -> PassM Exp2) -> DataCon -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ DataCon
"removeCopies: No indirection constructor found for: " DataCon -> DataCon -> DataCon
forall a. [a] -> [a] -> [a]
++ DataCon -> DataCon
forall a. Out a => a -> DataCon
sdoc DataCon
tycon
        [DataCon
dcon] -> do
          Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$
            [(Var, [Var], UrTy Var, Exp2)] -> Exp2 -> Exp2
forall loc dec (ext :: * -> * -> *).
[(Var, [loc], dec, PreExp ext loc dec)]
-> PreExp ext loc dec -> PreExp ext loc dec
mkLets ([(Var
indirection,[],DataCon -> Var -> UrTy Var
forall loc. DataCon -> loc -> UrTy loc
PackedTy DataCon
tycon Var
lout,
                      E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var (UrTy Var) -> Exp2) -> E2Ext Var (UrTy Var) -> Exp2
forall a b. (a -> b) -> a -> b
$ DataCon
-> DataCon
-> (Var, Var)
-> (Var, Var)
-> Exp2
-> E2Ext Var (UrTy Var)
forall loc dec.
DataCon
-> DataCon
-> (loc, loc)
-> (loc, loc)
-> E2 loc dec
-> E2Ext loc dec
IndirectionE DataCon
tycon DataCon
dcon (Var
lout , LocEnv
lenv LocEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
lout) (Var
lin, LocEnv
lenv LocEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
lin) Exp2
arg)])
            (Var -> Exp2
forall (ext :: * -> * -> *) loc dec. Var -> PreExp ext loc dec
VarE Var
indirection)
        [DataCon]
oth -> DataCon -> PassM Exp2
forall a. HasCallStack => DataCon -> a
error (DataCon -> PassM Exp2) -> DataCon -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ DataCon
"removeCopies: Multiple indirection constructors: " DataCon -> DataCon -> DataCon
forall a. [a] -> [a] -> [a]
++ [DataCon] -> DataCon
forall a. Out a => a -> DataCon
sdoc [DataCon]
oth

    LetE (Var
v,[Var]
locs,ty :: UrTy Var
ty@(PackedTy DataCon
tycon Var
_), (AppE Var
f [Var
lin,Var
lout] [Exp2
arg])) Exp2
bod | Var -> IsBoxed
isCopyFunName Var
f -> do
      -- Get the indirection datacon for this type
      let indrDcon :: [DataCon]
indrDcon = (DataCon -> IsBoxed) -> [DataCon] -> [DataCon]
forall a. (a -> IsBoxed) -> [a] -> [a]
filter DataCon -> IsBoxed
isIndirectionTag ([DataCon] -> [DataCon]) -> [DataCon] -> [DataCon]
forall a b. (a -> b) -> a -> b
$ Map Var (DDef (UrTy Var)) -> DataCon -> [DataCon]
forall a. Out a => DDefs a -> DataCon -> [DataCon]
getConOrdering Map Var (DDef (UrTy Var))
ddefs DataCon
tycon
      case [DataCon]
indrDcon of
        [] -> DataCon -> PassM Exp2
forall a. HasCallStack => DataCon -> a
error (DataCon -> PassM Exp2) -> DataCon -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ DataCon
"removeCopies: No indirection constructor found for: " DataCon -> DataCon -> DataCon
forall a. [a] -> [a] -> [a]
++ DataCon -> DataCon
forall a. Out a => a -> DataCon
sdoc DataCon
tycon
        [DataCon
dcon] -> do
          (Var, [Var], UrTy Var, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(Var, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (Var
v,[Var]
locs,UrTy Var
ty, E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var (UrTy Var) -> Exp2) -> E2Ext Var (UrTy Var) -> Exp2
forall a b. (a -> b) -> a -> b
$ DataCon
-> DataCon
-> (Var, Var)
-> (Var, Var)
-> Exp2
-> E2Ext Var (UrTy Var)
forall loc dec.
DataCon
-> DataCon
-> (loc, loc)
-> (loc, loc)
-> E2 loc dec
-> E2Ext loc dec
IndirectionE DataCon
tycon DataCon
dcon (Var
lout , LocEnv
lenv LocEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
lout) (Var
lin, LocEnv
lenv LocEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
lin) Exp2
arg) (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
            Map Var (DDef (UrTy Var))
-> FunDefs Exp2 -> LocEnv -> Env2 (UrTy Var) -> Exp2 -> PassM Exp2
removeCopiesExp Map Var (DDef (UrTy Var))
ddefs FunDefs Exp2
fundefs LocEnv
lenv (Var -> UrTy Var -> Env2 (UrTy Var) -> Env2 (UrTy Var)
forall a. Var -> a -> Env2 a -> Env2 a
extendVEnv Var
v UrTy Var
ty Env2 (UrTy Var)
env2) Exp2
bod
        [DataCon]
oth -> DataCon -> PassM Exp2
forall a. HasCallStack => DataCon -> a
error (DataCon -> PassM Exp2) -> DataCon -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ DataCon
"removeCopies: Multiple indirection constructors: " DataCon -> DataCon -> DataCon
forall a. [a] -> [a] -> [a]
++ [DataCon] -> DataCon
forall a. Out a => a -> DataCon
sdoc [DataCon]
oth

    Ext E2Ext Var (UrTy Var)
ext ->
      case E2Ext Var (UrTy Var)
ext of
        -- Update lenv with a binding for loc
        LetLocE Var
loc PreLocExp Var
FreeLE Exp2
bod -> do
          E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var (UrTy Var) -> Exp2)
-> (Exp2 -> E2Ext Var (UrTy Var)) -> Exp2 -> Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Var -> PreLocExp Var -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec. Var -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE Var
loc PreLocExp Var
forall loc. PreLocExp loc
FreeLE (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
            Map Var (DDef (UrTy Var))
-> FunDefs Exp2 -> LocEnv -> Env2 (UrTy Var) -> Exp2 -> PassM Exp2
removeCopiesExp Map Var (DDef (UrTy Var))
ddefs FunDefs Exp2
fundefs LocEnv
lenv Env2 (UrTy Var)
env2 Exp2
bod
        StartOfPkdCursor Var
cur -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var (UrTy Var) -> Exp2) -> E2Ext Var (UrTy Var) -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> E2Ext Var (UrTy Var)
forall loc dec. Var -> E2Ext loc dec
StartOfPkdCursor Var
cur
        TagCursor Var
a Var
b -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var (UrTy Var) -> Exp2) -> E2Ext Var (UrTy Var) -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> Var -> E2Ext Var (UrTy Var)
forall loc dec. Var -> Var -> E2Ext loc dec
TagCursor Var
a Var
b
        LetLocE Var
loc PreLocExp Var
rhs Exp2
bod -> do
          let reg :: Var
reg = case PreLocExp Var
rhs of
                      StartOfRegionLE Region
r  -> Region -> Var
regionToVar Region
r
                      InRegionLE Region
r -> Region -> Var
regionToVar Region
r
                      AfterConstantLE Int
_ Var
lc   -> LocEnv
lenv LocEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
lc
                      AfterVariableLE Var
_ Var
lc IsBoxed
_ -> LocEnv
lenv LocEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
lc
                      FromEndLE Var
lc           -> LocEnv
lenv LocEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
lc -- TODO: This needs to be fixed
          E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var (UrTy Var) -> Exp2)
-> (Exp2 -> E2Ext Var (UrTy Var)) -> Exp2 -> Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Var -> PreLocExp Var -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec. Var -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE Var
loc PreLocExp Var
rhs (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
            Map Var (DDef (UrTy Var))
-> FunDefs Exp2 -> LocEnv -> Env2 (UrTy Var) -> Exp2 -> PassM Exp2
removeCopiesExp Map Var (DDef (UrTy Var))
ddefs FunDefs Exp2
fundefs (Var -> Var -> LocEnv -> LocEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
loc Var
reg LocEnv
lenv) Env2 (UrTy Var)
env2 Exp2
bod
       -- Straightforward recursion
        RetE{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
        AddFixed{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
        LetRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
bod -> E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var (UrTy Var) -> Exp2)
-> (Exp2 -> E2Ext Var (UrTy Var)) -> Exp2 -> Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Region
-> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetRegionE Region
r RegionSize
sz Maybe RegionType
ty (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp2 -> PassM Exp2
go Exp2
bod
        LetParRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
bod -> E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var (UrTy Var) -> Exp2)
-> (Exp2 -> E2Ext Var (UrTy Var)) -> Exp2 -> Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Region
-> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetParRegionE Region
r RegionSize
sz Maybe RegionType
ty (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp2 -> PassM Exp2
go Exp2
bod
        FromEndE{}       -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
        BoundsCheck{}    -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
        IndirectionE{}   -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
        E2Ext Var (UrTy Var)
GetCilkWorkerNum -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
        LetAvail [Var]
vs Exp2
bod -> E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var (UrTy Var) -> Exp2)
-> (Exp2 -> E2Ext Var (UrTy Var)) -> Exp2 -> Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Var] -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec. [Var] -> E2 loc dec -> E2Ext loc dec
LetAvail [Var]
vs (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp2 -> PassM Exp2
go Exp2
bod
        AllocateTagHere{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
        AllocateScalarsHere{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
        SSPush{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
        SSPop{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex

    -- Straightforward recursion
    VarE{}     -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
    LitE{}     -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
    CharE{}    -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
    FloatE{}   -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
    LitSymE{}  -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
    AppE{}     -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
    PrimAppE{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
    DataConE{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
    ProjE Int
i Exp2
e  -> Int -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
Int -> PreExp ext loc dec -> PreExp ext loc dec
ProjE Int
i (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp2 -> PassM Exp2
go Exp2
e
    IfE Exp2
a Exp2
b Exp2
c  -> Exp2 -> Exp2 -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> PreExp ext loc dec -> PreExp ext loc dec -> PreExp ext loc dec
IfE (Exp2 -> Exp2 -> Exp2 -> Exp2)
-> PassM Exp2 -> PassM (Exp2 -> Exp2 -> Exp2)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp2 -> PassM Exp2
go Exp2
a PassM (Exp2 -> Exp2 -> Exp2) -> PassM Exp2 -> PassM (Exp2 -> Exp2)
forall a b. PassM (a -> b) -> PassM a -> PassM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp2 -> PassM Exp2
go Exp2
b PassM (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall a b. PassM (a -> b) -> PassM a -> PassM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp2 -> PassM Exp2
go Exp2
c
    MkProdE [Exp2]
ls -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
[PreExp ext loc dec] -> PreExp ext loc dec
MkProdE ([Exp2] -> Exp2) -> PassM [Exp2] -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp2 -> PassM Exp2) -> [Exp2] -> PassM [Exp2]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Exp2 -> PassM Exp2
go [Exp2]
ls
    LetE (Var
v,[Var]
locs,UrTy Var
ty, Exp2
rhs) Exp2
bod ->
      (Var, [Var], UrTy Var, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(Var, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE ((Var, [Var], UrTy Var, Exp2) -> Exp2 -> Exp2)
-> (Exp2 -> (Var, [Var], UrTy Var, Exp2)) -> Exp2 -> Exp2 -> Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Var
v,[Var]
locs,UrTy Var
ty,) (Exp2 -> Exp2 -> Exp2) -> PassM Exp2 -> PassM (Exp2 -> Exp2)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp2 -> PassM Exp2
go Exp2
rhs PassM (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall a b. PassM (a -> b) -> PassM a -> PassM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
        Map Var (DDef (UrTy Var))
-> FunDefs Exp2 -> LocEnv -> Env2 (UrTy Var) -> Exp2 -> PassM Exp2
removeCopiesExp Map Var (DDef (UrTy Var))
ddefs FunDefs Exp2
fundefs LocEnv
lenv (Var -> UrTy Var -> Env2 (UrTy Var) -> Env2 (UrTy Var)
forall a. Var -> a -> Env2 a -> Env2 a
extendVEnv Var
v UrTy Var
ty Env2 (UrTy Var)
env2) Exp2
bod
    CaseE Exp2
scrt [(DataCon, [(Var, Var)], Exp2)]
mp -> do
      let (VarE Var
v) = Exp2
scrt
          PackedTy DataCon
_ Var
tyloc = Var -> Env2 (UrTy Var) -> UrTy Var
forall a. Out a => Var -> Env2 a -> a
lookupVEnv Var
v Env2 (UrTy Var)
env2
          reg :: Var
reg = LocEnv
lenv LocEnv -> Var -> Var
forall k a. Ord k => Map k a -> k -> a
M.! Var
tyloc
      Exp2 -> [(DataCon, [(Var, Var)], Exp2)] -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> [(DataCon, [(Var, loc)], PreExp ext loc dec)]
-> PreExp ext loc dec
CaseE Exp2
scrt ([(DataCon, [(Var, Var)], Exp2)] -> Exp2)
-> PassM [(DataCon, [(Var, Var)], Exp2)] -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((DataCon, [(Var, Var)], Exp2)
 -> PassM (DataCon, [(Var, Var)], Exp2))
-> [(DataCon, [(Var, Var)], Exp2)]
-> PassM [(DataCon, [(Var, Var)], Exp2)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Var
-> LocEnv
-> Env2 (UrTy Var)
-> (DataCon, [(Var, Var)], Exp2)
-> PassM (DataCon, [(Var, Var)], Exp2)
docase Var
reg LocEnv
lenv Env2 (UrTy Var)
env2) [(DataCon, [(Var, Var)], Exp2)]
mp
    TimeIt Exp2
e UrTy Var
ty IsBoxed
b -> do
      Exp2
e' <- Exp2 -> PassM Exp2
go Exp2
e
      Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ Exp2 -> UrTy Var -> IsBoxed -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec -> dec -> IsBoxed -> PreExp ext loc dec
TimeIt Exp2
e' UrTy Var
ty IsBoxed
b
    WithArenaE Var
v Exp2
e -> do
      Exp2
e' <- Exp2 -> PassM Exp2
go Exp2
e
      Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp2 -> PassM Exp2) -> Exp2 -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ Var -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> PreExp ext loc dec -> PreExp ext loc dec
WithArenaE Var
v Exp2
e'
    SpawnE{}-> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
    Exp2
SyncE   -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
    MapE{}  -> DataCon -> PassM Exp2
forall a. HasCallStack => DataCon -> a
error (DataCon -> PassM Exp2) -> DataCon -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ DataCon
"go: TODO MapE"
    FoldE{} -> DataCon -> PassM Exp2
forall a. HasCallStack => DataCon -> a
error (DataCon -> PassM Exp2) -> DataCon -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ DataCon
"go: TODO FoldE"
  where
    go :: Exp2 -> PassM Exp2
go = Map Var (DDef (UrTy Var))
-> FunDefs Exp2 -> LocEnv -> Env2 (UrTy Var) -> Exp2 -> PassM Exp2
removeCopiesExp Map Var (DDef (UrTy Var))
ddefs FunDefs Exp2
fundefs LocEnv
lenv Env2 (UrTy Var)
env2
    docase :: Var
-> LocEnv
-> Env2 (UrTy Var)
-> (DataCon, [(Var, Var)], Exp2)
-> PassM (DataCon, [(Var, Var)], Exp2)
docase Var
reg LocEnv
lenv1 Env2 (UrTy Var)
env21 (DataCon
dcon,[(Var, Var)]
vlocs,Exp2
bod) = do
      -- Update the envs with bindings for pattern matched variables and locations.
      -- The locations point to the same region as the scrutinee.
      let ([Var]
vars,[Var]
locs) = [(Var, Var)] -> ([Var], [Var])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Var, Var)]
vlocs
          lenv1' :: LocEnv
lenv1' = (Var -> LocEnv -> LocEnv) -> LocEnv -> [Var] -> LocEnv
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Var
lc LocEnv
acc -> Var -> Var -> LocEnv -> LocEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
lc Var
reg LocEnv
acc) LocEnv
lenv1 [Var]
locs
          env21' :: Env2 (UrTy Var)
env21' = HasCallStack =>
DataCon
-> Map Var (DDef (UrTy Var))
-> [Var]
-> [Var]
-> Env2 (UrTy Var)
-> Env2 (UrTy Var)
DataCon
-> Map Var (DDef (UrTy Var))
-> [Var]
-> [Var]
-> Env2 (UrTy Var)
-> Env2 (UrTy Var)
extendPatternMatchEnv DataCon
dcon Map Var (DDef (UrTy Var))
ddefs [Var]
vars [Var]
locs Env2 (UrTy Var)
env21
      (DataCon
dcon,[(Var, Var)]
vlocs,) (Exp2 -> (DataCon, [(Var, Var)], Exp2))
-> PassM Exp2 -> PassM (DataCon, [(Var, Var)], Exp2)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Map Var (DDef (UrTy Var))
-> FunDefs Exp2 -> LocEnv -> Env2 (UrTy Var) -> Exp2 -> PassM Exp2
removeCopiesExp Map Var (DDef (UrTy Var))
ddefs FunDefs Exp2
fundefs LocEnv
lenv1' Env2 (UrTy Var)
env21' Exp2
bod)