{-# LANGUAGE Safe #-}
{-# LANGUAGE FlexibleContexts #-}
-- | Substitution and binder-refreshing for Eidos (doc/eidos.md §2, G2).
--
--   The uniqueness discipline makes substitution an environment map — no
--   capture is possible while the invariant holds — and concentrates the
--   invariant's maintenance in ONE primitive: 'refreshExp'/'refreshDefn',
--   the audited clone. Every pass that duplicates a term (inlining,
--   specialization, beta reduction, case-of-known-constructor) must route
--   the duplicated copy through a refresh; the linter's uniqueness rule
--   re-checks the invariant globally under --debug-lint.
--
--   Contracts:
--
--   * 'substVars' inserts each payload AS IS: sound only when each payload
--     lands at most once (or is binder-free). For the general case use
--     'substVarsRefreshing', which refreshes every inserted copy.
--   * 'refreshExp' freshens every binder in the term (term binders, join
--     labels, and — via 'refreshDefn' — signature type variables,
--     propagating the renaming through every type in the body) and leaves
--     free names untouched.
--   * Supplies: passes obtain fresh uniques from a state seeded above the
--     program's maximum ('nextUniq').
module ReWire.Eidos.Subst
      ( maxUniq, nextUniq
      , refreshExp, refreshDefn, instantiateDefn
      , substVars, substVarsRefreshing
      , occCounts, freeUniqs, occIds
      ) where

import ReWire.Eidos.Syntax
import ReWire.Eidos.Types (substTv)
import ReWire.SYB (queryWith)

import Control.Monad.State.Strict (MonadState, get, put)
import Data.Data (Data)
import Data.HashMap.Strict (HashMap)
import Data.Text (Text)

import qualified Data.HashMap.Strict as Map
import qualified Data.IntMap.Strict  as IM
import qualified Data.IntSet         as IS

-- | The largest unique occurring anywhere in a program — or any other
--   'Data' value holding 'Id's and 'TyVar's — (binders and occurrences;
--   the primitive basis' negative uniques never win).
maxUniq :: Data a => a -> Uniq
maxUniq :: forall a. Data a => a -> Int
maxUniq a
p = [Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ Int
0 Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
ids [Int] -> [Int] -> [Int]
forall a. Semigroup a => a -> a -> a
<> [Int]
tvs
      where ids :: [Uniq]
            ids :: [Int]
ids = (Id -> [Int]) -> a -> [Int]
forall a b r. (Data a, Typeable b) => (b -> [r]) -> a -> [r]
queryWith (\ Id
x -> [Id -> Int
idUniq Id
x]) a
p

            tvs :: [Uniq]
            tvs :: [Int]
tvs = (TyVar -> [Int]) -> a -> [Int]
forall a b r. (Data a, Typeable b) => (b -> [r]) -> a -> [r]
queryWith (\ TyVar
v -> [TyVar -> Int
tvUniq TyVar
v]) a
p

-- | A safe starting value for a pass's unique supply.
nextUniq :: Data a => a -> Uniq
nextUniq :: forall a. Data a => a -> Int
nextUniq = Int -> Int
forall a. Enum a => a -> a
succ (Int -> Int) -> (a -> Int) -> a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Int
forall a. Data a => a -> Int
maxUniq

freshU :: MonadState Uniq m => m Uniq
freshU :: forall (m :: * -> *). MonadState Int m => m Int
freshU = do
      u <- m Int
forall s (m :: * -> *). MonadState s m => m s
get
      put $ u + 1
      pure u

---
--- Refreshing (the audited clone primitive).
---

-- | The renaming environment: term binders, join labels, and signature
--   type variables (as a type substitution).
data Renames = Renames
      { Renames -> IntMap Id
rnVars  :: !(IM.IntMap Id)
      , Renames -> IntMap JoinId
rnJoins :: !(IM.IntMap JoinId)
      , Renames -> HashMap TyVar Ty
rnTvs   :: !(HashMap TyVar Ty)
      }

rn0 :: Renames
rn0 :: Renames
rn0 = IntMap Id -> IntMap JoinId -> HashMap TyVar Ty -> Renames
Renames IntMap Id
forall a. Monoid a => a
mempty IntMap JoinId
forall a. Monoid a => a
mempty HashMap TyVar Ty
forall a. Monoid a => a
mempty

-- | Freshen every binder in an expression; free names (and all types, when
--   no signature variables are being renamed) are untouched.
refreshExp :: MonadState Uniq m => Exp -> m Exp
refreshExp :: forall (m :: * -> *). MonadState Int m => Exp -> m Exp
refreshExp = Renames -> Exp -> m Exp
forall (m :: * -> *). MonadState Int m => Renames -> Exp -> m Exp
rex Renames
rn0

-- | Clone a definition with fresh binders throughout: a fresh definition
--   name (self-references in the body follow it), fresh signature type
--   variables (the renaming propagates through every type in the
--   parameters and body), and fresh parameter and local binders.
refreshDefn :: MonadState Uniq m => Defn -> m Defn
refreshDefn :: forall (m :: * -> *). MonadState Int m => Defn -> m Defn
refreshDefn (Defn Annote
an Id
x [Id]
params Exp
body Maybe DefnAttr
attr Maybe SpecOrigin
orig) = do
      let Sig [TyVar]
tvs Ty
t = Id -> Sig
idSig Id
x
      tvs' <- (TyVar -> m TyVar) -> [TyVar] -> m [TyVar]
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 TyVar -> m TyVar
forall (m :: * -> *). MonadState Int m => TyVar -> m TyVar
freshTv [TyVar]
tvs
      let rtv  = [(TyVar, Ty)] -> HashMap TyVar Ty
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
Map.fromList ([(TyVar, Ty)] -> HashMap TyVar Ty)
-> [(TyVar, Ty)] -> HashMap TyVar Ty
forall a b. (a -> b) -> a -> b
$ [TyVar] -> [Ty] -> [(TyVar, Ty)]
forall a b. [a] -> [b] -> [(a, b)]
zip [TyVar]
tvs ([Ty] -> [(TyVar, Ty)]) -> [Ty] -> [(TyVar, Ty)]
forall a b. (a -> b) -> a -> b
$ (TyVar -> Ty) -> [TyVar] -> [Ty]
forall a b. (a -> b) -> [a] -> [b]
map (Annote -> TyVar -> Ty
TyVarT Annote
an) [TyVar]
tvs'
          sig' = [TyVar] -> Ty -> Sig
Sig [TyVar]
tvs' (Ty -> Sig) -> Ty -> Sig
forall a b. (a -> b) -> a -> b
$ HashMap TyVar Ty -> Ty -> Ty
substTv HashMap TyVar Ty
rtv Ty
t
      x'   <- freshLike rtv x { idSig = sig' }
      let r0 = Renames
rn0 { rnVars = IM.singleton (idUniq x) x', rnTvs = rtv }
      (r1, params') <- bindIds r0 params
      body' <- rex r1 body
      pure $ Defn an x' params' body' attr orig
      where freshTv :: MonadState Uniq m => TyVar -> m TyVar
            freshTv :: forall (m :: * -> *). MonadState Int m => TyVar -> m TyVar
freshTv TyVar
v = do
                  u <- m Int
forall (m :: * -> *). MonadState Int m => m Int
freshU
                  pure v { tvUniq = u }

-- | Clone a definition at a type instantiation (the specializer's flavor
--   of the audited clone): the signature's variables are substituted away
--   by the given type arguments (the clone is monomorphic when they are
--   closed), every type in the parameters and body follows, and every
--   binder is refreshed. The clone is named by the given occurrence text
--   with a fresh unique; self-references in the body are NOT remapped —
--   they still name the origin at its instantiated type arguments, for
--   the caller's spine rewrite to resolve.
instantiateDefn :: MonadState Uniq m => Text -> [Ty] -> Defn -> m Defn
instantiateDefn :: forall (m :: * -> *).
MonadState Int m =>
Text -> [Ty] -> Defn -> m Defn
instantiateDefn Text
occ [Ty]
ts (Defn Annote
an Id
x [Id]
params Exp
body Maybe DefnAttr
attr Maybe SpecOrigin
orig) = do
      let Sig [TyVar]
tvs Ty
t = Id -> Sig
idSig Id
x
          rtv :: HashMap TyVar Ty
rtv       = [(TyVar, Ty)] -> HashMap TyVar Ty
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
Map.fromList ([(TyVar, Ty)] -> HashMap TyVar Ty)
-> [(TyVar, Ty)] -> HashMap TyVar Ty
forall a b. (a -> b) -> a -> b
$ [TyVar] -> [Ty] -> [(TyVar, Ty)]
forall a b. [a] -> [b] -> [(a, b)]
zip [TyVar]
tvs [Ty]
ts
      u <- m Int
forall (m :: * -> *). MonadState Int m => m Int
freshU
      let x' = Id
x { idOcc = occ, idUniq = u, idSig = Sig [] $ substTv rtv t }
      (r1, params') <- bindIds rn0 { rnTvs = rtv } params
      body' <- rex r1 body
      pure $ Defn an x' params' body' attr orig

-- | A fresh Id with the same occurrence text and a renamed signature.
freshLike :: MonadState Uniq m => HashMap TyVar Ty -> Id -> m Id
freshLike :: forall (m :: * -> *).
MonadState Int m =>
HashMap TyVar Ty -> Id -> m Id
freshLike HashMap TyVar Ty
rtv Id
x = do
      u <- m Int
forall (m :: * -> *). MonadState Int m => m Int
freshU
      let Sig tvs t = idSig x
      pure x { idUniq = u, idSig = Sig tvs $ substTv rtv t }

bindId :: MonadState Uniq m => Renames -> Id -> m (Renames, Id)
bindId :: forall (m :: * -> *).
MonadState Int m =>
Renames -> Id -> m (Renames, Id)
bindId Renames
r Id
x = do
      x' <- HashMap TyVar Ty -> Id -> m Id
forall (m :: * -> *).
MonadState Int m =>
HashMap TyVar Ty -> Id -> m Id
freshLike (Renames -> HashMap TyVar Ty
rnTvs Renames
r) Id
x
      pure (r { rnVars = IM.insert (idUniq x) x' $ rnVars r }, x')

bindIds :: MonadState Uniq m => Renames -> [Id] -> m (Renames, [Id])
bindIds :: forall (m :: * -> *).
MonadState Int m =>
Renames -> [Id] -> m (Renames, [Id])
bindIds = [Id] -> Renames -> [Id] -> m (Renames, [Id])
forall {f :: * -> *}.
MonadState Int f =>
[Id] -> Renames -> [Id] -> f (Renames, [Id])
go []
      where go :: [Id] -> Renames -> [Id] -> f (Renames, [Id])
go [Id]
acc Renames
r []       = (Renames, [Id]) -> f (Renames, [Id])
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Renames
r, [Id] -> [Id]
forall a. [a] -> [a]
reverse [Id]
acc)
            go [Id]
acc Renames
r (Id
x : [Id]
xs) = do
                  (r', x') <- Renames -> Id -> f (Renames, Id)
forall (m :: * -> *).
MonadState Int m =>
Renames -> Id -> m (Renames, Id)
bindId Renames
r Id
x
                  go (x' : acc) r' xs

rex :: MonadState Uniq m => Renames -> Exp -> m Exp
rex :: forall (m :: * -> *). MonadState Int m => Renames -> Exp -> m Exp
rex Renames
r Exp
e = case Exp
e of
      Var Annote
an Id
x        -> Exp -> m Exp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> m Exp) -> Exp -> m Exp
forall a b. (a -> b) -> a -> b
$ Annote -> Id -> Exp
Var Annote
an (Id -> Exp) -> Id -> Exp
forall a b. (a -> b) -> a -> b
$ Id -> Id
occ Id
x
      Con Annote
an Ty
t Text
c      -> Exp -> m Exp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> m Exp) -> Exp -> m Exp
forall a b. (a -> b) -> a -> b
$ Annote -> Ty -> Text -> Exp
Con Annote
an (Ty -> Ty
ty Ty
t) Text
c
      Prim Annote
an Ty
t Builtin
p     -> Exp -> m Exp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> m Exp) -> Exp -> m Exp
forall a b. (a -> b) -> a -> b
$ Annote -> Ty -> Builtin -> Exp
Prim Annote
an (Ty -> Ty
ty Ty
t) Builtin
p
      LitInt Annote
an Ty
t Integer
n   -> Exp -> m Exp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> m Exp) -> Exp -> m Exp
forall a b. (a -> b) -> a -> b
$ Annote -> Ty -> Integer -> Exp
LitInt Annote
an (Ty -> Ty
ty Ty
t) Integer
n
      LitStr {}       -> Exp -> m Exp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
      LitList Annote
an Ty
t [Exp]
es -> Annote -> Ty -> [Exp] -> Exp
LitList Annote
an (Ty -> Ty
ty Ty
t) ([Exp] -> Exp) -> m [Exp] -> m Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> m Exp) -> [Exp] -> m [Exp]
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 (Renames -> Exp -> m Exp
forall (m :: * -> *). MonadState Int m => Renames -> Exp -> m Exp
rex Renames
r) [Exp]
es
      LitVec Annote
an Ty
t [Exp]
es  -> Annote -> Ty -> [Exp] -> Exp
LitVec Annote
an (Ty -> Ty
ty Ty
t) ([Exp] -> Exp) -> m [Exp] -> m Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> m Exp) -> [Exp] -> m [Exp]
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 (Renames -> Exp -> m Exp
forall (m :: * -> *). MonadState Int m => Renames -> Exp -> m Exp
rex Renames
r) [Exp]
es
      App Annote
an Exp
f Arg
a      -> Annote -> Exp -> Arg -> Exp
App Annote
an (Exp -> Arg -> Exp) -> m Exp -> m (Arg -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Renames -> Exp -> m Exp
forall (m :: * -> *). MonadState Int m => Renames -> Exp -> m Exp
rex Renames
r Exp
f m (Arg -> Exp) -> m Arg -> m Exp
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Arg -> m Arg
forall (m :: * -> *). MonadState Int m => Arg -> m Arg
arg Arg
a
      Lam Annote
an Id
x Exp
b      -> do
            (r', x') <- Renames -> Id -> m (Renames, Id)
forall (m :: * -> *).
MonadState Int m =>
Renames -> Id -> m (Renames, Id)
bindId Renames
r Id
x
            Lam an x' <$> rex r' b
      Let Annote
an Bind
b Exp
body   -> do
            (r', b') <- Bind -> m (Renames, Bind)
forall (m :: * -> *). MonadState Int m => Bind -> m (Renames, Bind)
bnd Bind
b
            Let an b' <$> rex r' body
      Jump Annote
an JoinId
j [Exp]
es    -> Annote -> JoinId -> [Exp] -> Exp
Jump Annote
an (JoinId -> JoinId
jn JoinId
j) ([Exp] -> Exp) -> m [Exp] -> m Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> m Exp) -> [Exp] -> m [Exp]
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 (Renames -> Exp -> m Exp
forall (m :: * -> *). MonadState Int m => Renames -> Exp -> m Exp
rex Renames
r) [Exp]
es
      Case Annote
an Ty
t Exp
s Id
cb [Alt]
alts -> do
            s'        <- Renames -> Exp -> m Exp
forall (m :: * -> *). MonadState Int m => Renames -> Exp -> m Exp
rex Renames
r Exp
s
            (r', cb') <- bindId r cb
            alts'     <- mapM (alt r') alts
            pure $ Case an (ty t) s' cb' alts'
      where occ :: Id -> Id
            occ :: Id -> Id
