module Gibbon.Passes.RegionsInwards (regionsInwards) where

import qualified Data.Map as M
import qualified Data.Set as S
import qualified Data.List as L
import Data.Foldable as F
import Text.PrettyPrint.GenericPretty

import Gibbon.Common
import Gibbon.L2.Syntax
import Data.Maybe ()
import qualified Data.Maybe as S


data DelayedBind = DelayRegion Region RegionSize (Maybe RegionType)                                            --define data type that can be Region, Loc, LocExp to store the delayed bindings
                 | DelayLoc LocVar LocExp | DelayParRegion Region RegionSize (Maybe RegionType)
  deriving (Int -> DelayedBind -> ShowS
[DelayedBind] -> ShowS
DelayedBind -> String
(Int -> DelayedBind -> ShowS)
-> (DelayedBind -> String)
-> ([DelayedBind] -> ShowS)
-> Show DelayedBind
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DelayedBind -> ShowS
showsPrec :: Int -> DelayedBind -> ShowS
$cshow :: DelayedBind -> String
show :: DelayedBind -> String
$cshowList :: [DelayedBind] -> ShowS
showList :: [DelayedBind] -> ShowS
Show, (forall x. DelayedBind -> Rep DelayedBind x)
-> (forall x. Rep DelayedBind x -> DelayedBind)
-> Generic DelayedBind
forall x. Rep DelayedBind x -> DelayedBind
forall x. DelayedBind -> Rep DelayedBind x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. DelayedBind -> Rep DelayedBind x
from :: forall x. DelayedBind -> Rep DelayedBind x
$cto :: forall x. Rep DelayedBind x -> DelayedBind
to :: forall x. Rep DelayedBind x -> DelayedBind
Generic)

instance Out DelayedBind

type DelayedBindEnv = M.Map (S.Set LocVar) [DelayedBind]                         --define a Map from set to the DelayedBind data type

regionsInwards :: Prog2 -> PassM Prog2
regionsInwards :: Prog2 -> PassM Prog2
regionsInwards 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
    let scopeSetMain :: Set LocVar
