module Gibbon.Passes.ThreadRegions where

import qualified Data.List as L
import Data.Maybe ( fromJust )
import qualified Data.Map as M
import qualified Data.Set as S
import Data.Foldable ( foldrM )

import Gibbon.Common
import Gibbon.DynFlags
-- import Gibbon.NewL2.Syntax as L2
import Gibbon.NewL2.Syntax as NewL2

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

{-

Threading regions
~~~~~~~~~~~~~~~~~

Functions need end-of-regions cursors for various purposes. The output region
cursors are used for bounds checking (See [Infinite regions] in BoundsCheck).
The input region cursors are useful for garbage collection -- if there's an
indirection from R1 to R2 (input), we need to bump R2's refcount and therefore
need R2's cursor. This pass updates all call-sites to also pass region cursors.
They are prepended to the locations that AppE forms accept.
N.B. for output regions, we actually use end-of-chunk cursors, not
end-of-region cursors.

    AppE add1 [lin, lout] arg

becomes

    AppE add1 [regin, regout, lin, lout] arg


Moreover, functions must also return region cursors, at least for the output
regions. Consider this example:

    ...
    let (x, lout1) = AppE add1 [regin1, regout, lin1, lout] a1 in
    let (y, lout2) = AppE add1 [regin2, regout, lin2, lout1] a2 in
    ...

This is not correct. Because of bounds checking, the first call to add1 might
start using a new output chunk. And we shouldn't use regout in the second call
to add1 -- 'regout' is already full! So we have to thread these output regions,
just like we do the output locations.

    ...
    let (x, lout1, regout1) = AppE add1 [regin1, regout, lin1, lout] a1 in
    let (y, lout2, regout2) = AppE add1 [regin2, regout1, lin2, lout1] a2 in
    ...


-}


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

-- Maps the LHS of a constructor to the region of it's last field. Because of
-- parallelism the last field of constructor may not be in the same region as
-- it's tag. The region of the last field represents the "finished writing
-- output here region", so that's the region that should be threaded.
type RightmostRegEnv = M.Map LocVar RegVar

-- Location arguments
type FnLocArgs = [LREM]

-- Allocation env
type AllocEnv = M.Map Var TyCon

-- Regions of packed values
type PkdEnv = M.Map LocVar RegVar

-- A ordered list of locations within each region.
type OrderedLocsEnv = M.Map RegVar [LocVar]

-- Bound variables that map to their corresponding shortcut pointers
type RanEnv = M.Map Var Var

threadRegions :: NewL2.Prog2 -> PassM NewL2.Prog2
threadRegions :: Prog2 -> PassM Prog2
threadRegions 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
  [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 (DDefs Ty2 -> FunDefs Exp2 -> FunDef2 -> PassM FunDef2
threadRegionsFn DDefs (TyOf Exp2)
DDefs Ty2
ddefs FunDefs Exp2
fundefs) ([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 (m :: * -> *) a. Monad m => a -> m a
return 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
<$>
                  DDefs Ty2
-> FunDefs Exp2
-> [LREM]
-> RegEnv
-> Env2 Ty2
-> RegEnv
-> AllocEnv
-> AllocEnv
-> RegEnv
-> OrderedLocsEnv
-> RegEnv
-> Set Var
-> Set Var
-> Exp2
-> PassM Exp2
threadRegionsExp DDefs (TyOf Exp2)
DDefs Ty2
ddefs FunDefs Exp2
fundefs [] RegEnv
forall k a. Map k a
M.empty Env2 Ty2
env2 RegEnv
forall k a. Map k a
M.empty AllocEnv
forall k a. Map k a
M.empty AllocEnv
forall k a. Map k a
M.empty RegEnv
forall k a. Map k a
M.empty OrderedLocsEnv
forall k a. Map k a
M.empty RegEnv
forall k a. Map k a
M.empty Set Var
forall a. Set a
S.empty Set Var
forall a. Set a
S.empty 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)
ddefs FunDefs Exp2
fundefs' Maybe (Exp2, TyOf Exp2)
Maybe (Exp2, Ty2)
mainExp'

threadRegionsFn :: DDefs Ty2 -> FunDefs2 -> NewL2.FunDef2 -> PassM NewL2.FunDef2
threadRegionsFn :: DDefs Ty2 -> FunDefs Exp2 -> FunDef2 -> PassM FunDef2
threadRegionsFn DDefs Ty2
ddefs FunDefs Exp2
fundefs f :: FunDef2
f@FunDef{Var
funName :: forall ex. FunDef ex -> Var
funName :: Var
funName,[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,FunMeta
funMeta :: FunMeta
funMeta :: forall ex. FunDef ex -> FunMeta
funMeta,Exp2
funBody :: Exp2
funBody :: forall ex. FunDef ex -> ex
funBody} = do
  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)
      fn :: Ty2 -> M.Map Var TyCon -> M.Map Var TyCon
      fn :: Ty2 -> AllocEnv -> AllocEnv
fn = (\Ty2
ty AllocEnv
acc -> case Ty2 -> UrTy Var
unTy2 Ty2
ty of
                         PackedTy TyCon
tycon Var
loc -> Var -> TyCon -> AllocEnv -> AllocEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
loc TyCon
tycon AllocEnv
acc
                         ProdTy [UrTy Var]
tys -> (Ty2 -> AllocEnv -> AllocEnv) -> AllocEnv -> [Ty2] -> AllocEnv
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Ty2 -> AllocEnv -> AllocEnv
fn AllocEnv
acc ((UrTy Var -> Ty2) -> [UrTy Var] -> [Ty2]
forall a b. (a -> b) -> [a] -> [b]
map UrTy Var -> Ty2
MkTy2 [UrTy Var]
tys)
                         UrTy Var
_ -> AllocEnv
acc)
      rlocs_env :: AllocEnv
rlocs_env = (Ty2 -> AllocEnv -> AllocEnv) -> AllocEnv -> [Ty2] -> AllocEnv
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Ty2 -> AllocEnv -> AllocEnv
fn AllocEnv
forall k a. Map k a
M.empty (ArrowTy2 Ty2 -> [Ty2]
forall ty2. ArrowTy2 ty2 -> [ty2]
arrIns ArrowTy (TyOf Exp2)
ArrowTy2 Ty2
funTy)
      wlocs_env :: AllocEnv
wlocs_env = Ty2 -> AllocEnv -> AllocEnv
fn (ArrowTy2 Ty2 -> Ty2
forall ty2. ArrowTy2 ty2 -> ty2
arrOut ArrowTy (TyOf Exp2)
ArrowTy2 Ty2
funTy) AllocEnv
forall k a. Map k a
M.empty
      fnlocargs :: [LREM]
fnlocargs = (LRM -> LREM) -> [LRM] -> [LREM]
forall a b. (a -> b) -> [a] -> [b]
map LRM -> LREM
fromLRM (ArrowTy2 Ty2 -> [LRM]
forall ty2. ArrowTy2 ty2 -> [LRM]
locVars ArrowTy (TyOf Exp2)
ArrowTy2 Ty2
funTy)
      region_locs :: OrderedLocsEnv
region_locs = [(Var, [Var])] -> OrderedLocsEnv
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Var, [Var])] -> OrderedLocsEnv)
-> [(Var, [Var])] -> OrderedLocsEnv
forall a b. (a -> b) -> a -> b
$ (LRM -> (Var, [Var])) -> [LRM] -> [(Var, [Var])]
forall a b. (a -> b) -> [a] -> [b]
map (\(LRM Var
l Region
r Modality
_m) -> (Region -> Var
regionToVar Region
r, [Var
l])) (ArrowTy2 Ty2 -> [LRM]
forall ty2. ArrowTy2 ty2 -> [LRM]
locVars ArrowTy (TyOf Exp2)
ArrowTy2 Ty2
funTy)
  Exp2
bod' <- DDefs Ty2
-> FunDefs Exp2
-> [LREM]
-> RegEnv
-> Env2 Ty2
-> RegEnv
-> AllocEnv
-> AllocEnv
-> RegEnv
-> OrderedLocsEnv
-> RegEnv
-> Set Var
-> Set Var
-> Exp2
-> PassM Exp2
threadRegionsExp DDefs Ty2
ddefs FunDefs Exp2
fundefs [LREM]
fnlocargs RegEnv
initRegEnv Env2 Ty2
env2 RegEnv
forall k a. Map k a
M.empty AllocEnv
rlocs_env AllocEnv
wlocs_env RegEnv
forall k a. Map k a
M.empty OrderedLocsEnv
region_locs RegEnv
forall k a. Map k a
M.empty Set Var
forall a. Set a
S.empty Set Var
forall a. Set a
S.empty Exp2
funBody
  -- Boundschecking
  DynFlags
dflags <- PassM DynFlags
forall (m :: * -> *). MonadReader Config m => m DynFlags
getDynFlags
  let free_wlocs :: Set Var
free_wlocs = [Var] -> Set Var
forall a. Ord a => [a] -> Set a
S.fromList (ArrowTy2 Ty2 -> [Var]
forall ty2. ArrowTy2 ty2 -> [Var]
outLocVars ArrowTy (TyOf Exp2)
ArrowTy2 Ty2
funTy)
  let free_rlocs :: Set Var
free_rlocs = [Var] -> Set Var
forall a. Ord a => [a] -> Set a
S.fromList (ArrowTy2 Ty2 -> [Var]
forall ty2. ArrowTy2 ty2 -> [Var]
inLocVars ArrowTy (TyOf Exp2)
ArrowTy2 Ty2
funTy)
  let free_rlocs' :: Set (Maybe Var, Var)
free_rlocs' = let tmp :: [(Maybe Var, Var)]
tmp = ((Var, Ty2) -> [(Maybe Var, Var)])
-> [(Var, Ty2)] -> [(Maybe Var, Var)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(Var
x,Ty2
ty) -> case Ty2 -> UrTy Var
unTy2 Ty2
ty of
                                            PackedTy TyCon
_ Var
loc -> [(Var -> Maybe Var
forall a. a -> Maybe a
Just Var
x,Var
loc)]
                                            UrTy Var
_ -> []) ([(Var, Ty2)] -> [(Maybe Var, Var)])
-> [(Var, Ty2)] -> [(Maybe Var, Var)]
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)
                        tmp2 :: [(Maybe Var, Var)]
tmp2 = (Var -> (Maybe Var, Var)) -> [Var] -> [(Maybe Var, Var)]
forall a b. (a -> b) -> [a] -> [b]
map (\Var
x -> (Maybe Var
forall a. Maybe a
Nothing, Var
x)) ([Var] -> [(Maybe Var, Var)]) -> [Var] -> [(Maybe Var, Var)]
forall a b. (a -> b) -> a -> b
$ (Set Var -> [Var]
forall a. Set a -> [a]
S.toList Set Var
free_rlocs) [Var] -> [Var] -> [Var]
forall a. Eq a => [a] -> [a] -> [a]
L.\\ (((Maybe Var, Var) -> Var) -> [(Maybe Var, Var)] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe Var, Var) -> Var
forall a b. (a, b) -> b
snd [(Maybe Var, Var)]
tmp)
                    in [(Maybe Var, Var)] -> Set (Maybe Var, Var)
forall a. Ord a => [a] -> Set a
S.fromList ([(Maybe Var, Var)] -> Set (Maybe Var, Var))
-> [(Maybe Var, Var)] -> Set (Maybe Var, Var)
forall a b. (a -> b) -> a -> b
$ [(Maybe Var, Var)]
tmp [(Maybe Var, Var)] -> [(Maybe Var, Var)] -> [(Maybe Var, Var)]
forall a. [a] -> [a] -> [a]
++ [(Maybe Var, Var)]
tmp2
  ([(Var, [LocArg], Ty2, Exp2)]
rpush,[(Var, [LocArg], Ty2, Exp2)]
wpush,[(Var, [LocArg], Ty2, Exp2)]
rpop,[(Var, [LocArg], Ty2, Exp2)]
wpop) <- Set (Maybe Var, Var)
-> Set Var
-> AllocEnv
-> AllocEnv
-> RegEnv
-> PassM
     ([(Var, [LocArg], Ty2, Exp2)], [(Var, [LocArg], Ty2, Exp2)],
      [(Var, [LocArg], Ty2, Exp2)], [(Var, [LocArg], Ty2, Exp2)])
ss_ops Set (Maybe Var, Var)
free_rlocs' Set Var
free_wlocs AllocEnv
rlocs_env AllocEnv
wlocs_env RegEnv
initRegEnv
  let no_eager_promote :: Bool
no_eager_promote = GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_NoEagerPromote DynFlags
dflags
  let bod'' :: Exp2
bod'' = -- This function is always given a BigInfinite region.
              if GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_BigInfiniteRegions DynFlags
dflags Bool -> Bool -> Bool
|| Var -> Bool
isCopySansPtrsFunName Var
funName
              then Exp2
bod'
              else
                let packed_outs :: [UrTy Var]
packed_outs = UrTy Var -> [UrTy Var]
forall a. Show a => UrTy a -> [UrTy a]
getPackedTys (Ty2 -> UrTy Var
unTy2 (ArrowTy2 Ty2 -> Ty2
forall ty2. ArrowTy2 ty2 -> ty2
arrOut ArrowTy (TyOf Exp2)
ArrowTy2 Ty2
funTy))
                    locs_tycons :: AllocEnv
locs_tycons = (UrTy Var -> AllocEnv -> AllocEnv)
-> AllocEnv -> [UrTy Var] -> AllocEnv
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr
                                    (\UrTy Var
ty AllocEnv
acc ->
                                         case UrTy Var
ty of
                                           PackedTy TyCon
t Var
loc ->  Var -> TyCon -> AllocEnv -> AllocEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
loc TyCon
t AllocEnv
acc
                                           UrTy Var
_ -> AllocEnv
acc)
                                    AllocEnv
forall k a. Map k a
M.empty
                                    [UrTy Var]
packed_outs
                    boundschecks :: [(Var, [LocArg], Ty2, Exp2)]
boundschecks = (LRM -> [(Var, [LocArg], Ty2, Exp2)])
-> [LRM] -> [(Var, [LocArg], Ty2, Exp2)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap
                                     (\(LRM Var
loc Region
reg Modality
mode) ->
                                        if Modality
mode Modality -> Modality -> Bool
forall a. Eq a => a -> a -> Bool
== Modality
Output
                                        then let rv :: Var
rv = Region -> Var
regionToVar Region
reg
                                                 end_rv :: Var
end_rv = Var -> Var
toEndV Var
rv
                                                 -- rv = end_reg
                                                 bc :: Int
bc = DDefs Ty2 -> TyCon -> Int
boundsCheck DDefs Ty2
ddefs (AllocEnv
locs_tycons AllocEnv -> Var -> TyCon
forall k a. Ord k => Map k a -> k -> a
M.! Var
loc)
                                                 locarg :: LocArg
locarg = LREM -> LocArg
NewL2.Loc (Var -> Var -> Var -> Modality -> LREM
LREM Var
loc Var
rv Var
end_rv Modality
mode)
                                                 regarg :: LocArg
regarg = Var -> Modality -> Var -> LocArg
NewL2.EndOfReg Var
rv Modality
mode Var
end_rv
                                             in -- dbgTraceIt ("boundscheck" ++ sdoc ((locs_tycons M.! loc), bc)) $
                                                -- maintain shadowstack in no eager promotion mode
                                                [(Var
"_",[],UrTy Var -> Ty2
MkTy2 UrTy Var
forall loc. UrTy loc
IntTy, E2Ext LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext(E2Ext LocArg Ty2 -> Exp2) -> E2Ext LocArg Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Int -> LocArg -> LocArg -> E2Ext LocArg Ty2
forall loc dec. Int -> loc -> loc -> E2Ext loc dec
BoundsCheck Int
bc LocArg
regarg LocArg
locarg)]
                                        else [])
                                     (ArrowTy2 Ty2 -> [LRM]
forall ty2. ArrowTy2 ty2 -> [LRM]
locVars ArrowTy (TyOf Exp2)
ArrowTy2 Ty2
funTy)
                in
                   -- If eager promotion is disabled, growing a region can also trigger a GC.
                   if Bool
no_eager_promote Bool -> Bool -> Bool
&& FunMeta -> Bool
funCanTriggerGC FunMeta
funMeta
                   then [(Var, [LocArg], 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, [LocArg], Ty2, Exp2)]
rpush [(Var, [LocArg], Ty2, Exp2)]
-> [(Var, [LocArg], Ty2, Exp2)] -> [(Var, [LocArg], Ty2, Exp2)]
forall a. [a] -> [a] -> [a]
++ [(Var, [LocArg], Ty2, Exp2)]
wpush [(Var, [LocArg], Ty2, Exp2)]
-> [(Var, [LocArg], Ty2, Exp2)] -> [(Var, [LocArg], Ty2, Exp2)]
forall a. [a] -> [a] -> [a]
++ [(Var, [LocArg], Ty2, Exp2)]
boundschecks [(Var, [LocArg], Ty2, Exp2)]
-> [(Var, [LocArg], Ty2, Exp2)] -> [(Var, [LocArg], Ty2, Exp2)]
forall a. [a] -> [a] -> [a]
++ [(Var, [LocArg], Ty2, Exp2)]
wpop [(Var, [LocArg], Ty2, Exp2)]
-> [(Var, [LocArg], Ty2, Exp2)] -> [(Var, [LocArg], Ty2, Exp2)]
forall a. [a] -> [a] -> [a]
++ [(Var, [LocArg], Ty2, Exp2)]
rpop) Exp2
bod'
                   else [(Var, [LocArg], 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, [LocArg], Ty2, Exp2)]
boundschecks Exp2
bod'

  FunDef2 -> PassM FunDef2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (FunDef2 -> PassM FunDef2) -> FunDef2 -> PassM FunDef2
forall a b. (a -> b) -> a -> b
$ FunDef2
f {funBody :: Exp2
funBody = Exp2
bod''}



