{-| Do all things necessary to compile parallel allocations to a single region.

In the sequential semantics, (letloc-after x) can only run after x is written to
the store. In the parallel version, we relax this requirement. Every letloc-after
leads to creation of a new region, and we later tie things together with pointers.

    let x     = spawn (foo [l1])
    letloc l2 = after x
    let y     = foo [l2]
    _         = sync

will be transformed to:

    let x     = spawn (foo [l1])
    letregion r3
    letloc l3 = start r3
    let y     = foo [l3]
    _         = sync
    letloc l2 = after x
    tie l2 l3

Need a better name for this pass.

-}

module Gibbon.Passes.ParAlloc (parAlloc) where

import           Control.Monad ( when )
import           Data.Foldable ( foldrM )
import qualified Data.Map as M
import qualified Data.Set as S

import           Gibbon.L2.Syntax
import           Gibbon.Common
import           Gibbon.DynFlags

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

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

-- If there's a (letloc loc2 = after-variable x loc1) binding s.t. x is a
-- spawned variable (LHS of a SpawnE), this binding is swallowed into a
-- PAfter (loc2, (x, loc1)), and emitted back into the ast after a SyncE.
-- PVar is used to accomplish the same thing for let bindings.
data PendingBind = PVar   (Var,[LocVar],Ty2,Exp2)
                 | PAfter (LocVar, (Var, LocVar))
  deriving Int -> PendingBind -> ShowS
[PendingBind] -> ShowS
PendingBind -> String
(Int -> PendingBind -> ShowS)
-> (PendingBind -> String)
-> ([PendingBind] -> ShowS)
-> Show PendingBind
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PendingBind -> ShowS
showsPrec :: Int -> PendingBind -> ShowS
$cshow :: PendingBind -> String
show :: PendingBind -> String
$cshowList :: [PendingBind] -> ShowS
showList :: [PendingBind] -> ShowS
Show

-- The keys in this map are locations bound by a letloc-after, and they
-- map to the fresh locations that point to the start of a new region.
type AfterEnv = M.Map LocVar LocVar

parAlloc :: Prog2 -> PassM Prog2
parAlloc :: Prog2 -> PassM Prog2
parAlloc 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
  Bool
region_on_spawn <- GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_RegionOnSpawn (DynFlags -> Bool) -> PassM DynFlags -> PassM Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PassM DynFlags
forall (m :: * -> *). MonadReader Config m => m DynFlags
getDynFlags
  [FunDef2]
fds' <- (FunDef2 -> PassM FunDef2) -> [FunDef2] -> PassM [FunDef2]
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 FunDef2 -> PassM FunDef2
parAllocFn ([FunDef2] -> PassM [FunDef2]) -> [FunDef2] -> PassM [FunDef2]
forall a b. (a -> b) -> a -> b
$ FunDefs Exp2 -> [FunDef2]
forall k a. Map k a -> [a]
M.elems FunDefs Exp2
fundefs
  let fundefs' :: FunDefs Exp2
fundefs' = [(Var, FunDef2)] -> FunDefs Exp2
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Var, FunDef2)] -> FunDefs Exp2)
-> [(Var, FunDef2)] -> FunDefs Exp2
forall a b. (a -> b) -> a -> b
$ (FunDef2 -> (Var, FunDef2)) -> [FunDef2] -> [(Var, FunDef2)]
forall a b. (a -> b) -> [a] -> [b]
map (\FunDef2
f -> (FunDef2 -> Var
forall ex. FunDef ex -> Var
funName FunDef2
f,FunDef2
f)) [FunDef2]
fds'
      env2 :: Env2 Ty2
env2 = TyEnv Ty2 -> TyEnv (ArrowTy Ty2) -> Env2 Ty2
forall a. TyEnv a -> TyEnv (ArrowTy a) -> Env2 a
Env2 TyEnv Ty2
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, Ty2)
mainExp' <- case Maybe (Exp2, TyOf Exp2)
mainExp of
                Maybe (Exp2, TyOf Exp2)
Nothing -> Maybe (Exp2, Ty2) -> PassM (Maybe (Exp2, Ty2))
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Exp2, Ty2)
forall a. Maybe a
Nothing
                Just (Exp2
mn, TyOf Exp2
ty) -> (Exp2, Ty2) -> Maybe (Exp2, Ty2)
forall a. a -> Maybe a
Just ((Exp2, Ty2) -> Maybe (Exp2, Ty2))
-> (Exp2 -> (Exp2, Ty2)) -> Exp2 -> Maybe (Exp2, Ty2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (,TyOf Exp2
Ty2
ty) (Exp2 -> Maybe (Exp2, Ty2))
-> PassM Exp2 -> PassM (Maybe (Exp2, Ty2))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
                  DDefs2
-> FunDefs Exp2
-> Env2 Ty2
-> RegEnv
-> RegEnv
-> Maybe Var
-> [PendingBind]
-> Set Var
-> Set Var
-> Bool
-> Exp2
-> PassM Exp2
parAllocExp DDefs (TyOf Exp2)
DDefs2
ddefs FunDefs Exp2
fundefs Env2 Ty2
env2 RegEnv
forall k a. Map k a
M.empty RegEnv
forall k a. Map k a
M.empty Maybe Var
forall a. Maybe a
Nothing [] Set Var
forall a. Set a
S.empty Set Var
forall a. Set a
S.empty Bool
region_on_spawn Exp2
mn
  Prog2 -> PassM Prog2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (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)
ddefs FunDefs Exp2
fundefs' Maybe (Exp2, TyOf Exp2)
Maybe (Exp2, Ty2)
mainExp'
  where
    parAllocFn :: FunDef2 -> PassM FunDef2
    parAllocFn :: FunDef2 -> PassM FunDef2
parAllocFn f :: FunDef2
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
      -- if hasParallelism funTy
      -- then do
        Bool
region_on_spawn <- GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_RegionOnSpawn (DynFlags -> Bool) -> PassM DynFlags -> PassM Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PassM DynFlags
forall (m :: * -> *). MonadReader Config m => m DynFlags
getDynFlags
        DynFlags
dflags <- PassM DynFlags
forall (m :: * -> *). MonadReader Config m => m DynFlags
getDynFlags
        let ret_ty :: Ty2
ret_ty = ArrowTy2 Ty2 -> Ty2
forall ty2. ArrowTy2 ty2 -> ty2
arrOut ArrowTy (TyOf Exp2)
ArrowTy2 Ty2
funTy
        Bool -> PassM () -> PassM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (ArrowTy2 Ty2 -> Bool
forall ty2. ArrowTy2 ty2 -> Bool
hasParallelism ArrowTy (TyOf Exp2)
ArrowTy2 Ty2
funTy Bool -> Bool -> Bool
&& Ty2 -> Bool
forall a. Show a => UrTy a -> Bool
hasPacked Ty2
ret_ty Bool -> Bool -> Bool
&& GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_Gibbon1 DynFlags
dflags) (PassM () -> PassM ()) -> PassM () -> PassM ()
forall a b. (a -> b) -> a -> b
$
          String -> PassM ()
forall a. HasCallStack => String -> a
error String
"gibbon: Cannot compile parallel allocations in Gibbon1 mode."

        let initRegEnv :: RegEnv
initRegEnv = [(Var, Var)] -> RegEnv
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Var, Var)] -> RegEnv) -> [(Var, Var)] -> RegEnv
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 Ty2 -> [LRM]
forall ty2. ArrowTy2 ty2 -> [LRM]
locVars ArrowTy (TyOf Exp2)
ArrowTy2 Ty2
funTy)
            initTyEnv :: TyEnv Ty2