scopeSetMain = [LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList ([LocVar] -> Set LocVar) -> [LocVar] -> Set LocVar
forall a b. (a -> b) -> a -> b
$ (FunDef Exp2 -> LocVar) -> [FunDef Exp2] -> [LocVar]
forall a b. (a -> b) -> [a] -> [b]
map FunDef Exp2 -> LocVar
forall ex. FunDef ex -> LocVar
funName (FunDefs Exp2 -> [FunDef Exp2]
forall k a. Map k a -> [a]
M.elems FunDefs Exp2
fundefs)                --Init scopeSet with all the function names
        functionArgs :: Set LocVar
functionArgs = [LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList ([LocVar] -> Set LocVar) -> [LocVar] -> Set LocVar
forall a b. (a -> b) -> a -> b
$ (FunDef Exp2 -> [LocVar]) -> [FunDef Exp2] -> [LocVar]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap FunDef Exp2 -> [LocVar]
forall ex. FunDef ex -> [LocVar]
funArgs (FunDefs Exp2 -> [FunDef Exp2]
forall k a. Map k a -> [a]
M.elems FunDefs Exp2
fundefs)          --Init functionArgs with all the function arguments, concatenate into one list
        scopeSetFun :: Set LocVar
scopeSetFun  = Set LocVar
scopeSetMain Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union` Set LocVar
functionArgs                       --scope set for function body is the union of function args and the function names
    [FunDef Exp2]
fds' <- (FunDef Exp2 -> PassM (FunDef Exp2))
-> [FunDef Exp2] -> PassM [FunDef Exp2]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Set LocVar -> FunDef Exp2 -> PassM (FunDef Exp2)
placeRegionsInwardsFunBody Set LocVar
scopeSetFun) (FunDefs Exp2 -> [FunDef Exp2]
forall k a. Map k a -> [a]
M.elems FunDefs Exp2
fundefs)      --Delay Regions for the function body
    let fundefs' :: FunDefs Exp2
fundefs' = [(LocVar, FunDef Exp2)] -> FunDefs Exp2
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(LocVar, FunDef Exp2)] -> FunDefs Exp2)
-> [(LocVar, FunDef Exp2)] -> FunDefs Exp2
forall a b. (a -> b) -> a -> b
$ (FunDef Exp2 -> (LocVar, FunDef Exp2))
-> [FunDef Exp2] -> [(LocVar, FunDef Exp2)]
forall a b. (a -> b) -> [a] -> [b]
map (\FunDef Exp2
f -> (FunDef Exp2 -> LocVar
forall ex. FunDef ex -> LocVar
funName FunDef Exp2
f,FunDef Exp2
f)) [FunDef Exp2]
fds'
    Maybe (Exp2, UrTy LocVar)
mainExp' <- case Maybe (Exp2, TyOf Exp2)
mainExp of
        Maybe (Exp2, TyOf Exp2)
Nothing -> Maybe (Exp2, UrTy LocVar) -> PassM (Maybe (Exp2, UrTy LocVar))
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Exp2, UrTy LocVar)
forall a. Maybe a
Nothing
        Just (Exp2
mn, TyOf Exp2
ty)-> do
          let env :: Map k a
env = Map k a
forall k a. Map k a
M.empty                                                      --Use M.empty for creating the empty env
              in (Exp2, UrTy LocVar) -> Maybe (Exp2, UrTy LocVar)
forall a. a -> Maybe a
Just ((Exp2, UrTy LocVar) -> Maybe (Exp2, UrTy LocVar))
-> (Exp2 -> (Exp2, UrTy LocVar))
-> Exp2
-> Maybe (Exp2, UrTy LocVar)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (,TyOf Exp2
UrTy LocVar
ty) (Exp2 -> Maybe (Exp2, UrTy LocVar))
-> PassM Exp2 -> PassM (Maybe (Exp2, UrTy LocVar))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DelayedBindEnv -> Set LocVar -> Exp2 -> PassM Exp2
placeRegionInwards DelayedBindEnv
forall k a. Map k a
env Set LocVar
scopeSetMain Exp2
mn         --Delay Regions for the main function

    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 LocVar)
mainExp'                                        --return new ddefs, fundefs and mainExpressions



placeRegionsInwardsFunBody :: S.Set Var -> FunDef2  -> PassM FunDef2
placeRegionsInwardsFunBody :: Set LocVar -> FunDef Exp2 -> PassM (FunDef Exp2)
placeRegionsInwardsFunBody Set LocVar
scopeSet f :: FunDef Exp2
f@FunDef{Exp2
funBody :: Exp2
funBody :: forall ex. FunDef ex -> ex
funBody}  = do
  let env :: Map k a
env     = Map k a
forall k a. Map k a
M.empty                                                          --Create empty environment
  Exp2
funBody' <- DelayedBindEnv -> Set LocVar -> Exp2 -> PassM Exp2
placeRegionInwards DelayedBindEnv
forall k a. Map k a
env Set LocVar
scopeSet Exp2
funBody                            --Recursively delay regions for function body
  FunDef Exp2 -> PassM (FunDef Exp2)
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (FunDef Exp2 -> PassM (FunDef Exp2))
-> FunDef Exp2 -> PassM (FunDef Exp2)
forall a b. (a -> b) -> a -> b
$ FunDef Exp2
f {funBody :: Exp2
funBody = Exp2
funBody'}


placeRegionInwards :: DelayedBindEnv -> S.Set Var -> Exp2 -> PassM Exp2          --Recursive funtion that will move the regions inwards
placeRegionInwards :: DelayedBindEnv -> Set LocVar -> Exp2 -> PassM Exp2
placeRegionInwards DelayedBindEnv
env Set LocVar
scopeSet Exp2
ex  =
  case Exp2
ex of
    Ext E2Ext LocVar (UrTy LocVar)
ext ->
      case E2Ext LocVar (UrTy LocVar)
ext of

        LetRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
rhs -> do                                                   --take care of regions
          let key' :: Set LocVar
key' = LocVar -> Set LocVar
forall a. a -> Set a
S.singleton (Region -> LocVar
regionToVar Region
r)
              val' :: [DelayedBind]
val' = [Region -> RegionSize -> Maybe RegionType -> DelayedBind
DelayRegion Region
r RegionSize
sz Maybe RegionType
ty]
              env' :: DelayedBindEnv
env' = Set LocVar -> [DelayedBind] -> DelayedBindEnv -> DelayedBindEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Set LocVar
key' [DelayedBind]
val' DelayedBindEnv
env
              in DelayedBindEnv -> Set LocVar -> Exp2 -> PassM Exp2
placeRegionInwards DelayedBindEnv
env' Set LocVar
scopeSet Exp2
rhs

        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

        LetLocE LocVar
loc LocExp
phs Exp2
rhs -> do                                                --take care of locations
          case LocExp
phs of

            StartOfRegionLE Region
r -> do
              let keyList' :: [Set LocVar]
keyList' = DelayedBindEnv -> [Set LocVar]
forall k a. Map k a -> [k]
M.keys DelayedBindEnv
env
                  key' :: Maybe (Set LocVar)
key'     = (Set LocVar -> Bool) -> [Set LocVar] -> Maybe (Set LocVar)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
F.find (LocVar -> Set LocVar -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member (Region -> LocVar
regionToVar Region
r)) [Set LocVar]
keyList'
                  in case Maybe (Set LocVar)
key' of
                    Maybe (Set LocVar)
Nothing -> do
                      let key'' :: Set LocVar
key'' = LocVar -> Set LocVar
forall a. a -> Set a
S.singleton LocVar
loc
                          val' :: [DelayedBind]
val' = [LocVar -> LocExp -> DelayedBind
DelayLoc LocVar
loc LocExp
phs]
                          env' :: DelayedBindEnv
env' = Set LocVar -> [DelayedBind] -> DelayedBindEnv -> DelayedBindEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Set LocVar
key'' [DelayedBind]
val' DelayedBindEnv
env
                       in DelayedBindEnv -> Set LocVar -> Exp2 -> PassM Exp2
placeRegionInwards DelayedBindEnv
env' Set LocVar
scopeSet Exp2
rhs
                    Just Set LocVar
myKey -> do
                      let valList :: [DelayedBind]
valList  = [DelayedBind] -> Set LocVar -> DelayedBindEnv -> [DelayedBind]
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault [] Set LocVar
myKey DelayedBindEnv
env
                          myKey' :: Set LocVar
myKey'   = LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => a -> Set a -> Set a
S.insert LocVar
loc Set LocVar
myKey
                          valList' :: [DelayedBind]
valList' = [DelayedBind]
valList [DelayedBind] -> [DelayedBind] -> [DelayedBind]
forall a. [a] -> [a] -> [a]
++ [LocVar -> LocExp -> DelayedBind
DelayLoc LocVar
loc LocExp
phs]
                          tempDict :: DelayedBindEnv
tempDict = Set LocVar -> DelayedBindEnv -> DelayedBindEnv
forall k a. Ord k => k -> Map k a -> Map k a
M.delete Set LocVar
myKey DelayedBindEnv
env
                          newEnv :: DelayedBindEnv
newEnv   = Set LocVar -> [DelayedBind] -> DelayedBindEnv -> DelayedBindEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Set LocVar
myKey' [DelayedBind]
valList' DelayedBindEnv
tempDict
                          in DelayedBindEnv -> Set LocVar -> Exp2 -> PassM Exp2
placeRegionInwards DelayedBindEnv
newEnv Set LocVar
scopeSet Exp2
rhs         --recurse on rhs using the newenv

            AfterConstantLE Int
_ LocVar
loc' -> do                                    --In case statement, actual match = AfterConstantLE integralVal loc'
              let keyList' :: [Set LocVar]
keyList' = DelayedBindEnv -> [Set LocVar]
forall k a. Map k a -> [k]
M.keys DelayedBindEnv
env
                  key' :: Maybe (Set LocVar)
key'     = (Set LocVar -> Bool) -> [Set LocVar] -> Maybe (Set LocVar)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
F.find (LocVar -> Set LocVar -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member LocVar
loc') [Set LocVar]
keyList'
                  in case Maybe (Set LocVar)
key' of
                    Maybe (Set LocVar)
Nothing -> do
                      let key'' :: Set LocVar
key'' = LocVar -> Set LocVar
forall a. a -> Set a
S.singleton LocVar
loc
                          val' :: [DelayedBind]
val' = [LocVar -> LocExp -> DelayedBind
DelayLoc LocVar
loc LocExp
phs]
                          env' :: DelayedBindEnv
env' = Set LocVar -> [DelayedBind] -> DelayedBindEnv -> DelayedBindEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Set LocVar
key'' [DelayedBind]
val' DelayedBindEnv
env
                        in DelayedBindEnv -> Set LocVar -> Exp2 -> PassM Exp2
placeRegionInwards DelayedBindEnv
env' Set LocVar
scopeSet Exp2
rhs
                    Just Set LocVar
myKey -> do
                      let valList :: [DelayedBind]
valList  = [DelayedBind] -> Set LocVar -> DelayedBindEnv -> [DelayedBind]
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault [] Set LocVar
myKey DelayedBindEnv
env
                          myKey' :: Set LocVar
myKey'   = LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => a -> Set a -> Set a
S.insert LocVar
loc Set LocVar
myKey
                          valList' :: [DelayedBind]
valList' = [DelayedBind]
valList [DelayedBind] -> [DelayedBind] -> [DelayedBind]
forall a. [a] -> [a] -> [a]
++ [LocVar -> LocExp -> DelayedBind
DelayLoc LocVar
loc LocExp
phs]
                          tempDict :: DelayedBindEnv
tempDict = Set LocVar -> DelayedBindEnv -> DelayedBindEnv
forall k a. Ord k => k -> Map k a -> Map k a
M.delete Set LocVar
myKey DelayedBindEnv
env
                          newEnv :: DelayedBindEnv
newEnv   = Set LocVar -> [DelayedBind] -> DelayedBindEnv -> DelayedBindEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Set LocVar
myKey' [DelayedBind]
valList' DelayedBindEnv
tempDict
                          in DelayedBindEnv -> Set LocVar -> Exp2 -> PassM Exp2
placeRegionInwards DelayedBindEnv
newEnv Set LocVar
scopeSet Exp2
rhs

            AfterVariableLE LocVar
_ LocVar
loc' Bool
_ -> do                                  --In case statement, actual match = AfterVariableLE variable loc' boolVal
              let keyList' :: [Set LocVar]
keyList' = DelayedBindEnv -> [Set LocVar]
forall k a. Map k a -> [k]
M.keys DelayedBindEnv
env
                  key' :: Maybe (Set LocVar)
key'     = (Set LocVar -> Bool) -> [Set LocVar] -> Maybe (Set LocVar)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
F.find (LocVar -> Set LocVar -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member LocVar
loc') [Set LocVar]
keyList'
                  in case Maybe (Set LocVar)
key' of
                    Maybe (Set LocVar)
Nothing -> do
                        let key'' :: Set LocVar
key'' = LocVar -> Set LocVar
forall a. a -> Set a
S.singleton LocVar
loc
                            val' :: [DelayedBind]
val' = [LocVar -> LocExp -> DelayedBind
DelayLoc LocVar
loc LocExp
phs]
                            env' :: DelayedBindEnv
env' = Set LocVar -> [DelayedBind] -> DelayedBindEnv -> DelayedBindEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Set LocVar
key'' [DelayedBind]
val' DelayedBindEnv
env
                         in DelayedBindEnv -> Set LocVar -> Exp2 -> PassM Exp2
placeRegionInwards DelayedBindEnv
env' Set LocVar
scopeSet Exp2
rhs
                    Just Set LocVar
myKey -> do
                      let valList :: [DelayedBind]
valList  = [DelayedBind] -> Set LocVar -> DelayedBindEnv -> [DelayedBind]
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault [] Set LocVar
myKey DelayedBindEnv
env
                          myKey' :: Set LocVar
myKey'   = LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => a -> Set a -> Set a
S.insert LocVar
loc Set LocVar
myKey
                          valList' :: [DelayedBind]
valList' = [DelayedBind]
valList [DelayedBind] -> [DelayedBind] -> [DelayedBind]
forall a. [a] -> [a] -> [a]
++ [LocVar -> LocExp -> DelayedBind
DelayLoc LocVar
loc LocExp
phs]
                          tempDict :: DelayedBindEnv
tempDict = Set LocVar -> DelayedBindEnv -> DelayedBindEnv
forall k a. Ord k => k -> Map k a -> Map k a
M.delete Set LocVar
myKey DelayedBindEnv
env
                          newEnv :: DelayedBindEnv
newEnv   = Set LocVar -> [DelayedBind] -> DelayedBindEnv -> DelayedBindEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Set LocVar
myKey' [DelayedBind]
valList' DelayedBindEnv
tempDict
                          in DelayedBindEnv -> Set LocVar -> Exp2 -> PassM Exp2
placeRegionInwards DelayedBindEnv
newEnv Set LocVar
scopeSet Exp2
rhs

            InRegionLE Region
r -> do
              let keyList' :: [Set LocVar]
keyList' = DelayedBindEnv -> [Set LocVar]
forall k a. Map k a -> [k]
M.keys DelayedBindEnv
env
                  key' :: Maybe (Set LocVar)
key'     = (Set LocVar -> Bool) -> [Set LocVar] -> Maybe (Set LocVar)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
F.find (LocVar -> Set LocVar -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member (Region -> LocVar
regionToVar Region
r) ) [Set LocVar]
keyList'
                  in case Maybe (Set LocVar)
key' of
                    Maybe (Set LocVar)
Nothing -> String -> PassM Exp2
forall a. HasCallStack => String -> a
error String
"No existing region found for this Location in case InRegionLE"
                    Just Set LocVar
myKey -> do
                      let valList :: [DelayedBind]
valList  = [DelayedBind] -> Set LocVar -> DelayedBindEnv -> [DelayedBind]
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault [] Set LocVar
myKey DelayedBindEnv
env
                          myKey' :: Set LocVar
myKey'   = LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => a -> Set a -> Set a
S.insert LocVar
loc Set LocVar
myKey
                          valList' :: [DelayedBind]
valList' = [DelayedBind]
valList [DelayedBind] -> [DelayedBind] -> [DelayedBind]
forall a. [a] -> [a] -> [a]
++ [LocVar -> LocExp -> DelayedBind
DelayLoc LocVar
loc LocExp
phs]
                          tempDict :: DelayedBindEnv
tempDict = Set LocVar -> DelayedBindEnv -> DelayedBindEnv
forall k a. Ord k => k -> Map k a -> Map k a
M.delete Set LocVar
myKey DelayedBindEnv
env
                          newEnv :: DelayedBindEnv
newEnv   = Set LocVar -> [DelayedBind] -> DelayedBindEnv -> DelayedBindEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Set LocVar
myKey' [DelayedBind]
valList' DelayedBindEnv
tempDict
                          in DelayedBindEnv -> Set LocVar -> Exp2 -> PassM Exp2
placeRegionInwards DelayedBindEnv
newEnv Set LocVar
scopeSet Exp2
rhs

            FromEndLE LocVar
loc' -> do
              let keyList' :: [Set LocVar]
keyList' = DelayedBindEnv -> [Set LocVar]
forall k a. Map k a -> [k]
M.keys DelayedBindEnv
env
                  key' :: Maybe (Set LocVar)
key'     = (Set LocVar -> Bool) -> [Set LocVar] -> Maybe (Set LocVar)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
F.find (LocVar -> Set LocVar -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member LocVar
loc') [Set LocVar]
keyList'
                  in case Maybe (Set LocVar)
key' of
                    Maybe (Set LocVar)
Nothing -> String -> PassM Exp2
forall a. HasCallStack => String -> a
error String
"No existing variable found for this Location in case FromEndLE"
                    Just Set LocVar
myKey -> do
                      let valList :: [DelayedBind]
valList  = [DelayedBind] -> Set LocVar -> DelayedBindEnv -> [DelayedBind]
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault [] Set LocVar
myKey DelayedBindEnv
env
                          myKey' :: Set LocVar
myKey'   = LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => a -> Set a -> Set a
S.insert LocVar
loc Set LocVar
myKey
                          valList' :: [DelayedBind]
valList' = [DelayedBind]
valList [DelayedBind] -> [DelayedBind] -> [DelayedBind]
forall a. [a] -> [a] -> [a]
++ [LocVar -> LocExp -> DelayedBind
DelayLoc LocVar
loc LocExp
phs]
                          tempDict :: DelayedBindEnv
tempDict = Set LocVar -> DelayedBindEnv -> DelayedBindEnv
forall k a. Ord k => k -> Map k a -> Map k a
M.delete Set LocVar
myKey DelayedBindEnv
env
                          newEnv :: DelayedBindEnv
newEnv   = Set LocVar -> [DelayedBind] -> DelayedBindEnv -> DelayedBindEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Set LocVar
myKey' [DelayedBind]
valList' DelayedBindEnv
tempDict
                          in DelayedBindEnv -> Set LocVar -> Exp2 -> PassM Exp2
placeRegionInwards DelayedBindEnv
newEnv Set LocVar
scopeSet Exp2
rhs

            LocExp
FreeLE -> String -> PassM Exp2
forall a. HasCallStack => String -> a
error String
"Free LE not implemented yet!"                       --For FreeLE we need to figure out how to handle this?

        LetParRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
rhs -> do                                                --Handle a parallel LetRegion
          let key' :: Set LocVar
key' = LocVar -> Set LocVar
forall a. a -> Set a
S.singleton (Region -> LocVar
regionToVar Region
r)
              val' :: [DelayedBind]
val' = [Region -> RegionSize -> Maybe RegionType -> DelayedBind
DelayParRegion Region
r RegionSize
sz Maybe RegionType
ty]
              env' :: DelayedBindEnv
env' = Set LocVar -> [DelayedBind] -> DelayedBindEnv -> DelayedBindEnv
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Set LocVar
key' [DelayedBind]
val' DelayedBindEnv
env
              in DelayedBindEnv -> Set LocVar -> Exp2 -> PassM Exp2
placeRegionInwards DelayedBindEnv
env' Set LocVar
scopeSet Exp2
rhs


        RetE [LocVar]
locList LocVar
_                                  -> do              {- Look at the locList to see which variables are alive in env and codegen them before -}
                                                            let (DelayedBindEnv
_, Exp2
ex') = DelayedBindEnv -> Set LocVar -> Exp2 -> (DelayedBindEnv, Exp2)
dischargeBinds' DelayedBindEnv
env ([LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList [LocVar]
locList) Exp2
ex
                                                              in Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex'

        FromEndE LocVar
_                                     -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex        {- Actual type is FromEndE loc, Don't need to worry about it will appear later in the pipeline, Just return the expression -}
        BoundsCheck{}                                  -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex        {- Actual type is BoundsCheck integer l1 l2, Don't need to worry about it will appear later in the pipeline, Just return the expression -}
        AddFixed{}                                     -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex        {- Actual type is AddFixed variable integer, Return the expression -}
        IndirectionE{}                                 -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex        {- Actual type: IndirectionE tyCon dataCon (l1,v1) (l2,v2) rhs, skip the recursion, IndirectionE doesn't appear until later in the IR language, return the expression -}
        E2Ext LocVar (UrTy LocVar)
GetCilkWorkerNum                               -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex                   {- Just return the expression, there is no recusrion to do here -}
        LetAvail [LocVar]
vs Exp2
e                                  -> E2Ext LocVar (UrTy LocVar) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocVar (UrTy LocVar) -> Exp2)
-> (Exp2 -> E2Ext LocVar (UrTy LocVar)) -> Exp2 -> Exp2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [LocVar] -> Exp2 -> E2Ext LocVar (UrTy LocVar)
forall loc dec. [LocVar] -> E2 loc dec -> E2Ext loc dec
LetAvail [LocVar]
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  {- Recurse on the rhs directly -}
        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        -- Just return Nothing special here
    LitE{}                 -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex        -- Just return Nothing special here
    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        -- Just return Nothing special here
    LitSymE{}              -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex        -- Just return Nothing special here
    AppE LocVar
f [LocVar]
locVars [Exp2]
ls      -> do
                              let allKeys :: [Set LocVar]
allKeys  =  DelayedBindEnv -> [Set LocVar]
forall k a. Map k a -> [k]
M.keys DelayedBindEnv
env                                                             -- List of all keys from env
                                  keyList :: [Maybe (Set LocVar)]
keyList  = (LocVar -> Maybe (Set LocVar)) -> [LocVar] -> [Maybe (Set LocVar)]
forall a b. (a -> b) -> [a] -> [b]
map (\LocVar
variable -> (Set LocVar -> Bool) -> [Set LocVar] -> Maybe (Set LocVar)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
F.find (LocVar -> Set LocVar -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member LocVar
variable) [Set LocVar]
allKeys) [LocVar]
locVars           -- For each var in the input set find its corresponding key
                                  keyList' :: [Set LocVar]
keyList' = [Maybe (Set LocVar)] -> [Set LocVar]
forall a. [Maybe a] -> [a]
S.catMaybes [Maybe (Set LocVar)]
keyList                                                     -- Filter all the Nothing values from the list and let only Just values in the list
                                  newKeys :: [Set LocVar]
newKeys   = Set (Set LocVar) -> [Set LocVar]
forall a. Set a -> [a]
S.toList (Set (Set LocVar) -> [Set LocVar])
-> Set (Set LocVar) -> [Set LocVar]
forall a b. (a -> b) -> a -> b
$ [Set LocVar] -> Set (Set LocVar)
forall a. Ord a => [a] -> Set a
S.fromList [Set LocVar]
allKeys Set (Set LocVar) -> Set (Set LocVar) -> Set (Set LocVar)
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` [Set LocVar] -> Set (Set LocVar)
forall a. Ord a => [a] -> Set a
S.fromList [Set LocVar]
keyList'       -- Filter all the Nothing values from the list and let only Just values in the list
                                  newVals :: [[DelayedBind]]
newVals   = (Set LocVar -> [DelayedBind]) -> [Set LocVar] -> [[DelayedBind]]
forall a b. (a -> b) -> [a] -> [b]
map (\Set LocVar
key -> [DelayedBind] -> Set LocVar -> DelayedBindEnv -> [DelayedBind]
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault [] Set LocVar
key DelayedBindEnv
env) [Set LocVar]
newKeys
                                  tupleList :: [(Set LocVar, [DelayedBind])]
tupleList = [Set LocVar] -> [[DelayedBind]] -> [(Set LocVar, [DelayedBind])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Set LocVar]
newKeys [[DelayedBind]]
newVals
                                  newEnv' :: DelayedBindEnv
newEnv'   = [(Set LocVar, [DelayedBind])] -> DelayedBindEnv
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Set LocVar, [DelayedBind])]
tupleList
                               in do [Exp2]
