-- | An intermediate language with an effect system that captures traversals.
--
-- ASSUMES that the flatten pass has run, and thus we have trivial AppE operands.
--

module Gibbon.Passes.InferEffects
  (inferEffects, inferExp) where

import qualified Data.List as L
import Data.Set as S
import Data.Map as M

import Gibbon.Common
import Gibbon.L2.Syntax

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

-- | Chatter level for this module:
lvl :: Int
lvl :: Int
lvl = Int
5

type FunEnv2 = M.Map Var (ArrowTy2 Ty2)

type Deps = M.Map LocVar LocVar

updateDeps :: S.Set Effect -> Deps -> Deps
updateDeps :: Set Effect -> Deps -> Deps
updateDeps Set Effect
s = (LocVar -> Maybe LocVar) -> Deps -> Deps
forall a b k. (a -> Maybe b) -> Map k a -> Map k b
M.mapMaybe (\LocVar
lv -> if LocVar -> Set LocVar -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member LocVar
lv Set LocVar
ls then Maybe LocVar
forall a. Maybe a
Nothing else LocVar -> Maybe LocVar
forall a. a -> Maybe a
Just LocVar
lv)
  where ls :: Set LocVar
ls = (Effect -> LocVar) -> Set Effect -> Set LocVar
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map (\(Traverse LocVar
lv) -> LocVar
lv) Set Effect
s
-- idea: remove entries in map when they are satisfied by effect

metDep :: Deps -> S.Set Effect -> S.Set Effect
metDep :: Deps -> Set Effect -> Set Effect
metDep Deps
dps = (Effect -> Bool) -> Set Effect -> Set Effect
forall a. (a -> Bool) -> Set a -> Set a
S.filter (\(Traverse LocVar
lv) -> Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ LocVar -> Deps -> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.member LocVar
lv Deps
dps)
-- idea: remove effects that have (unmet) entries in map

locsEffect :: [LocVar] -> Set Effect
locsEffect :: [LocVar] -> Set Effect
locsEffect = [Effect] -> Set Effect
forall a. Ord a => [a] -> Set a
S.fromList ([Effect] -> Set Effect)
-> ([LocVar] -> [Effect]) -> [LocVar] -> Set Effect
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (LocVar -> Effect) -> [LocVar] -> [Effect]
forall a b. (a -> b) -> [a] -> [b]
L.map LocVar -> Effect
Traverse

-- | We initially populate all functions with MAXIMUM effect signatures.
--   Subsequently, these monotonically SHRINK until a fixpoint.
--   We also associate fresh location variables with packed types.
initialEnv :: FunDefs2 -> FunEnv2
initialEnv :: FunDefs2 -> FunEnv2
initialEnv FunDefs2
mp = (FunDef2 -> ArrowTy2 Ty2) -> FunDefs2 -> FunEnv2
forall a b k. (a -> b) -> Map k a -> Map k b
M.map FunDef2 -> ArrowTy2 Ty2
go FunDefs2
mp
  where
    go :: FunDef2 -> (ArrowTy2 Ty2)
    go :: FunDef2 -> ArrowTy2 Ty2
go FunDef{ArrowTy (TyOf Exp2)
funTy :: ArrowTy (TyOf Exp2)
funTy :: forall ex. FunDef ex -> ArrowTy (TyOf ex)
funTy} =
      let locs :: [LocVar]
locs       = ArrowTy2 Ty2 -> [LocVar]
forall ty2. ArrowTy2 ty2 -> [LocVar]
allLocVars ArrowTy (TyOf Exp2)
ArrowTy2 Ty2
funTy
          maxEffects :: Set Effect
maxEffects = [LocVar] -> Set Effect
locsEffect [LocVar]
locs
      in ArrowTy (TyOf Exp2)
ArrowTy2 Ty2
funTy { arrEffs :: Set Effect
arrEffs = Set Effect
maxEffects }


inferEffects :: Prog2 -> PassM Prog2
inferEffects :: Prog2 -> PassM Prog2
inferEffects prg :: Prog2
prg@Prog{DDefs (TyOf Exp2)
ddefs :: DDefs (TyOf Exp2)
ddefs :: forall ex. Prog ex -> DDefs (TyOf ex)
ddefs,FunDefs2
fundefs :: FunDefs2
fundefs :: forall ex. Prog ex -> FunDefs ex
fundefs} = do
  let finalFunTys :: FunEnv2
finalFunTys = Int -> FunDefs2 -> FunEnv2 -> FunEnv2
fixpoint Int
1 FunDefs2
fundefs (FunDefs2 -> FunEnv2
initialEnv FunDefs2
fundefs)
      funs :: FunDefs2
