-- | Infer region multiplicities
--
--   During inference, regions are merely annotated with a region metavariable.
--   InferRegionScope takes the next step and decides the region scope (global/dynamic)
--   and also assigns it a multiplicity.

module Gibbon.Passes.InferRegionScope
  (inferRegScope, inferRegScopeExp) where

import Data.Graph
import qualified Data.Map as M

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

-- All regions are "infinite" right now

-- | Infer region scope and multiplicity for a program annotated with regions
-- and locations.
inferRegScope :: Prog2 -> PassM Prog2
inferRegScope :: Prog2 -> PassM Prog2
inferRegScope 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 FunDef2 -> PassM FunDef2
inferRegScopeFun ([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'
  Maybe (Exp2, UrTy Var)
mainExp' <- case Maybe (Exp2, TyOf Exp2)
mainExp of
                Maybe (Exp2, TyOf Exp2)
Nothing -> Maybe (Exp2, UrTy Var) -> PassM (Maybe (Exp2, UrTy Var))
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Exp2, UrTy Var)
forall a. Maybe a
Nothing
                Just (Exp2
mn, TyOf Exp2
ty) -> (Exp2, UrTy Var) -> Maybe (Exp2, UrTy Var)
forall a. a -> Maybe a
Just ((Exp2, UrTy Var) -> Maybe (Exp2, UrTy Var))
-> (Exp2 -> (Exp2, UrTy Var)) -> Exp2 -> Maybe (Exp2, UrTy Var)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (,TyOf Exp2
UrTy Var
ty) (Exp2 -> Maybe (Exp2, UrTy Var))
-> PassM Exp2 -> PassM (Maybe (Exp2, UrTy Var))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Region Region -> Exp2 -> PassM Exp2
inferRegScopeExp Map Region Region
forall k a. Map k a
M.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, UrTy Var)
mainExp'

inferRegScopeFun :: FunDef2 -> PassM FunDef2
inferRegScopeFun :: FunDef2 -> PassM FunDef2
inferRegScopeFun f :: FunDef2
f@FunDef{Exp2
funBody :: Exp2
funBody :: forall ex. FunDef ex -> ex
funBody} = do
  Exp2
funBody' <- Map Region Region -> Exp2 -> PassM Exp2
inferRegScopeExp Map Region Region
forall k a. Map k a
M.empty Exp2
funBody
  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
funBody'}

{- Region scoping rules:
~~~~~~~~~~~~~~~~~~~~~~~~

A region in Gibbon can either be global or dynamic (local). Global regions have a
lifetime equal to that of the whole program. Whereas, dynamic regions can be created &
destroyed on the stack. The scope of the values in the region dictate it's scope.
If the value in the region escapes the context it was created in, it has be global.
Otherwise, a dynamic region would be sufficient. Consider two functions of the form:

    fnA = letregion ra in
          letloc la in
          let x = (Leaf 1) at la
          in x

    fnB = letregion rb in
          letloc lb in
          let y = (Leaf 1) at lb
          in 1

The return value of fnA resides in the region `ra`. Therefore `ra` has to be global.
Whereas, it's safe to mark the region `rb` as dynamic. The fn `inferRegScopeExp` traverses
the AST and transforms all region metavariables to fully scoped region variables.

Whenever it encounters a `letregion`, it builds a graph of dependencies of the
relevant part of the AST. And if there's a path between the region and
the end of the AST, the return value so to speak, the region is marked global.
For example, in fnA, it discovers the path `ra -> la -> x` and hence `ra` is global.
In fnB, there's no path from `rb` to 1.

-}

-- | Decide if a region should be global or local (dynamic).
--
--  Dynamic regions are stack allocated and automatically freed
inferRegScopeExp :: M.Map Region Region -> Exp2 -> PassM Exp2
inferRegScopeExp :: Map Region Region -> Exp2 -> PassM Exp2
inferRegScopeExp Map Region Region
env Exp2
ex =
  case Exp2
ex of
    Ext E2Ext Var (UrTy Var)
ext ->
      case E2Ext Var (UrTy Var)
ext of
        AddFixed{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
        LetRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
rhs ->
          case Region
r of
            MMapR{} -> E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var (UrTy Var) -> Exp2)