ls' <- (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 (DelayedBindEnv -> Set LocVar -> Exp2 -> PassM Exp2
placeRegionInwards DelayedBindEnv
newEnv' Set LocVar
scopeSet) [Exp2]
ls
                                     let (DelayedBindEnv
_, Exp2
ex') = DelayedBindEnv -> Set LocVar -> Exp2 -> (DelayedBindEnv, Exp2)
dischargeBinds' DelayedBindEnv
env ([LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList [LocVar]
locVars) (LocVar -> [LocVar] -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
LocVar -> [loc] -> [PreExp ext loc dec] -> PreExp ext loc dec
AppE LocVar
f [LocVar]
locVars [Exp2]
ls')
                                      in 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                                                                                   -- Just return, Nothing special here

    DataConE LocVar
loc String
dataCons [Exp2]
args      -> do
                                       let allKeys :: [Set LocVar]
allKeys  =  DelayedBindEnv -> [Set LocVar]
forall k a. Map k a -> [k]
M.keys DelayedBindEnv
env                                                         -- List of all keys from env
                                           freelist :: [Set LocVar]
freelist = (Exp2 -> Set LocVar) -> [Exp2] -> [Set LocVar]
forall a b. (a -> b) -> [a] -> [b]
map Exp2 -> Set LocVar
freeVars [Exp2]
args 
                                           freevars :: Set LocVar
freevars = (Set LocVar -> Set LocVar -> Set LocVar)
-> Set LocVar -> [Set LocVar] -> Set LocVar
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\Set LocVar
s1 Set LocVar
s2 -> Set LocVar
s1 Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union` Set LocVar
s2) (Set LocVar
forall a. Set a
S.empty) [Set LocVar]
freelist  
                                           keyList :: [Maybe (Set LocVar)]
keyList  = (LocVar -> Maybe (Set LocVar)) -> [LocVar] -> [Maybe (Set LocVar)]
forall a b. (a -> b) -> [a] -> [b]
map (\LocVar
variable -> (Set LocVar -> Bool) -> [Set LocVar] -> Maybe (Set LocVar)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
F.find (LocVar -> Set LocVar -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member LocVar
variable) [Set LocVar]
allKeys) ((Set LocVar -> [LocVar]
forall a. Set a -> [a]
S.toList Set LocVar
freevars) [LocVar] -> [LocVar] -> [LocVar]
forall a. [a] -> [a] -> [a]
++ [LocVar
loc])       -- For each var in the input set find its corresponding key
                                           keyList' :: [Set LocVar]
keyList' = [Maybe (Set LocVar)] -> [Set LocVar]
forall a. [Maybe a] -> [a]
S.catMaybes [Maybe (Set LocVar)]
keyList                                                 -- Filter all the Nothing values from the list and let only Just values in the list
                                           newKeys :: [Set LocVar]
newKeys   = Set (Set LocVar) -> [Set LocVar]
forall a. Set a -> [a]
S.toList (Set (Set LocVar) -> [Set LocVar])
-> Set (Set LocVar) -> [Set LocVar]
forall a b. (a -> b) -> a -> b
$ [Set LocVar] -> Set (Set LocVar)
forall a. Ord a => [a] -> Set a
S.fromList [Set LocVar]
allKeys Set (Set LocVar) -> Set (Set LocVar) -> Set (Set LocVar)
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` [Set LocVar] -> Set (Set LocVar)
forall a. Ord a => [a] -> Set a
S.fromList [Set LocVar]
keyList'   -- Filter all the Nothing values from the list and let only Just values in the list
                                           newVals :: [[DelayedBind]]
newVals   = (Set LocVar -> [DelayedBind]) -> [Set LocVar] -> [[DelayedBind]]
forall a b. (a -> b) -> [a] -> [b]
map (\Set LocVar
key -> [DelayedBind] -> Set LocVar -> DelayedBindEnv -> [DelayedBind]
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault [] Set LocVar
key DelayedBindEnv
env) [Set LocVar]
newKeys
                                           tupleList :: [(Set LocVar, [DelayedBind])]
tupleList = [Set LocVar] -> [[DelayedBind]] -> [(Set LocVar, [DelayedBind])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Set LocVar]
newKeys [[DelayedBind]]
newVals
                                           newEnv' :: DelayedBindEnv
newEnv'   = [(Set LocVar, [DelayedBind])] -> DelayedBindEnv
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Set LocVar, [DelayedBind])]
tupleList
                                           in do [Exp2]
args' <- (Exp2 -> PassM Exp2) -> [Exp2] -> PassM [Exp2]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (DelayedBindEnv -> Set LocVar -> Exp2 -> PassM Exp2
placeRegionInwards DelayedBindEnv
newEnv' Set LocVar
scopeSet) [Exp2]
args
                                                 let (DelayedBindEnv
_, Exp2
ex') = DelayedBindEnv -> Set LocVar -> Exp2 -> (DelayedBindEnv, Exp2)
dischargeBinds' DelayedBindEnv
env (Set LocVar
freevars Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union` (LocVar -> Set LocVar
forall a. a -> Set a
S.singleton LocVar
loc)) (LocVar -> String -> [Exp2] -> Exp2
forall (ext :: * -> * -> *) loc dec.
loc -> String -> [PreExp ext loc dec] -> PreExp ext loc dec
DataConE LocVar
loc String
dataCons [Exp2]
args')
                                                  in 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    {- Simple recursion on e -}

    IfE Exp2
a Exp2
b Exp2
c              -> do     -- Optimization for IF statements check the freeVariables in b and c, intersect them to avoid generating duplicate bindings at in then and else part
                                     let freeVarsB :: Set LocVar
freeVarsB  = Exp2 -> Set LocVar
freeVars Exp2
b
                                         freeVarsC :: Set LocVar
freeVarsC  = Exp2 -> Set LocVar
freeVars Exp2
c
                                         commonVars :: Set LocVar
commonVars = Set LocVar
freeVarsB Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.intersection` Set LocVar
freeVarsC
                                         --_          = dbgTraceIt (sdoc commonVars)
                                         allKeys :: [Set LocVar]
allKeys    = DelayedBindEnv -> [Set LocVar]
forall k a. Map k a -> [k]
M.keys DelayedBindEnv
env
                                         keyList :: [Maybe (Set LocVar)]
keyList    = (LocVar -> Maybe (Set LocVar)) -> [LocVar] -> [Maybe (Set LocVar)]
forall a b. (a -> b) -> [a] -> [b]
map (\LocVar
variable -> (Set LocVar -> Bool) -> [Set LocVar] -> Maybe (Set LocVar)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
F.find (LocVar -> Set LocVar -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member LocVar
variable) [Set LocVar]
allKeys) (Set LocVar -> [LocVar]
forall a. Set a -> [a]
S.toList Set LocVar
commonVars)
                                         keyList' :: [Set LocVar]
keyList'   = [Maybe (Set LocVar)] -> [Set LocVar]
forall a. [Maybe a] -> [a]
S.catMaybes [Maybe (Set LocVar)]
keyList
                                         newKeys :: [Set LocVar]
newKeys   = Set (Set LocVar) -> [Set LocVar]
forall a. Set a -> [a]
S.toList (Set (Set LocVar) -> [Set LocVar])
-> Set (Set LocVar) -> [Set LocVar]
forall a b. (a -> b) -> a -> b
$ [Set LocVar] -> Set (Set LocVar)
forall a. Ord a => [a] -> Set a
S.fromList [Set LocVar]
allKeys Set (Set LocVar) -> Set (Set LocVar) -> Set (Set LocVar)
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` [Set LocVar] -> Set (Set LocVar)
forall a. Ord a => [a] -> Set a
S.fromList [Set LocVar]
keyList'
                                         newVals :: [[DelayedBind]]
newVals    = (Set LocVar -> [DelayedBind]) -> [Set LocVar] -> [[DelayedBind]]
forall a b. (a -> b) -> [a] -> [b]
map (\Set LocVar
key -> [DelayedBind] -> Set LocVar -> DelayedBindEnv -> [DelayedBind]
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault [] Set LocVar
key DelayedBindEnv
env) [Set LocVar]
newKeys
                                         tupleList :: [(Set LocVar, [DelayedBind])]
tupleList  = [Set LocVar] -> [[DelayedBind]] -> [(Set LocVar, [DelayedBind])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Set LocVar]
newKeys [[DelayedBind]]
newVals
                                         newEnv' :: DelayedBindEnv
newEnv'    = [(Set LocVar, [DelayedBind])] -> DelayedBindEnv
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Set LocVar, [DelayedBind])]
tupleList
                                     Exp2
b' <- DelayedBindEnv -> Set LocVar -> Exp2 -> PassM Exp2
placeRegionInwards DelayedBindEnv
env Set LocVar
scopeSet Exp2
b       -- Recurse on b (Then part)
                                     Exp2
c' <- DelayedBindEnv -> Set LocVar -> Exp2 -> PassM Exp2
placeRegionInwards DelayedBindEnv
env Set LocVar
scopeSet Exp2
c       -- Recurse on c (Else part)
                                     let (DelayedBindEnv
_, Exp2
a') = DelayedBindEnv -> Set LocVar -> Exp2 -> (DelayedBindEnv, Exp2)
dischargeBinds' DelayedBindEnv
newEnv' Set LocVar
commonVars Exp2
a
                                     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 -> 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
b' Exp2
c'                             -- Return the new IfE expression {-dbgTraceIt (sdoc (commonVars, keyList, env, newEnv'))-}

    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                            {- Recurse over all expression in the tuple in the expression ls -}

    LetE (LocVar
v,[LocVar]
locs,UrTy LocVar
ty,Exp2
rhs) Exp2
bod      -> do
                                    let newScope :: Set LocVar
newScope = LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => a -> Set a -> Set a
S.insert LocVar
v Set LocVar
scopeSet                                                     {- The locs will be empty at this point, so just update scope set and recurse -}
                                        allKeys :: [Set LocVar]
allKeys  =  DelayedBindEnv -> [Set LocVar]
forall k a. Map k a -> [k]
M.keys DelayedBindEnv
env
                                        free_vars :: [LocVar]
free_vars =   UrTy LocVar -> [LocVar]
locsInTy UrTy LocVar
ty                                                          -- List of all keys from env
                                        keyList :: [Maybe (Set LocVar)]
keyList  = (LocVar -> Maybe (Set LocVar)) -> [LocVar] -> [Maybe (Set LocVar)]
forall a b. (a -> b) -> [a] -> [b]
map (\LocVar
variable -> (Set LocVar -> Bool) -> [Set LocVar] -> Maybe (Set LocVar)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
F.find (LocVar -> Set LocVar -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member LocVar
variable) [Set LocVar]
allKeys) [LocVar]
free_vars         -- For each var in the input set find its corresponding key
                                        keyList' :: [Set LocVar]
keyList' = [Maybe (Set LocVar)] -> [Set LocVar]
forall a. [Maybe a] -> [a]
S.catMaybes [Maybe (Set LocVar)]
keyList
                                        newKeys :: [Set LocVar]
newKeys   = Set (Set LocVar) -> [Set LocVar]
forall a. Set a -> [a]
S.toList (Set (Set LocVar) -> [Set LocVar])
-> Set (Set LocVar) -> [Set LocVar]
forall a b. (a -> b) -> a -> b
$ [Set LocVar] -> Set (Set LocVar)
forall a. Ord a => [a] -> Set a
S.fromList [Set LocVar]
allKeys Set (Set LocVar) -> Set (Set LocVar) -> Set (Set LocVar)
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` [Set LocVar] -> Set (Set LocVar)
forall a. Ord a => [a] -> Set a
S.fromList [Set LocVar]
keyList'       -- Filter all the Nothing values from the list and let only Just values in the list
                                        newVals :: [[DelayedBind]]
newVals   = (Set LocVar -> [DelayedBind]) -> [Set LocVar] -> [[DelayedBind]]
forall a b. (a -> b) -> [a] -> [b]
map (\Set LocVar
key -> [DelayedBind] -> Set LocVar -> DelayedBindEnv -> [DelayedBind]
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault [] Set LocVar
key DelayedBindEnv
env) [Set LocVar]
newKeys
                                        tupleList :: [(Set LocVar, [DelayedBind])]
tupleList = [Set LocVar] -> [[DelayedBind]] -> [(Set LocVar, [DelayedBind])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Set LocVar]
newKeys [[DelayedBind]]
newVals
                                        newEnv' :: DelayedBindEnv
newEnv'   = [(Set LocVar, [DelayedBind])] -> DelayedBindEnv
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Set LocVar, [DelayedBind])]
tupleList
                                        in do Exp2
