{-# LANGUAGE TemplateHaskell #-}
-- | This module performs limited common subexpression elimination
module Language.PureScript.CoreFn.CSE (optimizeCommonSubexpressions) where

import Protolude hiding (pass)

import Control.Lens (At(..), makeLenses, non, view, (%~), (.=), (.~), (<>~), (^.))
import Control.Monad.Supply (Supply)
import Control.Monad.Supply.Class (MonadSupply)
import Control.Monad.RWS (MonadWriter, RWST, censor, evalRWST, listen, pass, tell)
import Data.Bitraversable (bitraverse)
import Data.Functor.Compose (Compose(..))
import Data.IntMap.Monoidal qualified as IM
import Data.IntSet qualified as IS
import Data.Map qualified as M
import Data.Maybe (fromJust)
import Data.Semigroup (Min(..))
import Data.Semigroup.Generic (GenericSemigroupMonoid(..))

import Language.PureScript.AST.Literals (Literal(..))
import Language.PureScript.AST.SourcePos (nullSourceSpan)
import Language.PureScript.Constants.Libs qualified as C
import Language.PureScript.CoreFn.Ann (Ann)
import Language.PureScript.CoreFn.Binders (Binder(..))
import Language.PureScript.CoreFn.Expr (Bind(..), CaseAlternative(..), Expr(..))
import Language.PureScript.CoreFn.Meta (Meta(IsSyntheticApp))
import Language.PureScript.CoreFn.Traversals (everywhereOnValues, traverseCoreFn)
import Language.PureScript.Environment (dictTypeName)
import Language.PureScript.Names (pattern ByNullSourcePos, Ident(..), ModuleName, ProperName(..), Qualified(..), QualifiedBy(..), freshIdent, runIdent, toMaybeModuleName)
import Language.PureScript.PSString (decodeString)

-- |
-- `discuss f m` is an action that listens to the output of `m`, passes that
-- and its value through `f`, and uses (only) the value of the result to set
-- the new value and output. (Any output produced via the monad in `f` is
-- ignored, though other monadic effects will hold.)
--
discuss :: MonadWriter w m => ((a, w) -> m (b, w)) -> m a -> m b
discuss :: forall w (m :: * -> *) a b.
MonadWriter w m =>
((a, w) -> m (b, w)) -> m a -> m b
discuss (a, w) -> m (b, w)
f = forall w (m :: * -> *) a. MonadWriter w m => m (a, w -> w) -> m a
pass forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second forall a b. a -> b -> a
const) forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, w) -> m (b, w)
f forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< forall w (m :: * -> *) a. MonadWriter w m => m a -> m (a, w)
listen)

-- |
-- Modify the target of an optic in the state with a monadic computation that
-- returns some extra information of type `r` in a tuple.
--
-- I would prefer that this be a named function, but I don't know what to name
-- it. I went with symbols instead because the function that this operator most
-- resembles is `(%%=)`, which doesn't have a textual name as far as I know.
-- Compare the following (approximate) types:
--
-- @
-- (%%=)  :: MonadState s m => Lens s s a b -> (a ->   (r, b)) -> m r
-- (%%<~) :: MonadState s m => Lens s s a b -> (a -> m (r, b)) -> m r
-- @
--
-- Replacing the `=` with `<~` was inspired by analogy with the following pair:
--
-- @
-- (.=) :: MonadState s m => Lens s s a b ->   b -> m ()
-- (<~) :: MonadState s m => Lens s s a b -> m b -> m ()
-- @
--
-- I regret any confusion that ensues.
--
-- Note that there are two interpretations that could reasonably be expected
-- for this type.
--
-- @
-- (%%<~) :: MonadState s m => Lens s s a b -> (a -> m (r, b)) -> m r
-- @
--
-- One is:
-- * Get the focused `a` value from the monad
-- * Run the computation
-- * Get the new state from the returned monad
-- * Take the returned `b` value and set it in the new state
--
-- The other is:
-- * Get the focused `a` value from the monad
-- * Run the computation
-- * Take the returned `b` value and set it in the *original* state
-- * Put the result into the returned monad
--
-- This operator corresponds to the second interpretation. The purpose of this,
-- and part of the purpose of having this operator at all instead of composing
-- simpler operators, is to enable using the lens only once (on the original
-- state) instead of twice (for a get and a set on different states).
--
(%%<~)
  :: MonadState s m
  => ((a -> Compose m ((,) r) b) -> s -> Compose m ((,) r) s)
     -- ^ please read as Lens s s a b
  -> (a -> m (r, b))
  -> m r
(a -> Compose m ((,) r) b) -> s -> Compose m ((,) r) s
l %%<~ :: forall s (m :: * -> *) a r b.
MonadState s m =>
((a -> Compose m ((,) r) b) -> s -> Compose m ((,) r) s)
-> (a -> m (r, b)) -> m r
%%<~ a -> m (r, b)
f = forall s (m :: * -> *). MonadState s m => m s
get forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall {k1} {k2} (f :: k1 -> *) (g :: k2 -> k1) (a :: k2).
Compose f g a -> f (g a)
getCompose forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Compose m ((,) r) b) -> s -> Compose m ((,) r) s
l (forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> m (r, b)
f) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. a -> b -> a
const
infix 4 %%<~

-- |
-- A PluralityMap is like a weaker multiset: like a multiset, it can hold
-- several of the same value, but instead of keeping track of their exact
-- counts, it only records whether there is one (False) or more than one
-- (True).
--
newtype PluralityMap k = PluralityMap { forall k. PluralityMap k -> Map k Bool
getPluralityMap :: M.Map k Bool }