occ Id
x = case Int -> IntMap Id -> Maybe Id
forall a. Int -> IntMap a -> Maybe a
IM.lookup (Id -> Int
idUniq Id
x) (IntMap Id -> Maybe Id) -> IntMap Id -> Maybe Id
forall a b. (a -> b) -> a -> b
$ Renames -> IntMap Id
rnVars Renames
r of
                  Just Id
x'                       -> Id
x'
                  Maybe Id
Nothing | HashMap TyVar Ty -> Bool
forall k v. HashMap k v -> Bool
Map.null (Renames -> HashMap TyVar Ty
rnTvs Renames
r)  -> Id
x
                          | Bool
otherwise           -> Id
x { idSig = Sig tvs $ substTv (rnTvs r) t }
                        where Sig [TyVar]
tvs Ty
t = Id -> Sig
idSig Id
x

            jn :: JoinId -> JoinId
            jn :: JoinId -> JoinId
jn JoinId
j = case Int -> IntMap JoinId -> Maybe JoinId
forall a. Int -> IntMap a -> Maybe a
IM.lookup (Id -> Int
idUniq (Id -> Int) -> Id -> Int
forall a b. (a -> b) -> a -> b
$ JoinId -> Id
jpId JoinId
j) (IntMap JoinId -> Maybe JoinId) -> IntMap JoinId -> Maybe JoinId
forall a b. (a -> b) -> a -> b
$ Renames -> IntMap JoinId
rnJoins Renames
r of
                  Just JoinId