-> (Exp2 -> E2Ext Var (UrTy Var)) -> Exp2 -> Exp2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Region
-> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetRegionE Region
r RegionSize
sz Maybe RegionType
ty (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp2 -> PassM Exp2
go Exp2
rhs
            Region
_ ->
              let deps :: [(Var, Var, [Var])]
deps = Exp2 -> [(Var, Var, [Var])]
depList Exp2
ex
              in case [(Var, Var, [Var])]
deps of
                   ((Var
retVar,Var
_,[Var]
_):[(Var, Var, [Var])]
_) ->
                     let (Graph
g,Vertex -> (Var, Var, [Var])
_,Var -> Maybe Vertex
vtxF) = [(Var, Var, [Var])]
-> (Graph, Vertex -> (Var, Var, [Var]), Var -> Maybe Vertex)
forall key node.
Ord key =>
[(node, key, [key])]
-> (Graph, Vertex -> (node, key, [key]), key -> Maybe Vertex)
graphFromEdges [(Var, Var, [Var])]
deps
                         regV :: Var
regV = Region -> Var
regionToVar Region
r
                         -- Vertex of the region variable
                         regVertex :: Vertex
regVertex =
                           case Var -> Maybe Vertex
vtxF Var
regV of
                             Just Vertex
x  -> Vertex
x
                             Maybe Vertex
Nothing -> [Char] -> Vertex
forall a. HasCallStack => [Char] -> a
error ([Char] -> Vertex) -> [Char] -> Vertex
forall a b. (a -> b) -> a -> b
$ [Char]
"No vertex for:" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Region -> [Char]
forall a. Out a => a -> [Char]
sdoc Region
r
                         -- Vertex of the return value
                         retVertex :: Vertex
retVertex =
                           case Var -> Maybe Vertex
vtxF Var
retVar of
                             Just Vertex
x  -> Vertex
x
                             Maybe Vertex
Nothing -> [Char] -> Vertex
forall a. HasCallStack => [Char] -> a
error ([Char] -> Vertex) -> [Char] -> Vertex
forall a b. (a -> b) -> a -> b
$ [Char]
"No vertex for:" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Var -> [Char]
forall a. Out a => a -> [Char]
sdoc Var
retVar
                         -- The value in the region  escapes the current scope if there's
                         -- a path between the region variable and the thing returned.
                         -- TODO: Warn the user when this happens in a fn ?
                     in do DynFlags
dflags <- PassM DynFlags
forall (m :: * -> *). MonadReader Config m => m DynFlags
getDynFlags
                           let defaultMul :: Multiplicity
defaultMul = if (GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_BigInfiniteRegions DynFlags
dflags) Bool -> Bool -> Bool
||
                                               (GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_Gibbon1 DynFlags
dflags)
                                            then Multiplicity
BigInfinite
                                            else Multiplicity
Infinite
                           let scoped_reg :: Region
scoped_reg = if Graph -> Vertex -> Vertex -> Bool
path Graph
g Vertex
retVertex Vertex
regVertex
                                            then (Var -> Multiplicity -> Region
GlobR Var
regV Multiplicity
defaultMul)
                                            -- [2018.03.30] - TEMP: Turning off scoped buffers.
                                            -- else Ext$ LetRegionE (DynR regV mul) (inferRegScopeExp rhs)
                                            -- else (DynR regV mul)
                                            else (Var -> Multiplicity -> Region
GlobR Var
regV Multiplicity
defaultMul)
                           E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var (UrTy Var) -> Exp2)
-> (Exp2 -> E2Ext Var (UrTy Var)) -> Exp2 -> Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
                             Region
-> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetRegionE Region
scoped_reg RegionSize
Undefined Maybe RegionType
forall a. Maybe a
Nothing (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
                             Map Region Region -> Exp2 -> PassM Exp2
inferRegScopeExp (Region -> Region -> Map Region Region -> Map Region Region
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Region
r Region
scoped_reg Map Region Region
env) Exp2
rhs
                   [] -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex

        LetParRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
rhs ->
          case Region
r of
            MMapR{} -> E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var (UrTy Var) -> Exp2)