instance Ord k => Semigroup (PluralityMap k) where
  PluralityMap Map k Bool
l <> :: PluralityMap k -> PluralityMap k -> PluralityMap k
<> PluralityMap Map k Bool
r =
    let
      l' :: Map k Bool
l' = forall k a b. (k -> a -> b) -> Map k a -> Map k b
M.mapWithKey (\k
k -> (Bool -> Bool -> Bool
|| k
k forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map k Bool
r)) Map k Bool
l
    in forall k. Map k Bool -> PluralityMap k
PluralityMap forall a b. (a -> b) -> a -> b
$ Map k Bool
l' forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union` Map k Bool
r

instance Ord k => Monoid (PluralityMap k) where
  mempty :: PluralityMap k
mempty = forall k. Map k Bool -> PluralityMap k
PluralityMap forall k a. Map k a
M.empty

data BindingType = NonRecursive | Recursive deriving BindingType -> BindingType -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: BindingType -> BindingType -> Bool
$c/= :: BindingType -> BindingType -> Bool
== :: BindingType -> BindingType -> Bool
$c== :: BindingType -> BindingType -> Bool
Eq

-- |
-- Record summary data about an expression.
--
data CSESummary = CSESummary
  { CSESummary -> IntSet
_scopesUsed    :: IS.IntSet
    -- ^ set of the scope numbers used in this expression
  , CSESummary -> Maybe (Min Int)
_noFloatWithin :: Maybe (Min Int)
    -- ^ optionally a scope within which this expression is not to be floated
    -- (because the expression uses an identifier bound recursively in that
    -- scope)
  , CSESummary -> PluralityMap Ident
_plurality     :: PluralityMap Ident
    -- ^ which floated identifiers are used more than once in this expression
    -- (note that a single use inside an Abs will be considered multiple uses,
    -- as this pass doesn't know when/how many times an Abs will be executed)
  , CSESummary
-> MonoidalIntMap [(Ident, (PluralityMap Ident, Expr Ann))]
_newBindings   :: IM.MonoidalIntMap [(Ident, (PluralityMap Ident, Expr Ann))]
    -- ^ floated bindings, organized by scope number
  , CSESummary -> Map Ident (Expr Ann)
_toBeReinlined :: M.Map Ident (Expr Ann)
    -- ^ a map of floated identifiers that did not end up getting bound and
    -- will need to be reinlined at the end of the pass
  }
  deriving forall x. Rep CSESummary x -> CSESummary
forall x. CSESummary -> Rep CSESummary x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep CSESummary x -> CSESummary
$cfrom :: forall x. CSESummary -> Rep CSESummary x
Generic
  deriving (NonEmpty CSESummary -> CSESummary
CSESummary -> CSESummary -> CSESummary
forall b. Integral b => b -> CSESummary -> CSESummary
forall a.
(a -> a -> a)
-> (NonEmpty a -> a)
-> (forall b. Integral b => b -> a -> a)
-> Semigroup a
stimes :: forall b. Integral b => b -> CSESummary -> CSESummary
$cstimes :: forall b. Integral b => b -> CSESummary -> CSESummary
sconcat :: NonEmpty CSESummary -> CSESummary
$csconcat :: NonEmpty CSESummary -> CSESummary
<> :: CSESummary -> CSESummary -> CSESummary
$c<> :: CSESummary -> CSESummary -> CSESummary
Semigroup, Semigroup CSESummary
CSESummary
[CSESummary] -> CSESummary
CSESummary -> CSESummary -> CSESummary
forall a.
Semigroup a -> a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
mconcat :: [CSESummary] -> CSESummary
$cmconcat :: [CSESummary] -> CSESummary
mappend :: CSESummary -> CSESummary -> CSESummary
$cmappend :: CSESummary -> CSESummary -> CSESummary
mempty :: CSESummary
$cmempty :: CSESummary
Monoid) via GenericSemigroupMonoid CSESummary

-- |
-- Append a value at a given scope depth.
--
addToScope :: Semigroup v => Int -> v -> IM.MonoidalIntMap v -> IM.MonoidalIntMap v
addToScope :: forall v.
Semigroup v =>
Int -> v -> MonoidalIntMap v -> MonoidalIntMap v
addToScope Int
depth v
v
  = forall a.
(Maybe a -> Maybe a) -> Int -> MonoidalIntMap a -> MonoidalIntMap a
IM.alter (forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall b a. b -> (a -> b) -> Maybe a -> b
maybe v
v (forall a. Semigroup a => a -> a -> a
<> v
v)) Int
depth

-- |
-- Remove and return an entire scope from a map of bindings.
--
popScope :: Monoid v => Int -> IM.MonoidalIntMap v -> (v, IM.MonoidalIntMap v)
popScope :: forall v.
Monoid v =>
Int -> MonoidalIntMap v -> (v, MonoidalIntMap v)
popScope Int
depth
  = forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a.
(Int -> a -> Maybe a)
-> Int -> MonoidalIntMap a -> (Maybe a, MonoidalIntMap a)
IM.updateLookupWithKey (\Int
_ v
_ -> forall a. Maybe a
Nothing) Int
depth

-- |
-- Describe the context of an expression.
--
data CSEEnvironment = CSEEnvironment
  { CSEEnvironment -> Int
_depth :: Int
    -- ^ number of enclosing binding scopes (this includes not only Abs, but
    -- Let and CaseAlternative bindings)
  , CSEEnvironment -> Int
_deepestTopLevelScope :: Int
    -- ^ number of enclosing binding scopes outside the first Abs; used to
    -- decide whether to qualify floated identifiers
  , CSEEnvironment -> Map Ident (Int, BindingType)
_bound :: M.Map Ident (Int, BindingType)
    -- ^ map from identifiers to depth in which they are bound and whether
    -- or not the binding is recursive
  }

makeLenses ''CSESummary
makeLenses ''CSEEnvironment

-- |
-- Map from the shape of an expression to an identifier created to represent
-- that expression, organized by scope depth.
--
type CSEState = IM.MonoidalIntMap (M.Map (Expr ()) Ident)

-- |
-- The monad in which CSE takes place.
--
type CSEMonad a = RWST CSEEnvironment CSESummary CSEState Supply a

type HasCSEReader = MonadReader CSEEnvironment
type HasCSEWriter = MonadWriter CSESummary
type HasCSEState = MonadState CSEState

-- |
-- Run a CSEMonad computation; the return value is augmented with a map of
-- identifiers that should be replaced in the final expression because they
-- didn't end up needing to be floated.
--
runCSEMonad :: CSEMonad a -> Supply (a, M.Map Ident (Expr Ann))
runCSEMonad :: forall a. CSEMonad a -> Supply (a, Map Ident (Expr Ann))
runCSEMonad CSEMonad a
x = forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (forall s a. s -> Getting a s a -> a
^. Lens' CSESummary (Map Ident (Expr Ann))
toBeReinlined) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) r w s a.
Monad m =>
RWST r w s m a -> r -> s -> m (a, w)
evalRWST CSEMonad a
x (Int -> Int -> Map Ident (Int, BindingType) -> CSEEnvironment
CSEEnvironment Int
0 Int
0 forall k a. Map k a
M.empty) forall a. MonoidalIntMap a
IM.empty