funs = (FunDef2 -> FunDef2) -> FunDefs2 -> FunDefs2
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (\fn :: FunDef2
fn@FunDef{LocVar
funName :: LocVar
funName :: forall ex. FunDef ex -> LocVar
funName} ->
                       FunDef2
fn{ funTy :: ArrowTy (TyOf Exp2)
funTy = FunEnv2
finalFunTys FunEnv2 -> LocVar -> ArrowTy2 Ty2
forall k a. Ord k => Map k a -> k -> a
! LocVar
funName })
             FunDefs2
fundefs
  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
$ Prog2
prg { fundefs :: FunDefs2
fundefs = FunDefs2
funs }
  where
    fixpoint :: Int -> FunDefs2 -> FunEnv2 -> FunEnv2
    fixpoint :: Int -> FunDefs2 -> FunEnv2 -> FunEnv2
fixpoint Int
iter FunDefs2
funs FunEnv2
fenv =
       let funtys :: FunEnv2
funtys = (FunDef2 -> ArrowTy2 Ty2) -> FunDefs2 -> FunEnv2
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (DDefs Ty2 -> FunEnv2 -> FunDef2 -> ArrowTy2 Ty2
inferFunDef DDefs (TyOf Exp2)
DDefs Ty2
ddefs FunEnv2
fenv) FunDefs2
funs
       in
         if FunEnv2
fenv FunEnv2 -> FunEnv2 -> Bool
forall a. Eq a => a -> a -> Bool
== FunEnv2
funtys
         then Int -> [Char] -> FunEnv2 -> FunEnv2
forall a. Int -> [Char] -> a -> a
dbgTrace Int
lvl ([Char]
"\n<== Fixpoint completed after iteration "[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++Int -> [Char]
forall a. Show a => a -> [Char]
show Int
iter[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++[Char]
" ==>") (FunEnv2 -> FunEnv2) -> FunEnv2 -> FunEnv2
forall a b. (a -> b) -> a -> b
$ FunEnv2
fenv
         else Int -> FunDefs2 -> FunEnv2 -> FunEnv2
fixpoint (Int
iterInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) FunDefs2
funs FunEnv2
funtys


inferFunDef :: DDefs Ty2 -> FunEnv2 -> FunDef2 -> ArrowTy2 Ty2
inferFunDef :: DDefs Ty2 -> FunEnv2 -> FunDef2 -> ArrowTy2 Ty2
inferFunDef DDefs Ty2
ddfs FunEnv2
fenv FunDef{[LocVar]
funArgs :: [LocVar]
funArgs :: forall ex. FunDef ex -> [LocVar]
funArgs,Exp2
funBody :: Exp2
funBody :: forall ex. FunDef ex -> ex
funBody,ArrowTy (TyOf Exp2)
funTy :: forall ex. FunDef ex -> ArrowTy (TyOf ex)
funTy :: ArrowTy (TyOf Exp2)
funTy} = ArrowTy (TyOf Exp2)
ArrowTy2 Ty2
funTy { arrEffs :: Set Effect
arrEffs = Set Effect -> Set Effect -> Set Effect
forall a. Ord a => Set a -> Set a -> Set a
S.intersection Set Effect
travs Set Effect
eff }
  where
    env0 :: Map LocVar Ty2
env0  = [(LocVar, Ty2)] -> Map LocVar Ty2
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(LocVar, Ty2)] -> Map LocVar Ty2)
-> [(LocVar, Ty2)] -> Map LocVar Ty2
forall a b. (a -> b) -> a -> b
$ [LocVar] -> [Ty2] -> [(LocVar, Ty2)]
forall a b. [a] -> [b] -> [(a, b)]
zip [LocVar]
funArgs (ArrowTy2 Ty2 -> [Ty2]
forall ty2. ArrowTy2 ty2 -> [ty2]
arrIns ArrowTy (TyOf Exp2)
ArrowTy2 Ty2
funTy)
    travs :: Set Effect