initTyEnv  = [(Var, Ty2)] -> TyEnv Ty2
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Var, Ty2)] -> TyEnv Ty2) -> [(Var, Ty2)] -> TyEnv Ty2
forall a b. (a -> b) -> a -> b
$ [Var] -> [Ty2] -> [(Var, Ty2)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Var]
funArgs (ArrowTy2 Ty2 -> [Ty2]
forall ty2. ArrowTy2 ty2 -> [ty2]
arrIns ArrowTy (TyOf Exp2)
ArrowTy2 Ty2
funTy)
            env2 :: Env2 Ty2
env2 = TyEnv Ty2 -> TyEnv (ArrowTy Ty2) -> Env2 Ty2
forall a. TyEnv a -> TyEnv (ArrowTy a) -> Env2 a
Env2 TyEnv Ty2
initTyEnv (FunDefs Exp2 -> TyEnv (ArrowTy (TyOf Exp2))
forall a. FunDefs a -> TyEnv (ArrowTy (TyOf a))
initFunEnv FunDefs Exp2
fundefs)
            boundlocs :: Set Var
boundlocs = [Var] -> Set Var
forall a. Ord a => [a] -> Set a
S.fromList ([Var]
funArgs [Var] -> [Var] -> [Var]
forall a. [a] -> [a] -> [a]
++ ArrowTy2 Ty2 -> [Var]
forall ty2. ArrowTy2 ty2 -> [Var]
allLocVars ArrowTy (TyOf Exp2)
ArrowTy2 Ty2
funTy [Var] -> [Var] -> [Var]
forall a. [a] -> [a] -> [a]
++ ArrowTy2 Ty2 -> [Var]
forall ty2. ArrowTy2 ty2 -> [Var]
allRegVars ArrowTy (TyOf Exp2)
ArrowTy2 Ty2
funTy)
        Exp2
bod' <- DDefs2
-> FunDefs Exp2
-> Env2 Ty2
-> RegEnv
-> RegEnv
-> Maybe Var
-> [PendingBind]
-> Set Var
-> Set Var
-> Bool
-> Exp2
-> PassM Exp2
parAllocExp DDefs (TyOf Exp2)
DDefs2
ddefs FunDefs Exp2
fundefs Env2 Ty2
env2 RegEnv
initRegEnv RegEnv
forall k a. Map k a
M.empty Maybe Var
forall a. Maybe a
Nothing [] Set Var
forall a. Set a
S.empty Set Var
boundlocs Bool
region_on_spawn Exp2
funBody
        FunDef2 -> PassM FunDef2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FunDef2 -> PassM FunDef2) -> FunDef2 -> PassM FunDef2
forall a b. (a -> b) -> a -> b
$ FunDef2
f {funBody :: Exp2
funBody = Exp2
bod'}
      -- else pure f

parAllocExp :: DDefs2 -> FunDefs2 -> Env2 Ty2 -> RegEnv -> AfterEnv -> Maybe Var
            -> [PendingBind] -> S.Set Var -> S.Set LocVar -> Bool -> Exp2
            -> PassM Exp2
parAllocExp :: DDefs2
-> FunDefs Exp2
-> Env2 Ty2
-> RegEnv
-> RegEnv
-> Maybe Var
-> [PendingBind]
-> Set Var
-> Set Var
-> Bool
-> Exp2
-> PassM Exp2
parAllocExp DDefs2
ddefs FunDefs Exp2
fundefs Env2 Ty2
env2 RegEnv
reg_env RegEnv
after_env Maybe Var
mb_parent_id [PendingBind]
pending_binds Set Var
spawned Set Var
boundlocs Bool
region_on_spawn Exp2
ex =
  case Exp2
ex of
    LetE (Var
v, [Var]
endlocs, Ty2
ty, (SpawnE Var
f [Var]
locs [Exp2]
args)) Exp2
bod -> do
      let env2' :: Env2 Ty2
env2' = Var -> Ty2 -> Env2 Ty2 -> Env2 Ty2
forall a. Var -> a -> Env2 a -> Env2 a
extendVEnv Var
v Ty2
ty Env2 Ty2
env2
          spawned' :: Set Var
spawned' = Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.insert Var
v Set Var
spawned
          newlocs :: [Var]
newlocs = (Var -> Var) -> [Var] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map (\Var
loc -> Var -> Var -> RegEnv -> Var
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Var
loc Var
loc RegEnv
after_env) [Var]
locs
          ty' :: Ty2
ty' = RegEnv -> Ty2 -> Ty2
substLoc RegEnv
after_env Ty2
ty
          pending_binds' :: [PendingBind]
pending_binds' = (PendingBind -> PendingBind) -> [PendingBind] -> [PendingBind]
forall a b. (a -> b) -> [a] -> [b]
map
                      (\PendingBind
b -> case PendingBind
b of
                                PVar{} -> PendingBind
b
                                PAfter (Var
loc1, (Var
w, Var
loc2)) -> (Var, (Var, Var)) -> PendingBind
PAfter (Var
loc1, (Var
w, Var -> Var -> RegEnv -> Var
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Var
loc2 Var
loc2 RegEnv
after_env)))
                      [PendingBind]
pending_binds
          reg_env' :: RegEnv
reg_env' = (PendingBind -> RegEnv -> RegEnv)
-> RegEnv -> [PendingBind] -> RegEnv
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\PendingBind
b RegEnv
acc -> case PendingBind
b of
                                    PVar{} -> RegEnv
acc
                                    PAfter (Var
loc1, (Var
_, Var
loc2)) ->
                                      case Var -> RegEnv -> Maybe Var
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Var
loc1 RegEnv
reg_env of
                                        Maybe Var
Nothing -> RegEnv
acc
                                        Just{}  -> Var -> Var -> RegEnv -> RegEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
loc1 (RegEnv
reg_env RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
loc2) RegEnv
acc)
                       RegEnv
reg_env [PendingBind]
pending_binds'
      Var
parent_id <- Var -> PassM Var
forall (m :: * -> *). MonadState Int m => Var -> m Var
gensym Var
"parent_id"
      [Exp2]
args' <- (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]
args
      Exp2
bod'  <- DDefs2
-> FunDefs Exp2
-> Env2 Ty2
-> RegEnv
-> RegEnv
-> Maybe Var
-> [PendingBind]
-> Set Var
-> Set Var
-> Bool
-> Exp2
-> PassM Exp2
parAllocExp DDefs2
ddefs FunDefs Exp2
fundefs Env2 Ty2
env2' RegEnv
reg_env' RegEnv
after_env (Var -> Maybe Var
forall a. a -> Maybe a
Just Var
parent_id) [PendingBind]
pending_binds' Set Var
spawned' Set Var
boundlocs Bool
region_on_spawn Exp2
bod
      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