-- |
-- Mark all expressions floated out of this computation as "plural". This pass
-- assumes that any given Abs may be invoked multiple times, so any expressions
-- inside the Abs but floated out of it also count as having multiple uses,
-- even if they only appear once within the Abs. Consequently, any expressions
-- that can be floated out of an Abs won't be reinlined at the end.
--
enterAbs :: HasCSEWriter m => m a -> m a
enterAbs :: forall (m :: * -> *) a. HasCSEWriter m => m a -> m a
enterAbs = forall w (m :: * -> *) a. MonadWriter w m => (w -> w) -> m a -> m a
censor forall a b. (a -> b) -> a -> b
$ Lens' CSESummary (PluralityMap Ident)
plurality forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ forall k. Map k Bool -> PluralityMap k
PluralityMap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. a -> b -> a
const Bool
True) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k. PluralityMap k -> Map k Bool
getPluralityMap

-- |
-- Run the provided computation in a new scope.
--
newScope :: (HasCSEReader m, HasCSEWriter m) => Bool -> (Int -> m a) -> m a
newScope :: forall (m :: * -> *) a.
(HasCSEReader m, HasCSEWriter m) =>
Bool -> (Int -> m a) -> m a
newScope Bool
isTopLevel Int -> m a
body = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local CSEEnvironment -> CSEEnvironment
goDeeper forall a b. (a -> b) -> a -> b
$ do
  Int
d <- forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Lens' CSEEnvironment Int
depth
  forall w (m :: * -> *) a. MonadWriter w m => (w -> w) -> m a -> m a
censor (Int -> CSESummary -> CSESummary
filterToDepth Int
d) (Int -> m a
body Int
d)
  where
  filterToDepth :: Int -> CSESummary -> CSESummary
filterToDepth Int
d
    = (Lens' CSESummary IntSet
scopesUsed forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ (Int -> Bool) -> IntSet -> IntSet
IS.filter (forall a. Ord a => a -> a -> Bool
< Int
d))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Lens' CSESummary (Maybe (Min Int))
noFloatWithin forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (forall a. Ord a => a -> a -> Bool
< forall a. a -> Min a
Min Int
d))
  goDeeper :: CSEEnvironment -> CSEEnvironment
goDeeper env :: CSEEnvironment
env@CSEEnvironment{Int
Map Ident (Int, BindingType)
_bound :: Map Ident (Int, BindingType)
_deepestTopLevelScope :: Int
_depth :: Int
_bound :: CSEEnvironment -> Map Ident (Int, BindingType)
_deepestTopLevelScope :: CSEEnvironment -> Int
_depth :: CSEEnvironment -> Int
..} =
    if Bool
isTopLevel
    then CSEEnvironment
env{ _depth :: Int
_depth = Int
depth', _deepestTopLevelScope :: Int
_deepestTopLevelScope = Int
depth' }
    else CSEEnvironment
env{ _depth :: Int
_depth = Int
depth' }
    where 
    depth' :: Int
depth' = forall a. Enum a => a -> a
succ Int
_depth

-- |
-- Record a list of identifiers as being bound in the given scope.
--
withBoundIdents :: HasCSEReader m => [Ident] -> (Int, BindingType) -> m a -> m a
withBoundIdents :: forall (m :: * -> *) a.
HasCSEReader m =>
[Ident] -> (Int, BindingType) -> m a -> m a
withBoundIdents [Ident]
idents (Int, BindingType)
t = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (Lens' CSEEnvironment (Map Ident (Int, BindingType))
bound forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (Int, BindingType)
t))) [Ident]
idents)

-- |
-- Run the provided computation in a new scope in which the provided
-- identifiers are bound non-recursively.
--
newScopeWithIdents :: (HasCSEReader m, HasCSEWriter m) => Bool -> [Ident] -> m a -> m a
newScopeWithIdents :: forall (m :: * -> *) a.
(HasCSEReader m, HasCSEWriter m) =>
Bool -> [Ident] -> m a -> m a
newScopeWithIdents Bool
isTopLevel [Ident]
idents = forall (m :: * -> *) a.
(HasCSEReader m, HasCSEWriter m) =>
Bool -> (Int -> m a) -> m a
newScope Bool
isTopLevel forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall (m :: * -> *) a.
HasCSEReader m =>
[Ident] -> (Int, BindingType) -> m a -> m a
withBoundIdents [Ident]
idents forall b c a. (b -> c) -> (a -> b) -> a -> c
. (, BindingType
NonRecursive))