travs = [Effect] -> Set Effect
forall a. Ord a => [a] -> Set a
S.fromList ([Effect] -> Set Effect) -> [Effect] -> Set Effect
forall a b. (a -> b) -> a -> b
$ (LocVar -> Effect) -> [LocVar] -> [Effect]
forall a b. (a -> b) -> [a] -> [b]
L.map LocVar -> Effect
Traverse ([LocVar] -> [Effect]) -> [LocVar] -> [Effect]
forall a b. (a -> b) -> a -> b
$ ArrowTy2 Ty2 -> [LocVar]
forall ty2. ArrowTy2 ty2 -> [LocVar]
inLocVars ArrowTy (TyOf Exp2)
ArrowTy2 Ty2
funTy
    (Set Effect
eff,Maybe LocVar
_outLoc) = DDefs Ty2
-> FunEnv2
-> Map LocVar Ty2
-> Deps
-> Exp2
-> (Set Effect, Maybe LocVar)
inferExp DDefs Ty2
ddfs FunEnv2
fenv Map LocVar Ty2
env0 Deps
forall k a. Map k a
M.empty Exp2
funBody


inferExp :: DDefs Ty2 -> FunEnv2 -> TyEnv Ty2 -> Deps -> Exp2 -> (Set Effect, Maybe LocVar)
inferExp :: DDefs Ty2
-> FunEnv2
-> Map LocVar Ty2
-> Deps
-> Exp2
-> (Set Effect, Maybe LocVar)
inferExp DDefs Ty2
ddfs FunEnv2
fenv Map LocVar Ty2
env Deps
dps Exp2
expr =
  case Exp2
expr of
    -- QUESTION: does a variable reference count as traversing to the end?
    -- If so, the identity function has the traverse effect.
    -- I'd prefer that the identity function get type (Tree_p -{}-> Tree_p).
    VarE LocVar
v -> case LocVar -> Map LocVar Ty2 -> Maybe Ty2
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup LocVar
v Map LocVar Ty2
env of
                Just Ty2
ty -> (Set Effect
forall a. Set a
S.empty, Ty2 -> Maybe LocVar
packedLoc Ty2
ty)
                Maybe Ty2
Nothing -> [Char] -> (Set Effect, Maybe LocVar)
forall a. HasCallStack => [Char] -> a
error ([Char] -> (Set Effect, Maybe LocVar))
-> [Char] -> (Set Effect, Maybe LocVar)
forall a b. (a -> b) -> a -> b
$ [Char]
"Unknown var: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ LocVar -> [Char]
forall a. Out a => a -> [Char]
sdoc LocVar
v

    LitE Int
_    -> (Set Effect
forall a. Set a
S.empty, Maybe LocVar
forall a. Maybe a
Nothing)
    CharE Char
_   -> (Set Effect
forall a. Set a
S.empty, Maybe LocVar
forall a. Maybe a
Nothing)
    FloatE{}  -> (Set Effect
forall a. Set a
S.empty, Maybe LocVar
forall a. Maybe a
Nothing)
    LitSymE LocVar
_ -> (Set Effect
forall a. Set a
S.empty, Maybe LocVar
forall a. Maybe a
Nothing)

    AppE LocVar
v [LocVar]
locs [Exp2]
_e ->
      -- Substitue locations used at this particular call-site in the function
      -- effects computed so far
      let orgLocs :: [LocVar]
orgLocs = ArrowTy2 Ty2 -> [LocVar]
forall ty2. ArrowTy2 ty2 -> [LocVar]
allLocVars (FunEnv2
fenv FunEnv2 -> LocVar -> ArrowTy2 Ty2
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# LocVar
v)
          locMap :: Deps
locMap  = [(LocVar, LocVar)] -> Deps
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(LocVar, LocVar)] -> Deps) -> [(LocVar, LocVar)] -> Deps
forall a b. (a -> b) -> a -> b
$ [LocVar] -> [LocVar] -> [(LocVar, LocVar)]
forall a b. [a] -> [b] -> [(a, b)]
zip [LocVar]
orgLocs [LocVar]
locs
          eff :: Set Effect
eff     = ArrowTy2 Ty2 -> Set Effect
forall ty2. ArrowTy2 ty2 -> Set Effect
arrEffs (FunEnv2
fenv FunEnv2 -> LocVar -> ArrowTy2 Ty2
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# LocVar
v)
      in (Deps -> Set Effect -> Set Effect
metDep Deps
dps (Set Effect -> Set Effect) -> Set Effect -> Set Effect
forall a b. (a -> b) -> a -> b
$ Deps -> Set Effect -> Set Effect
substEffs Deps
locMap Set Effect
eff, Maybe LocVar
forall a. Maybe a
Nothing)

    -- If rands are already trivial, no traversal effects can occur here.
    -- All primitives operate on non-packed data.
    PrimAppE Prim Ty2