ex' <- (LocVar, [LocVar], UrTy LocVar, Exp2) -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
(LocVar, [loc], dec, PreExp ext loc dec)
-> PreExp ext loc dec -> PreExp ext loc dec
LetE ((LocVar, [LocVar], UrTy LocVar, Exp2) -> Exp2 -> Exp2)
-> (Exp2 -> (LocVar, [LocVar], UrTy LocVar, Exp2))
-> Exp2
-> Exp2
-> Exp2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (LocVar
v,[LocVar]
locs,UrTy LocVar
ty,) (Exp2 -> Exp2 -> Exp2) -> PassM Exp2 -> PassM (Exp2 -> Exp2)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DelayedBindEnv -> Set LocVar -> Exp2 -> PassM Exp2
placeRegionInwards DelayedBindEnv
newEnv' Set LocVar
newScope 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
<*> DelayedBindEnv -> Set LocVar -> Exp2 -> PassM Exp2
placeRegionInwards DelayedBindEnv
newEnv' Set LocVar
newScope Exp2
bod
                                              let (DelayedBindEnv
_, Exp2
ex'') = DelayedBindEnv -> Set LocVar -> Exp2 -> (DelayedBindEnv, Exp2)
dischargeBinds' DelayedBindEnv
env ([LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList [LocVar]
free_vars) Exp2
ex'
                                               in Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex''

    CaseE Exp2
scrt [(String, [(LocVar, LocVar)], Exp2)]
brs                -> do
      [(String, [(LocVar, LocVar)], Exp2)]
brs' <- ((String, [(LocVar, LocVar)], Exp2)
 -> PassM (String, [(LocVar, LocVar)], Exp2))