j' -> JoinId
j'
                  Maybe JoinId
Nothing -> JoinId
j

            ty :: Ty -> Ty
            ty :: Ty -> Ty
ty | HashMap TyVar Ty -> Bool
forall k v. HashMap k v -> Bool
Map.null (Renames -> HashMap TyVar Ty
rnTvs Renames
r) = Ty -> Ty
forall a. a -> a
id
               | Bool
otherwise          = HashMap TyVar Ty -> Ty -> Ty
substTv (Renames -> HashMap TyVar Ty
rnTvs Renames
r)

            arg :: MonadState Uniq m => Arg -> m Arg
            arg :: forall (m :: * -> *). MonadState Int m => Arg -> m Arg
arg = \ case
                  EArg Exp
x -> Exp -> Arg
EArg (Exp -> Arg) -> m Exp -> m Arg
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Renames -> Exp -> m Exp
forall (m :: * -> *). MonadState Int m => Renames -> Exp -> m Exp
rex Renames
r Exp
x
                  TArg Ty
t -> Arg -> m Arg
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Arg -> m Arg) -> Arg -> m Arg
forall a b. (a -> b) -> a -> b
$ Ty -> Arg
TArg (Ty -> Arg) -> Ty -> Arg
forall a b. (a -> b) -> a -> b
$ Ty -> Ty
ty Ty
t

            bnd :: MonadState Uniq m => Bind -> m (Renames, Bind)
            bnd :: forall (m :: * -> *). MonadState Int m => Bind -> m (Renames, Bind)