_ [Exp2]
rands -> [Exp2] -> (Set Effect, Maybe LocVar) -> (Set Effect, Maybe LocVar)
forall e a. (HasCallStack, Expression e) => [e] -> a -> a
assertTrivs [Exp2]
rands (Set Effect
forall a. Set a
S.empty, Maybe LocVar
forall a. Maybe a
Nothing)

    -- TODO: what would _locs have here ?
    LetE (LocVar
v,[LocVar]
_locs,Ty2
ty,Exp2
rhs) Exp2
bod ->
      let (Set Effect
effRhs,Maybe LocVar
_rhsLoc) = DDefs Ty2
-> FunEnv2
-> Map LocVar Ty2
-> Deps
-> Exp2
-> (Set Effect, Maybe LocVar)
inferExp DDefs Ty2
ddfs FunEnv2
fenv Map LocVar Ty2
env Deps
dps Exp2
rhs
          dps' :: Deps
dps' = Set Effect -> Deps -> Deps
updateDeps Set Effect
effRhs Deps
dps
          -- TODO: extend env with rhsLoc ? or _locs ?
          (Set Effect
effBod,Maybe LocVar
bLoc) = DDefs Ty2
-> FunEnv2
-> Map LocVar Ty2
-> Deps
-> Exp2
-> (Set Effect, Maybe LocVar)
inferExp DDefs Ty2
ddfs FunEnv2
fenv (LocVar -> Ty2 -> Map LocVar Ty2 -> Map LocVar Ty2
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert LocVar
v Ty2
ty Map LocVar Ty2
env) Deps
dps' Exp2
bod
      in (Set Effect -> Set Effect -> Set Effect
forall a. Ord a => Set a -> Set a -> Set a
S.union Set Effect
effRhs Set Effect
effBod, Maybe LocVar
bLoc)

    -- TODO: do we need to join locC and locA
    IfE Exp2
tst Exp2
consq Exp2
alt ->
      let (Set Effect
effT,Maybe LocVar
_locT) = DDefs Ty2
-> FunEnv2
-> Map LocVar Ty2
-> Deps
-> Exp2
-> (Set Effect, Maybe LocVar)
inferExp DDefs Ty2
ddfs FunEnv2
fenv Map LocVar Ty2
env Deps
dps Exp2
tst
          (Set Effect
effC,Maybe LocVar
locC) = DDefs Ty2
-> FunEnv2
-> Map LocVar Ty2
-> Deps
-> Exp2
-> (Set Effect, Maybe LocVar)
inferExp DDefs Ty2
ddfs FunEnv2
fenv Map LocVar Ty2
env Deps
dps Exp2
consq
          (Set Effect
effA,Maybe LocVar
locA) = DDefs Ty2
-> FunEnv2
-> Map LocVar Ty2
-> Deps
-> Exp2
-> (Set Effect, Maybe LocVar)
inferExp DDefs Ty2
ddfs FunEnv2
fenv Map LocVar Ty2
env Deps
dps Exp2
alt
          loc :: Maybe LocVar
loc = case (Maybe LocVar
locC,Maybe LocVar
locA) of
                  (Maybe LocVar
Nothing,Maybe LocVar
Nothing)  -> Maybe LocVar
forall a. Maybe a
Nothing
                  (Just LocVar
l', Maybe LocVar
Nothing) -> LocVar -> Maybe LocVar
forall a. a -> Maybe a
Just LocVar
l'
                  (Maybe LocVar
Nothing, Just LocVar
l') -> LocVar -> Maybe LocVar
forall a. a -> Maybe a
Just LocVar
l'
                  (Just LocVar
l', Just LocVar
_m) -> LocVar -> Maybe LocVar
forall a. a -> Maybe a
Just LocVar
l'
      in (Set Effect -> Set Effect -> Set Effect
forall a. Ord a => Set a -> Set a -> Set a
S.union Set Effect
effT (Set Effect -> Set Effect -> Set Effect
forall a. Ord a => Set a -> Set a -> Set a
S.intersection Set Effect
effC Set Effect
effA), Maybe LocVar
loc)

    -- ignore locations, won't have any effect on the generated effects. ??
    MkProdE [Exp2]
ls ->
      let ([Set Effect]
effs, [Maybe LocVar]
_locs) = [(Set Effect, Maybe LocVar)] -> ([Set Effect], [Maybe LocVar])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Set Effect, Maybe LocVar)] -> ([Set Effect], [Maybe LocVar]))
-> [(Set Effect, Maybe LocVar)] -> ([Set Effect], [Maybe LocVar])
forall a b. (a -> b) -> a -> b
$ (Exp2 -> (Set Effect, Maybe LocVar))
-> [Exp2] -> [(Set Effect, Maybe LocVar)]
forall a b. (a -> b) -> [a] -> [b]
L.map (DDefs Ty2
-> FunEnv2
-> Map LocVar Ty2
-> Deps
-> Exp2
-> (Set Effect, Maybe LocVar)
inferExp DDefs Ty2
ddfs FunEnv2
fenv Map LocVar Ty2
env Deps
dps) [Exp2]
ls
      in ([Set Effect] -> Set Effect
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions [Set Effect]
effs, Maybe LocVar
forall a. Maybe a
Nothing)

    SpawnE{} -> (Set Effect
forall a. Set a
S.empty, Maybe LocVar
forall a. Maybe a
Nothing)

    Exp2