-- |
-- Produce, or retrieve from the state, an identifier for referencing the given
-- expression, at and below the given depth.
--
generateIdentFor :: (HasCSEState m, MonadSupply m) => Int -> Expr () -> m (Bool, Ident)
generateIdentFor :: forall (m :: * -> *).
(HasCSEState m, MonadSupply m) =>
Int -> Expr () -> m (Bool, Ident)
generateIdentFor Int
d Expr ()
e = forall m. At m => Index m -> Lens' m (Maybe (IxValue m))
at Int
d forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Eq a => a -> Iso' (Maybe a) a
non forall a. Monoid a => a
mempty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall m. At m => Index m -> Lens' m (Maybe (IxValue m))
at Expr ()
e forall s (m :: * -> *) a r b.
MonadState s m =>
((a -> Compose m ((,) r) b) -> s -> Compose m ((,) r) s)
-> (a -> m (r, b)) -> m r
%%<~ \case
  Maybe Ident
Nothing    -> forall (m :: * -> *). MonadSupply m => Text -> m Ident
freshIdent (forall {a}. Expr a -> Text
nameHint Expr ()
e) forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \Ident
ident -> ((Bool
True, Ident
ident), forall a. a -> Maybe a
Just Ident
ident)
  Just Ident
ident -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Bool
False, Ident
ident), forall a. a -> Maybe a
Just Ident
ident)
  -- A reminder: as with %%=, the first element of the returned pair is the
  -- final result of the expression, and the second element is the value to
  -- stuff back through the lens into the state. (The difference is that %%<~
  -- enables doing monadic work in the RHS, namely `freshIdent` here.)
  where
  nameHint :: Expr a -> Text
nameHint = \case
    App a
_ Expr a
v1 Expr a
v2
      | Var a
_ Qualified Ident
n <- Expr a
v1
      , forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (a :: ProperNameType). Text -> ProperName a