-> [(String, [(LocVar, LocVar)], Exp2)]
-> PassM [(String, [(LocVar, LocVar)], 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
        (\(String
a,[(LocVar, LocVar)]
b,Exp2
c) -> do let varList :: [LocVar]
varList = ((LocVar, LocVar) -> LocVar) -> [(LocVar, LocVar)] -> [LocVar]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (LocVar, LocVar) -> LocVar
forall a b. (a, b) -> a
fst [(LocVar, LocVar)]
b                                                                       -- Get all the variables from the tuple list
                            newScope :: Set LocVar
newScope  = Set LocVar
scopeSet Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union` [LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList [LocVar]
varList                                          -- Make the newScope set by unioning the old one with the varList
                            allKeys :: [Set LocVar]
allKeys   =  DelayedBindEnv -> [Set LocVar]
forall k a. Map k a -> [k]
M.keys DelayedBindEnv
env
                            free_vars :: Set LocVar
free_vars = Exp2 -> Set LocVar
freeVars Exp2
c Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union` Set LocVar
newScope                                                  -- List of all keys from env
                            keyList :: [Maybe (Set LocVar)]
keyList   = (LocVar -> Maybe (Set LocVar)) -> [LocVar] -> [Maybe (Set LocVar)]
forall a b. (a -> b) -> [a] -> [b]
map (\LocVar
variable -> (Set LocVar -> Bool) -> [Set LocVar] -> Maybe (Set LocVar)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
F.find (LocVar -> Set LocVar -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member LocVar
variable) [Set LocVar]
allKeys) (Set LocVar -> [LocVar]
forall a. Set a -> [a]
S.toList Set LocVar
free_vars)     -- For each var in the input set find its corresponding key
                            keyList' :: [Set LocVar]
keyList'  = [Maybe (Set LocVar)] -> [Set LocVar]
forall a. [Maybe a] -> [a]
S.catMaybes [Maybe (Set LocVar)]
keyList
                            newKeys :: [Set LocVar]
newKeys   = Set (Set LocVar) -> [Set LocVar]
forall a. Set a -> [a]
S.toList (Set (Set LocVar) -> [Set LocVar])
-> Set (Set LocVar) -> [Set LocVar]
forall a b. (a -> b) -> a -> b
$ [Set LocVar] -> Set (Set LocVar)
forall a. Ord a => [a] -> Set a
S.fromList [Set LocVar]
allKeys Set (Set LocVar) -> Set (Set LocVar) -> Set (Set LocVar)
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` [Set LocVar] -> Set (Set LocVar)
forall a. Ord a => [a] -> Set a
S.fromList [Set LocVar]
keyList'               -- Filter all the Nothing values from the list and let only Just values in the list
                            newVals :: [[DelayedBind]]
newVals   = (Set LocVar -> [DelayedBind]) -> [Set LocVar] -> [[DelayedBind]]
forall a b. (a -> b) -> [a] -> [b]
map (\Set LocVar
key -> [DelayedBind] -> Set LocVar -> DelayedBindEnv -> [DelayedBind]
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault [] Set LocVar
key DelayedBindEnv
env) [Set LocVar]
newKeys
                            tupleList :: [(Set LocVar, [DelayedBind])]
tupleList = [Set LocVar] -> [[DelayedBind]] -> [(Set LocVar, [DelayedBind])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Set LocVar]
newKeys [[DelayedBind]]
newVals
                            newEnv' :: DelayedBindEnv
newEnv'   = [(Set LocVar, [DelayedBind])] -> DelayedBindEnv
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Set LocVar, [DelayedBind])]
tupleList
                        Exp2
c' <- DelayedBindEnv -> Set LocVar -> Exp2 -> PassM Exp2
placeRegionInwards DelayedBindEnv
env Set LocVar
newScope Exp2
c
                        let (DelayedBindEnv
_, Exp2
c'') = DelayedBindEnv -> Set LocVar -> Exp2 -> (DelayedBindEnv, Exp2)
dischargeBinds' DelayedBindEnv
newEnv' Set LocVar
free_vars Exp2
c'                                                -- Discharge the binds using the newScope and the dictionary
                         in (String, [(LocVar, LocVar)], Exp2)
-> PassM (String, [(LocVar, LocVar)], Exp2)
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return (String
a,[(LocVar, LocVar)]
b,Exp2
c'')) [(String, [(LocVar, LocVar)], Exp2)]
brs                                                                          -- dbgTraceIt (sdoc (free_vars, keyList, env, newEnv'))

      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 -> [(String, [(LocVar, LocVar)], Exp2)] -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec
-> [(String, [(LocVar, loc)], PreExp ext loc dec)]
-> PreExp ext loc dec
CaseE Exp2
scrt [(String, [(LocVar, LocVar)], Exp2)]
brs'
    TimeIt Exp2
e UrTy LocVar
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 LocVar -> Bool -> Exp2
forall (ext :: * -> * -> *) loc dec.
PreExp ext loc dec -> dec -> Bool -> PreExp ext loc dec
TimeIt Exp2
e' UrTy LocVar
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 LocVar
v Exp2
e                -> LocVar -> Exp2 -> Exp2
forall (ext :: * -> * -> *) loc dec.
LocVar -> PreExp ext loc dec -> PreExp ext loc dec
WithArenaE LocVar
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{}                        -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex                        -- Is there a recursion element to this?
    FoldE{}                       -> Exp2 -> PassM Exp2
forall a. a -> PassM a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp2
ex                        -- Is there a recursion element to this?
  where
    go :: Exp2 -> PassM Exp2
go = DelayedBindEnv -> Set LocVar -> Exp2 -> PassM Exp2
placeRegionInwards DelayedBindEnv
env Set LocVar
scopeSet

-- This is a function to discharge binds given a dictionary, scopeSet and expression where free variables might exist
dischargeBinds :: DelayedBindEnv -> S.Set Var -> Exp2 -> (DelayedBindEnv, Exp2)
dischargeBinds :: DelayedBindEnv -> Set LocVar -> Exp2 -> (DelayedBindEnv, Exp2)
dischargeBinds DelayedBindEnv
env Set LocVar
scopeSet Exp2
exp2 =
  let free_vars :: Set LocVar
free_vars        = Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
S.difference (Exp2 -> Set LocVar
freeVars Exp2
exp2) Set LocVar
scopeSet                         -- Take the difference of the scopeSet with the set that freeVar gives.
      (DelayedBindEnv
newEnv, Exp2
newExp) = Set LocVar -> DelayedBindEnv -> Exp2 -> (DelayedBindEnv, Exp2)
codeGen Set LocVar
free_vars DelayedBindEnv
env Exp2
exp2
  in  (DelayedBindEnv
newEnv, Exp2
newExp)

-- This is a duplicate function to the one above but instead it takes a Set of LocVar to codeGen directly instead of the expression and scopeSet.
dischargeBinds' :: DelayedBindEnv -> S.Set LocVar -> Exp2 -> (DelayedBindEnv, Exp2)
dischargeBinds' :: DelayedBindEnv -> Set LocVar -> Exp2 -> (DelayedBindEnv, Exp2)
dischargeBinds' DelayedBindEnv
env Set LocVar
free_vars Exp2
exp2 = do Set LocVar -> DelayedBindEnv -> Exp2 -> (DelayedBindEnv, Exp2)
codeGen Set LocVar
free_vars DelayedBindEnv
env Exp2
exp2

-- Use this function to codegen from the env by giving a set of variables you want to codegen from
codeGen :: S.Set LocVar -> DelayedBindEnv -> Exp2 -> (DelayedBindEnv, Exp2)
codeGen :: Set LocVar -> DelayedBindEnv -> Exp2 -> (DelayedBindEnv, Exp2)
codeGen Set LocVar
set DelayedBindEnv
env Exp2
body =
  let allKeys :: [Set LocVar]
allKeys   =  DelayedBindEnv -> [Set LocVar]
forall k a. Map k a -> [k]
M.keys DelayedBindEnv
env                                                          -- List of all keys from env
      keyList :: [Maybe (Set LocVar)]
keyList   = (LocVar -> Maybe (Set LocVar)) -> [LocVar] -> [Maybe (Set LocVar)]
forall a b. (a -> b) -> [a] -> [b]
map (\LocVar
variable -> (Set LocVar -> Bool) -> [Set LocVar] -> Maybe (Set LocVar)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
F.find (LocVar -> Set LocVar -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member LocVar
variable) [Set LocVar]
allKeys ) (Set LocVar -> [LocVar]
forall a. Set a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Set LocVar
set)  -- For each var in the input set find its corresponding key
      keyList' :: [Set LocVar]
keyList'  = Set (Set LocVar) -> [Set LocVar]
forall a. Set a -> [a]
S.toList (Set (Set LocVar) -> [Set LocVar])
-> Set (Set LocVar) -> [Set LocVar]
forall a b. (a -> b) -> a -> b
$ [Set LocVar] -> Set (Set LocVar)
forall a. Ord a => [a] -> Set a
S.fromList ([Set LocVar] -> Set (Set LocVar))
-> [Set LocVar] -> Set (Set LocVar)
forall a b. (a -> b) -> a -> b
$ [Maybe (Set LocVar)] -> [Set LocVar]
forall a. [Maybe a] -> [a]
S.catMaybes [Maybe (Set LocVar)]
keyList                          -- Filter out all the Nothing values from the list and let only Just values in the list
      valList :: [DelayedBind]
valList   = (Set LocVar -> [DelayedBind]) -> [Set LocVar] -> [DelayedBind]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\Set LocVar
key -> [DelayedBind] -> Set LocVar -> DelayedBindEnv -> [DelayedBind]
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault [] Set LocVar
key DelayedBindEnv
env) [Set LocVar]
keyList'            -- For each key in the keyList from before find the value associated with the key
      newKeys :: [Set LocVar]
newKeys   = Set (Set LocVar) -> [Set LocVar]
forall a. Set a -> [a]
S.toList (Set (Set LocVar) -> [Set LocVar])
-> Set (Set LocVar) -> [Set LocVar]
forall a b. (a -> b) -> a -> b
$ [Set LocVar] -> Set (Set LocVar)
forall a. Ord a => [a] -> Set a
S.fromList [Set LocVar]
allKeys Set (Set LocVar) -> Set (Set LocVar) -> Set (Set LocVar)
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` [Set LocVar] -> Set (Set LocVar)
forall a. Ord a => [a] -> Set a
S.fromList [Set LocVar]
keyList'     -- Filter all the Nothing values from the list and let only Just values in the list
      newVals :: [[DelayedBind]]
newVals   = (Set LocVar -> [DelayedBind]) -> [Set LocVar] -> [[DelayedBind]]
forall a b. (a -> b) -> [a] -> [b]
map (\Set LocVar
key -> [DelayedBind] -> Set LocVar -> DelayedBindEnv -> [DelayedBind]
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault [] Set LocVar
key DelayedBindEnv
env) [Set LocVar]
newKeys
      tupleList :: [(Set LocVar, [DelayedBind])]
tupleList = [Set LocVar] -> [[DelayedBind]] -> [(Set LocVar, [DelayedBind])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Set LocVar]
newKeys [[DelayedBind]]
newVals
      newEnv' :: DelayedBindEnv
newEnv'   = [(Set LocVar, [DelayedBind])] -> DelayedBindEnv
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Set LocVar, [DelayedBind])]
tupleList
      exps :: Exp2
exps      = (DelayedBind -> Exp2 -> Exp2) -> Exp2 -> [DelayedBind] -> Exp2
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr DelayedBind -> Exp2 -> Exp2
bindDelayedBind Exp2
body [DelayedBind]
valList                                   -- Get all the bindings for all the expressions in the key
   in {-dbgTraceIt (sdoc set)-} (DelayedBindEnv
newEnv', Exp2
exps)


bindDelayedBind :: DelayedBind -> Exp2 -> Exp2
bindDelayedBind :: DelayedBind -> Exp2 -> Exp2
bindDelayedBind DelayedBind
delayed Exp2
body =
  case DelayedBind
delayed of
    DelayRegion Region
r RegionSize
sz Maybe RegionType
ty -> E2Ext LocVar (UrTy LocVar) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocVar (UrTy LocVar) -> Exp2)
-> E2Ext LocVar (UrTy LocVar) -> Exp2
forall a b. (a -> b) -> a -> b
$ Region
-> RegionSize
-> Maybe RegionType
-> Exp2
-> E2Ext LocVar (UrTy LocVar)
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
body
    DelayParRegion Region