SyncE -> (Set Effect
forall a. Set a
S.empty, Maybe LocVar
forall a. Maybe a
Nothing)

    ProjE Int
_n Exp2
e ->
      let (Set Effect
eff, Maybe LocVar
_loc) = DDefs Ty2
-> FunEnv2
-> Map LocVar Ty2
-> Deps
-> Exp2
-> (Set Effect, Maybe LocVar)
inferExp DDefs Ty2
ddfs FunEnv2
fenv Map LocVar Ty2
env Deps
dps Exp2
e
      in (Set Effect
eff, Maybe LocVar
forall a. Maybe a
Nothing)

    CaseE Exp2
e [([Char], [(LocVar, LocVar)], Exp2)]
mp ->
      let (Set Effect
eff,Maybe LocVar
loc1) = DDefs Ty2
-> FunEnv2
-> Map LocVar Ty2
-> Deps
-> Exp2
-> (Set Effect, Maybe LocVar)
inferExp DDefs Ty2
ddfs FunEnv2
fenv Map LocVar Ty2
env Deps
dps Exp2
e
          ([Bool]
bools,[(Set Effect, Maybe LocVar)]
effsLocs) = [(Bool, (Set Effect, Maybe LocVar))]
-> ([Bool], [(Set Effect, Maybe LocVar)])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Bool, (Set Effect, Maybe LocVar))]
 -> ([Bool], [(Set Effect, Maybe LocVar)]))
-> [(Bool, (Set Effect, Maybe LocVar))]
-> ([Bool], [(Set Effect, Maybe LocVar)])
forall a b. (a -> b) -> a -> b
$ (([Char], [(LocVar, LocVar)], Exp2)
 -> (Bool, (Set Effect, Maybe LocVar)))
-> [([Char], [(LocVar, LocVar)], Exp2)]
-> [(Bool, (Set Effect, Maybe LocVar))]
forall a b. (a -> b) -> [a] -> [b]
L.map ([Char], [(LocVar, LocVar)], Exp2)
-> (Bool, (Set Effect, Maybe LocVar))
caserhs [([Char], [(LocVar, LocVar)], Exp2)]
mp
          ([Set Effect]
effs,[Maybe LocVar]
_) = [(Set Effect, Maybe LocVar)] -> ([Set Effect], [Maybe LocVar])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Set Effect, Maybe LocVar)]
effsLocs

          -- Should we check that we actually _have_ all cases ? Or are incomplete case
          -- matches enough for traversal ?

          -- Critical policy point!  We only get to the end if ALL
          -- branches get to the end.
          end :: Set Effect
end = if (Bool -> Bool) -> [Bool] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Bool -> Bool
forall a. a -> a
id [Bool]
bools
                then case Maybe LocVar
loc1 of
                       Just LocVar
v  -> Effect -> Set Effect
forall a. a -> Set a
S.singleton (LocVar -> Effect
Traverse LocVar
v)
                       Maybe LocVar
Nothing -> Set Effect
forall a. Set a
S.empty
                else Set Effect
forall a. Set a
S.empty
          ret :: Set Effect
ret = Set Effect -> Set Effect -> Set Effect
forall a. Ord a => Set a -> Set a -> Set a
S.union (Set Effect -> Set Effect -> Set Effect
forall a. Ord a => Set a -> Set a -> Set a
S.union Set Effect
eff Set Effect
end)
                        ((Set Effect -> Set Effect -> Set Effect)
-> [Set Effect] -> Set Effect
forall a. (a -> a -> a) -> [a] -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
L.foldl1 Set Effect -> Set Effect -> Set Effect
forall a. Ord a => Set a -> Set a -> Set a
S.intersection [Set Effect]
effs)
      in (Set Effect
ret, Maybe LocVar
forall a. Maybe a
Nothing)

    -- Construct output packed data.  We will always "scroll to the end" of
    -- output values, so they are not interesting for this effect analysis.
    DataConE LocVar