threadRegionsExp :: DDefs Ty2 -> FunDefs2 -> [LREM] -> RegEnv -> Env2 Ty2
                 -> RightmostRegEnv -> AllocEnv -> AllocEnv -> PkdEnv
                 -> OrderedLocsEnv -> RanEnv -> S.Set LocVar -> S.Set LocVar
                 -> NewL2.Exp2 -> PassM NewL2.Exp2
threadRegionsExp :: DDefs Ty2
-> FunDefs Exp2
-> [LREM]
-> RegEnv
-> Env2 Ty2
-> RegEnv
-> AllocEnv
-> AllocEnv
-> RegEnv
-> OrderedLocsEnv
-> RegEnv
-> Set Var
-> Set Var
-> Exp2
-> PassM Exp2
threadRegionsExp DDefs Ty2
ddefs FunDefs Exp2
fundefs [LREM]
fnLocArgs RegEnv
renv Env2 Ty2
env2 RegEnv
lfenv AllocEnv
rlocs_env AllocEnv
wlocs_env RegEnv
pkd_env OrderedLocsEnv
region_locs RegEnv
ran_env Set Var
indirs Set Var
redirs Exp2
ex =
  case Exp2
ex of
    AppE Var
f [LocArg]
applocs [Exp2]
args -> do
      let ty :: TyOf Exp2
ty = 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)
DDefs Ty2
ddefs Env2 (TyOf Exp2)
Env2 Ty2
env2 Exp2
ex
          argtys :: [Ty2]
argtys = (Exp2 -> Ty2) -> [Exp2] -> [Ty2]
forall a b. (a -> b) -> [a] -> [b]
map (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)
DDefs Ty2
ddefs Env2 (TyOf Exp2)
Env2 Ty2
env2) [Exp2]
args
          argtylocs :: [Var]
argtylocs = (Ty2 -> [Var]) -> [Ty2] -> [Var]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Ty2 -> [Var]
locsInTy [Ty2]
argtys
          in_regs :: [LocArg]
in_regs = (Var -> [LocArg] -> [LocArg]) -> [LocArg] -> [Var] -> [LocArg]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Var
x [LocArg]
acc -> if Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Var
x Set Var
indirs Bool -> Bool -> Bool
|| Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Var
x Set Var
redirs
                                     then (Var -> LocArg
EndOfReg_Tagged Var
x) LocArg -> [LocArg] -> [LocArg]
forall a. a -> [a] -> [a]
: [LocArg]
acc
                                     else case Var -> RegEnv -> Maybe Var
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Var
x RegEnv
ran_env of
                                            Just Var
ran -> (Var -> LocArg
EndOfReg_Tagged Var
ran) LocArg -> [LocArg] -> [LocArg]
forall a. a -> [a] -> [a]
: [LocArg]
acc
                                            Maybe Var
Nothing -> case Var -> RegEnv -> Maybe Var
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Var
x RegEnv
renv of
                                                         Just Var
r -> (Var -> Modality -> Var -> LocArg
NewL2.EndOfReg Var
r Modality
Input (Var -> Var
toEndV Var
r)) LocArg -> [LocArg] -> [LocArg]
forall a. a -> [a] -> [a]
: [LocArg]
acc
                                                         Maybe Var
Nothing -> [LocArg]
acc)
                    [] [Var]
argtylocs
      let applocs' :: [LocArg]
applocs' = (LocArg -> LocArg) -> [LocArg] -> [LocArg]
forall a b. (a -> b) -> [a] -> [b]
map (\LocArg
loc -> case LocArg
loc of
                                    NewL2.Loc LREM
lrem ->
                                      let x :: Var
x = LREM -> Var
lremLoc LREM
lrem in
                                        if Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Var
x Set Var
indirs Bool -> Bool -> Bool
|| Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Var
x Set Var
redirs
                                        then LREM -> LocArg
NewL2.Loc (LREM
lrem { lremEndReg :: Var
lremEndReg = Var -> Var
toEndFromTaggedV Var
x })
                                        else LocArg
loc
                                    LocArg
_ -> LocArg
loc)
                         [LocArg]
applocs
      -- If this function returns a Packed type, it'll have input and output
      -- locations and therefore, input and output regions.
      if UrTy Var -> Bool
forall a. Show a => UrTy a -> Bool
hasPacked (Ty2 -> UrTy Var
unTy2 TyOf Exp2
Ty2
ty)
      then do
        let out_tylocs :: [Var]
out_tylocs = Ty2 -> [Var]
locsInTy TyOf Exp2
Ty2
ty
        let out_regs :: [LocArg]
out_regs = (Var -> LocArg) -> [Var] -> [LocArg]
forall a b. (a -> b) -> [a] -> [b]
map (\Var
l -> let r :: Var
r = (RegEnv
renv RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
l) in Var -> Modality -> Var -> LocArg
NewL2.EndOfReg Var
r Modality
Output (Var -> Var
toEndV Var
r)) [Var]
out_tylocs
        let newapplocs :: [LocArg]
newapplocs = [LocArg]
in_regs [LocArg] -> [LocArg] -> [LocArg]
forall a. [a] -> [a] -> [a]
++ [LocArg]
out_regs [LocArg] -> [LocArg] -> [LocArg]
forall a. [a] -> [a] -> [a]
++ [LocArg]
applocs'
        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 -> [LocArg] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE Var
f [LocArg]
newapplocs [Exp2]
args
      -- Otherwise, only input regions.
      else do
        let newapplocs :: [LocArg]
newapplocs = [LocArg]
in_regs [LocArg] -> [LocArg] -> [LocArg]
forall a. [a] -> [a] -> [a]
++ [LocArg]
applocs
        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 -> [LocArg] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE Var
f [LocArg]
newapplocs [Exp2]
args

    LetE (Var
v,[LocArg]
locs,Ty2
ty, (AppE Var
f [LocArg]
applocs [Exp2]
args)) Exp2
bod -> do
        let argtylocs :: [Var]
argtylocs = (Exp2 -> [Var]) -> [Exp2] -> [Var]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap
                        (\Exp2
arg ->
                             let argty :: TyOf Exp2
argty = 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)
DDefs Ty2
ddefs Env2 (TyOf Exp2)
Env2 Ty2
env2 Exp2
arg in
                             case Exp2
arg of
                               VarE Var
w ->
                                 case Ty2 -> UrTy Var
unTy2 TyOf Exp2
Ty2
argty of
                                   -- Indirection or redirection cursor.
                                   UrTy Var
CursorTy -> [Var
w]
                                   UrTy Var
_ -> Ty2 -> [Var]
locsInTy TyOf Exp2
Ty2
argty
                               Exp2
_ -> Ty2 -> [Var]
locsInTy TyOf Exp2
Ty2
argty)
                        [Exp2]
args
        let in_regargs :: [LocArg]
in_regargs =
              (Var -> [LocArg] -> [LocArg]) -> [LocArg] -> [Var] -> [LocArg]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Var
x [LocArg]
acc -> if Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Var
x Set Var
indirs Bool -> Bool -> Bool
|| Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Var
x Set Var
redirs
                               then (Var -> LocArg
EndOfReg_Tagged Var
x) LocArg -> [LocArg] -> [LocArg]
forall a. a -> [a] -> [a]
: [LocArg]
acc
                               else case Var -> RegEnv -> Maybe Var
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Var
x RegEnv
ran_env of
                                      Just Var
ran -> (Var -> LocArg
EndOfReg_Tagged Var
ran) LocArg -> [LocArg] -> [LocArg]
forall a. a -> [a] -> [a]
: [LocArg]
acc
                                      Maybe Var
Nothing ->
                                        case Var -> RegEnv -> Maybe Var
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Var
x RegEnv
renv of
                                          Just Var
r -> (Var -> Modality -> Var -> LocArg
NewL2.EndOfReg Var
r Modality
Input (Var -> Var
toEndV Var
r)) LocArg -> [LocArg] -> [LocArg]
forall a. a -> [a] -> [a]
: [LocArg]
acc
                                          Maybe Var
Nothing -> [LocArg]
acc)
              [] [Var]
argtylocs
        --------------------
        let outretlocs :: [Var]
outretlocs = if UrTy Var -> Bool
forall a. Show a => UrTy a -> Bool
hasPacked (Ty2 -> UrTy Var
unTy2 Ty2
ty) then Ty2 -> [Var]
locsInTy Ty2
ty else []
            out_regvars :: [Var]