r RegionSize
sz Maybe RegionType
ty -> E2Ext LocVar (UrTy LocVar) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocVar (UrTy LocVar) -> Exp2)
-> E2Ext LocVar (UrTy LocVar) -> Exp2
forall a b. (a -> b) -> a -> b
$ Region
-> RegionSize
-> Maybe RegionType
-> Exp2
-> E2Ext LocVar (UrTy LocVar)
forall loc dec.
Region
-> RegionSize -> Maybe RegionType -> E2 loc dec -> E2Ext loc dec
LetParRegionE Region
r RegionSize
sz Maybe RegionType
ty Exp2
body
    DelayLoc LocVar
loc LocExp
locexp -> E2Ext LocVar (UrTy LocVar) -> Exp2
forall (ext :: * -> * -> *) loc dec.
ext loc dec -> PreExp ext loc dec
Ext (E2Ext LocVar (UrTy LocVar) -> Exp2)
-> E2Ext LocVar (UrTy LocVar) -> Exp2
forall a b. (a -> b) -> a -> b
$ LocVar -> LocExp -> Exp2 -> E2Ext LocVar (UrTy LocVar)
forall loc dec.
LocVar -> PreLocExp loc -> E2 loc dec -> E2Ext loc dec
LetLocE LocVar
loc LocExp
locexp Exp2
body