_loc [Char]
_dcon [Exp2]
es -> [Exp2] -> (Set Effect, Maybe LocVar) -> (Set Effect, Maybe LocVar)
forall e a. (HasCallStack, Expression e) => [e] -> a -> a
assertTrivs [Exp2]
es (Set Effect
forall a. Set a
S.empty, Maybe LocVar
forall a. Maybe a
Nothing)

    TimeIt Exp2
e Ty2
_ Bool
_ -> DDefs Ty2
-> FunEnv2
-> Map LocVar Ty2
-> Deps
-> Exp2
-> (Set Effect, Maybe LocVar)
inferExp DDefs Ty2
ddfs FunEnv2
fenv Map LocVar Ty2
env Deps
dps Exp2
e

    WithArenaE LocVar
_v Exp2
e -> DDefs Ty2
-> FunEnv2
-> Map LocVar Ty2
-> Deps
-> Exp2
-> (Set Effect, Maybe LocVar)
inferExp DDefs Ty2
ddfs FunEnv2
fenv Map LocVar Ty2
env Deps
dps Exp2
e

    MapE{} -> [Char] -> (Set Effect, Maybe LocVar)
forall a. HasCallStack => [Char] -> a
error [Char]
"inferEffects: MapE not handled."
    FoldE{} -> [Char] -> (Set Effect, Maybe LocVar)
forall a. HasCallStack => [Char] -> a
error [Char]
"inferEffects: FoldE not handled."

    Ext (LetRegionE Region
_ RegionSize
_ Maybe RegionType
_ Exp2
rhs) -> DDefs Ty2
-> FunEnv2
-> Map LocVar Ty2
-> Deps
-> Exp2
-> (Set Effect, Maybe LocVar)
inferExp DDefs Ty2
ddfs FunEnv2
fenv Map LocVar Ty2
env Deps
dps Exp2
rhs
    Ext (LetParRegionE Region
_ RegionSize
_ Maybe RegionType
_ Exp2
rhs) -> DDefs Ty2
-> FunEnv2
-> Map LocVar Ty2
-> Deps
-> Exp2
-> (Set Effect, Maybe LocVar)
inferExp DDefs Ty2
ddfs FunEnv2
fenv Map LocVar Ty2
env Deps
dps Exp2
rhs
    Ext (LetLocE LocVar
_ PreLocExp LocVar
_ Exp2
rhs)  -> DDefs Ty2
-> FunEnv2
-> Map LocVar Ty2
-> Deps
-> Exp2
-> (Set Effect, Maybe LocVar)
inferExp DDefs Ty2
ddfs FunEnv2
fenv Map LocVar Ty2
env Deps
dps Exp2
rhs
    Ext (StartOfPkdCursor{}) -> (Set Effect
forall a. Set a
S.empty, Maybe LocVar
forall a. Maybe a
Nothing)
    Ext (TagCursor{})   -> (Set Effect
forall a. Set a
S.empty, Maybe LocVar
forall a. Maybe a
Nothing)
    Ext (RetE [LocVar]
_ LocVar
_)         -> (Set Effect
forall a. Set a
S.empty, Maybe LocVar
forall a. Maybe a
Nothing)
    Ext (FromEndE LocVar
_ )      -> (Set Effect
forall a. Set a
S.empty, Maybe LocVar
forall a. Maybe a
Nothing)
    Ext (IndirectionE{})   -> (Set Effect
forall a. Set a
S.empty, Maybe LocVar
forall a. Maybe a
Nothing)
    Ext (BoundsCheck{})    -> (Set Effect
forall a. Set a
S.empty, Maybe LocVar
forall a. Maybe a
Nothing)
    Ext (AddFixed{})       -> [Char] -> (Set Effect, Maybe LocVar)
forall a. HasCallStack => [Char] -> a
error [Char]
"inferEffects: AddFixed not handled."
    Ext (E2Ext LocVar Ty2
GetCilkWorkerNum) -> (Set Effect
forall a. Set a
S.empty, Maybe LocVar
forall a. Maybe a
Nothing)
    Ext (LetAvail [LocVar]
_ Exp2
e)     -> DDefs Ty2
-> FunEnv2
-> Map LocVar Ty2
-> Deps
-> Exp2
-> (Set Effect, Maybe LocVar)
inferExp DDefs Ty2
ddfs FunEnv2
fenv Map LocVar Ty2
env Deps
dps Exp2
e
    Ext (AllocateTagHere{}) -> (Set Effect
forall a. Set a
S.empty, Maybe LocVar
forall a. Maybe a
Nothing)
    Ext (AllocateScalarsHere{}) -> (Set Effect
forall a. Set a
S.empty, Maybe LocVar
forall a. Maybe a
Nothing)
    Ext (SSPush{}) -> (Set Effect
forall a. Set a
S.empty, Maybe LocVar
forall a. Maybe a
Nothing)
    Ext (SSPop{}) -> (Set Effect
forall a. Set a
S.empty, Maybe LocVar
forall a. Maybe a
Nothing)

  where
    packedLoc :: Ty2 -> Maybe LocVar
    packedLoc :: Ty2 -> Maybe LocVar