out_regvars = (Var -> Var) -> [Var] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map (RegEnv
renv RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
#) [Var]
outretlocs
        [Var]
out_regvars' <- (Var -> PassM Var) -> [Var] -> PassM [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) -> [a] -> m [b]
mapM (\Var
r -> Var -> PassM Var
forall (m :: * -> *). MonadState Int m => Var -> m Var
gensym Var
r) [Var]
out_regvars
        let out_regargs :: [LocArg]
out_regargs = (Var -> LocArg) -> [Var] -> [LocArg]
forall a b. (a -> b) -> [a] -> [b]
map (\Var
r -> Var -> Modality -> Var -> LocArg
NewL2.EndOfReg Var
r Modality
Output (Var -> Var
toEndV Var
r)) [Var]
out_regvars
        let out_regargs' :: [LocArg]
out_regargs' = (Var -> LocArg) -> [Var] -> [LocArg]
forall a b. (a -> b) -> [a] -> [b]
map (\Var
r -> Var -> Modality -> Var -> LocArg
NewL2.EndOfReg Var
r Modality
Output (Var -> Var
toEndV Var
r)) [Var]
out_regvars'

        -- Indirections will return end-of-input-region cursor of the region
        -- where they're written, and not of their target.
        let in_regvars :: [Var]
in_regvars = (Var -> Var) -> [Var] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map (RegEnv
renv RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
#) [Var]
argtylocs
        [Var]
in_regvars' <- (Var -> PassM Var) -> [Var] -> PassM [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) -> [a] -> m [b]
mapM (\Var
r -> Var -> PassM Var
forall (m :: * -> *). MonadState Int m => Var -> m Var
gensym Var
r) [Var]
in_regvars
        let in_regargs' :: [LocArg]
in_regargs' = (Var -> LocArg) -> [Var] -> [LocArg]
forall a b. (a -> b) -> [a] -> [b]
map (\Var
r -> Var -> Modality -> Var -> LocArg
NewL2.EndOfReg Var
r Modality
Input (Var -> Var
toEndV Var
r)) [Var]
in_regvars'
        --------------------
        let ran_endofregs :: [(Var, Var)]
ran_endofregs = (Var -> (Var, Var)) -> [Var] -> [(Var, Var)]
forall a b. (a -> b) -> [a] -> [b]
map (\Var
loc -> (Var
loc,RegEnv
renv RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
loc)) ([Var] -> [(Var, Var)]) -> [Var] -> [(Var, Var)]
forall a b. (a -> b) -> a -> b
$
                            (UrTy Var -> Var) -> [UrTy Var] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map (\(PackedTy TyCon
_ Var
loc) -> Var
loc) ([UrTy Var] -> [Var]) -> [UrTy Var] -> [Var]
forall a b. (a -> b) -> a -> b
$
                            UrTy Var -> [UrTy Var]
forall a. Show a => UrTy a -> [UrTy a]
getPackedTys (Ty2 -> UrTy Var
unTy2 Ty2
ty)
        let pkd_env1 :: RegEnv
pkd_env1 = RegEnv
pkd_env RegEnv -> RegEnv -> RegEnv
forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union` ([(Var, Var)] -> RegEnv
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Var, Var)]
ran_endofregs)
        --------------------
        let applocs' :: [LocArg]
applocs' = (LocArg -> LocArg) -> [LocArg] -> [LocArg]
forall a b. (a -> b) -> [a] -> [b]
map (\LocArg
loc -> case LocArg
loc of
                                      NewL2.Loc LREM
lrem ->
                                        let x :: Var
x = LREM -> Var
lremLoc LREM
lrem in
                                          if Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Var
x Set Var
indirs Bool -> Bool -> Bool
|| Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Var
x Set Var
redirs
                                          then LREM -> LocArg
NewL2.Loc (LREM
lrem { lremEndReg :: Var
lremEndReg = Var -> Var
toEndFromTaggedV Var
x })
                                          else LocArg
loc
                                      LocArg
_ -> LocArg
loc)
                           [LocArg]
applocs
        let newapplocs :: [LocArg]
newapplocs = [LocArg]
in_regargs [LocArg] -> [LocArg] -> [LocArg]
forall a. [a] -> [a] -> [a]
++ [LocArg]
out_regargs [LocArg] -> [LocArg] -> [LocArg]
forall a. [a] -> [a] -> [a]
++ [LocArg]
applocs'
        let newretlocs :: [LocArg]
newretlocs = [LocArg]
in_regargs' [LocArg] -> [LocArg] -> [LocArg]
forall a. [a] -> [a] -> [a]
++ [LocArg]
out_regargs' [LocArg] -> [LocArg] -> [LocArg]
forall a. [a] -> [a] -> [a]
++ [LocArg]
locs
        --------------------
        -- 'locs' only has end-witnesses up to this pass. Make their regions
        -- same as regions of the locations that the function traverses.
        let traversed_indices :: [Int]
traversed_indices = let fnty :: ArrowTy (TyOf Exp2)
fnty = FunDef2 -> ArrowTy (TyOf Exp2)
forall ex. FunDef ex -> ArrowTy (TyOf ex)
funTy (FunDefs Exp2
fundefs FunDefs Exp2 -> Var -> FunDef2
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
f) in
                                      ((Int, Var) -> Int) -> [(Int, Var)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Var) -> Int
forall a b. (a, b) -> a
fst ([(Int, Var)] -> [Int]) -> [(Int, Var)] -> [Int]
forall a b. (a -> b) -> a -> b
$
                                      ((Int, Var) -> Bool) -> [(Int, Var)] -> [(Int, Var)]
forall a. (a -> Bool) -> [a] -> [a]
filter (\(Int
_i,Var
loc) -> Var -> Effect
Traverse Var
loc Effect -> Set Effect -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` (ArrowTy2 Ty2 -> Set Effect
forall ty2. ArrowTy2 ty2 -> Set Effect
arrEffs ArrowTy (TyOf Exp2)
ArrowTy2 Ty2
fnty)) ([(Int, Var)] -> [(Int, Var)]) -> [(Int, Var)] -> [(Int, Var)]
forall a b. (a -> b) -> a -> b
$
                                      [Int] -> [Var] -> [(Int, Var)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] (ArrowTy2 Ty2 -> [Var]
forall ty2. ArrowTy2 ty2 -> [Var]
allLocVars ArrowTy (TyOf Exp2)
ArrowTy2 Ty2
fnty)
        let renv1 :: RegEnv
renv1 = [(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
$ [Var] -> [Var] -> [(Var, Var)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((LocArg -> Var) -> [LocArg] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map LocArg -> Var
toLocVar [LocArg]
locs)
                                     ((Int -> Var) -> [Int] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
i -> let zipped :: [(Var, Var)]
zipped = [Var] -> [Var] -> [(Var, Var)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Var]
in_regvars [Var]
in_regvars'
                                                     r :: Var
r = RegEnv
renv RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# ([Var]
argtylocs [Var] -> Int -> Var
forall a. HasCallStack => [a] -> Int -> a
!! Int
i)
                                                     Just Var
r' = Var -> [(Var, Var)] -> Maybe Var
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Var
r [(Var, Var)]
zipped
                                                 in Var
r')
                                          [Int]
traversed_indices)
        let renv2 :: RegEnv
renv2 = RegEnv -> RegEnv -> RegEnv
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union RegEnv
renv1 RegEnv
renv

        -- Update input and returned locations to point to the fresh regions
        --

        let region_locs1 :: OrderedLocsEnv
region_locs1 = ((Var, Var) -> OrderedLocsEnv -> OrderedLocsEnv)
-> OrderedLocsEnv -> [(Var, Var)] -> OrderedLocsEnv
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\(Var
r,Var
r') OrderedLocsEnv
acc -> Var -> [Var] -> OrderedLocsEnv -> OrderedLocsEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
r' (OrderedLocsEnv
acc OrderedLocsEnv -> Var -> [Var]
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
r) OrderedLocsEnv
acc)
                             OrderedLocsEnv
region_locs
                             ([Var] -> [Var] -> [(Var, Var)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Var]
in_regvars [Var]
in_regvars')
        let region_locs2 :: OrderedLocsEnv
region_locs2 = ((Var, Var) -> OrderedLocsEnv -> OrderedLocsEnv)
-> OrderedLocsEnv -> [(Var, Var)] -> OrderedLocsEnv
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\(Var
r,Var
r') OrderedLocsEnv
acc -> Var -> [Var] -> OrderedLocsEnv -> OrderedLocsEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
r' (OrderedLocsEnv
acc OrderedLocsEnv -> Var -> [Var]
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
r) OrderedLocsEnv
acc)
                             OrderedLocsEnv
region_locs1
                             ([Var] -> [Var] -> [(Var, Var)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Var]
out_regvars [Var]
out_regvars')
        -- [2022.10.04]: do outregvars need to updated like inregvars below?
        let (RegEnv
renv3, Exp2
bod1) =
              ((Var, Var, Var) -> (RegEnv, Exp2) -> (RegEnv, Exp2))
-> (RegEnv, Exp2) -> [(Var, Var, Var)] -> (RegEnv, Exp2)
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,Var
r,Var
r') (RegEnv
acc, Exp2
bod_acc) ->
                       ( (Var -> Var -> RegEnv -> RegEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
lc Var
r' (RegEnv -> RegEnv) -> RegEnv -> RegEnv
forall a b. (a -> b) -> a -> b
$ (Var -> Var) -> RegEnv -> RegEnv
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (\Var
w -> if Var
w Var -> Var -> Bool
forall a. Eq a => a -> a -> Bool
== Var
r then Var
r' else Var
w) RegEnv
acc)
                       ,  Either Var Var -> Var -> Exp2 -> Exp2
substEndReg (Var -> Either Var Var
forall a b. b -> Either a b
Right Var
r) (Var -> Var
toEndV Var
r') Exp2
bod_acc))
              (RegEnv
renv2, Exp2
bod)
              ([Var] -> [Var] -> [Var] -> [(Var, Var, Var)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
L.zip3 [Var]
outretlocs [Var]
out_regvars [Var]
out_regvars')
        let (RegEnv
renv4, OrderedLocsEnv
region_locs3, Exp2
bod2) =
              ((Var, Var, Var)
 -> (RegEnv, OrderedLocsEnv, Exp2)
 -> (RegEnv, OrderedLocsEnv, Exp2))
-> (RegEnv, OrderedLocsEnv, Exp2)
-> [(Var, Var, Var)]
-> (RegEnv, OrderedLocsEnv, Exp2)
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,Var
r,Var
r') (RegEnv
acc1,OrderedLocsEnv
acc2,Exp2
bod_acc) ->
                       -- Keep the old mapping for lc in renv because at the end of a
                       -- branch for indirections we want to return the end-of-input-region
                       -- cursor this function was called with, as opposed to what
                       -- the function call using the indirection pointer returns.
                       if Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Var
lc Set Var
indirs then (RegEnv
acc1,OrderedLocsEnv
acc2,Exp2
bod_acc) else
                         let locs_in_r :: [Var]
locs_in_r = (OrderedLocsEnv
region_locs2 OrderedLocsEnv -> Var -> [Var]
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
r) in
                           case Var -> [Var] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
L.elemIndex Var
lc [Var]
locs_in_r of
                             Just Int
idx ->
                               if Int
idx Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== ([Var] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Var]
locs_in_r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
                               then let fake_last_loc :: Var
fake_last_loc = TyCon -> Var
toVar TyCon
"fake_" Var -> Var -> Var
`varAppend` Var
lc
                                        acc1' :: RegEnv
acc1' = Var -> Var -> RegEnv -> RegEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
fake_last_loc Var
r' RegEnv
acc1
                                        acc2' :: OrderedLocsEnv
acc2' = ([Var] -> [Var]) -> Var -> OrderedLocsEnv -> OrderedLocsEnv
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (\[Var]
ls -> [Var]
ls [Var] -> [Var] -> [Var]
forall a. [a] -> [a] -> [a]
++ [Var
fake_last_loc]) Var
r OrderedLocsEnv
acc2
                                        acc2'' :: OrderedLocsEnv
acc2'' = Var -> [Var] -> OrderedLocsEnv -> OrderedLocsEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
r' (OrderedLocsEnv
acc2' OrderedLocsEnv -> Var -> [Var]
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
r) OrderedLocsEnv
acc2'
                                        acc2''' :: OrderedLocsEnv
acc2''' = (Var -> OrderedLocsEnv -> OrderedLocsEnv)
-> OrderedLocsEnv -> [Var] -> OrderedLocsEnv
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Var
lc2 OrderedLocsEnv
acc3 -> ([Var] -> [Var]) -> Var -> OrderedLocsEnv -> OrderedLocsEnv
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (\[Var]
ls -> [Var]
ls [Var] -> [Var] -> [Var]
forall a. [a] -> [a] -> [a]
++ [Var
fake_last_loc]) (RegEnv
acc1' RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
lc2) OrderedLocsEnv
acc3) OrderedLocsEnv
acc2'' [Var]
locs_in_r
                                    in (RegEnv
acc1', OrderedLocsEnv
acc2''', Exp2
bod_acc)
                               else let ([Var]
_, [Var]
to_update) = Int -> [Var] -> ([Var], [Var])
forall a. Int -> [a] -> ([a], [a])
splitAt (Int
idxInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [Var]
locs_in_r
                                        updated :: RegEnv
updated = (Var -> Var -> Var) -> RegEnv -> RegEnv
forall k a b. (k -> a -> b) -> Map k a -> Map k b
M.mapWithKey (\Var
key Var
val -> if Var
key Var -> [Var] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Var]
to_update then Var
r' else Var
val) RegEnv
acc1
                                        bod_acc' :: Exp2
bod_acc' = (Var -> Exp2 -> Exp2) -> Exp2 -> [Var] -> Exp2
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr
                                                     (\Var
l Exp2
b -> Either Var Var -> Var -> Exp2 -> Exp2
substEndReg (Var -> Either Var Var
forall a b. a -> Either a b
Left Var
l) (Var -> Var
toEndV Var
r') Exp2
b)
                                                     Exp2
bod_acc
                                                     (Set Var -> [Var]
forall a. Set a -> [a]
S.toList (RegEnv -> Set Var
forall k a. Map k a -> Set k
M.keysSet RegEnv
acc1 Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.intersection` ([Var] -> Set Var
forall a. Ord a => [a] -> Set a
S.fromList [Var]
to_update)))
                                    in (RegEnv
updated, OrderedLocsEnv
acc2, Exp2
bod_acc')
                             Maybe Int
Nothing -> TyCon -> (RegEnv, OrderedLocsEnv, Exp2)
forall a. HasCallStack => TyCon -> a
error (TyCon -> (RegEnv, OrderedLocsEnv, Exp2))
-> TyCon -> (RegEnv, OrderedLocsEnv, Exp2)
forall a b. (a -> b) -> a -> b
$ TyCon
"threadRegionsExp: unbound loc " TyCon -> TyCon -> TyCon
forall a. [a] -> [a] -> [a]
++ (Var, [Var], Set Var, OrderedLocsEnv, Var) -> TyCon
forall a. Out a => a -> TyCon
sdoc (Var
lc,[Var]
locs_in_r,Set Var
indirs,OrderedLocsEnv
region_locs1,Var
r))
              (RegEnv
renv3, OrderedLocsEnv
region_locs2, Exp2
bod1)
              ([Var] -> [Var] -> [Var] -> [(Var, Var, Var)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
L.zip3 [Var]
argtylocs [Var]
in_regvars [Var]
in_regvars')
        -- TODO: only keep the rightmost end-of-input-region cursor in renv.
        --------------------
        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
            rlocs_env' :: AllocEnv
rlocs_env' = UrTy Var -> AllocEnv -> AllocEnv
forall {k}. Ord k => UrTy k -> Map k TyCon -> Map k TyCon
updRLocsEnv (Ty2 -> UrTy Var
unTy2 Ty2
ty) AllocEnv
rlocs_env
            wlocs_env' :: AllocEnv
wlocs_env' = (Var -> AllocEnv -> AllocEnv) -> AllocEnv -> [Var] -> AllocEnv
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Var
loc AllocEnv
acc -> Var -> AllocEnv -> AllocEnv
forall k a. Ord k => k -> Map k a -> Map k a
M.delete Var
loc AllocEnv
acc) AllocEnv
wlocs_env (Ty2 -> [Var]
locsInTy Ty2
ty)
        Exp2
bod3 <- DDefs Ty2
-> FunDefs Exp2
-> [LREM]
-> RegEnv
-> Env2 Ty2
-> RegEnv
-> AllocEnv
-> AllocEnv
-> RegEnv
-> OrderedLocsEnv
-> RegEnv
-> Set Var
-> Set Var
-> Exp2
-> PassM Exp2
threadRegionsExp DDefs Ty2
ddefs FunDefs Exp2
fundefs [LREM]
fnLocArgs RegEnv
renv4 Env2 Ty2
env2' RegEnv
lfenv AllocEnv
rlocs_env' AllocEnv
wlocs_env' RegEnv
pkd_env1 OrderedLocsEnv
region_locs3 RegEnv
ran_env Set Var
indirs Set Var
redirs Exp2
bod2

        -- shadowstack  ops
        --------------------
        let -- free = S.fromList $ freeLocVars bod
            free :: Set Var
free = Set Var -> Env2 Ty2 -> Exp2 -> Set Var
ss_free_locs ([Var] -> Set Var
forall a. Ord a => [a] -> Set a
S.fromList (Var
v Var -> [Var] -> [Var]
forall a. a -> [a] -> [a]
: Ty2 -> [Var]
locsInTy Ty2
ty [Var] -> [Var] -> [Var]
forall a. [a] -> [a] -> [a]
++ ((LocArg -> Var) -> [LocArg] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map LocArg -> Var
toLocVar [LocArg]
locs))) Env2 Ty2
env2' Exp2
bod
            free_wlocs :: Set Var
free_wlocs = Set Var
free Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.intersection` (AllocEnv -> Set Var
forall k a. Map k a -> Set k
M.keysSet AllocEnv
wlocs_env')
            free_rlocs :: Set Var
free_rlocs = Set Var
free Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.intersection` (AllocEnv -> Set Var
forall k a. Map k a -> Set k
M.keysSet AllocEnv
rlocs_env')
            free_rlocs' :: Set (Maybe Var, Var)
free_rlocs' = let tmp :: [(Maybe Var, Var)]
tmp = ((Var, Ty2) -> (Maybe Var, Var))
-> [(Var, Ty2)] -> [(Maybe Var, Var)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Var
x,(MkTy2 (PackedTy TyCon
_ Var
loc))) -> (Var -> Maybe Var
forall a. a -> Maybe a
Just Var
x,Var
loc)) ([(Var, Ty2)] -> [(Maybe Var, Var)])
-> [(Var, Ty2)] -> [(Maybe Var, Var)]
forall a b. (a -> b) -> a -> b
$
                                    ((Var, Ty2) -> Bool) -> [(Var, Ty2)] -> [(Var, Ty2)]
forall a. (a -> Bool) -> [a] -> [a]
filter (\(Var
_x,_y :: Ty2
_y@(MkTy2 (PackedTy TyCon
tycon Var
loc))) -> Var
loc Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set Var
free_rlocs Bool -> Bool -> Bool
&& TyCon
tycon TyCon -> TyCon -> Bool
forall a. Eq a => a -> a -> Bool
/= TyCon
hole_tycon)
                                           (TyEnv Ty2 -> [(Var, Ty2)]
forall k a. Map k a -> [(k, a)]
M.toList (TyEnv Ty2 -> [(Var, Ty2)]) -> TyEnv Ty2 -> [(Var, Ty2)]
forall a b. (a -> b) -> a -> b
$ (Ty2 -> Bool) -> TyEnv Ty2 -> TyEnv Ty2
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (UrTy Var -> Bool
forall a. UrTy a -> Bool
isPackedTy (UrTy Var -> Bool) -> (Ty2 -> UrTy Var) -> Ty2 -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ty2 -> UrTy Var
unTy2) (Env2 Ty2 -> TyEnv Ty2
forall a. Env2 a -> TyEnv a
vEnv Env2 Ty2
env2))
                              tmp2 :: [(Maybe Var, Var)]
tmp2 = (Var -> (Maybe Var, Var)) -> [Var] -> [(Maybe Var, Var)]
forall a b. (a -> b) -> [a] -> [b]
map (\Var
x -> (Maybe Var
forall a. Maybe a
Nothing, Var
x)) ([Var] -> [(Maybe Var, Var)]) -> [Var] -> [(Maybe Var, Var)]
forall a b. (a -> b) -> a -> b
$ (Set Var -> [Var]
forall a. Set a -> [a]
S.toList Set Var
free_rlocs) [Var] -> [Var] -> [Var]
forall a. Eq a => [a] -> [a] -> [a]
L.\\ (((Maybe Var, Var) -> Var) -> [(Maybe Var, Var)] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe Var, Var) -> Var
forall a b. (a, b) -> b
snd [(Maybe Var, Var)]
tmp)
                          in [(Maybe Var, Var)] -> Set (Maybe Var, Var)
forall a. Ord a => [a] -> Set a
S.fromList ([(Maybe Var, Var)] -> Set (Maybe Var, Var))
-> [(Maybe Var, Var)] -> Set (Maybe Var, Var)
forall a b. (a -> b) -> a -> b
$ [(Maybe Var, Var)]
tmp [(Maybe Var, Var)] -> [(Maybe Var, Var)] -> [(Maybe Var, Var)]
forall a. [a] -> [a] -> [a]
++ [(Maybe Var, Var)]
tmp2
        ([(Var, [LocArg], Ty2, Exp2)]
rpush,[(Var, [LocArg], Ty2, Exp2)]
wpush,[(Var, [LocArg], Ty2, Exp2)]
rpop,[(Var, [LocArg], Ty2, Exp2)]
wpop) <- Set (Maybe Var, Var)
-> Set Var
-> AllocEnv
-> AllocEnv
-> RegEnv
-> PassM
     ([(Var, [LocArg], Ty2, Exp2)], [(Var, [LocArg], Ty2, Exp2)],
      [(Var, [LocArg], Ty2, Exp2)], [(Var, [LocArg], Ty2, Exp2)])
ss_ops Set (Maybe Var, Var)
free_rlocs' Set Var
free_wlocs AllocEnv
rlocs_env AllocEnv
wlocs_env RegEnv
renv
        Bool
emit_ss <- PassM Bool
emit_ss_instrs
        if Bool
emit_ss Bool -> Bool -> Bool
&& FunMeta -> Bool
funCanTriggerGC (FunDef2 -> FunMeta
forall ex. FunDef ex -> FunMeta
funMeta (FunDefs Exp2
fundefs FunDefs Exp2 -> Var -> FunDef2
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
f))
          then do let binds :: [(Var, [LocArg], Ty2, Exp2)]
binds = [(Var, [LocArg], Ty2, Exp2)]
rpush [(Var, [LocArg], Ty2, Exp2)]
-> [(Var, [LocArg], Ty2, Exp2)] -> [(Var, [LocArg], Ty2, Exp2)]
forall a. [a] -> [a] -> [a]
++ [(Var, [LocArg], Ty2, Exp2)]
wpush [(Var, [LocArg], Ty2, Exp2)]
-> [(Var, [LocArg], Ty2, Exp2)] -> [(Var, [LocArg], Ty2, Exp2)]
forall a. [a] -> [a] -> [a]
++ [(Var
v, [LocArg]
newretlocs, Ty2
ty, Var -> [LocArg] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE Var
f [LocArg]
newapplocs [Exp2]
args)] [(Var, [LocArg], Ty2, Exp2)]
-> [(Var, [LocArg], Ty2, Exp2)] -> [(Var, [LocArg], Ty2, Exp2)]
forall a. [a] -> [a] -> [a]
++ [(Var, [LocArg], Ty2, Exp2)]
wpop [(Var, [LocArg], Ty2, Exp2)]
-> [(Var, [LocArg], Ty2, Exp2)] -> [(Var, [LocArg], Ty2, Exp2)]
forall a. [a] -> [a] -> [a]
++ [(Var, [LocArg], Ty2, Exp2)]
rpop
                  (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, [LocArg], 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, [LocArg], Ty2, Exp2)]
binds Exp2
bod3)
          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
$ [(Var, [LocArg], 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
v, [LocArg]
newretlocs, Ty2
ty, Var -> [LocArg] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE Var
f [LocArg]
newapplocs [Exp2]
args)] Exp2
bod3


    LetE (Var
v,[LocArg]
locs,Ty2
ty, (SpawnE Var
f [LocArg]
applocs [Exp2]
args)) Exp2
bod -> do
      let e' :: Exp2
e' = (Var, [LocArg], 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,[LocArg]
locs,Ty2
ty, (Var -> [LocArg] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE Var
f [LocArg]
applocs [Exp2]
args)) Exp2
bod
      Exp2
e'' <- DDefs Ty2
-> FunDefs Exp2
-> [LREM]
-> RegEnv
-> Env2 Ty2
-> RegEnv
-> AllocEnv
-> AllocEnv
-> RegEnv
-> OrderedLocsEnv
-> RegEnv
-> Set Var
-> Set Var
-> Exp2
-> PassM Exp2
threadRegionsExp DDefs Ty2
ddefs FunDefs Exp2
fundefs [LREM]
fnLocArgs RegEnv
renv Env2 Ty2
env2 RegEnv
lfenv AllocEnv
rlocs_env AllocEnv
wlocs_env RegEnv
pkd_env OrderedLocsEnv
region_locs RegEnv
ran_env Set Var
indirs Set Var
redirs 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
$ Var -> [Exp2] -> Exp2 -> Exp2
forall loc dec.
(Eq loc, Eq dec) =>
Var
-> [PreExp E2Ext loc dec]
-> PreExp E2Ext loc dec
-> PreExp E2Ext loc dec
changeAppToSpawn Var
f [Exp2]
args Exp2
e''

    -- AUDITME: this causes all all DataConE's to return an additional cursor.
    LetE (Var
v,[LocArg]
locs,ty :: Ty2
ty@(MkTy2 (PackedTy TyCon
_ Var
loc)), rhs :: Exp2
rhs@(DataConE LocArg
_ TyCon
_ [Exp2]
args)) Exp2
bod -> do
      let reg_of_tag :: Var
reg_of_tag = RegEnv
renv RegEnv -> Var -> Var
forall k a. Ord k => Map k a -> k -> a
M.! Var
loc
          lfenv' :: RegEnv
lfenv' = case [Exp2]
args of
                     [] -> RegEnv
lfenv
                     [Exp2]
_  ->
                       let last_ty :: TyOf Exp2
last_ty = 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)
DDefs Ty2
ddefs Env2 (TyOf Exp2)
Env2 Ty2
env2 ([Exp2] -> Exp2
forall a. HasCallStack => [a] -> a
last [Exp2]
args) in
                       case Ty2 -> UrTy Var
unTy2 TyOf Exp2
Ty2
last_ty of
                          PackedTy TyCon
_ Var
last_loc -> do
                            let reg_of_last_arg :: Var
reg_of_last_arg = RegEnv
renv RegEnv -> Var -> Var
forall k a. Ord k => Map k a -> k -> a
M.! Var
last_loc
                            if Var
reg_of_tag Var -> Var -> Bool
forall a. Eq a => a -> a -> Bool
/= Var
reg_of_last_arg
                            then Var -> Var -> RegEnv -> RegEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
loc Var
reg_of_last_arg RegEnv
lfenv
                            else RegEnv
lfenv
                          UrTy Var
_ -> RegEnv
lfenv
      let pkd_env1 :: RegEnv
pkd_env1 = Var -> Var -> RegEnv -> RegEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
loc (RegEnv
renv RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
loc) RegEnv
pkd_env
      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
          rlocs_env' :: AllocEnv
rlocs_env' = UrTy Var -> AllocEnv -> AllocEnv
forall {k}. Ord k => UrTy k -> Map k TyCon -> Map k TyCon
updRLocsEnv (Ty2 -> UrTy Var
unTy2 Ty2
ty) AllocEnv
rlocs_env
          wlocs_env' :: AllocEnv
wlocs_env' = (Var -> AllocEnv -> AllocEnv) -> AllocEnv -> [Var] -> AllocEnv
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Var
loc2 AllocEnv
acc -> Var -> AllocEnv -> AllocEnv
forall k a. Ord k => k -> Map k a -> Map k a
M.delete Var
loc2 AllocEnv
acc) AllocEnv
wlocs_env (Ty2 -> [Var]
locsInTy Ty2
ty)
      (Var, [LocArg], 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, [LocArg], Ty2, Exp2) -> Exp2 -> Exp2)
-> (Exp2 -> (Var, [LocArg], Ty2, Exp2)) -> Exp2 -> Exp2 -> Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Var
v,[LocArg]
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
<*>
        DDefs Ty2
-> FunDefs Exp2
-> [LREM]
-> RegEnv
-> Env2 Ty2
-> RegEnv
-> AllocEnv
-> AllocEnv
-> RegEnv
-> OrderedLocsEnv
-> RegEnv
-> Set Var
-> Set Var
-> Exp2
-> PassM Exp2
threadRegionsExp DDefs Ty2
ddefs FunDefs Exp2
fundefs [LREM]
fnLocArgs RegEnv
renv Env2 Ty2
env2' RegEnv
lfenv' AllocEnv
rlocs_env' AllocEnv
wlocs_env' RegEnv
pkd_env1 OrderedLocsEnv
region_locs RegEnv
ran_env Set Var
indirs Set Var
redirs Exp2
bod

    LetE (Var
v,[LocArg]
locs,ty :: Ty2
ty@(MkTy2 (PackedTy TyCon
_ Var
loc)),(Ext (IndirectionE TyCon
tcon TyCon
dcon (LocArg
a,LocArg
_b) (LocArg
c,LocArg
_d) Exp2
cpy))) Exp2
bod -> do
      let fn :: Var -> Modality -> LocArg
fn Var
x Modality
mode = if Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Var
x Set Var
indirs Bool -> Bool -> Bool
|| Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Var
x Set Var
redirs
                      then (Var -> LocArg
EndOfReg_Tagged Var
x)
                      else case Var -> RegEnv -> Maybe Var
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Var
x RegEnv
ran_env of
                             Just Var
ran -> (Var -> LocArg
EndOfReg_Tagged Var
ran)
                             Maybe Var
Nothing -> case Var -> RegEnv -> Maybe Var
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Var
x RegEnv
renv of
                                          Just Var
r -> (Var -> Modality -> Var -> LocArg
NewL2.EndOfReg Var
r Modality
mode (Var -> Var
toEndV Var
r))
                                          Maybe Var
Nothing -> TyCon -> LocArg
forall a. HasCallStack => TyCon -> a
error (TyCon -> LocArg) -> TyCon -> LocArg
forall a b. (a -> b) -> a -> b
$ TyCon
"threadRegionsExp: unbound loc " TyCon -> TyCon -> TyCon
forall a. [a] -> [a] -> [a]
++ Var -> TyCon
forall a. Out a => a -> TyCon
sdoc Var
x
      let b' :: LocArg
b' = Var -> Modality -> LocArg
fn (LocArg -> Var
toLocVar LocArg
a) Modality
Output
      let d' :: LocArg
d' = Var -> Modality -> LocArg
fn (LocArg -> Var
toLocVar LocArg
c) Modality
Input
      let fn2 :: LocArg -> Var -> LocArg
fn2 (Loc LREM
lrem) Var
end = LREM -> LocArg
Loc (LREM
lrem { lremEndReg :: Var
lremEndReg = Var
end })
          fn2 LocArg
oth Var
_ = TyCon -> LocArg
forall a. HasCallStack => TyCon -> a
error (TyCon -> LocArg) -> TyCon -> LocArg
forall a b. (a -> b) -> a -> b
$ TyCon
"fn2: " TyCon -> TyCon -> TyCon
forall a. [a] -> [a] -> [a]
++ LocArg -> TyCon
forall a. Out a => a -> TyCon
sdoc LocArg
oth
      let a' :: LocArg
a' = LocArg -> Var -> LocArg
fn2 LocArg
a (LocArg -> Var
toLocVar LocArg
b')
          c' :: LocArg
c' = LocArg -> Var -> LocArg
fn2 LocArg
c (LocArg -> Var
toLocVar LocArg
d')
      let pkd_env' :: RegEnv
pkd_env' = Var -> Var -> RegEnv -> RegEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
loc (RegEnv
renv RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
loc) RegEnv
pkd_env
      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
          rlocs_env' :: AllocEnv
rlocs_env' = UrTy Var -> AllocEnv -> AllocEnv
forall {k}. Ord k => UrTy k -> Map k TyCon -> Map k TyCon
updRLocsEnv (Ty2 -> UrTy Var
unTy2 Ty2
ty) AllocEnv
rlocs_env
          wlocs_env' :: AllocEnv
wlocs_env' = (Var -> AllocEnv -> AllocEnv) -> AllocEnv -> [Var] -> AllocEnv
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Var
loc2 AllocEnv
acc -> Var -> AllocEnv -> AllocEnv
forall k a. Ord k => k -> Map k a -> Map k a
M.delete Var
loc2 AllocEnv
acc) AllocEnv
wlocs_env (Ty2 -> [Var]
locsInTy Ty2
ty)
      Exp2
bod' <- DDefs Ty2
-> FunDefs Exp2
-> [LREM]
-> RegEnv
-> Env2 Ty2
-> RegEnv
-> AllocEnv
-> AllocEnv
-> RegEnv
-> OrderedLocsEnv
-> RegEnv
-> Set Var
-> Set Var
-> Exp2
-> PassM Exp2
threadRegionsExp DDefs Ty2
ddefs FunDefs Exp2
fundefs [LREM]
fnLocArgs RegEnv
renv Env2 Ty2
env2' RegEnv
lfenv AllocEnv
rlocs_env' AllocEnv
wlocs_env' RegEnv
pkd_env' OrderedLocsEnv
region_locs RegEnv
ran_env Set Var
indirs Set Var
redirs Exp2
bod
      let boundscheck :: Exp2 -> Exp2
boundscheck = let locarg :: LocArg
locarg = LocArg
a'
                            regarg :: LocArg
regarg = LocArg
b'
                            -- bc = boundsCheck ddefs tcon
                            bc :: Int
bc = Int
18
                        in (Var, [LocArg], 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
"_",[],UrTy Var -> Ty2
MkTy2 UrTy Var
forall loc. UrTy loc
IntTy, E2Ext LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext(E2Ext LocArg Ty2 -> Exp2) -> E2Ext LocArg Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Int -> LocArg -> LocArg -> E2Ext LocArg Ty2
forall loc dec. Int -> loc -> loc -> E2Ext loc dec
BoundsCheck Int
bc LocArg
regarg LocArg
locarg)
      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 -> Exp2
boundscheck (Exp2 -> Exp2) -> Exp2 -> Exp2
forall a b. (a -> b) -> a -> b
$ (Var, [LocArg], 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,[LocArg]
locs,Ty2
ty,(E2Ext LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (TyCon
-> TyCon
-> (LocArg, LocArg)
-> (LocArg, LocArg)
-> Exp2
-> E2Ext LocArg Ty2
forall loc dec.
TyCon
-> TyCon -> (loc, loc) -> (loc, loc) -> E2 loc dec -> E2Ext loc dec
IndirectionE TyCon
tcon TyCon
dcon (LocArg
a',LocArg
b') (LocArg
c',LocArg
d') Exp2
cpy))) Exp2
bod'

    Ext (StartOfPkdCursor Var
cur) -> do
      let (PackedTy TyCon
_ Var
loc) = Ty2 -> UrTy Var
unTy2 (Var -> Env2 Ty2 -> Ty2
forall a. Out a => Var -> Env2 a -> a
lookupVEnv Var
cur Env2 Ty2
env2)
      case Var -> RegEnv -> Maybe Var
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Var
loc RegEnv
pkd_env of
        Just Var
reg -> 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
$ E2Ext LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocArg Ty2 -> Exp2) -> E2Ext LocArg Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> Var -> E2Ext LocArg Ty2
forall loc dec. Var -> Var -> E2Ext loc dec
TagCursor Var
loc (Var -> Var
toEndV Var
reg)
        Maybe Var
Nothing -> TyCon -> PassM Exp2
forall a. HasCallStack => TyCon -> a
error (TyCon -> PassM Exp2) -> TyCon -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ TyCon
"threadRegionsExp: unbound " TyCon -> TyCon -> TyCon
forall a. [a] -> [a] -> [a]
++ (Var, RegEnv) -> TyCon
forall a. Out a => a -> TyCon
sdoc (Var
loc, RegEnv
pkd_env)

    -- Sometimes, this expression can have RetE forms. We should collect and update
    -- the locs here appropriately.
    LetE (Var
v,[LocArg]
locs,Ty2
ty, rhs :: Exp2
rhs@(TimeIt{})) Exp2
bod -> do
       Exp2
rhs' <- Exp2 -> PassM Exp2
go Exp2
rhs
       let retlocs :: [LocArg]
retlocs = Exp2 -> [LocArg]
findRetLocs Exp2
rhs'
           newretlocs :: [LocArg]
newretlocs = [LocArg]
retlocs [LocArg] -> [LocArg] -> [LocArg]
forall a. [a] -> [a] -> [a]
++ [LocArg]
locs
       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
           rlocs_env' :: AllocEnv
rlocs_env' = UrTy Var -> AllocEnv -> AllocEnv
forall {k}. Ord k => UrTy k -> Map k TyCon -> Map k TyCon
updRLocsEnv (Ty2 -> UrTy Var
unTy2 Ty2
ty) AllocEnv
rlocs_env
           wlocs_env' :: AllocEnv
wlocs_env' = (Var -> AllocEnv -> AllocEnv) -> AllocEnv -> [Var] -> AllocEnv
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Var
loc AllocEnv
acc -> Var -> AllocEnv -> AllocEnv
forall k a. Ord k => k -> Map k a -> Map k a
M.delete Var
loc AllocEnv
acc) AllocEnv
wlocs_env (Ty2 -> [Var]
locsInTy Ty2
ty)
       Exp2
bod1 <- DDefs Ty2
-> FunDefs Exp2
-> [LREM]
-> RegEnv
-> Env2 Ty2
-> RegEnv
-> AllocEnv
-> AllocEnv
-> RegEnv
-> OrderedLocsEnv
-> RegEnv
-> Set Var
-> Set Var
-> Exp2
-> PassM Exp2
threadRegionsExp DDefs Ty2
ddefs FunDefs Exp2
fundefs [LREM]
fnLocArgs RegEnv
renv Env2 Ty2
env2' RegEnv
lfenv AllocEnv
rlocs_env' AllocEnv
wlocs_env' RegEnv
pkd_env OrderedLocsEnv
region_locs RegEnv
ran_env Set Var
indirs Set Var
redirs Exp2
bod

       -- shadowstack  ops
       --------------------
       let -- free = S.fromList $ freeLocVars bod
            free :: Set Var
free = Set Var -> Env2 Ty2 -> Exp2 -> Set Var
ss_free_locs ([Var] -> Set Var
forall a. Ord a => [a] -> Set a
S.fromList (Var
v Var -> [Var] -> [Var]
forall a. a -> [a] -> [a]
: Ty2 -> [Var]
locsInTy Ty2
ty [Var] -> [Var] -> [Var]
forall a. [a] -> [a] -> [a]
++ ((LocArg -> Var) -> [LocArg] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map LocArg -> Var
toLocVar [LocArg]
locs))) Env2 Ty2
env2' Exp2
bod
            free_wlocs :: Set Var
free_wlocs = Set Var
free Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.intersection` (AllocEnv -> Set Var
forall k a. Map k a -> Set k
M.keysSet AllocEnv
wlocs_env')
            free_rlocs :: Set Var
free_rlocs = Set Var
free Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.intersection` (AllocEnv -> Set Var
forall k a. Map k a -> Set k
M.keysSet AllocEnv
rlocs_env')
            free_rlocs' :: Set (Maybe Var, Var)
free_rlocs' = let tmp :: [(Maybe Var, Var)]
tmp = ((Var, Ty2) -> (Maybe Var, Var))
-> [(Var, Ty2)] -> [(Maybe Var, Var)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Var
x,(MkTy2 (PackedTy TyCon
_ Var
loc))) -> (Var -> Maybe Var
forall a. a -> Maybe a
Just Var
x,Var
loc)) ([(Var, Ty2)] -> [(Maybe Var, Var)])
-> [(Var, Ty2)] -> [(Maybe Var, Var)]
forall a b. (a -> b) -> a -> b
$
                                    ((Var, Ty2) -> Bool) -> [(Var, Ty2)] -> [(Var, Ty2)]
forall a. (a -> Bool) -> [a] -> [a]
filter (\(Var
_x,_y :: Ty2
_y@(MkTy2 (PackedTy TyCon
tycon Var
loc))) -> Var
loc Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set Var
free_rlocs Bool -> Bool -> Bool
&& TyCon
tycon TyCon -> TyCon -> Bool
forall a. Eq a => a -> a -> Bool
/= TyCon
hole_tycon)
                                           (TyEnv Ty2 -> [(Var, Ty2)]
forall k a. Map k a -> [(k, a)]
M.toList (TyEnv Ty2 -> [(Var, Ty2)]) -> TyEnv Ty2 -> [(Var, Ty2)]
forall a b. (a -> b) -> a -> b
$ (Ty2 -> Bool) -> TyEnv Ty2 -> TyEnv Ty2
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (UrTy Var -> Bool
forall a. UrTy a -> Bool
isPackedTy (UrTy Var -> Bool) -> (Ty2 -> UrTy Var) -> Ty2 -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ty2 -> UrTy Var
unTy2) (Env2 Ty2 -> TyEnv Ty2
forall a. Env2 a -> TyEnv a
vEnv Env2 Ty2
env2))
                              tmp2 :: [(Maybe Var, Var)]
tmp2 = (Var -> (Maybe Var, Var)) -> [Var] -> [(Maybe Var, Var)]
forall a b. (a -> b) -> [a] -> [b]
map (\Var
x -> (Maybe Var
forall a. Maybe a
Nothing, Var
x)) ([Var] -> [(Maybe Var, Var)]) -> [Var] -> [(Maybe Var, Var)]
forall a b. (a -> b) -> a -> b
$ (Set Var -> [Var]
forall a. Set a -> [a]
S.toList Set Var
free_rlocs) [Var] -> [Var] -> [Var]
forall a. Eq a => [a] -> [a] -> [a]
L.\\ (((Maybe Var, Var) -> Var) -> [(Maybe Var, Var)] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe Var, Var) -> Var
forall a b. (a, b) -> b
snd [(Maybe Var, Var)]
tmp)
                          in [(Maybe Var, Var)] -> Set (Maybe Var, Var)
forall a. Ord a => [a] -> Set a
S.fromList ([(Maybe Var, Var)] -> Set (Maybe Var, Var))
-> [(Maybe Var, Var)] -> Set (Maybe Var, Var)
forall a b. (a -> b) -> a -> b
$ [(Maybe Var, Var)]
tmp [(Maybe Var, Var)] -> [(Maybe Var, Var)] -> [(Maybe Var, Var)]
forall a. [a] -> [a] -> [a]
++ [(Maybe Var, Var)]
tmp2
       ([(Var, [LocArg], Ty2, Exp2)]
rpush,[(Var, [LocArg], Ty2, Exp2)]
wpush,[(Var, [LocArg], Ty2, Exp2)]
rpop,[(Var, [LocArg], Ty2, Exp2)]
wpop) <- Set (Maybe Var, Var)
-> Set Var
-> AllocEnv
-> AllocEnv
-> RegEnv
-> PassM
     ([(Var, [LocArg], Ty2, Exp2)], [(Var, [LocArg], Ty2, Exp2)],
      [(Var, [LocArg], Ty2, Exp2)], [(Var, [LocArg], Ty2, Exp2)])
ss_ops Set (Maybe Var, Var)
free_rlocs' Set Var
free_wlocs AllocEnv
rlocs_env AllocEnv
wlocs_env RegEnv
renv
       Bool
emit_ss <- PassM Bool
emit_ss_instrs
       if Bool
emit_ss
         then do let binds :: [(Var, [LocArg], Ty2, Exp2)]
binds = [(Var, [LocArg], Ty2, Exp2)]
rpush [(Var, [LocArg], Ty2, Exp2)]
-> [(Var, [LocArg], Ty2, Exp2)] -> [(Var, [LocArg], Ty2, Exp2)]
forall a. [a] -> [a] -> [a]
++ [(Var, [LocArg], Ty2, Exp2)]
wpush [(Var, [LocArg], Ty2, Exp2)]
-> [(Var, [LocArg], Ty2, Exp2)] -> [(Var, [LocArg], Ty2, Exp2)]
forall a. [a] -> [a] -> [a]
++ [(Var
v, [LocArg]
newretlocs, Ty2
ty, Exp2
rhs')] [(Var, [LocArg], Ty2, Exp2)]
-> [(Var, [LocArg], Ty2, Exp2)] -> [(Var, [LocArg], Ty2, Exp2)]
forall a. [a] -> [a] -> [a]
++ [(Var, [LocArg], Ty2, Exp2)]
wpop [(Var, [LocArg], Ty2, Exp2)]
-> [(Var, [LocArg], Ty2, Exp2)] -> [(Var, [LocArg], Ty2, Exp2)]
forall a. [a] -> [a] -> [a]
++ [(Var, [LocArg], Ty2, Exp2)]
rpop
                 (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, [LocArg], 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, [LocArg], Ty2, Exp2)]
binds 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
$ [(Var, [LocArg], 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
v, [LocArg]
newretlocs, Ty2
ty, Exp2
rhs')] Exp2
bod1

    LetE (Var
v,[LocArg]
locs,Ty2
ty,rhs :: Exp2
rhs@(Ext (AllocateTagHere Var
x TyCon
x_tycon))) Exp2
bod -> do
      let -- x_tycon = (wlocs_env # x)
          rlocs_env' :: AllocEnv
rlocs_env' = Var -> TyCon -> AllocEnv -> AllocEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
x TyCon
x_tycon AllocEnv
rlocs_env
          wlocs_env' :: AllocEnv
wlocs_env' = Var -> AllocEnv -> AllocEnv
forall k a. Ord k => k -> Map k a -> Map k a
M.delete Var
x AllocEnv
wlocs_env
      ((Var, [LocArg], 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,[LocArg]
locs,Ty2
ty,Exp2
rhs)) (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
        DDefs Ty2
-> FunDefs Exp2
-> [LREM]
-> RegEnv
-> Env2 Ty2
-> RegEnv
-> AllocEnv
-> AllocEnv
-> RegEnv
-> OrderedLocsEnv
-> RegEnv
-> Set Var
-> Set Var
-> Exp2
-> PassM Exp2
threadRegionsExp DDefs Ty2
ddefs FunDefs Exp2
fundefs [LREM]
fnLocArgs RegEnv
renv (Var -> Ty2 -> Env2 Ty2 -> Env2 Ty2
forall a. Var -> a -> Env2 a -> Env2 a
extendVEnv Var
v Ty2
ty Env2 Ty2
env2) RegEnv
lfenv AllocEnv
rlocs_env' AllocEnv
wlocs_env' RegEnv
pkd_env OrderedLocsEnv
region_locs RegEnv
ran_env Set Var
indirs Set Var
redirs Exp2
bod

    LetE (Var
v,[LocArg]
locs,Ty2
ty, Exp2
rhs) Exp2
bod ->
      (Var, [LocArg], 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, [LocArg], Ty2, Exp2) -> Exp2 -> Exp2)
-> (Exp2 -> (Var, [LocArg], Ty2, Exp2)) -> Exp2 -> Exp2 -> Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Var
v,[LocArg]
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
<*>
        DDefs Ty2
-> FunDefs Exp2
-> [LREM]
-> RegEnv
-> Env2 Ty2
-> RegEnv
-> AllocEnv
-> AllocEnv
-> RegEnv
-> OrderedLocsEnv
-> RegEnv
-> Set Var
-> Set Var
-> Exp2
-> PassM Exp2
threadRegionsExp DDefs Ty2
ddefs FunDefs Exp2
fundefs [LREM]
fnLocArgs RegEnv
renv (Var -> Ty2 -> Env2 Ty2 -> Env2 Ty2
forall a. Var -> a -> Env2 a -> Env2 a
extendVEnv Var
v Ty2
ty Env2 Ty2
env2) RegEnv
lfenv AllocEnv
rlocs_env AllocEnv
wlocs_env RegEnv
pkd_env OrderedLocsEnv
region_locs RegEnv
ran_env Set Var
indirs Set Var
redirs Exp2
bod

    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
<$> DDefs Ty2
-> FunDefs Exp2
-> [LREM]
-> RegEnv
-> Env2 Ty2
-> RegEnv
-> AllocEnv
-> AllocEnv
-> RegEnv
-> OrderedLocsEnv
-> RegEnv
-> Set Var
-> Set Var
-> Exp2
-> PassM Exp2
threadRegionsExp DDefs Ty2
ddefs FunDefs Exp2
fundefs [LREM]
fnLocArgs RegEnv
renv (Var -> Ty2 -> Env2 Ty2 -> Env2 Ty2
forall a. Var -> a -> Env2 a -> Env2 a
extendVEnv Var
v (UrTy Var -> Ty2
MkTy2 UrTy Var
forall loc. UrTy loc
ArenaTy) Env2 Ty2
env2) RegEnv
lfenv AllocEnv
rlocs_env AllocEnv
wlocs_env RegEnv
pkd_env OrderedLocsEnv
region_locs RegEnv
ran_env Set Var
indirs Set Var
redirs Exp2
e

    Ext E2Ext LocArg Ty2
ext ->
      case E2Ext LocArg Ty2
ext of
        AddFixed{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
        LetLocE Var
loc PreLocExp LocArg
FreeLE Exp2
bod ->
          E2Ext LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocArg Ty2 -> Exp2)
-> (Exp2 -> E2Ext LocArg Ty2) -> Exp2 -> Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Var -> PreLocExp LocArg -> Exp2 -> E2Ext LocArg Ty2
forall loc dec. Var -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE Var
loc PreLocExp LocArg
forall loc. PreLocExp loc
FreeLE (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
            DDefs Ty2
-> FunDefs Exp2
-> [LREM]
-> RegEnv
-> Env2 Ty2
-> RegEnv
-> AllocEnv
-> AllocEnv
-> RegEnv
-> OrderedLocsEnv
-> RegEnv
-> Set Var
-> Set Var
-> Exp2
-> PassM Exp2
threadRegionsExp DDefs Ty2
ddefs FunDefs Exp2
fundefs [LREM]
fnLocArgs RegEnv
renv Env2 Ty2
env2 RegEnv
lfenv AllocEnv
rlocs_env AllocEnv
wlocs_env RegEnv
pkd_env OrderedLocsEnv
region_locs RegEnv
ran_env Set Var
indirs Set Var
redirs Exp2
bod
        -- Update renv with a binding for loc
        LetLocE Var
loc PreLocExp LocArg
rhs Exp2
bod -> do
          let reg :: Var
reg = case PreLocExp LocArg
rhs of
                      StartOfRegionLE Region
r      -> Region -> Var
regionToVar Region
r
                      InRegionLE Region
r           -> Region -> Var
regionToVar Region
r
                      AfterConstantLE Int
_ LocArg
lc   -> RegEnv
renv RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# (LocArg -> Var
toLocVar LocArg
lc)
                      AfterVariableLE Var
_ LocArg
lc Bool
_ -> RegEnv
renv RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# (LocArg -> Var
toLocVar LocArg
lc)
                      FromEndLE LocArg
lc           -> RegEnv
renv RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# (LocArg -> Var
toLocVar LocArg
lc)
              wlocs_env' :: AllocEnv
wlocs_env' = Var -> TyCon -> AllocEnv -> AllocEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
loc TyCon
hole_tycon AllocEnv
wlocs_env
              region_locs1 :: OrderedLocsEnv
region_locs1 = case PreLocExp LocArg
rhs of
                               AfterConstantLE{} -> ([Var] -> [Var]) -> Var -> OrderedLocsEnv -> OrderedLocsEnv
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (\[Var]
locs -> [Var]
locs [Var] -> [Var] -> [Var]
forall a. [a] -> [a] -> [a]
++ [Var
loc]) Var
reg OrderedLocsEnv
region_locs
                               AfterVariableLE{} -> ([Var] -> [Var]) -> Var -> OrderedLocsEnv -> OrderedLocsEnv
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (\[Var]
locs -> [Var]
locs [Var] -> [Var] -> [Var]
forall a. [a] -> [a] -> [a]
++ [Var
loc]) Var
reg OrderedLocsEnv
region_locs
                               StartOfRegionLE{} -> Var -> [Var] -> OrderedLocsEnv -> OrderedLocsEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
reg [Var
loc] OrderedLocsEnv
region_locs
                               PreLocExp LocArg
_ -> OrderedLocsEnv
region_locs
          E2Ext LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocArg Ty2 -> Exp2)
-> (Exp2 -> E2Ext LocArg Ty2) -> Exp2 -> Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Var -> PreLocExp LocArg -> Exp2 -> E2Ext LocArg Ty2
forall loc dec. Var -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE Var
loc PreLocExp LocArg
rhs (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
            DDefs Ty2
-> FunDefs Exp2
-> [LREM]
-> RegEnv
-> Env2 Ty2
-> RegEnv
-> AllocEnv
-> AllocEnv
-> RegEnv
-> OrderedLocsEnv
-> RegEnv
-> Set Var
-> Set Var
-> Exp2
-> PassM Exp2
threadRegionsExp DDefs Ty2
ddefs FunDefs Exp2
fundefs [LREM]
fnLocArgs (Var -> Var -> RegEnv -> RegEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
loc Var
reg RegEnv
renv) Env2 Ty2
env2 RegEnv
lfenv AllocEnv
rlocs_env AllocEnv
wlocs_env' RegEnv
pkd_env OrderedLocsEnv
region_locs1 RegEnv
ran_env Set Var
indirs Set Var
redirs Exp2
bod

        RetE [LocArg]
locs Var
v -> do
          let ty :: Ty2
ty = Var -> Env2 Ty2 -> Ty2
forall a. Out a => Var -> Env2 a -> a
lookupVEnv Var
v Env2 Ty2
env2
              fn :: Modality -> Var -> LocArg
fn Modality
m = (\Var
r -> Var -> Modality -> Var -> LocArg
NewL2.EndOfReg Var
r Modality
m (Var -> Var
toEndV Var
r))
              outtylocs :: [Var]
outtylocs = Ty2 -> [Var]
locsInTy Ty2
ty
              outtyregvars :: [Var]
outtyregvars =
                (Var -> [Var] -> [Var]) -> [Var] -> [Var] -> [Var]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Var
loc [Var]
acc -> case Var -> RegEnv -> Maybe Var
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Var
loc RegEnv
lfenv of
                                     Maybe Var
Nothing -> (RegEnv
renv RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
loc) Var -> [Var] -> [Var]
forall a. a -> [a] -> [a]
: [Var]
acc
                                     Just Var
r  -> Var
r Var -> [Var] -> [Var]
forall a. a -> [a] -> [a]
: [Var]
acc)
                [] [Var]
outtylocs
              outtyregargs :: [LocArg]
outtyregargs = (Var -> LocArg) -> [Var] -> [LocArg]
forall a b. (a -> b) -> [a] -> [b]
map (Modality -> Var -> LocArg
fn Modality
Output) [Var]
outtyregvars
              inregvars :: [Var]
inregvars = (LREM -> Var) -> [LREM] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map (\LREM
lrm -> let r :: Var
r = RegEnv
renv RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# (LREM -> Var
lremLoc LREM
lrm)
                                           last_loc :: Var
last_loc = [Var] -> Var
forall a. HasCallStack => [a] -> a
last (OrderedLocsEnv
region_locs OrderedLocsEnv -> Var -> [Var]
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
r)
                                           r' :: Var
r' = (RegEnv
renv RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
last_loc)
                                       in Var
r') ([LREM] -> [Var]) -> [LREM] -> [Var]
forall a b. (a -> b) -> a -> b
$
                          (LREM -> Bool) -> [LREM] -> [LREM]
forall a. (a -> Bool) -> [a] -> [a]
filter (\LREM
lrm -> LREM -> Modality
lremMode LREM
lrm Modality -> Modality -> Bool
forall a. Eq a => a -> a -> Bool
== Modality
Input) [LREM]
fnLocArgs
              inregargs :: [LocArg]
inregargs = (Var -> LocArg) -> [Var] -> [LocArg]
forall a b. (a -> b) -> [a] -> [b]
map (Modality -> Var -> LocArg
fn Modality
Input) [Var]
inregvars
              newlocs :: [LocArg]
newlocs = [LocArg]
inregargs [LocArg] -> [LocArg] -> [LocArg]
forall a. [a] -> [a] -> [a]
++ [LocArg]
outtyregargs
          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
$ E2Ext LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocArg Ty2 -> Exp2) -> E2Ext LocArg Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ [LocArg] -> Var -> E2Ext LocArg Ty2
forall loc dec. [loc] -> Var -> E2Ext loc dec
RetE ([LocArg]
newlocs [LocArg] -> [LocArg] -> [LocArg]
forall a. [a] -> [a] -> [a]
++ [LocArg]
locs) Var
v

        TagCursor Var
a Var
b -> 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
$ E2Ext LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocArg Ty2 -> Exp2) -> E2Ext LocArg Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> Var -> E2Ext LocArg Ty2
forall loc dec. Var -> Var -> E2Ext loc dec
TagCursor Var
a Var
b
        LetRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
bod -> do
          -- shadowstack  ops
          --------------------
          let -- free = S.fromList $ freeLocVars bod
              free :: Set Var
free = Set Var -> Env2 Ty2 -> Exp2 -> Set Var
ss_free_locs (Var -> Set Var
forall a. a -> Set a
S.singleton (Region -> Var
regionToVar Region
r)) Env2 Ty2
env2 Exp2
bod
              free_wlocs :: Set Var
free_wlocs = Set Var
free Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.intersection` (AllocEnv -> Set Var
forall k a. Map k a -> Set k
M.keysSet AllocEnv
wlocs_env)
              free_rlocs :: Set Var
free_rlocs = Set Var
free Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.intersection` (AllocEnv -> Set Var
forall k a. Map k a -> Set k
M.keysSet AllocEnv
rlocs_env)
              free_rlocs' :: Set (Maybe Var, Var)
free_rlocs' = let tmp :: [(Maybe Var, Var)]
tmp = ((Var, Ty2) -> (Maybe Var, Var))
-> [(Var, Ty2)] -> [(Maybe Var, Var)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Var
x,(MkTy2 (PackedTy TyCon
_ Var
loc))) -> (Var -> Maybe Var
forall a. a -> Maybe a
Just Var
x,Var
loc)) ([(Var, Ty2)] -> [(Maybe Var, Var)])
-> [(Var, Ty2)] -> [(Maybe Var, Var)]
forall a b. (a -> b) -> a -> b
$
                                      ((Var, Ty2) -> Bool) -> [(Var, Ty2)] -> [(Var, Ty2)]
forall a. (a -> Bool) -> [a] -> [a]
filter (\(Var
_x,_y :: Ty2
_y@(MkTy2 (PackedTy TyCon
tycon Var
loc))) -> Var
loc Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set Var
free_rlocs Bool -> Bool -> Bool
&& TyCon
tycon TyCon -> TyCon -> Bool
forall a. Eq a => a -> a -> Bool
/= TyCon
hole_tycon)
                                             (TyEnv Ty2 -> [(Var, Ty2)]
forall k a. Map k a -> [(k, a)]
M.toList (TyEnv Ty2 -> [(Var, Ty2)]) -> TyEnv Ty2 -> [(Var, Ty2)]
forall a b. (a -> b) -> a -> b
$ (Ty2 -> Bool) -> TyEnv Ty2 -> TyEnv Ty2
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (UrTy Var -> Bool
forall a. UrTy a -> Bool
isPackedTy (UrTy Var -> Bool) -> (Ty2 -> UrTy Var) -> Ty2 -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ty2 -> UrTy Var
unTy2) (Env2 Ty2 -> TyEnv Ty2
forall a. Env2 a -> TyEnv a
vEnv Env2 Ty2
env2))
                                tmp2 :: [(Maybe Var, Var)]
tmp2 = (Var -> (Maybe Var, Var)) -> [Var] -> [(Maybe Var, Var)]
forall a b. (a -> b) -> [a] -> [b]
map (\Var
x -> (Maybe Var
forall a. Maybe a
Nothing, Var
x)) ([Var] -> [(Maybe Var, Var)]) -> [Var] -> [(Maybe Var, Var)]
forall a b. (a -> b) -> a -> b
$ (Set Var -> [Var]
forall a. Set a -> [a]
S.toList Set Var
free_rlocs) [Var] -> [Var] -> [Var]
forall a. Eq a => [a] -> [a] -> [a]
L.\\ (((Maybe Var, Var) -> Var) -> [(Maybe Var, Var)] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe Var, Var) -> Var
forall a b. (a, b) -> b
snd [(Maybe Var, Var)]
tmp)
                            in [(Maybe Var, Var)] -> Set (Maybe Var, Var)
forall a. Ord a => [a] -> Set a
S.fromList ([(Maybe Var, Var)] -> Set (Maybe Var, Var))
-> [(Maybe Var, Var)] -> Set (Maybe Var, Var)
forall a b. (a -> b) -> a -> b
$ [(Maybe Var, Var)]
tmp [(Maybe Var, Var)] -> [(Maybe Var, Var)] -> [(Maybe Var, Var)]
forall a. [a] -> [a] -> [a]
++ [(Maybe Var, Var)]
tmp2
          ([(Var, [LocArg], Ty2, Exp2)]
rpush,[(Var, [LocArg], Ty2, Exp2)]
wpush,[(Var, [LocArg], Ty2, Exp2)]
rpop,[(Var, [LocArg], Ty2, Exp2)]
wpop) <- Set (Maybe Var, Var)
-> Set Var
-> AllocEnv
-> AllocEnv
-> RegEnv
-> PassM
     ([(Var, [LocArg], Ty2, Exp2)], [(Var, [LocArg], Ty2, Exp2)],
      [(Var, [LocArg], Ty2, Exp2)], [(Var, [LocArg], Ty2, Exp2)])
ss_ops Set (Maybe Var, Var)
free_rlocs' Set Var
free_wlocs AllocEnv
rlocs_env AllocEnv
wlocs_env RegEnv
renv
          Bool
emit_ss <- PassM Bool
emit_ss_instrs
          Exp2
bod' <- Exp2 -> PassM Exp2
go Exp2
bod
          if Bool
emit_ss
            then do let pre :: Exp2 -> Exp2
pre = [(Var, [LocArg], 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, [LocArg], Ty2, Exp2)]
rpush [(Var, [LocArg], Ty2, Exp2)]
-> [(Var, [LocArg], Ty2, Exp2)] -> [(Var, [LocArg], Ty2, Exp2)]
forall a. [a] -> [a] -> [a]
++ [(Var, [LocArg], Ty2, Exp2)]
wpush)
                        post :: Exp2
post = [(Var, [LocArg], 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, [LocArg], Ty2, Exp2)]
wpop [(Var, [LocArg], Ty2, Exp2)]
-> [(Var, [LocArg], Ty2, Exp2)] -> [(Var, [LocArg], Ty2, Exp2)]
forall a. [a] -> [a] -> [a]
++ [(Var, [LocArg], Ty2, Exp2)]
rpop) 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
$ Exp2 -> Exp2
pre (E2Ext LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocArg Ty2 -> Exp2) -> E2Ext LocArg Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Region
-> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext LocArg Ty2
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
post)
            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 LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocArg Ty2 -> Exp2) -> E2Ext LocArg Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Region
-> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext LocArg Ty2
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
bod'
        LetParRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
bod -> E2Ext LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocArg Ty2 -> Exp2)
-> (Exp2 -> E2Ext LocArg Ty2) -> Exp2 -> Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Region
-> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext LocArg 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
<$> 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 Int
sz LocArg
_bound LocArg
cur -> do
          let reg :: Var
reg = Var -> Var
toEndV (RegEnv
renv RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# (LocArg -> Var
toLocVar LocArg
cur))
          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
$ E2Ext LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocArg Ty2 -> Exp2) -> E2Ext LocArg Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Int -> LocArg -> LocArg -> E2Ext LocArg Ty2
forall loc dec. Int -> loc -> loc -> E2Ext loc dec
BoundsCheck Int
sz (Var -> Modality -> Var -> LocArg
NewL2.EndOfReg Var
reg Modality
Output (Var -> Var
toEndV Var
reg)) LocArg
cur
        IndirectionE{}   -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
        E2Ext LocArg Ty2
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 LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocArg Ty2 -> Exp2)
-> (Exp2 -> E2Ext LocArg Ty2) -> Exp2 -> Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Var] -> Exp2 -> E2Ext LocArg 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

    -- 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
    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
    CaseE Exp2
scrt [(TyCon, [(Var, LocArg)], Exp2)]
mp -> do
      let (VarE Var
v) = Exp2
scrt
          PackedTy TyCon
_ Var
tyloc = Ty2 -> UrTy Var
unTy2 (Var -> Env2 Ty2 -> Ty2
forall a. Out a => Var -> Env2 a -> a
lookupVEnv Var
v Env2 Ty2
env2)
          reg :: Var
reg = RegEnv
renv RegEnv -> Var -> Var
forall k a. Ord k => Map k a -> k -> a
M.! Var
tyloc
      Exp2 -> [(TyCon, [(Var, LocArg)], Exp2)] -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> [(TyCon, [(Var, loc)], PreExp ext loc dec)]
-> PreExp ext loc dec
CaseE Exp2
scrt ([(TyCon, [(Var, LocArg)], Exp2)] -> Exp2)
-> PassM [(TyCon, [(Var, LocArg)], Exp2)] -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((TyCon, [(Var, LocArg)], Exp2)
 -> PassM (TyCon, [(Var, LocArg)], Exp2))
-> [(TyCon, [(Var, LocArg)], Exp2)]
-> PassM [(TyCon, [(Var, LocArg)], 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
-> RegEnv
-> Env2 Ty2
-> RegEnv
-> AllocEnv
-> AllocEnv
-> RegEnv
-> OrderedLocsEnv
-> RegEnv
-> Set Var
-> Set Var
-> (TyCon, [(Var, LocArg)], Exp2)
-> PassM (TyCon, [(Var, LocArg)], Exp2)
docase Var
reg RegEnv
renv Env2 Ty2
env2 RegEnv
lfenv AllocEnv
rlocs_env AllocEnv
wlocs_env RegEnv
pkd_env OrderedLocsEnv
region_locs RegEnv
ran_env Set Var
indirs Set Var
redirs) [(TyCon, [(Var, LocArg)], 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 (m :: * -> *) a. Monad m => a -> m a
return (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
    SpawnE{} -> TyCon -> PassM Exp2
forall a. HasCallStack => TyCon -> a
error TyCon
"threadRegionsExp: Unbound SpawnE"
    Exp2
SyncE    -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
    MapE{}  -> TyCon -> PassM Exp2
forall a. HasCallStack => TyCon -> a
error (TyCon -> PassM Exp2) -> TyCon -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ TyCon
"threadRegionsExp: TODO MapE"
    FoldE{} -> TyCon -> PassM Exp2
forall a. HasCallStack => TyCon -> a
error (TyCon -> PassM Exp2) -> TyCon -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ TyCon
"threadRegionsExp: TODO FoldE"
  where
    emit_ss_instrs :: PassM Bool
emit_ss_instrs =
      do DynFlags
dflags <- PassM DynFlags
forall (m :: * -> *). MonadReader Config m => m DynFlags
getDynFlags
         Bool -> PassM Bool
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> PassM Bool) -> Bool -> PassM Bool
forall a b. (a -> b) -> a -> b
$ GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_GenGc DynFlags
dflags Bool -> Bool -> Bool
&& Bool -> Bool
not (GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_DisableGC DynFlags
dflags)

    go :: Exp2 -> PassM Exp2
go = DDefs Ty2
-> FunDefs Exp2
-> [LREM]
-> RegEnv
-> Env2 Ty2
-> RegEnv
-> AllocEnv
-> AllocEnv
-> RegEnv
-> OrderedLocsEnv
-> RegEnv
-> Set Var
-> Set Var
-> Exp2
-> PassM Exp2
threadRegionsExp DDefs Ty2
ddefs FunDefs Exp2
fundefs [LREM]
fnLocArgs RegEnv
renv Env2 Ty2
env2 RegEnv
lfenv AllocEnv
rlocs_env AllocEnv
wlocs_env RegEnv
pkd_env OrderedLocsEnv
region_locs RegEnv
ran_env Set Var
indirs Set Var
redirs

    docase :: Var
-> RegEnv
-> Env2 Ty2
-> RegEnv
-> AllocEnv
-> AllocEnv
-> RegEnv
-> OrderedLocsEnv
-> RegEnv
-> Set Var
-> Set Var
-> (TyCon, [(Var, LocArg)], Exp2)
-> PassM (TyCon, [(Var, LocArg)], Exp2)
docase Var
reg RegEnv
renv1 Env2 Ty2
env21 RegEnv
lfenv1 AllocEnv
rlocs_env1 AllocEnv
wlocs_env1 RegEnv
pkd_env1 OrderedLocsEnv
region_locs1 RegEnv
ran_env1 Set Var
indirs1 Set Var
redirs1 (TyCon
dcon,[(Var, LocArg)]
vlocargs,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,[LocArg]
locargs) = [(Var, LocArg)] -> ([Var], [LocArg])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Var, LocArg)]
vlocargs
          dcon_tys :: [Ty2]
dcon_tys = DDefs Ty2 -> TyCon -> [Ty2]
forall a. Out a => DDefs a -> TyCon -> [a]
lookupDataCon DDefs Ty2
ddefs TyCon
dcon
          locs :: [Var]
locs = (LocArg -> Var) -> [LocArg] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map LocArg -> Var
toLocVar [LocArg]
locargs
          renv0 :: RegEnv
renv0  = if TyCon -> Bool
isIndirectionTag TyCon
dcon Bool -> Bool -> Bool
|| TyCon -> Bool
isRedirectionTag TyCon
dcon
                   then (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
renv1 [Var]
vars
                   else RegEnv
renv1
          renv1' :: RegEnv
renv1' = (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
renv0 [Var]
locs
          env21' :: Env2 Ty2
env21' = HasCallStack =>
TyCon -> DDefs Ty2 -> [Var] -> [Var] -> Env2 Ty2 -> Env2 Ty2
TyCon -> DDefs Ty2 -> [Var] -> [Var] -> Env2 Ty2 -> Env2 Ty2
extendPatternMatchEnv TyCon
dcon DDefs Ty2
ddefs [Var]
vars [Var]
locs Env2 Ty2
env21
          rlocs_env1' :: AllocEnv
rlocs_env1' = ((Var, Ty2) -> AllocEnv -> AllocEnv)
-> AllocEnv -> [(Var, Ty2)] -> AllocEnv
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\(Var
loc,Ty2
ty) AllocEnv
acc ->
                                case Ty2 -> UrTy Var
unTy2 Ty2
ty of
                                  PackedTy TyCon
tycon Var
_ -> Var -> TyCon -> AllocEnv -> AllocEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
loc TyCon
tycon AllocEnv
acc
                                  UrTy Var
_ -> AllocEnv
acc)
                              AllocEnv
rlocs_env1
                              ([Var] -> [Ty2] -> [(Var, Ty2)]
forall a b.
(Show a, Show b, HasCallStack) =>
[a] -> [b] -> [(a, b)]
fragileZip [Var]
locs [Ty2]
dcon_tys)
          pkd_env1' :: RegEnv
pkd_env1' = ((Var, Ty2) -> RegEnv -> RegEnv)
-> RegEnv -> [(Var, Ty2)] -> 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
loc,Ty2
ty) RegEnv
acc ->
                                case Ty2 -> UrTy Var
unTy2 Ty2
ty of
                                  PackedTy{} -> Var -> Var -> RegEnv -> RegEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Var
loc Var
reg RegEnv
acc
                                  UrTy Var
_ -> RegEnv
acc)
                              RegEnv
pkd_env1
                              ([Var] -> [Ty2] -> [(Var, Ty2)]
forall a b.
(Show a, Show b, HasCallStack) =>
[a] -> [b] -> [(a, b)]
fragileZip [Var]
locs [Ty2]
dcon_tys)
          indirs1' :: Set Var
indirs1' = if TyCon -> Bool
isIndirectionTag TyCon
dcon
                     then Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.insert ([Var] -> Var
forall a. HasCallStack => [a] -> a
head [Var]
vars) Set Var
indirs1
                     else Set Var
indirs1
          redirs1' :: Set Var
redirs1' = if TyCon -> Bool
isRedirectionTag TyCon
dcon
                     then Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.insert ([Var] -> Var
forall a. HasCallStack => [a] -> a
head [Var]
vars) Set Var
redirs1
                     else Set Var
redirs1
          region_locs1' :: OrderedLocsEnv
region_locs1' = if TyCon -> Bool
isIndirectionTag TyCon
dcon Bool -> Bool -> Bool
|| TyCon -> Bool
isRedirectionTag TyCon
dcon
                          then ([Var] -> [Var]) -> Var -> OrderedLocsEnv -> OrderedLocsEnv
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (\[Var]
val -> [Var]
val [Var] -> [Var] -> [Var]
forall a. [a] -> [a] -> [a]
++ Int -> [Var] -> [Var]
forall a. Int -> [a] -> [a]
take Int
1 [Var]
vars) Var
reg OrderedLocsEnv
region_locs1
                          else ([Var] -> [Var]) -> Var -> OrderedLocsEnv -> OrderedLocsEnv
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (\[Var]
val -> [Var]
val [Var] -> [Var] -> [Var]
forall a. [a] -> [a] -> [a]
++ [Var]
locs) Var
reg OrderedLocsEnv
region_locs1
          num_cursor_tys :: Int
num_cursor_tys = [Ty2] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Ty2] -> Int) -> [Ty2] -> Int
forall a b. (a -> b) -> a -> b
$ (Ty2 -> Bool) -> [Ty2] -> [Ty2]
forall a. (a -> Bool) -> [a] -> [a]
filter ((UrTy Var -> UrTy Var -> Bool
forall a. Eq a => a -> a -> Bool
== UrTy Var
forall loc. UrTy loc
CursorTy) (UrTy Var -> Bool) -> (Ty2 -> UrTy Var) -> Ty2 -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ty2 -> UrTy Var
unTy2) [Ty2]
dcon_tys
          ran_env1' :: RegEnv
ran_env1' = RegEnv
ran_env1 RegEnv -> RegEnv -> RegEnv
forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union`
                      (if TyCon -> Bool
isIndirectionTag TyCon
dcon Bool -> Bool -> Bool
|| TyCon -> Bool
isRedirectionTag TyCon
dcon then RegEnv
forall k a. Map k a
M.empty else
                         [(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
$ [Var] -> [Var] -> [(Var, Var)]
forall a b. [a] -> [b] -> [(a, b)]
zip
                         ([Var] -> [Var]
forall a. [a] -> [a]
reverse ([Var] -> [Var]) -> [Var] -> [Var]
forall a b. (a -> b) -> a -> b
$ Int -> [Var] -> [Var]
forall a. Int -> [a] -> [a]
take Int
num_cursor_tys ([Var] -> [Var]) -> [Var] -> [Var]
forall a b. (a -> b) -> a -> b
$ [Var] -> [Var]
forall a. [a] -> [a]
reverse [Var]
locs)
                         (Int -> [Var] -> [Var]
forall a. Int -> [a] -> [a]
take Int
num_cursor_tys [Var]
vars))
      (TyCon
dcon,[(Var, LocArg)]
vlocargs,) (Exp2 -> (TyCon, [(Var, LocArg)], Exp2))
-> PassM Exp2 -> PassM (TyCon, [(Var, LocArg)], Exp2)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
         (DDefs Ty2
-> FunDefs Exp2
-> [LREM]
-> RegEnv
-> Env2 Ty2
-> RegEnv
-> AllocEnv
-> AllocEnv
-> RegEnv
-> OrderedLocsEnv
-> RegEnv
-> Set Var
-> Set Var
-> Exp2
-> PassM Exp2
threadRegionsExp DDefs Ty2
ddefs FunDefs Exp2
fundefs [LREM]
fnLocArgs RegEnv
renv1' Env2 Ty2
env21' RegEnv
lfenv1 AllocEnv
rlocs_env1' AllocEnv
wlocs_env1 RegEnv
pkd_env1' OrderedLocsEnv
region_locs1' RegEnv
ran_env1' Set Var
indirs1' Set Var
redirs1' Exp2
bod)

    ss_free_locs :: S.Set Var -> Env2 Ty2 -> Exp2 -> S.Set Var
    ss_free_locs :: Set Var -> Env2 Ty2 -> Exp2 -> Set Var
ss_free_locs Set Var
bound Env2 Ty2
env20 Exp2
ex0 =
                     (Var -> Var) -> Set Var -> Set Var
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map (\Var
w -> case Var -> TyEnv Ty2 -> Maybe Ty2
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Var
w (Env2 Ty2 -> TyEnv Ty2
forall a. Env2 a -> TyEnv a
vEnv Env2 Ty2
env20) of
                                       -- assumption: it's a location
                                       Maybe Ty2
Nothing -> Var
w
                                       Just (MkTy2 (PackedTy TyCon
_ Var
loc)) -> Var
loc
                                       Just Ty2
wty -> TyCon -> Var
forall a. HasCallStack => TyCon -> a
error (TyCon -> Var) -> TyCon -> Var
forall a b. (a -> b) -> a -> b
$ TyCon
"threadRegionsExp: unexpected type " TyCon -> TyCon -> TyCon
forall a. [a] -> [a] -> [a]
++ (Var, Ty2) -> TyCon
forall a. Show a => a -> TyCon
show (Var
w,Ty2
wty))
                           (Exp2 -> Set Var
allFreeVars_sans_datacon_args Exp2
ex0 Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.difference`
                            (Set Var
bound Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.union`
                             TyEnv Ty2 -> Set Var
forall k a. Map k a -> Set k
M.keysSet ((Ty2 -> Bool) -> TyEnv Ty2 -> TyEnv Ty2
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (Bool -> Bool
not (Bool -> Bool) -> (Ty2 -> Bool) -> Ty2 -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UrTy Var -> Bool
forall a. UrTy a -> Bool
isPackedTy (UrTy Var -> Bool) -> (Ty2 -> UrTy Var) -> Ty2 -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ty2 -> UrTy Var
unTy2) (Env2 Ty2 -> TyEnv Ty2
forall a. Env2 a -> TyEnv a
vEnv Env2 Ty2
env20)) Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.union`
                             Map Var (ArrowTy2 Ty2) -> Set Var
forall k a. Map k a -> Set k
M.keysSet (Env2 Ty2 -> TyEnv (ArrowTy Ty2)
forall a. Env2 a -> TyEnv (ArrowTy a)
fEnv Env2 Ty2
env20)))

    updRLocsEnv :: UrTy k -> Map k TyCon -> Map k TyCon
updRLocsEnv UrTy k
t Map k TyCon
acc =
                 case UrTy k
t of
                   PackedTy TyCon
tycon k
loc -> k -> TyCon -> Map k TyCon -> Map k TyCon
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
loc TyCon
tycon Map k TyCon
acc
                   ProdTy [UrTy k]
tys -> (UrTy k -> Map k TyCon -> Map k TyCon)
-> Map k TyCon -> [UrTy k] -> Map k TyCon
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr UrTy k -> Map k TyCon -> Map k TyCon
updRLocsEnv Map k TyCon
acc [UrTy k]
tys
                   UrTy k
_ -> Map k TyCon
acc


hole_tycon :: String
hole_tycon :: TyCon
hole_tycon = TyCon
"HOLE"

ss_ops :: S.Set (Maybe Var, LocVar) -> S.Set Var -> AllocEnv -> AllocEnv -> RegEnv ->
          PassM
          ([(Var, [LocArg], Ty2, Exp2)], [(Var, [LocArg], Ty2, Exp2)],
           [(Var, [LocArg], Ty2, Exp2)], [(Var, [LocArg], Ty2, Exp2)])
ss_ops :: Set (Maybe Var, Var)
-> Set Var
-> AllocEnv
-> AllocEnv
-> RegEnv
-> PassM
     ([(Var, [LocArg], Ty2, Exp2)], [(Var, [LocArg], Ty2, Exp2)],
      [(Var, [LocArg], Ty2, Exp2)], [(Var, [LocArg], Ty2, Exp2)])
ss_ops Set (Maybe Var, Var)
free_rlocs Set Var
free_wlocs AllocEnv
rlocs_env AllocEnv
wlocs_env RegEnv
renv = do
      [(Var, [LocArg], Ty2, Exp2)]
rpush <- (((Maybe Var, Var)
 -> [(Var, [LocArg], Ty2, Exp2)]
 -> PassM [(Var, [LocArg], Ty2, Exp2)])
-> [(Var, [LocArg], Ty2, Exp2)]
-> Set (Maybe Var, Var)
-> PassM [(Var, [LocArg], Ty2, Exp2)]
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> b -> m b) -> b -> t a -> m b
foldrM (\(Maybe Var
mb_x,Var
loc) [(Var, [LocArg], Ty2, Exp2)]
acc -> do
                         Var
push <- Var -> PassM Var
forall (m :: * -> *). MonadState Int m => Var -> m Var
gensym Var
"ss_push"
                         let tycon :: TyCon
tycon = AllocEnv
rlocs_env AllocEnv -> Var -> TyCon
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
loc
                         if TyCon
tycon TyCon -> TyCon -> Bool
forall a. Eq a => a -> a -> Bool
== TyCon
hole_tycon
                         then [(Var, [LocArg], Ty2, Exp2)] -> PassM [(Var, [LocArg], Ty2, Exp2)]
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Var, [LocArg], Ty2, Exp2)]
acc
                         else case Maybe Var
mb_x of
                                Maybe Var
Nothing -> [(Var, [LocArg], Ty2, Exp2)] -> PassM [(Var, [LocArg], Ty2, Exp2)]
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Var
push,[],UrTy Var -> Ty2
MkTy2 ([UrTy Var] -> UrTy Var
forall loc. [UrTy loc] -> UrTy loc
ProdTy []), E2Ext LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocArg Ty2 -> Exp2) -> E2Ext LocArg Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ SSModality -> Var -> Var -> TyCon -> E2Ext LocArg Ty2
forall loc dec. SSModality -> Var -> Var -> TyCon -> E2Ext loc dec
SSPush SSModality
Read Var
loc (Var -> Var
toEndV (RegEnv
renv RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
loc)) TyCon
tycon) (Var, [LocArg], Ty2, Exp2)
-> [(Var, [LocArg], Ty2, Exp2)] -> [(Var, [LocArg], Ty2, Exp2)]
forall a. a -> [a] -> [a]
: [(Var, [LocArg], Ty2, Exp2)]
acc)
                                Just Var
x  -> [(Var, [LocArg], Ty2, Exp2)] -> PassM [(Var, [LocArg], Ty2, Exp2)]
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Var
push,[],UrTy Var -> Ty2
MkTy2 ([UrTy Var] -> UrTy Var
forall loc. [UrTy loc] -> UrTy loc
ProdTy []), E2Ext LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocArg Ty2 -> Exp2) -> E2Ext LocArg Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ SSModality -> Var -> Var -> TyCon -> E2Ext LocArg Ty2
forall loc dec. SSModality -> Var -> Var -> TyCon -> E2Ext loc dec
SSPush SSModality
Read Var
x (Var -> Var
toEndV (RegEnv
renv RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
loc)) TyCon
tycon) (Var, [LocArg], Ty2, Exp2)
-> [(Var, [LocArg], Ty2, Exp2)] -> [(Var, [LocArg], Ty2, Exp2)]
forall a. a -> [a] -> [a]
: [(Var, [LocArg], Ty2, Exp2)]
acc))
                       []
                       Set (Maybe Var, Var)
free_rlocs) :: PassM [(Var, [LocArg], Ty2, Exp2)]
      [(Var, [LocArg], Ty2, Exp2)]
wpush <- ((Var
 -> [(Var, [LocArg], Ty2, Exp2)]
 -> PassM [(Var, [LocArg], Ty2, Exp2)])
-> [(Var, [LocArg], Ty2, Exp2)]
-> Set Var
-> PassM [(Var, [LocArg], Ty2, Exp2)]
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> b -> m b) -> b -> t a -> m b
foldrM (\Var
x [(Var, [LocArg], Ty2, Exp2)]
acc -> do
                         Var
push <- Var -> PassM Var
forall (m :: * -> *). MonadState Int m => Var -> m Var
gensym Var
"ss_push"
                         let tycon :: TyCon
tycon = AllocEnv
wlocs_env AllocEnv -> Var -> TyCon
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
x
                         if TyCon
tycon TyCon -> TyCon -> Bool
forall a. Eq a => a -> a -> Bool
== TyCon
hole_tycon
                         then [(Var, [LocArg], Ty2, Exp2)] -> PassM [(Var, [LocArg], Ty2, Exp2)]
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Var, [LocArg], Ty2, Exp2)]
acc
                         else [(Var, [LocArg], Ty2, Exp2)] -> PassM [(Var, [LocArg], Ty2, Exp2)]
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Var
push,[],UrTy Var -> Ty2
MkTy2 ([UrTy Var] -> UrTy Var
forall loc. [UrTy loc] -> UrTy loc
ProdTy []), E2Ext LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocArg Ty2 -> Exp2) -> E2Ext LocArg Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ SSModality -> Var -> Var -> TyCon -> E2Ext LocArg Ty2
forall loc dec. SSModality -> Var -> Var -> TyCon -> E2Ext loc dec
SSPush SSModality
Write Var
x (Var -> Var
toEndV (RegEnv
renv RegEnv -> Var -> Var
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Var
x)) TyCon
tycon) (Var, [LocArg], Ty2, Exp2)
-> [(Var, [LocArg], Ty2, Exp2)] -> [(Var, [LocArg], Ty2, Exp2)]
forall a. a -> [a] -> [a]
: [(Var, [LocArg], Ty2, Exp2)]
acc))
                       []
                       Set Var
free_wlocs) :: PassM [(Var, [LocArg], Ty2, Exp2)]
      let fn :: (a, b, c, PreExp E2Ext loc dec)
-> PassM (Var, b, c, PreExp E2Ext loc dec)
fn = (\(a
_x,b
locs,c
ty,Ext (SSPush SSModality
a Var
b Var
c TyCon
_)) -> Var -> PassM Var
forall (m :: * -> *). MonadState Int m => Var -> m Var
gensym Var
"ss_pop" PassM Var
-> (Var -> PassM (Var, b, c, PreExp E2Ext loc dec))
-> PassM (Var, b, c, PreExp E2Ext loc dec)
forall a b. PassM a -> (a -> PassM b) -> PassM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Var
y -> (Var, b, c, PreExp E2Ext loc dec)
-> PassM (Var, b, c, PreExp E2Ext loc dec)
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Var
y,b
locs,c
ty,E2Ext loc dec -> PreExp E2Ext loc dec
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (SSModality -> Var -> Var -> E2Ext loc dec
forall loc dec. SSModality -> Var -> Var -> E2Ext loc dec
SSPop SSModality
a Var
b Var
c)))
      [(Var, [LocArg], Ty2, Exp2)]
rpop <- ((Var, [LocArg], Ty2, Exp2) -> PassM (Var, [LocArg], Ty2, Exp2))
-> [(Var, [LocArg], Ty2, Exp2)]
-> PassM [(Var, [LocArg], Ty2, 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, [LocArg], Ty2, Exp2) -> PassM (Var, [LocArg], Ty2, Exp2)
forall {a} {b} {c} {loc} {dec} {loc} {dec}.
(a, b, c, PreExp E2Ext loc dec)
-> PassM (Var, b, c, PreExp E2Ext loc dec)
fn ([(Var, [LocArg], Ty2, Exp2)] -> [(Var, [LocArg], Ty2, Exp2)]
forall a. [a] -> [a]
reverse [(Var, [LocArg], Ty2, Exp2)]
rpush)
      [(Var, [LocArg], Ty2, Exp2)]
wpop <- ((Var, [LocArg], Ty2, Exp2) -> PassM (Var, [LocArg], Ty2, Exp2))
-> [(Var, [LocArg], Ty2, Exp2)]
-> PassM [(Var, [LocArg], Ty2, 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, [LocArg], Ty2, Exp2) -> PassM (Var, [LocArg], Ty2, Exp2)
forall {a} {b} {c} {loc} {dec} {loc} {dec}.
(a, b, c, PreExp E2Ext loc dec)
-> PassM (Var, b, c, PreExp E2Ext loc dec)
fn ([(Var, [LocArg], Ty2, Exp2)] -> [(Var, [LocArg], Ty2, Exp2)]
forall a. [a] -> [a]
reverse [(Var, [LocArg], Ty2, Exp2)]
wpush)
      ([(Var, [LocArg], Ty2, Exp2)], [(Var, [LocArg], Ty2, Exp2)],
 [(Var, [LocArg], Ty2, Exp2)], [(Var, [LocArg], Ty2, Exp2)])
-> PassM
     ([(Var, [LocArg], Ty2, Exp2)], [(Var, [LocArg], Ty2, Exp2)],
      [(Var, [LocArg], Ty2, Exp2)], [(Var, [LocArg], Ty2, Exp2)])
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([(Var, [LocArg], Ty2, Exp2)]
rpush,[(Var, [LocArg], Ty2, Exp2)]
wpush,[(Var, [LocArg], Ty2, Exp2)]
rpop,[(Var, [LocArg], Ty2, Exp2)]
wpop)


-- Inspect an AST and return locations in a RetE form.
findRetLocs :: Exp2 -> [LocArg]
findRetLocs :: Exp2 -> [LocArg]
findRetLocs Exp2
e0 = Exp2 -> [LocArg] -> [LocArg]
go Exp2
e0 []
  where
    go :: Exp2 -> [LocArg] -> [LocArg]
    go :: Exp2 -> [LocArg] -> [LocArg]
go Exp2
ex [LocArg]
acc =
      case Exp2
ex of
        VarE{}    -> [LocArg]
acc
        LitE{}    -> [LocArg]
acc
        CharE{}   -> [LocArg]
acc
        FloatE{}  -> [LocArg]
acc
        LitSymE{} -> [LocArg]
acc
        AppE Var
_ [LocArg]
_ [Exp2]
args   -> (Exp2 -> [LocArg] -> [LocArg]) -> [LocArg] -> [Exp2] -> [LocArg]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Exp2 -> [LocArg] -> [LocArg]
go [LocArg]
acc [Exp2]
args
        PrimAppE Prim Ty2
_ [Exp2]
args -> (Exp2 -> [LocArg] -> [LocArg]) -> [LocArg] -> [Exp2] -> [LocArg]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Exp2 -> [LocArg] -> [LocArg]
go [LocArg]
acc [Exp2]
args
        LetE (Var
_,[LocArg]
_,Ty2
_,Exp2
rhs) Exp2
bod -> do
          (Exp2 -> [LocArg] -> [LocArg]) -> [LocArg] -> [Exp2] -> [LocArg]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Exp2 -> [LocArg] -> [LocArg]
go [LocArg]
acc [Exp2
rhs,Exp2
bod]
        IfE Exp2
a Exp2
b Exp2
c  -> (Exp2 -> [LocArg] -> [LocArg]) -> [LocArg] -> [Exp2] -> [LocArg]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Exp2 -> [LocArg] -> [LocArg]
go [LocArg]
acc [Exp2
a,Exp2
b,Exp2
c]
        MkProdE [Exp2]
xs -> (Exp2 -> [LocArg] -> [LocArg]) -> [LocArg] -> [Exp2] -> [LocArg]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Exp2 -> [LocArg] -> [LocArg]
go [LocArg]
acc [Exp2]
xs
        ProjE Int
_ Exp2
e  -> Exp2 -> [LocArg] -> [LocArg]
go Exp2
e [LocArg]
acc
        DataConE LocArg
_ TyCon
_ [Exp2]
args -> (Exp2 -> [LocArg] -> [LocArg]) -> [LocArg] -> [Exp2] -> [LocArg]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Exp2 -> [LocArg] -> [LocArg]
go [LocArg]
acc [Exp2]
args
        CaseE Exp2
_ [(TyCon, [(Var, LocArg)], Exp2)]
mp ->
          ((TyCon, [(Var, LocArg)], Exp2) -> [LocArg] -> [LocArg])
-> [LocArg] -> [(TyCon, [(Var, LocArg)], Exp2)] -> [LocArg]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\(TyCon
_,[(Var, LocArg)]
_,Exp2
c) [LocArg]
acc2 -> Exp2 -> [LocArg] -> [LocArg]
go Exp2
c [LocArg]
acc2) [LocArg]
acc [(TyCon, [(Var, LocArg)], Exp2)]
mp
        TimeIt Exp2
e Ty2
_ty Bool
_b  -> Exp2 -> [LocArg] -> [LocArg]
go Exp2
e [LocArg]
acc
        WithArenaE Var
_v Exp2
e -> Exp2 -> [LocArg] -> [LocArg]
go Exp2
e [LocArg]
acc
        SpawnE{} -> [LocArg]
acc
        SyncE{}  -> [LocArg]
acc
        Ext E2Ext LocArg Ty2
ext ->
          case E2Ext LocArg Ty2
ext of
            LetRegionE Region
_ RegionSize
_ Maybe RegionType
_ Exp2
bod  -> Exp2 -> [LocArg] -> [LocArg]
go Exp2
bod [LocArg]
acc
            LetParRegionE Region
_ RegionSize
_ Maybe RegionType
_ Exp2
bod  -> Exp2 -> [LocArg] -> [LocArg]
go Exp2
bod [LocArg]
acc
            LetLocE Var
_ PreLocExp LocArg
_ Exp2
bod   -> Exp2 -> [LocArg] -> [LocArg]
go Exp2
bod [LocArg]
acc
            StartOfPkdCursor{} -> [LocArg]
acc
            TagCursor{}    -> [LocArg]
acc
            RetE [LocArg]
locs Var
_       -> [LocArg]
locs [LocArg] -> [LocArg] -> [LocArg]
forall a. [a] -> [a] -> [a]
++ [LocArg]
acc
            FromEndE{}        -> [LocArg]
acc
            BoundsCheck{}     -> [LocArg]
acc
            IndirectionE{}    -> [LocArg]
acc
            AddFixed{}        -> [LocArg]
acc
            E2Ext LocArg Ty2
GetCilkWorkerNum  -> [LocArg]
acc
            LetAvail [Var]
_ Exp2
bod    -> Exp2 -> [LocArg] -> [LocArg]
go Exp2
bod [LocArg]
acc
            AllocateTagHere{} -> [LocArg]
acc
            AllocateScalarsHere{} -> [LocArg]
acc
            SSPush{} -> [LocArg]
acc
            SSPop{} -> [LocArg]
acc
        MapE{}  -> TyCon -> [LocArg]
forall a. HasCallStack => TyCon -> a
error TyCon
"findRetLocs: TODO MapE"
        FoldE{}  -> TyCon -> [LocArg]
forall a. HasCallStack => TyCon -> a
error TyCon
"findRetLocs: TODO FoldE"

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

-- Maximal sum of sizes of scalars before the first packed thing in the
-- constructors of this type. The assumption is that whatever writes
-- that packed value will do a bounds check again. Note that only AppE's
-- do boundschecking, DataConE's dont. We should fix this.
boundsCheck :: DDefs2 -> TyCon -> Int
boundsCheck :: DDefs Ty2 -> TyCon -> Int
boundsCheck DDefs Ty2
ddefs TyCon
tycon =
  let dcons :: [TyCon]
dcons = DDefs Ty2 -> TyCon -> [TyCon]
forall a. Out a => DDefs a -> TyCon -> [TyCon]
getConOrdering DDefs Ty2
ddefs TyCon
tycon
      spaceReqd :: t Ty2 -> (Int, Bool)
spaceReqd t Ty2
tys = ((Int, Bool) -> Ty2 -> (Int, Bool))
-> (Int, Bool) -> t Ty2 -> (Int, Bool)
forall b a. (b -> a -> b) -> b -> t a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\(Int
bytes, Bool
seen_packed) Ty2
ty ->
                               if Bool
seen_packed
                               then ( Int
bytes, Bool
seen_packed )
                               else if UrTy Var -> Bool
forall a. Show a => UrTy a -> Bool
hasPacked (Ty2 -> UrTy Var
unTy2 Ty2
ty)
                               then ( Int
bytes, Bool
True )
                               else ( Int
bytes Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Maybe Int -> Int
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Int -> Int) -> Maybe Int -> Int
forall a b. (a -> b) -> a -> b
$ UrTy Var -> Maybe Int
forall a. UrTy a -> Maybe Int
sizeOfTy (Ty2 -> UrTy Var
unTy2 Ty2
ty)), Bool
False ))
                        (Int
0, Bool
False)
                        t Ty2
tys
      tyss :: [[Ty2]]
tyss = (TyCon -> [Ty2]) -> [TyCon] -> [[Ty2]]
forall a b. (a -> b) -> [a] -> [b]
map (DDefs Ty2 -> TyCon -> [Ty2]
forall a. Out a => DDefs a -> TyCon -> [a]
lookupDataCon DDefs Ty2
ddefs) [TyCon]
dcons
      vals :: [Int]
vals = ([Ty2] -> Int) -> [[Ty2]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ((Int, Bool) -> Int
forall a b. (a, b) -> a
fst ((Int, Bool) -> Int) -> ([Ty2] -> (Int, Bool)) -> [Ty2] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Ty2] -> (Int, Bool)
forall {t :: * -> *}. Foldable t => t Ty2 -> (Int, Bool)
spaceReqd) [[Ty2]]
tyss
      -- Add a byte for the tag.
      num_bytes :: Int
num_bytes = (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Int]
vals)
  -- Reserve additional space for a redirection node or a forwarding pointer.
  in Int
num_bytes Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
9

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

-- gFreeVars ++ locations ++ region variables - (args to datacons)
allFreeVars_sans_datacon_args :: Exp2 -> S.Set Var
allFreeVars_sans_datacon_args :: Exp2 -> Set Var
allFreeVars_sans_datacon_args Exp2
ex =
  case Exp2
ex of
    AppE Var
_ [LocArg]
locs [Exp2]
args -> [Var] -> Set Var
forall a. Ord a => [a] -> Set a
S.fromList ((LocArg -> Var) -> [LocArg] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map LocArg -> Var
toLocVar [LocArg]
locs) Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.union` ([Set Var] -> Set Var
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions ((Exp2 -> Set Var) -> [Exp2] -> [Set Var]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Set Var
allFreeVars_sans_datacon_args [Exp2]
args))
    PrimAppE Prim Ty2
_ [Exp2]
args -> ([Set Var] -> Set Var
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions ((Exp2 -> Set Var) -> [Exp2] -> [Set Var]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Set Var
allFreeVars_sans_datacon_args [Exp2]
args))
    LetE (Var
v,[LocArg]
locs,Ty2
_,Exp2
rhs) Exp2
bod -> ([Var] -> Set Var
forall a. Ord a => [a] -> Set a
S.fromList ((LocArg -> Var) -> [LocArg] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map LocArg -> Var
toLocVar [LocArg]
locs) Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.union` (Exp2 -> Set Var
allFreeVars_sans_datacon_args Exp2
rhs) Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.union` (Exp2 -> Set Var
allFreeVars_sans_datacon_args Exp2
bod))
                               Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` Var -> Set Var
forall a. a -> Set a
S.singleton Var
v
    IfE Exp2
a Exp2
b Exp2
c -> Exp2 -> Set Var
allFreeVars_sans_datacon_args Exp2
a Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.union` Exp2 -> Set Var
allFreeVars_sans_datacon_args Exp2
b Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.union` Exp2 -> Set Var
allFreeVars_sans_datacon_args Exp2
c
    MkProdE [Exp2]
args -> ([Set Var] -> Set Var
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions ((Exp2 -> Set Var) -> [Exp2] -> [Set Var]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Set Var
allFreeVars_sans_datacon_args [Exp2]
args))
    ProjE Int
_ Exp2
bod -> Exp2 -> Set Var
allFreeVars_sans_datacon_args Exp2
bod
    CaseE Exp2
scrt [(TyCon, [(Var, LocArg)], Exp2)]
brs -> (Exp2 -> Set Var
allFreeVars_sans_datacon_args Exp2
scrt) Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.union` ([Set Var] -> Set Var
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions (((TyCon, [(Var, LocArg)], Exp2) -> Set Var)
-> [(TyCon, [(Var, LocArg)], Exp2)] -> [Set Var]
forall a b. (a -> b) -> [a] -> [b]
map (\(TyCon
_,[(Var, LocArg)]
vlocs,Exp2
c) -> Exp2 -> Set Var
allFreeVars_sans_datacon_args Exp2
c Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.difference`
                                                                                   [Var] -> Set Var
forall a. Ord a => [a] -> Set a
S.fromList (((Var, LocArg) -> Var) -> [(Var, LocArg)] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map (Var, LocArg) -> Var
forall a b. (a, b) -> a
fst [(Var, LocArg)]
vlocs) Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.difference`
                                                                                   [Var] -> Set Var
forall a. Ord a => [a] -> Set a
S.fromList (((Var, LocArg) -> Var) -> [(Var, LocArg)] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map (LocArg -> Var
toLocVar (LocArg -> Var)
-> ((Var, LocArg) -> LocArg) -> (Var, LocArg) -> Var
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Var, LocArg) -> LocArg
forall a b. (a, b) -> b
snd) [(Var, LocArg)]
vlocs))
                                                                  [(TyCon, [(Var, LocArg)], Exp2)]
brs))
    DataConE LocArg
loc TyCon
_ [Exp2]
_args -> Var -> Set Var
forall a. a -> Set a
S.singleton (LocArg -> Var
toLocVar LocArg
loc)
    TimeIt Exp2
e Ty2
_ Bool
_ -> Exp2 -> Set Var
allFreeVars_sans_datacon_args Exp2
e
    WithArenaE Var
_ Exp2
e -> Exp2 -> Set Var
allFreeVars_sans_datacon_args Exp2
e
    SpawnE Var
_ [LocArg]
locs [Exp2]
args -> [Var] -> Set Var
forall a. Ord a => [a] -> Set a
S.fromList ((LocArg -> Var) -> [LocArg] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map LocArg -> Var
toLocVar [LocArg]
locs) Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.union` ([Set Var] -> Set Var
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions ((Exp2 -> Set Var) -> [Exp2] -> [Set Var]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Set Var
allFreeVars_sans_datacon_args [Exp2]
args))
    Ext E2Ext LocArg Ty2
ext ->
      case E2Ext LocArg Ty2
ext of
        LetRegionE Region
r RegionSize
_sz Maybe RegionType
_ty Exp2
bod -> Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.delete (Region -> Var
regionToVar Region
r) (Exp2 -> Set Var
allFreeVars_sans_datacon_args Exp2
bod)
        LetParRegionE Region
r RegionSize
_sz Maybe RegionType
_ty Exp2
bod -> Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.delete (Region -> Var
regionToVar Region
r) (Exp2 -> Set Var
allFreeVars_sans_datacon_args Exp2
bod)
        LetLocE Var
loc PreLocExp LocArg
locexp Exp2
bod -> Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.delete Var
loc (Exp2 -> Set Var
allFreeVars_sans_datacon_args Exp2
bod Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.union` PreLocExp LocArg -> Set Var
forall a. FreeVars a => a -> Set Var
gFreeVars PreLocExp LocArg
locexp)
        StartOfPkdCursor Var
cur   -> Var -> Set Var
forall a. a -> Set a
S.singleton Var
cur
        TagCursor Var
a Var
b-> [Var] -> Set Var
forall a. Ord a => [a] -> Set a
S.fromList [Var
a,Var
b]
        RetE [LocArg]
locs Var
v     -> Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.insert Var
v ([Var] -> Set Var
forall a. Ord a => [a] -> Set a
S.fromList ((LocArg -> Var) -> [LocArg] -> [Var]
forall a b. (a -> b) -> [a] -> [b]
map LocArg -> Var
toLocVar [LocArg]
locs))
        FromEndE LocArg
loc    -> Var -> Set Var
forall a. a -> Set a
S.singleton (LocArg -> Var
toLocVar LocArg
loc)
        BoundsCheck Int
_ LocArg
reg LocArg
cur -> [Var] -> Set Var
forall a. Ord a => [a] -> Set a
S.fromList [LocArg -> Var
toLocVar LocArg
reg, LocArg -> Var
toLocVar LocArg
cur]
        IndirectionE TyCon
_ TyCon
_ (LocArg
a,LocArg
b) (LocArg
c,LocArg
d) Exp2
_ -> [Var] -> Set Var
forall a. Ord a => [a] -> Set a
S.fromList ([Var] -> Set Var) -> [Var] -> Set Var
forall a b. (a -> b) -> a -> b
$ [LocArg -> Var
toLocVar LocArg
a, LocArg -> Var
toLocVar LocArg
b, LocArg -> Var
toLocVar LocArg
c, LocArg -> Var
toLocVar LocArg
d]
        AddFixed Var
v Int
_    -> Var -> Set Var
forall a. a -> Set a
S.singleton Var
v
        E2Ext LocArg Ty2
GetCilkWorkerNum-> Set Var
forall a. Set a
S.empty
        LetAvail [Var]
vs Exp2
bod -> [Var] -> Set Var
forall a. Ord a => [a] -> Set a
S.fromList [Var]
vs Set Var -> Set Var -> Set Var
forall a. Ord a => Set a -> Set a -> Set a
`S.union` Exp2 -> Set Var
forall a. FreeVars a => a -> Set Var
gFreeVars Exp2
bod
        AllocateTagHere Var
loc TyCon
_ -> Var -> Set Var
forall a. a -> Set a
S.singleton Var
loc
        AllocateScalarsHere Var
loc -> Var -> Set Var
forall a. a -> Set a
S.singleton Var
loc
        SSPush SSModality
_ Var
a Var
b TyCon
_ -> [Var] -> Set Var
forall a. Ord a => [a] -> Set a
S.fromList [Var
a,Var
b]
        SSPop SSModality
_ Var
a Var
b -> [Var] -> Set Var
forall a. Ord a => [a] -> Set a
S.fromList [Var
a,Var
b]
    Exp2
_ -> Exp2 -> Set Var
forall a. FreeVars a => a -> Set Var
gFreeVars Exp2
ex


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

substEndReg :: Either LocVar RegVar -> RegVar -> Exp2 -> Exp2
substEndReg :: Either Var Var -> Var -> Exp2 -> Exp2
substEndReg Either Var Var
loc_or_reg Var
end_reg Exp2
ex =
  case Exp2
ex of
    AppE Var
f [LocArg]
locs [Exp2]
args -> Var -> [LocArg] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE Var
f ((LocArg -> LocArg) -> [LocArg] -> [LocArg]
forall a b. (a -> b) -> [a] -> [b]
map LocArg -> LocArg
gosubst [LocArg]
locs) ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Exp2
go [Exp2]
args)
    PrimAppE Prim Ty2
pr [Exp2]
args -> Prim Ty2 -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Prim dec -> [PreExp ext loc dec] -> PreExp ext loc dec
PrimAppE Prim Ty2
pr ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Exp2
go [Exp2]
args)
    LetE (Var
v,[LocArg]
locs,Ty2
ty,Exp2
rhs) Exp2
bod -> (Var, [LocArg], 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,(LocArg -> LocArg) -> [LocArg] -> [LocArg]
forall a b. (a -> b) -> [a] -> [b]
map LocArg -> LocArg
gosubst [LocArg]
locs,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]
args             -> [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]
map Exp2 -> Exp2
go [Exp2]
args)
    ProjE Int
i Exp2
bod              -> Int -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
Int -> PreExp ext loc dec -> PreExp ext loc dec
ProjE Int
i (Exp2 -> Exp2
go Exp2
bod)
    CaseE Exp2
scrt [(TyCon, [(Var, LocArg)], Exp2)]
brs           -> Exp2 -> [(TyCon, [(Var, LocArg)], Exp2)] -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> [(TyCon, [(Var, loc)], PreExp ext loc dec)]
-> PreExp ext loc dec
CaseE (Exp2 -> Exp2
go Exp2
scrt) (((TyCon, [(Var, LocArg)], Exp2) -> (TyCon, [(Var, LocArg)], Exp2))
-> [(TyCon, [(Var, LocArg)], Exp2)]
-> [(TyCon, [(Var, LocArg)], Exp2)]
forall a b. (a -> b) -> [a] -> [b]
map (\(TyCon
dcon,[(Var, LocArg)]
vlocs,Exp2
c) -> (TyCon
dcon,[(Var, LocArg)]
vlocs,Exp2 -> Exp2
go Exp2
c)) [(TyCon, [(Var, LocArg)], Exp2)]
brs)
    DataConE LocArg
loc TyCon
dcon [Exp2]
args   -> LocArg -> TyCon -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
loc -> TyCon -> [PreExp ext loc dec] -> PreExp ext loc dec
DataConE (LocArg -> LocArg
gosubst LocArg
loc) TyCon
dcon ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Exp2
go [Exp2]
args)
    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 Var
f [LocArg]
locs [Exp2]
args -> Var -> [LocArg] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
SpawnE Var
f ((LocArg -> LocArg) -> [LocArg] -> [LocArg]
forall a b. (a -> b) -> [a] -> [b]
map LocArg -> LocArg
gosubst [LocArg]
locs) ((Exp2 -> Exp2) -> [Exp2] -> [Exp2]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Exp2
go [Exp2]
args)
    Ext E2Ext LocArg Ty2
ext ->
      case E2Ext LocArg Ty2
ext of
        LetRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
bod    -> E2Ext LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocArg Ty2 -> Exp2) -> E2Ext LocArg Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Region
-> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext LocArg 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
bod)
        LetParRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
bod -> E2Ext LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocArg Ty2 -> Exp2) -> E2Ext LocArg Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Region
-> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext LocArg 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
bod)
        LetLocE Var
loc PreLocExp LocArg
locexp Exp2
bod    -> E2Ext LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocArg Ty2 -> Exp2) -> E2Ext LocArg Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Var -> PreLocExp LocArg -> Exp2 -> E2Ext LocArg Ty2
forall loc dec. Var -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE Var
loc PreLocExp LocArg
locexp (Exp2 -> Exp2
go Exp2
bod)
        RetE [LocArg]
locs Var
v               -> E2Ext LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocArg Ty2 -> Exp2) -> E2Ext LocArg Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ [LocArg] -> Var -> E2Ext LocArg Ty2
forall loc dec. [loc] -> Var -> E2Ext loc dec
RetE ((LocArg -> LocArg) -> [LocArg] -> [LocArg]
forall a b. (a -> b) -> [a] -> [b]
map LocArg -> LocArg
gosubst [LocArg]
locs) Var
v
        FromEndE LocArg
loc              -> E2Ext LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocArg Ty2 -> Exp2) -> E2Ext LocArg Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ LocArg -> E2Ext LocArg Ty2
forall loc dec. loc -> E2Ext loc dec
FromEndE (LocArg -> LocArg
gosubst LocArg
loc)
        BoundsCheck Int
i LocArg
reg LocArg
cur     -> E2Ext LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocArg Ty2 -> Exp2) -> E2Ext LocArg Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ Int -> LocArg -> LocArg -> E2Ext LocArg Ty2
forall loc dec. Int -> loc -> loc -> E2Ext loc dec
BoundsCheck Int
i (LocArg -> LocArg
gosubst LocArg
reg) (LocArg -> LocArg
gosubst LocArg
cur)
        IndirectionE TyCon
tycon TyCon
dcon (LocArg
a,LocArg
b) (LocArg
c,LocArg
d) Exp2
e ->
          E2Ext LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocArg Ty2 -> Exp2) -> E2Ext LocArg Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ TyCon
-> TyCon
-> (LocArg, LocArg)
-> (LocArg, LocArg)
-> Exp2
-> E2Ext LocArg Ty2
forall loc dec.
TyCon
-> TyCon -> (loc, loc) -> (loc, loc) -> E2 loc dec -> E2Ext loc dec
IndirectionE TyCon
tycon TyCon
dcon (LocArg -> LocArg
gosubst LocArg
a, LocArg -> LocArg
gosubst LocArg
b) (LocArg -> LocArg
gosubst LocArg
c, LocArg -> LocArg
gosubst LocArg
d) Exp2
e
        LetAvail [Var]
vs Exp2
bod -> E2Ext LocArg Ty2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocArg Ty2 -> Exp2) -> E2Ext LocArg Ty2 -> Exp2
forall a b. (a -> b) -> a -> b
$ [Var] -> Exp2 -> E2Ext LocArg Ty2
forall loc dec. [Var] -> E2 loc dec -> E2Ext loc dec
LetAvail [Var]
vs (Exp2 -> Exp2
go Exp2
bod)
        StartOfPkdCursor{}    -> Exp2
ex
        TagCursor{}           -> Exp2
ex
        AddFixed{}            -> Exp2
ex
        E2Ext LocArg Ty2
GetCilkWorkerNum      -> Exp2
ex
        AllocateTagHere{}     -> Exp2
ex
        AllocateScalarsHere{} -> Exp2
ex
        SSPush{}              -> Exp2
ex
        SSPop{}               -> Exp2
ex
    Exp2
_ -> Exp2
ex
 where
  go :: Exp2 -> Exp2
go = Either Var Var -> Var -> Exp2 -> Exp2
substEndReg Either Var Var
loc_or_reg Var
end_reg
  gosubst :: LocArg -> LocArg
gosubst = Either Var Var -> Var -> LocArg -> LocArg
substEndReg_locarg Either Var Var
loc_or_reg Var
end_reg

substEndReg_locarg :: Either LocVar RegVar -> RegVar -> LocArg -> LocArg
substEndReg_locarg :: Either Var Var -> Var -> LocArg -> LocArg
substEndReg_locarg Either Var Var
loc_or_reg Var
end_reg LocArg
locarg =
  case LocArg
locarg of
    Loc LREM
lrem          -> case Either Var Var
loc_or_reg of
                           Left Var
loc0 ->  if LREM -> Var
lremLoc LREM
lrem Var -> Var -> Bool
forall a. Eq a => a -> a -> Bool
== Var
loc0
                                         then LREM -> LocArg
Loc (LREM
lrem { lremEndReg :: Var
lremEndReg = Var
end_reg })
                                         else LocArg
locarg
                           Right Var
reg0 -> if LREM -> Var
lremReg LREM
lrem Var -> Var -> Bool
forall a. Eq a => a -> a -> Bool
== Var
reg0
                                         then LREM -> LocArg
Loc (LREM
lrem { lremEndReg :: Var
lremEndReg = Var
end_reg })
                                         else LocArg
locarg
    EndWitness LREM
lrem Var
v -> case Either Var Var
loc_or_reg of
                           Left Var
loc0 ->  if LREM -> Var
lremLoc LREM
lrem Var -> Var -> Bool
forall a. Eq a => a -> a -> Bool
== Var
loc0
                                         then LREM -> Var -> LocArg
EndWitness (LREM
lrem { lremEndReg :: Var
lremEndReg = Var
end_reg }) Var
v
                                         else LocArg
locarg
                           Right Var
reg0 -> if LREM -> Var
lremReg LREM
lrem Var -> Var -> Bool
forall a. Eq a => a -> a -> Bool
== Var
reg0
                                         then LREM -> Var -> LocArg
EndWitness (LREM
lrem { lremEndReg :: Var
lremEndReg = Var
end_reg }) Var
v
                                         else LocArg
locarg
    Reg{}             -> LocArg
locarg
    EndOfReg{}        -> LocArg
locarg
    EndOfReg_Tagged{} -> LocArg
locarg