ProperName forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ident -> Text
runIdent) Qualified Ident
n forall a. Eq a => a -> a -> Bool
== forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (a :: ProperNameType). ProperName a -> ProperName a
dictTypeName Qualified (ProperName 'ClassName)
C.IsSymbol
      , Literal a
_ (ObjectLiteral [(PSString
_, Abs a
_ Ident
_ (Literal a
_ (StringLiteral PSString
str)))]) <- Expr a
v2
      , Just Text
decodedStr <- PSString -> Maybe Text
decodeString PSString
str
        -> Text
decodedStr forall a. Semigroup a => a -> a -> a
<> Text
"IsSymbol"
      | Bool
otherwise
        -> Expr a -> Text
nameHint Expr a
v1
    Var a
_ (Qualified QualifiedBy
_ Ident
ident)
      | Ident Text
name             <- Ident
ident -> Text
name
      | GenIdent (Just Text
name) Integer
_ <- Ident
ident -> Text
name
    Accessor a
_ PSString
prop Expr a
_
      | Just Text
decodedProp <- PSString -> Maybe Text
decodeString PSString
prop -> Text
decodedProp
    Expr a
_ -> Text
"ref"

nullAnn :: Ann
nullAnn :: Ann
nullAnn = (SourceSpan
nullSourceSpan, [], forall a. Maybe a
Nothing, forall a. Maybe a
Nothing)

-- |
-- Use a map to substitute local Vars in a list of Binds.
--
replaceLocals :: M.Map Ident (Expr Ann) -> [Bind Ann] -> [Bind Ann]
replaceLocals :: Map Ident (Expr Ann) -> [Bind Ann] -> [Bind Ann]
replaceLocals Map Ident (Expr Ann)
m = if forall k a. Map k a -> Bool
M.null Map Ident (Expr Ann)
m then forall a. a -> a
identity else forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
map Bind Ann -> Bind Ann
f' where
  (Bind Ann -> Bind Ann
f', Expr Ann -> Expr Ann
g', Binder Ann -> Binder Ann
_) = forall a.
(Bind a -> Bind a)
-> (Expr a -> Expr a)
-> (Binder a -> Binder a)
-> (Bind a -> Bind a, Expr a -> Expr a, Binder a -> Binder a)
everywhereOnValues forall a. a -> a
identity Expr Ann -> Expr Ann
f forall a. a -> a
identity
  f :: Expr Ann -> Expr Ann
f e :: Expr Ann
e@(Var Ann
_ (Qualified QualifiedBy
_ Ident
ident)) = forall b a. b -> (a -> b) -> Maybe a -> b
maybe Expr Ann
e Expr Ann -> Expr Ann
g' forall a b. (a -> b) -> a -> b
$ Ident
ident forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map Ident (Expr Ann)
m
  f Expr Ann
e = Expr Ann
e

-- |
-- Store in the monad a new binding for the given expression, returning a Var
-- referencing it. The provided CSESummary will be transformed to reflect the
-- replacement.
--
floatExpr
  :: (HasCSEReader m, HasCSEState m, MonadSupply m)
  => QualifiedBy
  -> (Expr Ann, CSESummary)
  -> m (Expr Ann, CSESummary)
floatExpr :: forall (m :: * -> *).
(HasCSEReader m, HasCSEState m, MonadSupply m) =>
QualifiedBy -> (Expr Ann, CSESummary) -> m (Expr Ann, CSESummary)
floatExpr QualifiedBy
topLevelQB = \case
  (Expr Ann
e, w :: CSESummary
w@CSESummary{ _noFloatWithin :: CSESummary -> Maybe (Min Int)
_noFloatWithin = Maybe (Min Int)
Nothing, Map Ident (Expr Ann)
IntSet
MonoidalIntMap [(Ident, (PluralityMap Ident, Expr Ann))]
PluralityMap Ident
_toBeReinlined :: Map Ident (Expr Ann)
_newBindings :: MonoidalIntMap [(Ident, (PluralityMap Ident, Expr Ann))]
_plurality :: PluralityMap Ident
_scopesUsed :: IntSet
_toBeReinlined :: CSESummary -> Map Ident (Expr Ann)
_newBindings :: CSESummary
-> MonoidalIntMap [(Ident, (PluralityMap Ident, Expr Ann))]
_plurality :: CSESummary -> PluralityMap Ident
_scopesUsed :: CSESummary -> IntSet
.. }) -> do
    let deepestScope :: Int
deepestScope = if IntSet -> Bool
IS.null IntSet
_scopesUsed then Int
0 else IntSet -> Int
IS.findMax IntSet
_scopesUsed
    (Bool
isNew, Ident
ident) <- forall (m :: * -> *).
(HasCSEState m, MonadSupply m) =>
Int -> Expr () -> m (Bool, Ident)
generateIdentFor Int
deepestScope (forall (f :: * -> *) a. Functor f => f a -> f ()
void Expr Ann
e)
    Int
topLevel <- forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Lens' CSEEnvironment Int
deepestTopLevelScope
    let qb :: QualifiedBy
qb = if Int
deepestScope forall a. Ord a => a -> a -> Bool
> Int
topLevel then QualifiedBy
ByNullSourcePos else QualifiedBy
topLevelQB
    let w' :: CSESummary
w' = CSESummary
w
          forall a b. a -> (a -> b) -> b
& (if Bool
isNew then Lens'
  CSESummary
  (MonoidalIntMap [(Ident, (PluralityMap Ident, Expr Ann))])
newBindings forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ forall v.
Semigroup v =>
Int -> v -> MonoidalIntMap v -> MonoidalIntMap v
addToScope Int
deepestScope [(Ident
ident, (PluralityMap Ident
_plurality, Expr Ann
e))] else forall a. a -> a
identity)
          forall a b. a -> (a -> b) -> b
& Lens' CSESummary (PluralityMap Ident)
plurality forall s t a b. ASetter s t a b -> b -> s -> t
.~ forall k. Map k Bool -> PluralityMap k
PluralityMap (forall k a. k -> a -> Map k a
M.singleton Ident
ident Bool
False)
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. a -> Qualified Ident -> Expr a
Var Ann
nullAnn (forall a. QualifiedBy -> a -> Qualified a
Qualified QualifiedBy
qb Ident
ident), CSESummary
w')
  (Expr Ann
e, CSESummary
w) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (Expr Ann
e, CSESummary
w)

-- |
-- Take possession of the Binds intended to be added to the current scope,
-- removing them from the state, and return the list of Binds along with
-- whatever value is returned by the provided computation.
--
getNewBinds
  :: (HasCSEReader m, HasCSEState m, HasCSEWriter m)
  => m a
  -> m ([Bind Ann], a)
getNewBinds :: forall (m :: * -> *) a.
(HasCSEReader m, HasCSEState m, HasCSEWriter m) =>
m a -> m ([Bind Ann], a)
getNewBinds =
  forall w (m :: * -> *) a b.
MonadWriter w m =>
((a, w) -> m (b, w)) -> m a -> m b
discuss forall a b. (a -> b) -> a -> b
$ \(a
a, CSESummary
w) -> do
    Int
d <- forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Lens' CSEEnvironment Int
depth
    forall m. At m => Index m -> Lens' m (Maybe (IxValue m))
at Int
d forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
.= forall a. Maybe a
Nothing
    let ([(Ident, (PluralityMap Ident, Expr Ann))]
floatedHere, CSESummary
w') = Lens'
  CSESummary
  (MonoidalIntMap [(Ident, (PluralityMap Ident, Expr Ann))])
newBindings (forall v.
Monoid v =>
Int -> MonoidalIntMap v -> (v, MonoidalIntMap v)
popScope Int
d) CSESummary
w
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (, a
a) forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Ident, (PluralityMap Ident, Expr Ann))
-> ([Bind Ann], CSESummary) -> ([Bind Ann], CSESummary)
handleFloat ([], CSESummary
w') [(Ident, (PluralityMap Ident, Expr Ann))]
floatedHere
  where
  handleFloat :: (Ident, (PluralityMap Ident, Expr Ann))
-> ([Bind Ann], CSESummary) -> ([Bind Ann], CSESummary)
handleFloat (Ident
ident, (PluralityMap Ident
p, Expr Ann
e)) ([Bind Ann]
bs, CSESummary
w) =
    if forall a. HasCallStack => Maybe a -> a
fromJust forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Ident
ident forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k. PluralityMap k -> Map k Bool
getPluralityMap forall a b. (a -> b) -> a -> b
$ CSESummary
w forall s a. s -> Getting a s a -> a
^. Lens' CSESummary (PluralityMap Ident)
plurality
    then (forall a. a -> Ident -> Expr a -> Bind a
NonRec Ann
nullAnn Ident
ident Expr Ann
e forall a. a -> [a] -> [a]
: [Bind Ann]
bs, CSESummary
w')
    else ([Bind Ann]
bs, CSESummary
w' forall a b. a -> (a -> b) -> b
& Lens' CSESummary (Map Ident (Expr Ann))
toBeReinlined forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Ident
ident Expr Ann
e)
    where w' :: CSESummary
w' = CSESummary
w forall a b. a -> (a -> b) -> b
& Lens' CSESummary (PluralityMap Ident)
plurality forall a s t. Semigroup a => ASetter s t a a -> a -> s -> t
<>~ PluralityMap Ident
p

-- |
-- Like getNewBinds, but also stores the Binds in a Let wrapping the provided
-- expression. If said expression is already a Let, adds these Binds to that
-- Let instead.
--
getNewBindsAsLet
  :: (HasCSEReader m, HasCSEWriter m, HasCSEState m)
  => m (Expr Ann)
  -> m (Expr Ann)
getNewBindsAsLet :: forall (m :: * -> *).
(HasCSEReader m, HasCSEWriter m, HasCSEState m) =>
m (Expr Ann) -> m (Expr Ann)
getNewBindsAsLet = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry [Bind Ann] -> Expr Ann -> Expr Ann
go) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a.
(HasCSEReader m, HasCSEState m, HasCSEWriter m) =>
m a -> m ([Bind Ann], a)
getNewBinds where
  go :: [Bind Ann] -> Expr Ann -> Expr Ann
go [Bind Ann]
bs = if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Bind Ann]
bs then forall a. a -> a
identity else \case
    Let Ann
a [Bind Ann]
bs' Expr Ann
e' -> forall a. a -> [Bind a] -> Expr a -> Expr a
Let Ann
a ([Bind Ann]
bs forall a. [a] -> [a] -> [a]
++ [Bind Ann]
bs') Expr Ann
e'
    Expr Ann
e'           -> forall a. a -> [Bind a] -> Expr a -> Expr a
Let Ann
nullAnn [Bind Ann]
bs Expr Ann
e'

-- |
-- Feed the Writer part of the monad with the requirements of this name.
--
summarizeName
  :: (HasCSEReader m, HasCSEWriter m)
  => ModuleName
  -> Qualified Ident
  -> m ()
summarizeName :: forall (m :: * -> *).
(HasCSEReader m, HasCSEWriter m) =>
ModuleName -> Qualified Ident -> m ()
summarizeName ModuleName
mn (Qualified QualifiedBy
mn' Ident
ident) = do
  Map Ident (Int, BindingType)
m <- forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Lens' CSEEnvironment (Map Ident (Int, BindingType))
bound
  let (Int
s, BindingType
bt) =
        forall a. a -> Maybe a -> a
fromMaybe (Int
0, BindingType
NonRecursive) forall a b. (a -> b) -> a -> b
$
          forall (f :: * -> *). Alternative f => Bool -> f ()
guard (forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (forall a. Eq a => a -> a -> Bool
== ModuleName
mn) (QualifiedBy -> Maybe ModuleName
toMaybeModuleName QualifiedBy
mn')) forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Ident
ident forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map Ident (Int, BindingType)
m
  forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell forall a b. (a -> b) -> a -> b
$ forall a. Monoid a => a
mempty
       forall a b. a -> (a -> b) -> b
& Lens' CSESummary IntSet
scopesUsed forall s t a b. ASetter s t a b -> b -> s -> t
.~ Int -> IntSet
IS.singleton Int
s
       forall a b. a -> (a -> b) -> b
& Lens' CSESummary (Maybe (Min Int))
noFloatWithin forall s t a b. ASetter s t a b -> b -> s -> t
.~ (forall (f :: * -> *). Alternative f => Bool -> f ()
guard (BindingType
bt forall a. Eq a => a -> a -> Bool
== BindingType
Recursive) forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> forall a. a -> Min a
Min Int
s)

-- |
-- Collect all the Idents put in scope by a list of Binders.
--
identsFromBinders :: [Binder a] -> [Ident]
identsFromBinders :: forall a. [Binder a] -> [Ident]
identsFromBinders = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall {a}. Binder a -> [Ident]
identsFromBinder where
  identsFromBinder :: Binder a -> [Ident]
identsFromBinder = \case
    LiteralBinder a
_ (ArrayLiteral [Binder a]
xs)  -> forall a. [Binder a] -> [Ident]
identsFromBinders [Binder a]
xs
    LiteralBinder a
_ (ObjectLiteral [(PSString, Binder a)]
xs) -> forall a. [Binder a] -> [Ident]
identsFromBinders (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
map forall a b. (a, b) -> b
snd [(PSString, Binder a)]
xs)
    VarBinder a
_ Ident
ident                  -> [Ident
ident]
    ConstructorBinder a
_ Qualified (ProperName 'TypeName)
_ Qualified (ProperName 'ConstructorName)
_ [Binder a]
xs         -> forall a. [Binder a] -> [Ident]
identsFromBinders [Binder a]
xs
    NamedBinder a
_ Ident
ident Binder a
x              -> Ident
ident forall a. a -> [a] -> [a]
: Binder a -> [Ident]
identsFromBinder Binder a
x
    LiteralBinder a
_ BooleanLiteral{}   -> []
    LiteralBinder a
_ CharLiteral{}      -> []
    LiteralBinder a
_ NumericLiteral{}   -> []
    LiteralBinder a
_ StringLiteral{}    -> []
    NullBinder{}                       -> []

-- |
-- Float synthetic Apps (right now, the only Apps marked as synthetic are type
-- class dictionaries being fed to functions with constraints, superclass
-- accessors, and instances of IsSymbol) to a new or existing Let as close to
-- the top level as possible.
--
optimizeCommonSubexpressions :: ModuleName -> [Bind Ann] -> Supply [Bind Ann]
optimizeCommonSubexpressions :: ModuleName -> [Bind Ann] -> Supply [Bind Ann]
optimizeCommonSubexpressions ModuleName
mn
  = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (forall a b c. (a -> b -> c) -> b -> a -> c
flip Map Ident (Expr Ann) -> [Bind Ann] -> [Bind Ann]
replaceLocals))
  forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. CSEMonad a -> Supply (a, Map Ident (Expr Ann))
runCSEMonad
  forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall a. [a] -> [a] -> [a]
(++))
  forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a.
(HasCSEReader m, HasCSEState m, HasCSEWriter m) =>
m a -> m ([Bind Ann], a)
getNewBinds
  forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst
  forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a.
Bool -> CSEMonad a -> [Bind Ann] -> CSEMonad ([Bind Ann], a)
handleBinds Bool
True (forall (f :: * -> *) a. Applicative f => a -> f a
pure ())

  where

  -- This is the one place (I think?) that keeps this from being a general
  -- common subexpression elimination pass.
  shouldFloatExpr :: Expr Ann -> Bool
  shouldFloatExpr :: Expr Ann -> Bool
shouldFloatExpr = \case
    App (SourceSpan
_, [Comment]
_, Maybe SourceType
_, Just Meta
IsSyntheticApp) Expr Ann
e Expr Ann
_ -> Expr Ann -> Bool
isSimple Expr Ann
e
    Expr Ann
_                                      -> Bool
False

  isSimple :: Expr Ann -> Bool
  isSimple :: Expr Ann -> Bool
isSimple = \case
    Var{}          -> Bool
True
    Accessor Ann
_ PSString
_ Expr Ann
e -> Expr Ann -> Bool
isSimple Expr Ann
e
    Expr Ann
_              -> Bool
False

  handleAndWrapExpr :: Expr Ann -> CSEMonad (Expr Ann)
  handleAndWrapExpr :: Expr Ann -> CSEMonad (Expr Ann)
handleAndWrapExpr = forall (m :: * -> *).
(HasCSEReader m, HasCSEWriter m, HasCSEState m) =>
m (Expr Ann) -> m (Expr Ann)
getNewBindsAsLet forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr Ann -> CSEMonad (Expr Ann)
handleExpr

  (Bind Ann
-> RWST CSEEnvironment CSESummary CSEState Supply (Bind Ann)
handleBind, Expr Ann -> CSEMonad (Expr Ann)
handleExprDefault, Binder Ann
-> RWST CSEEnvironment CSESummary CSEState Supply (Binder Ann)
handleBinder, CaseAlternative Ann
-> RWST
     CSEEnvironment CSESummary CSEState Supply (CaseAlternative Ann)
_) = forall (f :: * -> *) a.
Applicative f =>
(Bind a -> f (Bind a))
-> (Expr a -> f (Expr a))
-> (Binder a -> f (Binder a))
-> (CaseAlternative a -> f (CaseAlternative a))
-> (Bind a -> f (Bind a), Expr a -> f (Expr a),
    Binder a -> f (Binder a),
    CaseAlternative a -> f (CaseAlternative a))
traverseCoreFn Bind Ann
-> RWST CSEEnvironment CSESummary CSEState Supply (Bind Ann)
handleBind Expr Ann -> CSEMonad (Expr Ann)
handleExpr Binder Ann
-> RWST CSEEnvironment CSESummary CSEState Supply (Binder Ann)
handleBinder CaseAlternative Ann
-> RWST
     CSEEnvironment CSESummary CSEState Supply (CaseAlternative Ann)
handleCaseAlternative

  topLevelQB :: QualifiedBy
topLevelQB = ModuleName -> QualifiedBy
ByModuleName ModuleName
mn

  handleExpr :: Expr Ann -> CSEMonad (Expr Ann)
  handleExpr :: Expr Ann -> CSEMonad (Expr Ann)
handleExpr = forall w (m :: * -> *) a b.
MonadWriter w m =>
((a, w) -> m (b, w)) -> m a -> m b
discuss (forall (m :: * -> *) a. Monad m => m Bool -> m a -> m a -> m a
ifM (Expr Ann -> Bool
shouldFloatExpr forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) (forall (m :: * -> *).
(HasCSEReader m, HasCSEState m, MonadSupply m) =>
QualifiedBy -> (Expr Ann, CSESummary) -> m (Expr Ann, CSESummary)
floatExpr QualifiedBy
topLevelQB) forall (f :: * -> *) a. Applicative f => a -> f a
pure) forall b c a. (b -> c) -> (a -> b) -> a -> c
. \case
    Abs Ann
a Ident
ident Expr Ann
e   -> forall (m :: * -> *) a. HasCSEWriter m => m a -> m a
enterAbs forall a b. (a -> b) -> a -> b
$ forall a. a -> Ident -> Expr a -> Expr a
Abs Ann
a Ident
ident forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a.
(HasCSEReader m, HasCSEWriter m) =>
Bool -> [Ident] -> m a -> m a
newScopeWithIdents Bool
False [Ident
ident] (Expr Ann -> CSEMonad (Expr Ann)
handleAndWrapExpr Expr Ann
e)
    v :: Expr Ann
v@(Var Ann
_ Qualified Ident
qname) -> forall (m :: * -> *).
(HasCSEReader m, HasCSEWriter m) =>
ModuleName -> Qualified Ident -> m ()
summarizeName ModuleName
mn Qualified Ident
qname forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Expr Ann
v
    Let Ann
a [Bind Ann]
bs Expr Ann
e      -> forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (forall a. a -> [Bind a] -> Expr a -> Expr a
Let Ann
a) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a.
Bool -> CSEMonad a -> [Bind Ann] -> CSEMonad ([Bind Ann], a)
handleBinds Bool
False (Expr Ann -> CSEMonad (Expr Ann)
handleExpr Expr Ann
e) [Bind Ann]
bs
    Expr Ann
x               -> Expr Ann -> CSEMonad (Expr Ann)
handleExprDefault Expr Ann
x

  handleCaseAlternative :: CaseAlternative Ann -> CSEMonad (CaseAlternative Ann)
  handleCaseAlternative :: CaseAlternative Ann
-> RWST
     CSEEnvironment CSESummary CSEState Supply (CaseAlternative Ann)
handleCaseAlternative (CaseAlternative [Binder Ann]
bs Either [(Expr Ann, Expr Ann)] (Expr Ann)
x) = forall a.
[Binder a]
-> Either [(Guard a, Guard a)] (Guard a) -> CaseAlternative a
CaseAlternative [Binder Ann]
bs forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> do
    forall (m :: * -> *) a.
(HasCSEReader m, HasCSEWriter m) =>
Bool -> [Ident] -> m a -> m a
newScopeWithIdents Bool
False (forall a. [Binder a] -> [Ident]
identsFromBinders [Binder Ann]
bs) forall a b. (a -> b) -> a -> b
$
      forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall a b. (a -> b) -> a -> b
$ forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse Expr Ann -> CSEMonad (Expr Ann)
handleAndWrapExpr Expr Ann -> CSEMonad (Expr Ann)
handleAndWrapExpr) Expr Ann -> CSEMonad (Expr Ann)
handleAndWrapExpr Either [(Expr Ann, Expr Ann)] (Expr Ann)
x

  handleBinds :: forall a. Bool -> CSEMonad a -> [Bind Ann] -> CSEMonad ([Bind Ann], a)
  handleBinds :: forall a.
Bool -> CSEMonad a -> [Bind Ann] -> CSEMonad ([Bind Ann], a)
handleBinds Bool
isTopLevel = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Bind Ann -> CSEMonad ([Bind Ann], a) -> CSEMonad ([Bind Ann], a)
go forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (f :: * -> *) a. Applicative f => a -> f a
pure where
    go :: Bind Ann -> CSEMonad ([Bind Ann], a) -> CSEMonad ([Bind Ann], a)
    go :: Bind Ann -> CSEMonad ([Bind Ann], a) -> CSEMonad ([Bind Ann], a)
go Bind Ann
b CSEMonad ([Bind Ann], a)
inner = case Bind Ann
b of
      -- For a NonRec Bind, traverse the bound expression in the current scope
      -- and then create a new scope for any remaining Binds and/or whatever
      -- inner thing all these Binds are applied to.
      NonRec Ann
a Ident
ident Expr Ann
e -> do
        Expr Ann
e' <- Expr Ann -> CSEMonad (Expr Ann)
handleExpr Expr Ann
e
        forall (m :: * -> *) a.
(HasCSEReader m, HasCSEWriter m) =>
Bool -> [Ident] -> m a -> m a
newScopeWithIdents Bool
isTopLevel [Ident
ident] forall a b. (a -> b) -> a -> b
$
          Bind Ann -> CSEMonad ([Bind Ann], a)
prependToNewBindsFromInner forall a b. (a -> b) -> a -> b
$ forall a. a -> Ident -> Expr a -> Bind a
NonRec Ann
a Ident
ident Expr Ann
e'
      Rec [((Ann, Ident), Expr Ann)]
es ->
        -- For a Rec Bind, the bound expressions need a new scope in which all
        -- these identifiers are bound recursively; then the remaining Binds
        -- and the inner thing can be traversed in the same scope with the same
        -- identifiers now bound non-recursively.
        forall (m :: * -> *) a.
(HasCSEReader m, HasCSEWriter m) =>
Bool -> (Int -> m a) -> m a
newScope Bool
isTopLevel forall a b. (a -> b) -> a -> b
$ \Int
d -> do
          let idents :: [Ident]
idents = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
map (forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [((Ann, Ident), Expr Ann)]
es
          [((Ann, Ident), Expr Ann)]
es' <- forall (m :: * -> *) a.
HasCSEReader m =>
[Ident] -> (Int, BindingType) -> m a -> m a
withBoundIdents [Ident]
idents (Int
d, BindingType
Recursive) forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Expr Ann -> CSEMonad (Expr Ann)
handleExpr) [((Ann, Ident), Expr Ann)]
es
          forall (m :: * -> *) a.
HasCSEReader m =>
[Ident] -> (Int, BindingType) -> m a -> m a
withBoundIdents [Ident]
idents (Int
d, BindingType
NonRecursive) forall a b. (a -> b) -> a -> b
$
            Bind Ann -> CSEMonad ([Bind Ann], a)
prependToNewBindsFromInner forall a b. (a -> b) -> a -> b
$ forall a. [((a, Ident), Expr a)] -> Bind a
Rec [((Ann, Ident), Expr Ann)]
es'

      where

      prependToNewBindsFromInner :: Bind Ann -> CSEMonad ([Bind Ann], a)
      prependToNewBindsFromInner :: Bind Ann -> CSEMonad ([Bind Ann], a)
prependToNewBindsFromInner Bind Ann
hd = forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (Bind Ann
hd forall a. a -> [a] -> [a]
:) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => m (m a) -> m a
join forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a.
(HasCSEReader m, HasCSEState m, HasCSEWriter m) =>
m a -> m ([Bind Ann], a)
getNewBinds CSEMonad ([Bind Ann], a)
inner