-> (Exp2 -> E2Ext Var (UrTy Var)) -> Exp2 -> Exp2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Region
-> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetParRegionE Region
r RegionSize
sz Maybe RegionType
ty (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp2 -> PassM Exp2
go Exp2
rhs
            Region
_ ->
              let deps :: [(Var, Var, [Var])]
deps = Exp2 -> [(Var, Var, [Var])]
depList Exp2
ex
              in case [(Var, Var, [Var])]
deps of
                   ((Var
retVar,Var
_,[Var]
_):[(Var, Var, [Var])]
_) ->
                     let (Graph
g,Vertex -> (Var, Var, [Var])
_,Var -> Maybe Vertex
vtxF) = [(Var, Var, [Var])]
-> (Graph, Vertex -> (Var, Var, [Var]), Var -> Maybe Vertex)
forall key node.
Ord key =>
[(node, key, [key])]
-> (Graph, Vertex -> (node, key, [key]), key -> Maybe Vertex)
graphFromEdges [(Var, Var, [Var])]
deps
                         regV :: Var
regV = Region -> Var
regionToVar Region
r
                         -- Vertex of the region variable
                         regVertex :: Vertex
regVertex =
                           case Var -> Maybe Vertex
vtxF Var
regV of
                             Just Vertex
x  -> Vertex
x
                             Maybe Vertex
Nothing -> [Char] -> Vertex
forall a. HasCallStack => [Char] -> a
error ([Char] -> Vertex) -> [Char] -> Vertex
forall a b. (a -> b) -> a -> b
$ [Char]
"No vertex for:" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Region -> [Char]
forall a. Out a => a -> [Char]
sdoc Region
r
                         -- Vertex of the return value
                         retVertex :: Vertex
retVertex =
                           case Var -> Maybe Vertex
vtxF Var
retVar of
                             Just Vertex
x  -> Vertex
x
                             Maybe Vertex
Nothing -> [Char] -> Vertex
forall a. HasCallStack => [Char] -> a
error ([Char] -> Vertex) -> [Char] -> Vertex
forall a b. (a -> b) -> a -> b
$ [Char]
"No vertex for:" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Var -> [Char]
forall a. Out a => a -> [Char]
sdoc Var
retVar
                         -- The value in the region  escapes the current scope if there's
                         -- a path between the region variable and the thing returned.
                         -- TODO: Warn the user when this happens in a fn ?
                     in do DynFlags
dflags <- PassM DynFlags
forall (m :: * -> *). MonadReader Config m => m DynFlags
getDynFlags
                           let defaultMul :: Multiplicity
defaultMul = if (GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_BigInfiniteRegions DynFlags
dflags) Bool -> Bool -> Bool
||
                                               (GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_Gibbon1 DynFlags
dflags)
                                            then Multiplicity
BigInfinite
                                            else Multiplicity
Infinite
                           let scoped_reg :: Region
scoped_reg = if Graph -> Vertex -> Vertex -> Bool
path Graph
g Vertex
retVertex Vertex
regVertex
                                            then (Var -> Multiplicity -> Region
GlobR Var
regV Multiplicity
defaultMul)
                                            -- [2018.03.30] - TEMP: Turning off scoped buffers.
                                            -- else Ext$ LetRegionE (DynR regV mul) (inferRegScopeExp rhs)
                                            -- else (DynR regV mul)
                                            else (Var -> Multiplicity -> Region
GlobR Var
regV Multiplicity
defaultMul)
                           E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var (UrTy Var) -> Exp2)
-> (Exp2 -> E2Ext Var (UrTy Var)) -> Exp2 -> Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Region
-> RegionSize -> Maybe RegionType -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetParRegionE (Var -> Multiplicity -> Region
GlobR Var
regV Multiplicity
defaultMul) RegionSize
Undefined Maybe RegionType
forall a. Maybe a
Nothing (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
                                   (Map Region Region -> Exp2 -> PassM Exp2
inferRegScopeExp (Region -> Region -> Map Region Region -> Map Region Region
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Region
r Region
scoped_reg Map Region Region
env) Exp2
rhs)
                   [] -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex

        -- Straightforward recursion
        LetLocE Var
loc PreLocExp Var
le Exp2
bod -> do
          let le' :: PreLocExp Var
le' = case PreLocExp Var
le of
                      StartOfRegionLE Region
r -> Region -> PreLocExp Var
forall loc. Region -> PreLocExp loc
StartOfRegionLE (Map Region Region
env Map Region Region -> Region -> Region
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Region
r)
                      InRegionLE Region
r -> Region -> PreLocExp Var
forall loc. Region -> PreLocExp loc
InRegionLE (Map Region Region
env Map Region Region -> Region -> Region
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# Region
r)
                      PreLocExp Var
_ -> PreLocExp Var
le
          E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var (UrTy Var) -> Exp2)