-- A function for use specific to this pass which gives all the possible variables and local variables that are used in a particular expression
-- This pass was made speciic because other version in gibbon don't return location variables, this version also adds location variables to the
-- returned set

freeVars :: Exp2 -> S.Set Var
freeVars :: Exp2 -> Set LocVar
freeVars Exp2
ex = case Exp2
ex of
  Ext E2Ext LocVar (UrTy LocVar)
ext                           ->
    case E2Ext LocVar (UrTy LocVar)
ext of
      LetRegionE Region
_ RegionSize
_ Maybe RegionType
_ Exp2
rhs          -> Exp2 -> Set LocVar
freeVars Exp2
rhs
      LetLocE LocVar
_ LocExp
phs Exp2
rhs             ->
        case LocExp
phs of
        StartOfRegionLE Region
_                 -> Exp2 -> Set LocVar
freeVars Exp2
rhs
        AfterConstantLE Int
_ LocVar
_         -> Exp2 -> Set LocVar
freeVars Exp2
rhs
        AfterVariableLE{}           -> Exp2 -> Set LocVar
freeVars Exp2
rhs
        InRegionLE Region
_                -> Exp2 -> Set LocVar
freeVars Exp2
rhs
        FromEndLE LocVar
_                 -> Exp2 -> Set LocVar
freeVars Exp2
rhs
        LocExp