packedLoc Ty2
ty = case Ty2
ty of
                     PackedTy [Char]
_ LocVar
loc -> LocVar -> Maybe LocVar
forall a. a -> Maybe a
Just LocVar
loc
                     Ty2
_ -> Maybe LocVar
forall a. Maybe a
Nothing

    caserhs :: (DataCon, [(Var,LocVar)], Exp2) -> (Bool, (Set Effect, Maybe LocVar))
    -- We've gotten "to the end" of a nullary constructor just by matching it:
    caserhs :: ([Char], [(LocVar, LocVar)], Exp2)
-> (Bool, (Set Effect, Maybe LocVar))
caserhs ([Char]
_dcon,[],Exp2
e) = ( Bool
True , DDefs Ty2
-> FunEnv2
-> Map LocVar Ty2
-> Deps
-> Exp2
-> (Set Effect, Maybe LocVar)
inferExp DDefs Ty2
ddfs FunEnv2
fenv Map LocVar Ty2
env Deps
dps Exp2
e )
    caserhs ([Char]
dcon,[(LocVar, LocVar)]
patVs,Exp2
e) =
      let ([LocVar]
vars,[LocVar]
locs) = [(LocVar, LocVar)] -> ([LocVar], [LocVar])
forall a b. [(a, b)] -> ([a], [b])
L.unzip [(LocVar, LocVar)]
patVs
          tys :: [Ty2]
tys    = DDefs Ty2 -> [Char] -> [Ty2]
forall a. Out a => DDefs a -> [Char] -> [a]
lookupDataCon DDefs Ty2
ddfs [Char]
dcon
          zipped :: [(LocVar, Ty2)]