-> (Exp2 -> E2Ext Var (UrTy Var)) -> Exp2 -> Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Var -> PreLocExp Var -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec. Var -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE Var
loc PreLocExp Var
le' (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp2 -> PassM Exp2
go Exp2
bod)
        StartOfPkdCursor{}-> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
        TagCursor{}-> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
        RetE{}     -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
        FromEndE{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
        BoundsCheck{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
        IndirectionE{}-> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
        E2Ext Var (UrTy Var)
GetCilkWorkerNum -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
        LetAvail [Var]
vs Exp2
e    -> E2Ext Var (UrTy Var) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext Var (UrTy Var) -> Exp2)
-> (Exp2 -> E2Ext Var (UrTy Var)) -> Exp2 -> Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Var] -> Exp2 -> E2Ext Var (UrTy Var)
forall loc dec. [Var] -> E2 loc dec -> E2Ext loc dec
LetAvail [Var]
vs (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp2 -> PassM Exp2
go Exp2
e
        AllocateTagHere{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
        AllocateScalarsHere{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
        SSPush{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
        SSPop{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex

    -- Straightforward recursion ...
    VarE{}     -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
    LitE{}     -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
    CharE{}    -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
    FloatE{}   -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
    LitSymE{}  -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
    AppE{}     -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
    PrimAppE{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
    DataConE{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex
    ProjE Vertex
i Exp2
e  -> Vertex -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
Vertex -> PreExp ext loc dec -> PreExp ext loc dec
ProjE Vertex
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
a) (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
b PassM (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall a b. PassM (a -> b) -> PassM a -> PassM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp2 -> PassM Exp2
go Exp2
c
    MkProdE [Exp2]
ls -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
[PreExp ext loc dec] -> PreExp ext loc dec
MkProdE ([Exp2] -> Exp2) -> PassM [Exp2] -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp2 -> PassM Exp2) -> [Exp2] -> PassM [Exp2]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Exp2 -> PassM Exp2
go [Exp2]
ls
    LetE (Var
v,[Var]
locs,UrTy Var
ty,Exp2
rhs) Exp2
bod -> (Var, [Var], UrTy Var, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(Var, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE ((Var, [Var], UrTy Var, Exp2) -> Exp2 -> Exp2)
-> (Exp2 -> (Var, [Var], UrTy Var, Exp2)) -> Exp2 -> Exp2 -> Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Var
v,[Var]
locs,UrTy Var
ty,) (Exp2 -> Exp2 -> Exp2) -> PassM Exp2 -> PassM (Exp2 -> Exp2)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp2 -> PassM Exp2
go Exp2
rhs PassM (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall a b. PassM (a -> b) -> PassM a -> PassM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp2 -> PassM Exp2
go Exp2
bod
    CaseE Exp2
scrt [([Char], [(Var, Var)], Exp2)]
mp -> (Exp2 -> [([Char], [(Var, Var)], Exp2)] -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> [([Char], [(Var, loc)], PreExp ext loc dec)]
-> PreExp ext loc dec
CaseE Exp2
scrt) ([([Char], [(Var, Var)], Exp2)] -> Exp2)
-> PassM [([Char], [(Var, Var)], Exp2)] -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (([Char], [(Var, Var)], Exp2)
 -> PassM ([Char], [(Var, Var)], Exp2))
-> [([Char], [(Var, Var)], Exp2)]
-> PassM [([Char], [(Var, Var)], Exp2)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (\([Char]
a,[(Var, Var)]
b,Exp2
c) -> ([Char]
a,[(Var, Var)]
b,) (Exp2 -> ([Char], [(Var, Var)], Exp2))
-> PassM Exp2 -> PassM ([Char], [(Var, Var)], Exp2)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp2 -> PassM Exp2
go Exp2
c) [([Char], [(Var, Var)], Exp2)]
mp
    TimeIt Exp2
e UrTy Var
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 -> UrTy Var -> Bool -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec -> dec -> Bool -> PreExp ext loc dec
TimeIt Exp2
e' UrTy Var
ty Bool
b
    SpawnE{} -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
    SyncE{}  -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp2
ex
    WithArenaE Var
v Exp2
e -> Var -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
Var -> PreExp ext loc dec -> PreExp ext loc dec
WithArenaE Var
v (Exp2 -> Exp2) -> PassM Exp2 -> PassM Exp2
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp2 -> PassM Exp2
go Exp2
e
    MapE{}  -> [Char] -> PassM Exp2
forall a. HasCallStack => [Char] -> a
error ([Char] -> PassM Exp2) -> [Char] -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ [Char]
"inferRegScopeExp: TODO MapE"
    FoldE{} -> [Char] -> PassM Exp2
forall a. HasCallStack => [Char] -> a
error ([Char] -> PassM Exp2) -> [Char] -> PassM Exp2
forall a b. (a -> b) -> a -> b
$ [Char]
"inferRegScopeExp: TODO FoldE"
  where
    go :: Exp2 -> PassM Exp2
go = Map Region Region -> Exp2 -> PassM Exp2
inferRegScopeExp Map Region Region
env