Skip to content

Commit

Permalink
Better names for combinators that manipulate the state of OCCM actions.
Browse files Browse the repository at this point in the history
`censored` is inspired by `censor` from `MonadWriter`, but (i) can't
reuse that name, and (ii) adverbs are better here anyway because we
are modifying actions.
  • Loading branch information
axch committed Jul 15, 2023
1 parent bac03f8 commit c9ab62c
Showing 1 changed file with 25 additions and 17 deletions.
42 changes: 25 additions & 17 deletions src/lib/OccAnalysis.hs
Original file line number Diff line number Diff line change
Expand Up @@ -144,13 +144,27 @@ countFreeVarsAsOccurrencesB obj =
forM_ (freeAtomVarsList $ Abs obj UnitE) \name -> do
modify (<> FV (singletonNameMapE name $ AccessInfo One accessOnce))

-- Run the given action with its own FV state, and return the FVs it
-- accumulates for post-processing.
isolated :: OCCM n a -> OCCM n (a, FV n)
isolated action = do
-- Run the given action with its own FV state, and return the FVs it accumulates
-- for post-processing. Merging them back in is up to the caller.
separately :: OCCM n a -> OCCM n (a, FV n)
separately action = do
r <- ask
lift11 $ lift11 $ runStateT1 (runReaderT1 r action) mempty

-- Run the given action with its own FV state, and process its accumulated FVs
-- before merging.
censored :: (FV n -> FV n) -> OCCM n a -> OCCM n a
censored f act = do
(a, fvs) <- separately act
modify (<> f fvs)
return a

-- Run the given action with its own FV state, then merge its accumulated FVs
-- afterwards. (This is only meaningful if the action reads the FV state.)
isolated :: OCCM n a -> OCCM n a
isolated = censored id
{-# INLINE isolated #-}

-- Extend the IxExpr environment
extend :: (BindsOneName b (AtomNameC SimpIR))
=> b any n -> IxExpr n -> OCCM n a -> OCCM n a
Expand All @@ -171,8 +185,8 @@ ixExpr name = do
-- including statically.
inlinedLater :: (HoistableE e) => e n -> OCCM n (e n)
inlinedLater obj = do
(_, fvs) <- isolated $ countFreeVarsAsOccurrences obj
modify (<> useManyTimesStatic (useManyTimes fvs))
void $ censored (useManyTimesStatic . useManyTimes)
$ countFreeVarsAsOccurrences obj
return obj

-- === Computing IxExpr summaries ===
Expand Down Expand Up @@ -289,9 +303,7 @@ occNest a (Abs decls ans) = case decls of
\d'@(Let b' (DeclBinding _ expr')) rest -> do
exprIx <- summaryExpr $ sink expr'
extend b' exprIx do
(below, belowfvs) <- isolated do
occNest (sink a) rest >>= wrapWithCachedFVs
modify (<> belowfvs)
below <- isolated $ occNest (sink a) rest >>= wrapWithCachedFVs
accessInfo <- getAccessInfo $ binderName d'
let usage = usageInfo accessInfo
let dceAttempt = case isPureDecl of
Expand Down Expand Up @@ -350,7 +362,7 @@ instance HasOCC SExpr where
Case scrut alts (EffTy effs ty) -> do
scrut' <- occ accessOnce scrut
scrutIx <- summary scrut
(alts', innerFVs) <- unzip <$> mapM (isolated . occAlt a scrutIx) alts
(alts', innerFVs) <- unzip <$> mapM (separately . occAlt a scrutIx) alts
modify (<> foldl' Occ.max zero innerFVs)
ty' <- occTy ty
countFreeVarsAsOccurrences effs
Expand Down Expand Up @@ -399,14 +411,10 @@ instance HasOCC (Hof SimpIR) where
ixDict' <- inlinedLater ixDict
occWithBinder (Abs b body) \b' body' -> do
extend b' (Occ.Var $ binderName b') do
(body'', bodyFV) <- isolated (occNest accessOnce body')
modify (<> abstractFor b' bodyFV)
body'' <- censored (abstractFor b') (occNest accessOnce body')
return $ For ann ixDict' (UnaryLamExpr b' body'')
For _ _ _ -> error "For body should be a unary lambda expression"
While body -> While <$> do
(body', bodyFV) <- isolated $ occNest accessOnce body
modify (<> useManyTimes bodyFV)
return body'
While body -> While <$> censored useManyTimes (occNest accessOnce body)
RunReader ini bd -> do
iniIx <- summary ini
bd' <- oneShot a [Deterministic [], iniIx] bd
Expand Down Expand Up @@ -466,7 +474,7 @@ occWithBinder
-> (forall l. DExt n l => Binder SimpIR n l -> e l -> OCCM l a)
-> OCCM n a
occWithBinder (Abs (b:>ty) body) cont = do
(ty', fvs) <- isolated $ occTy ty
(ty', fvs) <- separately $ occTy ty
ans <- refreshAbs (Abs (b:>ty') body) cont
modify (<> fvs)
return ans
Expand Down

0 comments on commit c9ab62c

Please sign in to comment.