zipped = [LocVar] -> [Ty2] -> [Char] -> [(LocVar, Ty2)]
forall a b.
(Show a, Show b, HasCallStack) =>
[a] -> [b] -> [Char] -> [(a, b)]
fragileZip' [LocVar]
vars [Ty2]
tys ([Char]
"Error in "[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
dcon [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++[Char]
" case: "
                                         [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++[Char]
"pattern vars, "[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++[LocVar] -> [Char]
forall a. Show a => a -> [Char]
show [LocVar]
vars[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++
                                         [Char]
", do not match the number of types "
                                         [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++[Ty2] -> [Char]
forall a. Show a => a -> [Char]
show [Ty2]
tys)

          env' :: Map LocVar Ty2
env' = Map LocVar Ty2 -> Map LocVar Ty2 -> Map LocVar Ty2
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union Map LocVar Ty2
env ([(LocVar, Ty2)] -> Map LocVar Ty2
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(LocVar, Ty2)]
zipped)

          packedOnly :: [(LocVar, Ty2)]
packedOnly = ((LocVar, Ty2) -> Bool) -> [(LocVar, Ty2)] -> [(LocVar, Ty2)]
forall a. (a -> Bool) -> [a] -> [a]
L.filter (\(LocVar
_,Ty2
t) -> Ty2 -> Bool
forall a. Show a => UrTy a -> Bool
hasPacked Ty2
t) [(LocVar, Ty2)]
zipped

          makeDps :: [(LocVar, Ty2)] -> Deps
makeDps [] = Deps
dps
          makeDps [(LocVar, Ty2)
_] = Deps
dps
          makeDps ((LocVar
loc,Ty2
ty):[(LocVar, Ty2)]
lts)
            | Ty2 -> Bool
forall a. Show a => UrTy a -> Bool
hasPacked Ty2
ty = case ((LocVar, Ty2) -> Bool) -> [(LocVar, Ty2)] -> [(LocVar, Ty2)]
forall a. (a -> Bool) -> [a] -> [a]
L.filter (Ty2 -> Bool
forall a. Show a => UrTy a -> Bool
hasPacked (Ty2 -> Bool) -> ((LocVar, Ty2) -> Ty2) -> (LocVar, Ty2) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (LocVar, Ty2) -> Ty2
forall a b. (a, b) -> b
snd) [(LocVar, Ty2)]
lts of
                               (LocVar
x,Ty2
_):[(LocVar, Ty2)]
_ -> LocVar -> LocVar -> Deps -> Deps
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert LocVar
loc LocVar
x (Deps -> Deps) -> Deps -> Deps
forall a b. (a -> b) -> a -> b
$ [(LocVar, Ty2)] -> Deps
makeDps [(LocVar, Ty2)]
lts
                               [(LocVar, Ty2)]
_ -> [(LocVar, Ty2)] -> Deps
makeDps [(LocVar, Ty2)]
lts
            | Bool
otherwise = [(LocVar, Ty2)] -> Deps
makeDps [(LocVar, Ty2)]
lts

          dps' :: Deps
dps' = [(LocVar, Ty2)] -> Deps
makeDps ([(LocVar, Ty2)] -> [(LocVar, Ty2)]
forall a. [a] -> [a]
reverse ([(LocVar, Ty2)] -> [(LocVar, Ty2)])
-> [(LocVar, Ty2)] -> [(LocVar, Ty2)]
forall a b. (a -> b) -> a -> b
$ [LocVar] -> [Ty2] -> [(LocVar, Ty2)]
forall a b. [a] -> [b] -> [(a, b)]
zip [LocVar]
locs [Ty2]
tys)

          subst_mp :: Deps
subst_mp = [(LocVar, LocVar)] -> Deps
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(LocVar, LocVar)] -> Deps) -> [(LocVar, LocVar)] -> Deps
forall a b. (a -> b) -> a -> b
$ [LocVar] -> [LocVar] -> [(LocVar, LocVar)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Ty2 -> [LocVar]) -> [Ty2] -> [LocVar]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Ty2 -> [LocVar]
locsInTy [Ty2]
tys) [LocVar]
locs
          (Set Effect
eff,Maybe LocVar
_) = DDefs Ty2
-> FunEnv2
-> Map LocVar Ty2
-> Deps
-> Exp2
-> (Set Effect, Maybe LocVar)
inferExp DDefs Ty2
ddfs FunEnv2
fenv Map LocVar Ty2
env' Deps
dps' Exp2
e
          eff' :: Set Effect
eff' = Deps -> Set Effect -> Set Effect
substEffs Deps
subst_mp Set Effect
eff
          winner :: Bool
winner = -- If there is NO packed child data, then our object has static size:
                   ((Ty2 -> Bool) -> [Ty2] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
L.all (Bool -> Bool
not (Bool -> Bool) -> (Ty2 -> Bool) -> Ty2 -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ty2 -> Bool
forall a. Show a => UrTy a -> Bool
hasPacked) [Ty2]
tys) Bool -> Bool -> Bool
||

                   -- Or if all non-static items were traversed:
                   (case [(LocVar, Ty2)]
packedOnly of
                      [] -> Bool
False
                      [(LocVar, Ty2)]
ls -> let patVMap :: Deps
patVMap = [(LocVar, LocVar)] -> Deps
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(LocVar, LocVar)]
patVs
                                packedlocs :: [LocVar]
packedlocs = ((LocVar, Ty2) -> LocVar) -> [(LocVar, Ty2)] -> [LocVar]
forall a b. (a -> b) -> [a] -> [b]
L.map (\(LocVar
a,Ty2
_) -> Deps
patVMap Deps -> LocVar -> LocVar
forall a b.
(Ord a, Out a, Out b, Show a, HasCallStack) =>
Map a b -> a -> b
# LocVar
a) [(LocVar, Ty2)]
ls
                            in (LocVar -> Bool) -> [LocVar] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\LocVar
x -> Effect -> Set Effect -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member (LocVar -> Effect
Traverse LocVar
x) Set Effect
eff') [LocVar]
packedlocs)

                   -- Or maybe the last-use rule applies:
                   -- TODO

          -- Also, in any binding form we are obligated to not return
          -- our local bindings in traversal side effects:
          isNotLocal :: Effect -> Bool
isNotLocal (Traverse LocVar
v) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ LocVar -> [LocVar] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
L.elem LocVar
v [LocVar]
locs
          stripped :: Set Effect
stripped = (Effect -> Bool) -> Set Effect -> Set Effect
forall a. (a -> Bool) -> Set a -> Set a
S.filter Effect -> Bool
isNotLocal Set Effect
eff'
      in  ( Bool
winner, (Set Effect
stripped,Maybe LocVar
forall a. Maybe a
Nothing) )