$ (Var, [Var], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(Var, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (Var
parent_id, [], Ty2
forall loc. UrTy loc
IntTy, E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext E2Ext Var Ty2
forall loc dec. E2Ext loc dec
GetCilkWorkerNum) (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$
             (Var, [Var], Ty2, 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]
endlocs, Ty2
ty', (Var -> [Var] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
SpawnE Var
f [Var]
newlocs [Exp2]
args')) Exp2
bod'

    LetE (Var
v, [Var]
endlocs, Ty2
ty, Exp2
SyncE) Exp2
bod -> do
      let env2' :: Env2 Ty2
env2' = Var -> Ty2 -> Env2 Ty2 -> Env2 Ty2
forall a. Var -> a -> Env2 a -> Env2 a
extendVEnv Var
v Ty2
ty Env2 Ty2
env2
          boundlocs' :: Set Var
boundlocs' = [Set Var] -> Set Var
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions [Set Var
spawned, Set Var
boundlocs,(RegEnv -> Set Var
forall k a. Map k a -> Set k
M.keysSet RegEnv
after_env)] Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.union`
                       (PendingBind -> Set Var -> Set Var)
-> Set Var -> [PendingBind] -> Set Var
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\PendingBind
b Set Var
acc ->
                                case PendingBind
b of
                                  PVar (Var
a,[Var]
_,Ty2
_,Exp2
_) -> Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.insert Var
a Set Var
acc
                                  PAfter (Var
a,(Var, Var)
_)   -> Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.insert Var
a Set Var
acc)
                         Set Var
forall a. Set a
S.empty [PendingBind]
pending_binds
      Exp2
bod1 <- DDefs2
-> FunDefs Exp2
-> Env2 Ty2
-> RegEnv
-> RegEnv
-> Maybe Var
-> [PendingBind]
-> Set Var
-> Set Var
-> Bool
-> Exp2
-> PassM Exp2
parAllocExp DDefs2
ddefs FunDefs Exp2
fundefs Env2 Ty2
env2' RegEnv
reg_env RegEnv
after_env Maybe Var
forall a. Maybe a
Nothing [] Set Var
forall a. Set a
S.empty Set Var
boundlocs' Bool
region_on_spawn Exp2
bod
      Exp2
bod2 <- ((Var, Var) -> Exp2 -> PassM Exp2)
-> Exp2 -> [(Var, Var)] -> PassM Exp2
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> b -> m b) -> b -> t a -> m b
foldrM
                 (\(Var
from, Var
to) Exp2
acc -> do
                    Var
indr <- Var -> PassM Var
forall (m :: * -> *). MonadState Int m => Var -> m Var
gensym Var
"pindr"
                    let Just String
tycon = (Ty2 -> Maybe String -> Maybe String)
-> Maybe String -> [Ty2] -> Maybe String
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Ty2
ty2 Maybe String
acc2 ->
                                                case Ty2
ty2 of
                                                  PackedTy String
tycon2 Var
loc | Var
loc Var -> Var -> Bool
forall a. Eq a => a -> a -> Bool
== Var
from -> String -> Maybe String
forall a. a -> Maybe a
Just String
tycon2
                                                  Ty2
_ -> Maybe String
acc2)
                                       Maybe String
forall a. Maybe a
Nothing (TyEnv Ty2 -> [Ty2]
forall k a. Map k a -> [a]
M.elems (Env2 Ty2 -> TyEnv Ty2
forall a. Env2 a -> TyEnv a
vEnv Env2 Ty2
env2))
                        indr_dcon :: String
indr_dcon = [String] -> String
forall a. HasCallStack => [a] -> a
head ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ (String -> Bool) -> [String] -> [String]
forall a. (a -> Bool) -> [a] -> [a]
filter String -> Bool
isIndirectionTag ([String] -> [String]) -> [String] -> [String]
forall a b. (a -> b) -> a -> b
$ DDefs2 -> String -> [String]
forall a. Out a => DDefs a -> String -> [String]
getConOrdering DDefs2
ddefs String
tycon
                        rhs :: Exp2
rhs = E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ String
-> String -> (Var, Var) -> (Var, Var) -> Exp2 -> E2Ext Var Ty2
forall loc dec.
String
-> String
-> (loc, loc)
-> (loc, loc)
-> E2 loc dec
-> E2Ext loc dec
IndirectionE String
tycon String
indr_dcon (Var
from, RegEnv
reg_env RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
from) (Var
to, RegEnv
reg_env RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
to) (Var -> [Var] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE Var
"nocopy" [] [])
                    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
$ (Var, [Var], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(Var, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (Var
indr, [], String -> Var -> Ty2
forall loc. String -> loc -> UrTy loc
PackedTy String
tycon Var
from, Exp2
rhs) Exp2
acc)
                 Exp2
bod1 (RegEnv -> [(Var, Var)]
forall k a. Map k a -> [(k, a)]
M.toList RegEnv
after_env)
      let bod3 :: Exp2
bod3 = (Exp2 -> PendingBind -> Exp2) -> Exp2 -> [PendingBind] -> Exp2
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
                     (\Exp2
acc PendingBind
b ->
                        case PendingBind
b of
                          PVar (Var, [Var], Ty2, Exp2)
vbnd -> [(Var, [Var], Ty2, Exp2)] -> Exp2 -> Exp2
forall loc dec (ext :: * -> * -> *).
[(Var, [loc], dec, PreExp ext loc dec)]
-> PreExp ext loc dec -> PreExp ext loc dec
mkLets [(Var, [Var], Ty2, Exp2)
vbnd] Exp2
acc
                          PAfter (Var
loc1, (Var
w, Var
loc2)) -> E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> PreLocExp Var -> Exp2 -> E2Ext Var Ty2
forall loc dec. Var -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE Var
loc1 (Var -> Var -> Bool -> PreLocExp Var
forall loc. Var -> loc -> Bool -> PreLocExp loc
AfterVariableLE Var
w Var
loc2 Bool
False) (Exp2 -> E2Ext Var Ty2) -> Exp2 -> E2Ext Var Ty2
forall a b. (a -> b) -> a -> b
$ Exp2
acc)
                     Exp2
bod2 [PendingBind]
pending_binds
      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
$ (Var, [Var], Ty2, 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]
endlocs, Ty2
ty, Exp2
forall (ext :: * -> * -> *) loc dec. PreExp ext loc dec
SyncE) Exp2
bod3

    AppE Var
f [Var]
locs [Exp2]
args -> do
      let newlocs :: [Var]
newlocs = (Var -> Var) -> [Var] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map (\Var
loc -> Var -> Var -> RegEnv -> Var
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Var
loc Var
loc RegEnv
after_env) [Var]
locs
      [Exp2]
args' <- (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]
args
      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
$ Var -> [Var] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE Var
f [Var]
newlocs [Exp2]
args'

    DataConE Var
loc String
dcon [Exp2]
args  -> do
      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
$ Var -> String -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
loc -> String -> [PreExp ext loc dec] -> PreExp ext loc dec
DataConE (Var -> Var -> RegEnv -> Var
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Var
loc Var
loc RegEnv
after_env) String
dcon [Exp2]
args

    LetE (Var
v, [Var]
locs, Ty2
ty, Exp2
rhs) Exp2
bod -> do
      let ty' :: Ty2
ty' = RegEnv -> Ty2 -> Ty2
substLoc RegEnv
after_env Ty2
ty
          pending_binds' :: [PendingBind]
pending_binds' = (PendingBind -> PendingBind) -> [PendingBind] -> [PendingBind]
forall a b. (a -> b) -> [a] -> [b]
map
                      (\PendingBind
b -> case PendingBind
b of
                                PVar{} -> PendingBind
b
                                PAfter (Var
loc1, (Var
w, Var
loc2)) -> (Var, (Var, Var)) -> PendingBind
PAfter (Var
loc1, (Var
w, Var -> Var -> RegEnv -> Var
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Var
loc2 Var
loc2 RegEnv
after_env)))
                      [PendingBind]
pending_binds
          reg_env' :: RegEnv
reg_env' = (PendingBind -> RegEnv -> RegEnv)
-> RegEnv -> [PendingBind] -> RegEnv
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\PendingBind
b RegEnv
acc -> case PendingBind
b of
                                    PVar{} -> RegEnv
acc
                                    PAfter (Var
loc1, (Var
_, Var
loc2)) ->
                                      case Var -> RegEnv -> Maybe Var
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Var
loc1 RegEnv
reg_env of
                                        Maybe Var
Nothing -> RegEnv
acc
                                        Just{}  -> Var -> Var -> RegEnv -> RegEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
loc1 (RegEnv
reg_env RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
loc2) RegEnv
acc)
                       RegEnv
reg_env [PendingBind]
pending_binds'
          env2' :: Env2 Ty2
env2' = Var -> Ty2 -> Env2 Ty2 -> Env2 Ty2
forall a. Var -> a -> Env2 a -> Env2 a
extendVEnv Var
v Ty2
ty Env2 Ty2
env2

          vars :: Set Var
vars = Exp2 -> Set Var
forall a. FreeVars a => a -> Set Var
gFreeVars (RegEnv -> Exp2 -> Exp2
substLocInExp RegEnv
after_env Exp2
rhs) Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` (FunDefs Exp2 -> Set Var
forall k a. Map k a -> Set k
M.keysSet FunDefs Exp2
fundefs)
          used :: Set Var
used = (Exp2 -> Set Var
allFreeVars (RegEnv -> Exp2 -> Exp2
substLocInExp RegEnv
after_env Exp2
rhs)) Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` (FunDefs Exp2 -> Set Var
forall k a. Map k a -> Set k
M.keysSet FunDefs Exp2
fundefs)

      -- Swallow this binding, and add v to 'spawned'
      if Bool -> Bool
not (Set Var -> Set Var -> Bool
forall a. Ord a => Set a -> Set a -> Bool
S.disjoint Set Var
vars Set Var
spawned)
      then do
        Exp2
rhs' <- Exp2 -> PassM Exp2
go Exp2
rhs
        let pending_binds'' :: [PendingBind]
pending_binds'' = (Var, [Var], Ty2, Exp2) -> PendingBind
PVar (Var
v, [Var]
locs, Ty2
ty', Exp2
rhs') PendingBind -> [PendingBind] -> [PendingBind]
forall a. a -> [a] -> [a]
: [PendingBind]
pending_binds'
            spawned' :: Set Var
spawned' = Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.insert Var
v Set Var
spawned
        DDefs2
-> FunDefs Exp2
-> Env2 Ty2
-> RegEnv
-> RegEnv
-> Maybe Var
-> [PendingBind]
-> Set Var
-> Set Var
-> Bool
-> Exp2
-> PassM Exp2
parAllocExp DDefs2
ddefs FunDefs Exp2
fundefs Env2 Ty2
env2' RegEnv
reg_env' RegEnv
after_env Maybe Var
mb_parent_id [PendingBind]
pending_binds'' Set Var
spawned' Set Var
boundlocs Bool
region_on_spawn Exp2
bod
      -- Swallow this binding, and but don't add v to 'spawned'
      else if Bool -> Bool
not (Set Var -> Set Var -> Bool
forall a. Ord a => Set a -> Set a -> Bool
S.isSubsetOf Set Var
used Set Var
boundlocs)
      then do
        Exp2
rhs' <- Exp2 -> PassM Exp2
go Exp2
rhs
        let pending_binds'' :: [PendingBind]
pending_binds'' = (Var, [Var], Ty2, Exp2) -> PendingBind
PVar (Var
v, [Var]
locs, Ty2
ty', Exp2
rhs') PendingBind -> [PendingBind] -> [PendingBind]
forall a. a -> [a] -> [a]
: [PendingBind]
pending_binds'
        DDefs2
-> FunDefs Exp2
-> Env2 Ty2
-> RegEnv
-> RegEnv
-> Maybe Var
-> [PendingBind]
-> Set Var
-> Set Var
-> Bool
-> Exp2
-> PassM Exp2
parAllocExp DDefs2
ddefs FunDefs Exp2
fundefs Env2 Ty2
env2' RegEnv
reg_env' RegEnv
after_env Maybe Var
mb_parent_id [PendingBind]
pending_binds'' Set Var
spawned Set Var
boundlocs Bool
region_on_spawn Exp2
bod
      -- Emit this binding as usual
      else if Set Var -> Set Var -> Bool
forall a. Ord a => Set a -> Set a -> Bool
S.disjoint Set Var
vars Set Var
spawned Bool -> Bool -> Bool
&& Set Var -> Set Var -> Bool
forall a. Ord a => Set a -> Set a -> Bool
S.isSubsetOf Set Var
used Set Var
boundlocs
      then do
        let boundlocs' :: Set Var
boundlocs' = Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.insert Var
v Set Var
boundlocs Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.union` ([Var] -> Set Var
forall a. Ord a => [a] -> Set a
S.fromList [Var]
locs)
        (Var, [Var], Ty2, 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], Ty2, Exp2) -> Exp2 -> Exp2)
-> (Exp2 -> (Var, [Var], Ty2, Exp2)) -> Exp2 -> Exp2 -> Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Var
v,[Var]
locs,Ty2
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
<*> DDefs2
-> FunDefs Exp2
-> Env2 Ty2
-> RegEnv
-> RegEnv
-> Maybe Var
-> [PendingBind]
-> Set Var
-> Set Var
-> Bool
-> Exp2
-> PassM Exp2
parAllocExp DDefs2
ddefs FunDefs Exp2
fundefs Env2 Ty2
env2' RegEnv
reg_env' RegEnv
after_env Maybe Var
mb_parent_id [PendingBind]
pending_binds' Set Var
spawned Set Var
boundlocs' Bool
region_on_spawn Exp2
bod
      else String -> PassM Exp2
forall a. HasCallStack => String -> a
error String
"parAlloc: LetE"

    -- Straightforward recursion
    VarE{}     -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
    LitE{}     -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
    CharE{}    -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
    FloatE{}   -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
    LitSymE{}  -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
    PrimAppE{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure 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
    CaseE Exp2
scrt [(String, [(Var, Var)], Exp2)]
mp -> do
      let (VarE Var
v) = Exp2
scrt
          PackedTy String
_ Var
tyloc = Var -> Env2 Ty2 -> Ty2
forall a. Out a => Var -> Env2 a -> a
lookupVEnv Var
v Env2 Ty2
env2
          reg :: Var
reg = RegEnv
reg_env RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
tyloc
      Exp2 -> [(String, [(Var, Var)], Exp2)] -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> [(String, [(Var, loc)], PreExp ext loc dec)]
-> PreExp ext loc dec
CaseE Exp2
scrt ([(String, [(Var, Var)], Exp2)] -> Exp2)
-> PassM [(String, [(Var, Var)], Exp2)] -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((String, [(Var, Var)], Exp2)
 -> PassM (String, [(Var, Var)], Exp2))
-> [(String, [(Var, Var)], Exp2)]
-> PassM [(String, [(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
-> Env2 Ty2
-> RegEnv
-> RegEnv
-> Maybe Var
-> [PendingBind]
-> Set Var
-> Set Var
-> (String, [(Var, Var)], Exp2)
-> PassM (String, [(Var, Var)], Exp2)
docase Var
reg Env2 Ty2
env2 RegEnv
reg_env RegEnv
after_env Maybe Var
mb_parent_id [PendingBind]
pending_binds Set Var
spawned Set Var
boundlocs) [(String, [(Var, Var)], Exp2)]
mp
    TimeIt Exp2
e Ty2
ty Bool
b -> do
      Exp2
e' <- Exp2 -> PassM Exp2
go Exp2
e
      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
$ Exp2 -> Ty2 -> Bool -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec -> dec -> Bool -> PreExp ext loc dec
TimeIt Exp2
e' Ty2
ty Bool
b
    WithArenaE Var
v Exp2
e -> (Var -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> PreExp ext loc dec -> PreExp ext loc dec
WithArenaE Var
v) (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp2 -> PassM Exp2
go Exp2
e
    SpawnE{} -> String -> PassM Exp2
forall a. HasCallStack => String -> a
error String
"parAllocExp: unbound SpawnE"
    SyncE{}  -> String -> PassM Exp2
forall a. HasCallStack => String -> a
error String
"parAllocExp: unbound SyncE"
    Ext E2Ext Var Ty2
ext  ->
      case E2Ext Var Ty2
ext of
        LetRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
bod       -> E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> (Exp2 -> E2Ext Var Ty2) -> Exp2 -> Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Region -> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext Var Ty2
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
<$>
                                    DDefs2
-> FunDefs Exp2
-> Env2 Ty2
-> RegEnv
-> RegEnv
-> Maybe Var
-> [PendingBind]
-> Set Var
-> Set Var
-> Bool
-> Exp2
-> PassM Exp2
parAllocExp DDefs2
ddefs FunDefs Exp2
fundefs Env2 Ty2
env2 RegEnv
reg_env RegEnv
after_env Maybe Var
mb_parent_id [PendingBind]
pending_binds Set Var
spawned (Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.insert (Region -> Var
regionToVar Region
r) Set Var
boundlocs) Bool
region_on_spawn Exp2
bod
        LetParRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
bod    -> E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> (Exp2 -> E2Ext Var Ty2) -> Exp2 -> Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Region -> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext Var Ty2
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
<$>
                                    DDefs2
-> FunDefs Exp2
-> Env2 Ty2
-> RegEnv
-> RegEnv
-> Maybe Var
-> [PendingBind]
-> Set Var
-> Set Var
-> Bool
-> Exp2
-> PassM Exp2
parAllocExp DDefs2
ddefs FunDefs Exp2
fundefs Env2 Ty2
env2 RegEnv
reg_env RegEnv
after_env Maybe Var
mb_parent_id [PendingBind]
pending_binds Set Var
spawned (Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.insert (Region -> Var
regionToVar Region
r) Set Var
boundlocs) Bool
region_on_spawn 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 Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> E2Ext Var Ty2
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 Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> Var -> E2Ext Var Ty2
forall loc dec. Var -> Var -> E2Ext loc dec
TagCursor Var
a Var
b

        LetLocE Var
loc PreLocExp Var
locexp Exp2
bod -> do
          case PreLocExp Var
locexp of
            -- Binding is swallowed, and it's continuation allocates in a fresh region.
            AfterVariableLE Var
v Var
loc2 Bool
True | Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Var
v Set Var
spawned -> do
              let (Just Var
parent_id) = Maybe Var
mb_parent_id
              Var
cont_id <- Var -> PassM Var
forall (m :: * -> *). MonadState Int m => Var -> m Var
gensym Var
"cont_id"
              Var
r <- Var -> PassM Var
forall (m :: * -> *). MonadState Int m => Var -> m Var
gensym Var
"rafter"
              Var
newloc <- Var -> PassM Var
forall (m :: * -> *). MonadState Int m => Var -> m Var
gensym Var
"loc"
              let newreg :: Region
newreg = Var -> Region
VarR Var
r
                  reg :: Var
reg = RegEnv
reg_env RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
loc2
                  after_env' :: RegEnv
after_env' = Var -> Var -> RegEnv -> RegEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
loc Var
newloc RegEnv
after_env
                  pending_binds' :: [PendingBind]
pending_binds'   = (Var, (Var, Var)) -> PendingBind
PAfter (Var
loc, (Var
v, Var
loc2)) PendingBind -> [PendingBind] -> [PendingBind]
forall a. a -> [a] -> [a]
: [PendingBind]
pending_binds
                  reg_env' :: RegEnv
reg_env'   = Var -> Var -> RegEnv -> RegEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
loc Var
reg (RegEnv -> RegEnv) -> RegEnv -> RegEnv
forall a b. (a -> b) -> a -> b
$ Var -> Var -> RegEnv -> RegEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
newloc Var
r RegEnv
reg_env
                  boundlocs1 :: Set Var
boundlocs1 = Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.insert Var
newloc Set Var
boundlocs
                  boundlocs2 :: Set Var
boundlocs2 = Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.insert Var
loc (Set Var -> Set Var) -> Set Var -> Set Var
forall a b. (a -> b) -> a -> b
$ Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.insert Var
v Set Var
boundlocs
              -- If this continuation is stolen
              Exp2
bod1 <- DDefs2
-> FunDefs Exp2
-> Env2 Ty2
-> RegEnv
-> RegEnv
-> Maybe Var
-> [PendingBind]
-> Set Var
-> Set Var
-> Bool
-> Exp2
-> PassM Exp2
parAllocExp DDefs2
ddefs FunDefs Exp2
fundefs Env2 Ty2
env2 RegEnv
reg_env' RegEnv
after_env' (Var -> Maybe Var
forall a. a -> Maybe a
Just Var
cont_id) [PendingBind]
pending_binds' Set Var
spawned Set Var
boundlocs1 Bool
region_on_spawn (RegEnv -> Exp2 -> Exp2
substLocInExp RegEnv
after_env' Exp2
bod)
              -- If it's not stolen
              Exp2
bod2 <- DDefs2
-> FunDefs Exp2
-> Env2 Ty2
-> RegEnv
-> RegEnv
-> Maybe Var
-> [PendingBind]
-> Set Var
-> Set Var
-> Bool
-> Exp2
-> PassM Exp2
parAllocExp DDefs2
ddefs FunDefs Exp2
fundefs Env2 Ty2
env2 RegEnv
reg_env' RegEnv
after_env (Var -> Maybe Var
forall a. a -> Maybe a
Just Var
cont_id) [PendingBind]
pending_binds (Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.delete Var
v Set Var
spawned) Set Var
boundlocs2 Bool
region_on_spawn Exp2
bod
              Var
not_stolen  <- Var -> PassM Var
forall (m :: * -> *). MonadState Int m => Var -> m Var
gensym Var
"not_stolen"
              -- If we are given the --region_on_spawn flag, we disable the region-on-steal optimization.
              if Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Var
loc2 Set Var
boundlocs Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
region_on_spawn
              then
                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
$ (Var, [Var], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(Var, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (Var
cont_id, [], Ty2
forall loc. UrTy loc
IntTy, E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext E2Ext Var Ty2
forall loc dec. E2Ext loc dec
GetCilkWorkerNum) (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$
                       (Var, [Var], Ty2, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(Var, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE (Var
not_stolen, [], Ty2
forall loc. UrTy loc
BoolTy, Prim Ty2 -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Prim dec -> [PreExp ext loc dec] -> PreExp ext loc dec
PrimAppE Prim Ty2
forall ty. Prim ty
EqIntP [Var -> Exp2
forall (ext :: * -> * -> *) loc dec. Var -> PreExp ext loc dec
VarE Var
cont_id, Var -> Exp2
forall (ext :: * -> * -> *) loc dec. Var -> PreExp ext loc dec
VarE Var
parent_id]) (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$
                       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 (Var -> Exp2
forall (ext :: * -> * -> *) loc dec. Var -> PreExp ext loc dec
VarE Var
not_stolen)
                           (E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ [Var] -> Exp2 -> E2Ext Var Ty2
forall loc dec. [Var] -> E2 loc dec -> E2Ext loc dec
LetAvail [Var
v] (Exp2 -> E2Ext Var Ty2) -> Exp2 -> E2Ext Var Ty2
forall a b. (a -> b) -> a -> b
$
                            E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> PreLocExp Var -> Exp2 -> E2Ext Var Ty2
forall loc dec. Var -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE Var
loc (Var -> Var -> Bool -> PreLocExp Var
forall loc. Var -> loc -> Bool -> PreLocExp loc
AfterVariableLE Var
v Var
loc2 Bool
False) Exp2
bod2) -- don't allocate in a fresh region
                           (E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Region -> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext Var Ty2
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetParRegionE Region
newreg RegionSize
Undefined Maybe RegionType
forall a. Maybe a
Nothing (Exp2 -> E2Ext Var Ty2) -> Exp2 -> E2Ext Var Ty2
forall a b. (a -> b) -> a -> b
$ E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> PreLocExp Var -> Exp2 -> E2Ext Var Ty2
forall loc dec. Var -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE Var
newloc (Region -> PreLocExp Var
forall loc. Region -> PreLocExp loc
StartOfRegionLE Region
newreg) Exp2
bod1)
              else
                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 Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Region -> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext Var Ty2
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetParRegionE Region
newreg RegionSize
Undefined Maybe RegionType
forall a. Maybe a
Nothing (Exp2 -> E2Ext Var Ty2) -> Exp2 -> E2Ext Var Ty2
forall a b. (a -> b) -> a -> b
$ E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> PreLocExp Var -> Exp2 -> E2Ext Var Ty2
forall loc dec. Var -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE Var
newloc (Region -> PreLocExp Var
forall loc. Region -> PreLocExp loc
StartOfRegionLE Region
newreg) Exp2
bod1

            -- Binding is swallowed, but no fresh region is created. This can brought back safely after a sync.
            AfterVariableLE Var
v Var
loc2 Bool
True | Bool -> Bool
not (Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Var
loc2 Set Var
boundlocs) Bool -> Bool -> Bool
|| Bool -> Bool
not (Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Var
v Set Var
boundlocs) -> do
              let pending_binds' :: [PendingBind]
pending_binds'  = (Var, (Var, Var)) -> PendingBind
PAfter (Var
loc, (Var
v, Var
loc2)) PendingBind -> [PendingBind] -> [PendingBind]
forall a. a -> [a] -> [a]
: [PendingBind]
pending_binds
                  reg :: Var
reg      = RegEnv
reg_env RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
loc2
                  reg_env' :: RegEnv
reg_env' = Var -> Var -> RegEnv -> RegEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
loc Var
reg RegEnv
reg_env
              DDefs2
-> FunDefs Exp2
-> Env2 Ty2
-> RegEnv
-> RegEnv
-> Maybe Var
-> [PendingBind]
-> Set Var
-> Set Var
-> Bool
-> Exp2
-> PassM Exp2
parAllocExp DDefs2
ddefs FunDefs Exp2
fundefs Env2 Ty2
env2 RegEnv
reg_env' RegEnv
after_env Maybe Var
mb_parent_id [PendingBind]
pending_binds' Set Var
spawned Set Var
boundlocs Bool
region_on_spawn Exp2
bod

            AfterVariableLE Var
v Var
loc2 Bool
True | Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Var
loc2 Set Var
boundlocs Bool -> Bool -> Bool
&& Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Var
v Set Var
boundlocs -> do
              let reg :: Var
reg = RegEnv
reg_env RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
loc2
                  reg_env' :: RegEnv
reg_env'  = Var -> Var -> RegEnv -> RegEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
loc Var
reg RegEnv
reg_env
                  boundlocs' :: Set Var
boundlocs'= Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.insert Var
loc Set Var
boundlocs
              Exp2
bod' <- DDefs2
-> FunDefs Exp2
-> Env2 Ty2
-> RegEnv
-> RegEnv
-> Maybe Var
-> [PendingBind]
-> Set Var
-> Set Var
-> Bool
-> Exp2
-> PassM Exp2
parAllocExp DDefs2
ddefs FunDefs Exp2
fundefs Env2 Ty2
env2 RegEnv
reg_env' RegEnv
after_env Maybe Var
mb_parent_id [PendingBind]
pending_binds Set Var
spawned Set Var
boundlocs' Bool
region_on_spawn Exp2
bod
              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 Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> PreLocExp Var -> Exp2 -> E2Ext Var Ty2
forall loc dec. Var -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE Var
loc (Var -> Var -> Bool -> PreLocExp Var
forall loc. Var -> loc -> Bool -> PreLocExp loc
AfterVariableLE Var
v Var
loc2 Bool
False) Exp2
bod'

            PreLocExp Var
FreeLE -> do
              let boundlocs' :: Set Var
boundlocs'= Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.insert Var
loc Set Var
boundlocs
              Exp2
bod' <- DDefs2
-> FunDefs Exp2
-> Env2 Ty2
-> RegEnv
-> RegEnv
-> Maybe Var
-> [PendingBind]
-> Set Var
-> Set Var
-> Bool
-> Exp2
-> PassM Exp2
parAllocExp DDefs2
ddefs FunDefs Exp2
fundefs Env2 Ty2
env2 RegEnv
reg_env RegEnv
after_env Maybe Var
mb_parent_id [PendingBind]
pending_binds Set Var
spawned Set Var
boundlocs' Bool
region_on_spawn Exp2
bod
              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 Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> PreLocExp Var -> Exp2 -> E2Ext Var Ty2
forall loc dec. Var -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE Var
loc PreLocExp Var
locexp Exp2
bod'

            PreLocExp Var
_ -> do
              let reg :: Var
reg = case PreLocExp Var
locexp of
                          StartOfRegionLE Region
r  -> Region -> Var
regionToVar Region
r
                          InRegionLE Region
r -> Region -> Var
regionToVar Region
r
                          AfterConstantLE Int
_ Var
lc   -> RegEnv
reg_env RegEnv -> 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 Bool
_ -> RegEnv
reg_env RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
lc
                          FromEndLE Var
lc           -> RegEnv
reg_env RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
lc
                  reg_env' :: RegEnv
reg_env'  = Var -> Var -> RegEnv -> RegEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
loc Var
reg RegEnv
reg_env
                  boundlocs' :: Set Var
boundlocs'= Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.insert Var
loc Set Var
boundlocs
              Exp2
bod' <- DDefs2
-> FunDefs Exp2
-> Env2 Ty2
-> RegEnv
-> RegEnv
-> Maybe Var
-> [PendingBind]
-> Set Var
-> Set Var
-> Bool
-> Exp2
-> PassM Exp2
parAllocExp DDefs2
ddefs FunDefs Exp2
fundefs Env2 Ty2
env2 RegEnv
reg_env' RegEnv
after_env Maybe Var
mb_parent_id [PendingBind]
pending_binds Set Var
spawned Set Var
boundlocs' Bool
region_on_spawn Exp2
bod
              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 Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> PreLocExp Var -> Exp2 -> E2Ext Var Ty2
forall loc dec. Var -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE Var
loc PreLocExp Var
locexp Exp2
bod'
        RetE{}         -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
        FromEndE{}     -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
        BoundsCheck{}  -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
        IndirectionE{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
        AddFixed{}     -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
        E2Ext Var Ty2
GetCilkWorkerNum->Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
        LetAvail [Var]
vs Exp2
bod -> E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> (Exp2 -> E2Ext Var Ty2) -> Exp2 -> Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Var] -> Exp2 -> E2Ext Var Ty2
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 (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
        AllocateScalarsHere{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
        SSPush{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
        SSPop{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
    MapE{}  -> String -> PassM Exp2
forall a. HasCallStack => String -> a
error (String -> PassM Exp2) -> String -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ String
"parAllocExp: TODO MapE"
    FoldE{} -> String -> PassM Exp2
forall a. HasCallStack => String -> a
error (String -> PassM Exp2) -> String -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ String
"parAllocExp: TODO FoldE"
  where
    go :: Exp2 -> PassM Exp2
go = DDefs2
-> FunDefs Exp2
-> Env2 Ty2
-> RegEnv
-> RegEnv
-> Maybe Var
-> [PendingBind]
-> Set Var
-> Set Var
-> Bool
-> Exp2
-> PassM Exp2
parAllocExp DDefs2
ddefs FunDefs Exp2
fundefs Env2 Ty2
env2 RegEnv
reg_env RegEnv
after_env Maybe Var
mb_parent_id [PendingBind]
pending_binds Set Var
spawned Set Var
boundlocs Bool
region_on_spawn

    docase :: Var
-> Env2 Ty2
-> RegEnv
-> RegEnv
-> Maybe Var
-> [PendingBind]
-> Set Var
-> Set Var
-> (String, [(Var, Var)], Exp2)
-> PassM (String, [(Var, Var)], Exp2)
docase Var
reg Env2 Ty2
env21 RegEnv
reg_env2 RegEnv
after_env2 Maybe Var
mb_parent_id2 [PendingBind]
pending_binds2 Set Var
spawned2 Set Var
boundlocs2 (String
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
          reg_env2' :: RegEnv
reg_env2' = (Var -> RegEnv -> RegEnv) -> RegEnv -> [Var] -> RegEnv
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 RegEnv
acc -> Var -> Var -> RegEnv -> RegEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
lc Var
reg RegEnv
acc) RegEnv
reg_env2 [Var]
locs
          env21' :: Env2 Ty2
env21'    = HasCallStack =>
String -> DDefs2 -> [Var] -> [Var] -> Env2 Ty2 -> Env2 Ty2
String -> DDefs2 -> [Var] -> [Var] -> Env2 Ty2 -> Env2 Ty2
extendPatternMatchEnv String
dcon DDefs2
ddefs [Var]
vars [Var]
locs Env2 Ty2
env21
          boundlocs2' :: Set Var
boundlocs2' = Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
S.union ([Var] -> Set Var
forall a. Ord a => [a] -> Set a
S.fromList ([Var]
vars [Var] -> [Var] -> [Var]
forall a. [a] -> [a] -> [a]
++ [Var]
locs)) Set Var
boundlocs2
      (String
dcon,[(Var, Var)]
vlocs,) (Exp2 -> (String, [(Var, Var)], Exp2))
-> PassM Exp2 -> PassM (String, [(Var, Var)], Exp2)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DDefs2
-> FunDefs Exp2
-> Env2 Ty2
-> RegEnv
-> RegEnv
-> Maybe Var
-> [PendingBind]
-> Set Var
-> Set Var
-> Bool
-> Exp2
-> PassM Exp2
parAllocExp DDefs2
ddefs FunDefs Exp2
fundefs Env2 Ty2
env21' RegEnv
reg_env2' RegEnv
after_env2 Maybe Var
mb_parent_id2 [PendingBind]
pending_binds2 Set Var
spawned2 Set Var
boundlocs2' Bool
region_on_spawn Exp2
bod


substLocInExp :: M.Map LocVar LocVar -> Exp2 -> Exp2
substLocInExp :: RegEnv -> Exp2 -> Exp2
substLocInExp RegEnv
mp Exp2
ex1 =
  case Exp2
ex1 of
    VarE{}    -> Exp2
ex1
    LitE{}    -> Exp2
ex1
    CharE{}   -> Exp2
ex1
    FloatE{}  -> Exp2
ex1
    LitSymE{} -> Exp2
ex1
    AppE Var
f [Var]
locs [Exp2]
args -> Var -> [Var] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE Var
f ((Var -> Var) -> [Var] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map (\Var
l -> Var -> Var
sub Var
l) [Var]
locs) ([Exp2] -> Exp2) -> [Exp2] -> Exp2
forall a b. (a -> b) -> a -> b
$ (Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Exp2
go [Exp2]
args
    PrimAppE Prim Ty2
f [Exp2]
args  -> Prim Ty2 -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Prim dec -> [PreExp ext loc dec] -> PreExp ext loc dec
PrimAppE Prim Ty2
f ([Exp2] -> Exp2) -> [Exp2] -> Exp2
forall a b. (a -> b) -> a -> b
$ (Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Exp2
go [Exp2]
args
    LetE (Var
v,[Var]
loc,Ty2
ty,Exp2
rhs) Exp2
bod -> do
      (Var, [Var], Ty2, 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]
loc,Ty2
ty, Exp2 -> Exp2
go Exp2
rhs) (Exp2 -> Exp2
go Exp2
bod)
    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
go Exp2
a) (Exp2 -> Exp2
go Exp2
b) (Exp2 -> Exp2
go Exp2
c)
    MkProdE [Exp2]
xs -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
[PreExp ext loc dec] -> PreExp ext loc dec
MkProdE ([Exp2] -> Exp2) -> [Exp2] -> Exp2
forall a b. (a -> b) -> a -> b
$ (Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Exp2
go [Exp2]
xs
    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) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Exp2 -> Exp2
go Exp2
e
    DataConE Var
loc String
dcon [Exp2]
args -> Var -> String -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
loc -> String -> [PreExp ext loc dec] -> PreExp ext loc dec
DataConE (Var -> Var
sub Var
loc) String
dcon ([Exp2] -> Exp2) -> [Exp2] -> Exp2
forall a b. (a -> b) -> a -> b
$ (Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Exp2
go [Exp2]
args
    CaseE Exp2
scrt [(String, [(Var, Var)], Exp2)]
pats ->
      Exp2 -> [(String, [(Var, Var)], Exp2)] -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> [(String, [(Var, loc)], PreExp ext loc dec)]
-> PreExp ext loc dec
CaseE (Exp2 -> Exp2
go Exp2
scrt) ([(String, [(Var, Var)], Exp2)] -> Exp2)
-> [(String, [(Var, Var)], Exp2)] -> Exp2
forall a b. (a -> b) -> a -> b
$ ((String, [(Var, Var)], Exp2) -> (String, [(Var, Var)], Exp2))
-> [(String, [(Var, Var)], Exp2)] -> [(String, [(Var, Var)], Exp2)]
forall a b. (a -> b) -> [a] -> [b]
map (\(String
a,[(Var, Var)]
b,Exp2
c) -> (String
a,[(Var, Var)]
b, Exp2 -> Exp2
go Exp2
c)) [(String, [(Var, Var)], Exp2)]
pats
    TimeIt Exp2
e Ty2
ty Bool
b  -> Exp2 -> Ty2 -> Bool -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec -> dec -> Bool -> PreExp ext loc dec
TimeIt (Exp2 -> Exp2
go Exp2
e) Ty2
ty Bool
b
    WithArenaE Var
v Exp2
e -> Var -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> PreExp ext loc dec -> PreExp ext loc dec
WithArenaE Var
v (Exp2 -> Exp2
go Exp2
e)
    SpawnE{} -> Exp2
ex1
    SyncE{}  -> Exp2
ex1
    Ext E2Ext Var Ty2
ext ->
      case E2Ext Var Ty2
ext of
        LetRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
rhs  -> E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Region -> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext Var Ty2
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetRegionE Region
r RegionSize
sz Maybe RegionType
ty (Exp2 -> Exp2
go Exp2
rhs)
        LetParRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
rhs -> E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Region -> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext Var Ty2
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetParRegionE Region
r RegionSize
sz Maybe RegionType
ty (Exp2 -> Exp2
go Exp2
rhs)
        LetLocE Var
l PreLocExp Var
lhs Exp2
rhs -> E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> PreLocExp Var -> Exp2 -> E2Ext Var Ty2
forall loc dec. Var -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE Var
l (PreLocExp Var -> PreLocExp Var
go2 PreLocExp Var
lhs) (Exp2 -> Exp2
go Exp2
rhs)
        StartOfPkdCursor Var
v -> E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> E2Ext Var Ty2
forall loc dec. Var -> E2Ext loc dec
StartOfPkdCursor Var
v
        TagCursor Var
a Var
b  -> E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> Var -> E2Ext Var Ty2
forall loc dec. Var -> Var -> E2Ext loc dec
TagCursor Var
a Var
b
        RetE [Var]
locs Var
v       -> E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ [Var] -> Var -> E2Ext Var Ty2
forall loc dec. [loc] -> Var -> E2Ext loc dec
RetE ((Var -> Var) -> [Var] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map (\Var
l -> Var -> Var
sub Var
l) [Var]
locs) Var
v
        FromEndE Var
loc      -> E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> E2Ext Var Ty2
forall loc dec. loc -> E2Ext loc dec
FromEndE (Var -> Var
sub Var
loc)
        BoundsCheck Int
i Var
l1 Var
l2 -> E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Int -> Var -> Var -> E2Ext Var Ty2
forall loc dec. Int -> loc -> loc -> E2Ext loc dec
BoundsCheck Int
i (Var -> Var
sub Var
l1) (Var -> Var
sub Var
l2)
        IndirectionE String
tc String
dc (Var
l1,Var
v1) (Var
l2,Var
v2) Exp2
e -> E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ String
-> String -> (Var, Var) -> (Var, Var) -> Exp2 -> E2Ext Var Ty2
forall loc dec.
String
-> String
-> (loc, loc)
-> (loc, loc)
-> E2 loc dec
-> E2Ext loc dec
IndirectionE String
tc String
dc (Var -> Var
sub Var
l1, Var
v1) (Var -> Var
sub Var
l2, Var
v2) (Exp2 -> Exp2
go Exp2
e)
        AddFixed{}        -> Exp2
ex1
        E2Ext Var Ty2
GetCilkWorkerNum  -> Exp2
ex1
        LetAvail [Var]
vs Exp2
bod   -> E2Ext Var Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var Ty2 -> Exp2) -> E2Ext Var Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ [Var] -> Exp2 -> E2Ext Var Ty2
forall loc dec. [Var] -> E2 loc dec -> E2Ext loc dec
LetAvail [Var]
vs (Exp2 -> Exp2
go Exp2
bod)
        AllocateTagHere{} -> Exp2
ex1
        AllocateScalarsHere{} -> Exp2
ex1
        SSPush{} -> Exp2
ex1
        SSPop{} -> Exp2
ex1
    MapE{}  -> String -> Exp2
forall a. HasCallStack => String -> a
error String
"substLocInExpExp: TODO MapE"
    FoldE{}  -> String -> Exp2
forall a. HasCallStack => String -> a
error String
"substLocInExpExp: TODO FoldE"

  where go :: Exp2 -> Exp2
go = RegEnv -> Exp2 -> Exp2
substLocInExp RegEnv
mp
        sub :: Var -> Var
sub Var
loc = Var -> Var -> RegEnv -> Var
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Var
loc Var
loc RegEnv
mp

        go2 :: PreLocExp Var -> PreLocExp Var
go2 PreLocExp Var
lexp = case PreLocExp Var
lexp of
                     StartOfRegionLE{} -> PreLocExp Var
lexp
                     AfterConstantLE Int
i Var
loc -> Int -> Var -> PreLocExp Var
forall loc. Int -> loc -> PreLocExp loc
AfterConstantLE Int
i (Var -> Var
sub Var
loc)
                     AfterVariableLE Var
i Var
loc Bool
b -> Var -> Var -> Bool -> PreLocExp Var
forall loc. Var -> loc -> Bool -> PreLocExp loc
AfterVariableLE Var
i (Var -> Var
sub Var
loc) Bool
b
                     InRegionLE{} -> PreLocExp Var
lexp
                     PreLocExp Var
FreeLE -> PreLocExp Var
lexp
                     FromEndLE Var
loc -> Var -> PreLocExp Var
forall loc. loc -> PreLocExp loc
FromEndLE (Var -> Var
sub Var
loc)