_                           -> Set LocVar
forall a. Set a
S.empty
      E2Ext LocVar (UrTy LocVar)
_                             -> Set LocVar
forall a. Set a
S.empty

  LetE (LocVar
_,[LocVar]
locs, UrTy LocVar
ty,Exp2
rhs) Exp2
bod         -> [LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList [LocVar]
locs  Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union` [LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList (UrTy LocVar -> [LocVar]
locsInTy UrTy LocVar
ty) Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union` Exp2 -> Set LocVar
freeVars Exp2
rhs Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union` Exp2 -> Set LocVar
freeVars Exp2
bod
  LitE Int
_                            -> Set LocVar
forall a. Set a
S.empty
  LitSymE LocVar
_                         -> Set LocVar
forall a. Set a
S.empty
  VarE LocVar
v                            -> LocVar -> Set LocVar
forall a. a -> Set a
S.singleton LocVar
v
  AppE LocVar
v [LocVar]
locvarList [Exp2]
ls              -> [Set LocVar] -> Set LocVar
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions ((Exp2 -> Set LocVar) -> [Exp2] -> [Set LocVar]
forall a b. (a -> b) -> [a] -> [b]
L.map Exp2 -> Set LocVar
freeVars [Exp2]
ls) Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union` LocVar -> Set LocVar
forall a. a -> Set a
S.singleton LocVar
v Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union` [LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList [LocVar]
locvarList
  PrimAppE Prim (UrTy LocVar)
_ [Exp2]
ls                     -> [Set LocVar] -> Set LocVar
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions ((Exp2 -> Set LocVar) -> [Exp2] -> [Set LocVar]
forall a b. (a -> b) -> [a] -> [b]
L.map Exp2 -> Set LocVar
freeVars [Exp2]
ls)
  MkProdE [Exp2]
ls                        -> [Set LocVar] -> Set LocVar
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions ((Exp2 -> Set LocVar) -> [Exp2] -> [Set LocVar]
forall a b. (a -> b) -> [a] -> [b]
L.map Exp2 -> Set LocVar
freeVars [Exp2]
ls)
  DataConE LocVar
locVar String
_ [Exp2]
ls              -> LocVar -> Set LocVar
forall a. a -> Set a
S.singleton LocVar
locVar  Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union`  [Set LocVar] -> Set LocVar
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions ((Exp2 -> Set LocVar) -> [Exp2] -> [Set LocVar]
forall a b. (a -> b) -> [a] -> [b]
L.map Exp2 -> Set LocVar
freeVars [Exp2]
ls)
  ProjE Int
_ Exp2
e                         -> Exp2 -> Set LocVar
freeVars Exp2
e
  IfE Exp2
e1 Exp2
e2 Exp2
e3                      -> [Set LocVar] -> Set LocVar
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions [Exp2 -> Set LocVar
freeVars Exp2
e1, Exp2 -> Set LocVar
freeVars Exp2
e2, Exp2 -> Set LocVar
freeVars Exp2
e3]
  CaseE Exp2
e [(String, [(LocVar, LocVar)], Exp2)]
ls                        -> Exp2 -> Set LocVar
freeVars Exp2
e Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union`
                                        [Set LocVar] -> Set LocVar
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions (((String, [(LocVar, LocVar)], Exp2) -> Set LocVar)
-> [(String, [(LocVar, LocVar)], Exp2)] -> [Set LocVar]
forall a b. (a -> b) -> [a] -> [b]
L.map (\(String
_, [(LocVar, LocVar)]
vlocs, Exp2
ee) ->
                                           let ([LocVar]
vars, [LocVar]
locVars) = [(LocVar, LocVar)] -> ([LocVar], [LocVar])
forall a b. [(a, b)] -> ([a], [b])
unzip [(LocVar, LocVar)]
vlocs
                                           in Exp2 -> Set LocVar
freeVars Exp2
ee Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union` [LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList [LocVar]
vars Set LocVar -> Set LocVar -> Set LocVar
forall a. Ord a => Set a -> Set a -> Set a
`S.union` [LocVar] -> Set LocVar
forall a. Ord a => [a] -> Set a
S.fromList [LocVar]
locVars) [(String, [(LocVar, LocVar)], Exp2)]
ls)
  Exp2
_                                 -> Set LocVar
forall a. Set a
S.empty