bnd = \ case
                  NonRec Id
x Exp
rhs -> do
                        rhs'     <- Renames -> Exp -> m Exp
forall (m :: * -> *). MonadState Int m => Renames -> Exp -> m Exp
rex Renames
r Exp
rhs
                        (r', x') <- bindId r x
                        pure (r', NonRec x' rhs')
                  Rec [(Id, Exp)]
bs       -> do
                        (r', xs') <- Renames -> [Id] -> m (Renames, [Id])
forall (m :: * -> *).
MonadState Int m =>
Renames -> [Id] -> m (Renames, [Id])
bindIds Renames
r ([Id] -> m (Renames, [Id])) -> [Id] -> m (Renames, [Id])
forall a b. (a -> b) -> a -> b
$ ((Id, Exp) -> Id) -> [(Id, Exp)] -> [Id]
forall a b. (a -> b) -> [a] -> [b]
map (Id, Exp) -> Id
forall a b. (a, b) -> a
fst [(Id, Exp)]
bs
                        rhss'     <- mapM (rex r' . snd) bs
                        pure (r', Rec $ zip xs' rhss')
                  -- The label scopes over the body of the LET (joins are
                  -- non-recursive; the renaming is threaded through the
                  -- join's own body too, where it is inert on lint-clean
                  -- input); the params scope only over the join body.
                  Join JoinId
j [Id]
ps Exp
b  -> do
                        (rj, jx') <- Renames -> Id -> m (Renames, Id)
forall (m :: * -> *).
MonadState Int m =>
Renames -> Id -> m (Renames, Id)
bindId Renames
r (Id -> m (Renames, Id)) -> Id -> m (Renames, Id)
forall a b. (a -> b) -> a -> b
$ JoinId -> Id
jpId JoinId
j
                        let j'  = Id -> Int -> JoinId
JoinId Id
jx' (Int -> JoinId) -> Int -> JoinId
forall a b. (a -> b) -> a -> b
$ JoinId -> Int
jpArity JoinId
j
                            rj' = Renames
rj { rnJoins = IM.insert (idUniq $ jpId j) j' $ rnJoins rj }
                        (rp, ps') <- bindIds rj' ps
                        b'        <- rex rp b
                        pure (rj', Join j' ps' b')

            alt :: MonadState Uniq m => Renames -> Alt -> m Alt
            alt :: forall (m :: * -> *). MonadState Int m => Renames -> Alt -> m Alt
alt Renames
r' (Alt Annote
an AltCon
c [Id]
xs Exp
body) = do
                  (r'', xs') <- Renames -> [Id] -> m (Renames, [Id])
forall (m :: * -> *).
MonadState Int m =>
Renames -> [Id] -> m (Renames, [Id])
bindIds Renames
r' [Id]
xs
                  Alt an c xs' <$> rex r'' body

---
--- Substitution.
---

-- | Substitute expressions for variable occurrences (by unique). Each
--   payload is inserted as is: use only when each payload can land at most
--   once, or is binder-free; otherwise use 'substVarsRefreshing'.
substVars :: IM.IntMap Exp -> Exp -> Exp
substVars :: IntMap Exp -> Exp -> Exp
substVars IntMap Exp
s = Exp -> Exp
go
      where go :: Exp -> Exp
            go :: Exp -> Exp
go Exp
e = case Exp
e of
                  Var Annote
_ Id
x         -> Exp -> Int -> IntMap Exp -> Exp
forall a. a -> Int -> IntMap a -> a
IM.findWithDefault Exp
e (Id -> Int
idUniq Id
x) IntMap Exp
s
                  Con {}          -> Exp
e
                  Prim {}         -> Exp
e
                  LitInt {}       -> Exp
e
                  LitStr {}       -> Exp
e
                  LitList Annote
an Ty
t [Exp]
es -> Annote -> Ty -> [Exp] -> Exp
LitList Annote
an Ty
t ([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ (Exp -> Exp) -> [Exp] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> Exp
go [Exp]
es
                  LitVec Annote
an Ty
t [Exp]
es  -> Annote -> Ty -> [Exp] -> Exp
LitVec Annote
an Ty
t ([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ (Exp -> Exp) -> [Exp] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> Exp
go [Exp]
es
                  App Annote
an Exp
f Arg
a      -> Annote -> Exp -> Arg -> Exp
App Annote
an (Exp -> Exp
go Exp
f) (Arg -> Exp) -> Arg -> Exp
forall a b. (a -> b) -> a -> b
$ Arg -> Arg
goArg Arg
a
                  Lam Annote
an Id
x Exp
b      -> Annote -> Id -> Exp -> Exp
Lam Annote
an Id
x (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
go Exp
b
                  Let Annote
an Bind
b Exp
body   -> Annote -> Bind -> Exp -> Exp
Let Annote
an (Bind -> Bind
goBind Bind
b) (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
go Exp
body
                  Jump Annote
an JoinId
j [Exp]
es    -> Annote -> JoinId -> [Exp] -> Exp
Jump Annote
an JoinId
j ([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ (Exp -> Exp) -> [Exp] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> Exp
go [Exp]
es
                  Case Annote
an Ty
t Exp
sc Id
cb [Alt]
alts -> Annote -> Ty -> Exp -> Id -> [Alt] -> Exp
Case Annote
an Ty
t (Exp -> Exp
go Exp
sc) Id
cb [ Annote -> AltCon -> [Id] -> Exp -> Alt
Alt Annote
aan AltCon
c [Id]
xs (Exp -> Exp
go Exp
b) | Alt Annote
aan AltCon
c [Id]
xs Exp
b <- [Alt]
alts ]

            goArg :: Arg -> Arg
            goArg :: Arg -> Arg
goArg = \ case
                  EArg Exp
e -> Exp -> Arg
EArg (Exp -> Arg) -> Exp -> Arg
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
go Exp
e
                  Arg
t      -> Arg
t

            goBind :: Bind -> Bind
            goBind :: Bind -> Bind
goBind = \ case
                  NonRec Id
x Exp
rhs -> Id -> Exp -> Bind
NonRec Id
x (Exp -> Bind) -> Exp -> Bind
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
go Exp
rhs
                  Rec [(Id, Exp)]
bs       -> [(Id, Exp)] -> Bind
Rec [ (Id
x, Exp -> Exp
go Exp
rhs) | (Id
x, Exp
rhs) <- [(Id, Exp)]
bs ]
                  Join JoinId
j [Id]
ps Exp
b  -> JoinId -> [Id] -> Exp -> Bind
Join JoinId
j [Id]
ps (Exp -> Bind) -> Exp -> Bind
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
go Exp
b

-- | Substitute expressions for variable occurrences, refreshing every
--   inserted copy (the uniqueness-preserving form).
substVarsRefreshing :: MonadState Uniq m => IM.IntMap Exp -> Exp -> m Exp
substVarsRefreshing :: forall (m :: * -> *).
MonadState Int m =>
IntMap Exp -> Exp -> m Exp
substVarsRefreshing IntMap Exp
s = Exp -> m Exp
forall (m :: * -> *). MonadState Int m => Exp -> m Exp
go
      where go :: MonadState Uniq m => Exp -> m Exp
            go :: forall (m :: * -> *). MonadState Int m => Exp -> m Exp
go Exp
e = case Exp
e of
                  Var Annote
_ Id
x
                        | Just Exp
e' <- Int -> IntMap Exp -> Maybe Exp
forall a. Int -> IntMap a -> Maybe a
IM.lookup (Id -> Int
idUniq Id
x) IntMap Exp
s -> Exp -> m Exp
forall (m :: * -> *). MonadState Int m => Exp -> m Exp
refreshExp Exp
e'
                        | Bool
otherwise                         -> Exp -> m Exp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
                  Con {}          -> Exp -> m Exp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
                  Prim {}         -> Exp -> m Exp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
                  LitInt {}       -> Exp -> m Exp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
                  LitStr {}       -> Exp -> m Exp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
                  LitList Annote
an Ty
t [Exp]
es -> Annote -> Ty -> [Exp] -> Exp
LitList Annote
an Ty
t ([Exp] -> Exp) -> m [Exp] -> m Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> m Exp) -> [Exp] -> m [Exp]
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 Exp -> m Exp
forall (m :: * -> *). MonadState Int m => Exp -> m Exp
go [Exp]
es
                  LitVec Annote
an Ty
t [Exp]
es  -> Annote -> Ty -> [Exp] -> Exp
LitVec Annote
an Ty
t ([Exp] -> Exp) -> m [Exp] -> m Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> m Exp) -> [Exp] -> m [Exp]
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 Exp -> m Exp
forall (m :: * -> *). MonadState Int m => Exp -> m Exp
go [Exp]
es
                  App Annote
an Exp
f Arg
a      -> Annote -> Exp -> Arg -> Exp
App Annote
an (Exp -> Arg -> Exp) -> m Exp -> m (Arg -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> m Exp
forall (m :: * -> *). MonadState Int m => Exp -> m Exp
go Exp
f m (Arg -> Exp) -> m Arg -> m Exp
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Arg -> m Arg
forall (m :: * -> *). MonadState Int m => Arg -> m Arg
goArg Arg
a
                  Lam Annote
an Id
x Exp
b      -> Annote -> Id -> Exp -> Exp
Lam Annote
an Id
x (Exp -> Exp) -> m Exp -> m Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> m Exp
forall (m :: * -> *). MonadState Int m => Exp -> m Exp
go Exp
b
                  Let Annote
an Bind
b Exp
body   -> Annote -> Bind -> Exp -> Exp
Let Annote
an (Bind -> Exp -> Exp) -> m Bind -> m (Exp -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bind -> m Bind
forall (m :: * -> *). MonadState Int m => Bind -> m Bind
goBind Bind
b m (Exp -> Exp) -> m Exp -> m Exp
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> m Exp
forall (m :: * -> *). MonadState Int m => Exp -> m Exp
go Exp
body
                  Jump Annote
an JoinId
j [Exp]
es    -> Annote -> JoinId -> [Exp] -> Exp
Jump Annote
an JoinId
j ([Exp] -> Exp) -> m [Exp] -> m Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> m Exp) -> [Exp] -> m [Exp]
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 Exp -> m Exp
forall (m :: * -> *). MonadState Int m => Exp -> m Exp
go [Exp]
es
                  Case Annote
an Ty
t Exp
sc Id
cb [Alt]
alts -> do
                        sc'   <- Exp -> m Exp
forall (m :: * -> *). MonadState Int m => Exp -> m Exp
go Exp
sc
                        alts' <- mapM (\ (Alt Annote
aan AltCon
c [Id]
xs Exp
b) -> Annote -> AltCon -> [Id] -> Exp -> Alt
Alt Annote
aan AltCon
c [Id]
xs (Exp -> Alt) -> m Exp -> m Alt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> m Exp
forall (m :: * -> *). MonadState Int m => Exp -> m Exp
go Exp
b) alts
                        pure $ Case an t sc' cb alts'

            goArg :: MonadState Uniq m => Arg -> m Arg
            goArg :: forall (m :: * -> *). MonadState Int m => Arg -> m Arg
goArg = \ case
                  EArg Exp
e -> Exp -> Arg
EArg (Exp -> Arg) -> m Exp -> m Arg
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> m Exp
forall (m :: * -> *). MonadState Int m => Exp -> m Exp
go Exp
e
                  Arg
t      -> Arg -> m Arg
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Arg
t

            goBind :: MonadState Uniq m => Bind -> m Bind
            goBind :: forall (m :: * -> *). MonadState Int m => Bind -> m Bind
goBind = \ case
                  NonRec Id
x Exp
rhs -> Id -> Exp -> Bind
NonRec Id
x (Exp -> Bind) -> m Exp -> m Bind
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> m Exp
forall (m :: * -> *). MonadState Int m => Exp -> m Exp
go Exp
rhs
                  Rec [(Id, Exp)]
bs       -> [(Id, Exp)] -> Bind
Rec ([(Id, Exp)] -> Bind) -> m [(Id, Exp)] -> m Bind
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Id, Exp) -> m (Id, Exp)) -> [(Id, Exp)] -> m [(Id, Exp)]
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 (\ (Id
x, Exp
rhs) -> (Id
x, ) (Exp -> (Id, Exp)) -> m Exp -> m (Id, Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> m Exp
forall (m :: * -> *). MonadState Int m => Exp -> m Exp
go Exp
rhs) [(Id, Exp)]
bs
                  Join JoinId
j [Id]
ps Exp
b  -> JoinId -> [Id] -> Exp -> Bind
Join JoinId
j [Id]
ps (Exp -> Bind) -> m Exp -> m Bind
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> m Exp
forall (m :: * -> *). MonadState Int m => Exp -> m Exp
go Exp
b

---
--- Occurrence analysis.
---

-- | One representative occurrence Id per occurring unique. An occurrence
--   carries its binder's signature, so this recovers binder Ids (e.g. for
--   captured locals) without an enclosing environment.
occIds :: Exp -> IM.IntMap Id
occIds :: Exp -> IntMap Id
occIds = \ case
      Var Annote
_ Id
v         -> Int -> Id -> IntMap Id
forall a. Int -> a -> IntMap a
IM.singleton (Id -> Int
idUniq Id
v) Id
v
      App Annote
_ Exp
f Arg
a       -> Exp -> IntMap Id
occIds Exp
f IntMap Id -> IntMap Id -> IntMap Id
forall a. Semigroup a => a -> a -> a
<> Arg -> IntMap Id
argIds Arg
a
      Lam Annote
_ Id
_ Exp
b       -> Exp -> IntMap Id
occIds Exp
b
      Let Annote
_ Bind
b Exp
body    -> Bind -> IntMap Id
bindIds' Bind
b IntMap Id -> IntMap Id -> IntMap Id
forall a. Semigroup a => a -> a -> a
<> Exp -> IntMap Id
occIds Exp
body
      Jump Annote
_ JoinId
_ [Exp]
es     -> [IntMap Id] -> IntMap Id
forall (f :: * -> *) a. Foldable f => f (IntMap a) -> IntMap a
IM.unions ([IntMap Id] -> IntMap Id) -> [IntMap Id] -> IntMap Id
forall a b. (a -> b) -> a -> b
$ (Exp -> IntMap Id) -> [Exp] -> [IntMap Id]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> IntMap Id
occIds [Exp]
es
      Case Annote
_ Ty
_ Exp
s Id
_ [Alt]
as -> Exp -> IntMap Id
occIds Exp
s IntMap Id -> IntMap Id -> IntMap Id
forall a. Semigroup a => a -> a -> a
<> [IntMap Id] -> IntMap Id
forall (f :: * -> *) a. Foldable f => f (IntMap a) -> IntMap a
IM.unions [ Exp -> IntMap Id
occIds Exp
b | Alt Annote
_ AltCon
_ [Id]
_ Exp
b <- [Alt]
as ]
      LitList Annote
_ Ty
_ [Exp]
es  -> [IntMap Id] -> IntMap Id
forall (f :: * -> *) a. Foldable f => f (IntMap a) -> IntMap a
IM.unions ([IntMap Id] -> IntMap Id) -> [IntMap Id] -> IntMap Id
forall a b. (a -> b) -> a -> b
$ (Exp -> IntMap Id) -> [Exp] -> [IntMap Id]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> IntMap Id
occIds [Exp]
es
      LitVec Annote
_ Ty
_ [Exp]
es   -> [IntMap Id] -> IntMap Id
forall (f :: * -> *) a. Foldable f => f (IntMap a) -> IntMap a
IM.unions ([IntMap Id] -> IntMap Id) -> [IntMap Id] -> IntMap Id
forall a b. (a -> b) -> a -> b
$ (Exp -> IntMap Id) -> [Exp] -> [IntMap Id]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> IntMap Id
occIds [Exp]
es
      Exp
_               -> IntMap Id
forall a. Monoid a => a
mempty
      where argIds :: Arg -> IM.IntMap Id
            argIds :: Arg -> IntMap Id
argIds = \ case
                  EArg Exp
e -> Exp -> IntMap Id
occIds Exp
e
                  Arg
_      -> IntMap Id
forall a. Monoid a => a
mempty

            bindIds' :: Bind -> IM.IntMap Id
            bindIds' :: Bind -> IntMap Id
bindIds' = \ case
                  NonRec Id
_ Exp
rhs -> Exp -> IntMap Id
occIds Exp
rhs
                  Rec [(Id, Exp)]
bs       -> [IntMap Id] -> IntMap Id
forall (f :: * -> *) a. Foldable f => f (IntMap a) -> IntMap a
IM.unions ([IntMap Id] -> IntMap Id) -> [IntMap Id] -> IntMap Id
forall a b. (a -> b) -> a -> b
$ ((Id, Exp) -> IntMap Id) -> [(Id, Exp)] -> [IntMap Id]
forall a b. (a -> b) -> [a] -> [b]
map (Exp -> IntMap Id
occIds (Exp -> IntMap Id) -> ((Id, Exp) -> Exp) -> (Id, Exp) -> IntMap Id
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Id, Exp) -> Exp
forall a b. (a, b) -> b
snd) [(Id, Exp)]
bs
                  Join JoinId
_ [Id]
_ Exp
b   -> Exp -> IntMap Id
occIds Exp
b

-- | The free variables of an expression, by unique: occurrences whose
--   binding site is not within the expression. (Global binder uniqueness
--   makes this a plain set difference — no shadowing is possible.)
freeUniqs :: Exp -> IS.IntSet
freeUniqs :: Exp -> IntSet
freeUniqs Exp
e = IntMap Int -> IntSet
forall a. IntMap a -> IntSet
IM.keysSet (Exp -> IntMap Int
occCounts Exp
e) IntSet -> IntSet -> IntSet
IS.\\ Exp -> IntSet
binderUniqs Exp
e
      where binderUniqs :: Exp -> IS.IntSet
            binderUniqs :: Exp -> IntSet
binderUniqs = \ case
                  Lam Annote
_ Id
x Exp
b       -> Int -> IntSet -> IntSet
IS.insert (Id -> Int
idUniq Id
x) (IntSet -> IntSet) -> IntSet -> IntSet
forall a b. (a -> b) -> a -> b
$ Exp -> IntSet
binderUniqs Exp
b
                  Let Annote
_ Bind
b Exp
body    -> Bind -> IntSet
bnd Bind
b IntSet -> IntSet -> IntSet
forall a. Semigroup a => a -> a -> a
<> Exp -> IntSet
binderUniqs Exp
body
                  App Annote
_ Exp
f Arg
a       -> Exp -> IntSet
binderUniqs Exp
f IntSet -> IntSet -> IntSet
forall a. Semigroup a => a -> a -> a
<> Arg -> IntSet
arg Arg
a
                  Jump Annote
_ JoinId
_ [Exp]
es     -> [IntSet] -> IntSet
forall (f :: * -> *). Foldable f => f IntSet -> IntSet
IS.unions ([IntSet] -> IntSet) -> [IntSet] -> IntSet
forall a b. (a -> b) -> a -> b
$ (Exp -> IntSet) -> [Exp] -> [IntSet]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> IntSet
binderUniqs [Exp]
es
                  Case Annote
_ Ty
_ Exp
s Id
x [Alt]
as -> Int -> IntSet -> IntSet
IS.insert (Id -> Int
idUniq Id
x) (IntSet -> IntSet) -> IntSet -> IntSet
forall a b. (a -> b) -> a -> b
$ Exp -> IntSet
binderUniqs Exp
s IntSet -> IntSet -> IntSet
forall a. Semigroup a => a -> a -> a
<> [IntSet] -> IntSet
forall (f :: * -> *). Foldable f => f IntSet -> IntSet
IS.unions [ [Int] -> IntSet
IS.fromList ((Id -> Int) -> [Id] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Id -> Int
idUniq [Id]
xs) IntSet -> IntSet -> IntSet
forall a. Semigroup a => a -> a -> a
<> Exp -> IntSet
binderUniqs Exp
b | Alt Annote
_ AltCon
_ [Id]
xs Exp
b <- [Alt]
as ]
                  LitList Annote
_ Ty
_ [Exp]
es  -> [IntSet] -> IntSet
forall (f :: * -> *). Foldable f => f IntSet -> IntSet
IS.unions ([IntSet] -> IntSet) -> [IntSet] -> IntSet
forall a b. (a -> b) -> a -> b
$ (Exp -> IntSet) -> [Exp] -> [IntSet]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> IntSet
binderUniqs [Exp]
es
                  LitVec Annote
_ Ty
_ [Exp]
es   -> [IntSet] -> IntSet
forall (f :: * -> *). Foldable f => f IntSet -> IntSet
IS.unions ([IntSet] -> IntSet) -> [IntSet] -> IntSet
forall a b. (a -> b) -> a -> b
$ (Exp -> IntSet) -> [Exp] -> [IntSet]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> IntSet
binderUniqs [Exp]
es
                  Exp
_               -> IntSet
forall a. Monoid a => a
mempty

            bnd :: Bind -> IS.IntSet
            bnd :: Bind -> IntSet
bnd = \ case
                  NonRec Id
x Exp
rhs -> Int -> IntSet -> IntSet
IS.insert (Id -> Int
idUniq Id
x) (IntSet -> IntSet) -> IntSet -> IntSet
forall a b. (a -> b) -> a -> b
$ Exp -> IntSet
binderUniqs Exp
rhs
                  Rec [(Id, Exp)]
bs       -> [Int] -> IntSet
IS.fromList (((Id, Exp) -> Int) -> [(Id, Exp)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Id -> Int
idUniq (Id -> Int) -> ((Id, Exp) -> Id) -> (Id, Exp) -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Id, Exp) -> Id
forall a b. (a, b) -> a
fst) [(Id, Exp)]
bs) IntSet -> IntSet -> IntSet
forall a. Semigroup a => a -> a -> a
<> [IntSet] -> IntSet
forall (f :: * -> *). Foldable f => f IntSet -> IntSet
IS.unions (((Id, Exp) -> IntSet) -> [(Id, Exp)] -> [IntSet]
forall a b. (a -> b) -> [a] -> [b]
map (Exp -> IntSet
binderUniqs (Exp -> IntSet) -> ((Id, Exp) -> Exp) -> (Id, Exp) -> IntSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Id, Exp) -> Exp
forall a b. (a, b) -> b
snd) [(Id, Exp)]
bs)
                  Join JoinId
j [Id]
ps Exp
b  -> [Int] -> IntSet
IS.fromList (Id -> Int
idUniq (JoinId -> Id
jpId JoinId
j) Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: (Id -> Int) -> [Id] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Id -> Int
idUniq [Id]
ps) IntSet -> IntSet -> IntSet
forall a. Semigroup a => a -> a -> a
<> Exp -> IntSet
binderUniqs Exp
b

            arg :: Arg -> IS.IntSet
            arg :: Arg -> IntSet
arg = \ case
                  EArg Exp
x -> Exp -> IntSet
binderUniqs Exp
x
                  Arg
_      -> IntSet
forall a. Monoid a => a
mempty

-- | Variable- and jump-occurrence counts by unique (an occurrence of a
--   join label at a jump counts for the label's Id). Dead binders are
--   absent from the map.
occCounts :: Exp -> IM.IntMap Int
occCounts :: Exp -> IntMap Int
occCounts = Exp -> IntMap Int
go
      where go :: Exp -> IM.IntMap Int
            go :: Exp -> IntMap Int
go = \ case
                  Var Annote
_ Id
x         -> Int -> Int -> IntMap Int
forall a. Int -> a -> IntMap a
IM.singleton (Id -> Int
idUniq Id
x) Int
1
                  Jump Annote
_ JoinId
j [Exp]
es     -> (Int -> Int -> Int) -> [IntMap Int] -> IntMap Int
forall (f :: * -> *) a.
Foldable f =>
(a -> a -> a) -> f (IntMap a) -> IntMap a
IM.unionsWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) ([IntMap Int] -> IntMap Int) -> [IntMap Int] -> IntMap Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> IntMap Int
forall a. Int -> a -> IntMap a
IM.singleton (Id -> Int
idUniq (Id -> Int) -> Id -> Int
forall a b. (a -> b) -> a -> b
$ JoinId -> Id
jpId JoinId
j) Int
1 IntMap Int -> [IntMap Int] -> [IntMap Int]
forall a. a -> [a] -> [a]
: (Exp -> IntMap Int) -> [Exp] -> [IntMap Int]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> IntMap Int
go [Exp]
es
                  App Annote
_ Exp
f Arg
a       -> (Int -> Int -> Int) -> IntMap Int -> IntMap Int -> IntMap Int
forall a. (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
IM.unionWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) (Exp -> IntMap Int
go Exp
f) (IntMap Int -> IntMap Int) -> IntMap Int -> IntMap Int
forall a b. (a -> b) -> a -> b
$ Arg -> IntMap Int
goArg Arg
a
                  Lam Annote
_ Id
_ Exp
b       -> Exp -> IntMap Int
go Exp
b
                  Let Annote
_ Bind
b Exp
body    -> (Int -> Int -> Int) -> IntMap Int -> IntMap Int -> IntMap Int
forall a. (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
IM.unionWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) (Bind -> IntMap Int
goBind Bind
b) (IntMap Int -> IntMap Int) -> IntMap Int -> IntMap Int
forall a b. (a -> b) -> a -> b
$ Exp -> IntMap Int
go Exp
body
                  Case Annote
_ Ty
_ Exp
s Id
_ [Alt]
as -> (Int -> Int -> Int) -> [IntMap Int] -> IntMap Int
forall (f :: * -> *) a.
Foldable f =>
(a -> a -> a) -> f (IntMap a) -> IntMap a
IM.unionsWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) ([IntMap Int] -> IntMap Int) -> [IntMap Int] -> IntMap Int
forall a b. (a -> b) -> a -> b
$ Exp -> IntMap Int
go Exp
s IntMap Int -> [IntMap Int] -> [IntMap Int]
forall a. a -> [a] -> [a]
: [ Exp -> IntMap Int
go Exp
b | Alt Annote
_ AltCon
_ [Id]
_ Exp
b <- [Alt]
as ]
                  LitList Annote
_ Ty
_ [Exp]
es  -> (Int -> Int -> Int) -> [IntMap Int] -> IntMap Int
forall (f :: * -> *) a.
Foldable f =>
(a -> a -> a) -> f (IntMap a) -> IntMap a
IM.unionsWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) ([IntMap Int] -> IntMap Int) -> [IntMap Int] -> IntMap Int
forall a b. (a -> b) -> a -> b
$ (Exp -> IntMap Int) -> [Exp] -> [IntMap Int]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> IntMap Int
go [Exp]
es
                  LitVec Annote
_ Ty
_ [Exp]
es   -> (Int -> Int -> Int) -> [IntMap Int] -> IntMap Int
forall (f :: * -> *) a.
Foldable f =>
(a -> a -> a) -> f (IntMap a) -> IntMap a
IM.unionsWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) ([IntMap Int] -> IntMap Int) -> [IntMap Int] -> IntMap Int
forall a b. (a -> b) -> a -> b
$ (Exp -> IntMap Int) -> [Exp] -> [IntMap Int]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> IntMap Int
go [Exp]
es
                  Exp
_               -> IntMap Int
forall a. Monoid a => a
mempty

            goArg :: Arg -> IM.IntMap Int
            goArg :: Arg -> IntMap Int
goArg = \ case
                  EArg Exp
e -> Exp -> IntMap Int
go Exp
e
                  Arg
_      -> IntMap Int
forall a. Monoid a => a
mempty

            goBind :: Bind -> IM.IntMap Int
            goBind :: Bind -> IntMap Int
goBind = \ case
                  NonRec Id
_ Exp
rhs -> Exp -> IntMap Int
go Exp
rhs
                  Rec [(Id, Exp)]
bs       -> (Int -> Int -> Int) -> [IntMap Int] -> IntMap Int
forall (f :: * -> *) a.
Foldable f =>
(a -> a -> a) -> f (IntMap a) -> IntMap a
IM.unionsWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) ([IntMap Int] -> IntMap Int) -> [IntMap Int] -> IntMap Int
forall a b. (a -> b) -> a -> b
$ ((Id, Exp) -> IntMap Int) -> [(Id, Exp)] -> [IntMap Int]
forall a b. (a -> b) -> [a] -> [b]
map (Exp -> IntMap Int
go (Exp -> IntMap Int)
-> ((Id, Exp) -> Exp) -> (Id, Exp) -> IntMap Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Id, Exp) -> Exp
forall a b. (a, b) -> b
snd) [(Id, Exp)]
bs
                  Join JoinId
_ [Id]
_ Exp
b   -> Exp -> IntMap Int
go Exp
b