{-# LANGUAGE Safe #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
-- | The Eidos simplifier: occurrence-driven, let-preserving partial
--   evaluation (the retired reduce\/specialize\/purge loop's successor,
--   doc/eidos.md §8). 'simplify' repeats specialize-on-values, purge, and
--   reduce until every definition is synthable and dictionary-free,
--   bounded by @--depth@ with the retired loop's diagnostic.
--
--   * 'reduceExp' is let-preserving: beta redexes become lets, and a let
--     is inlined only when its binder is dead, used once, bound to an
--     atom, or dictionary-typed ('dictTyCons' — never representable, so
--     substituted at every occurrence); a representable multi-use binding
--     KEEPS its let (sharing).
--     The one lambda-lifting policy surviving from the retired pipeline
--     is explicit here (Clash's LiftNonRep): a multi-use *function-typed*
--     let cannot stay (it is not representable), so it is lifted to a
--     top-level definition over its captured locals, named with the
--     @$LL.@ compiler-lifted prefix (display-only: the machine fold
--     strips it when naming the Hyle definition). Top-level references
--     are never unfolded — the function hierarchy is preserved;
--     higher-orderness dies by argument baking, not call inlining.
--
--   * 'specialize' on values (the retired value-argument specializer):
--     a call to a top-level definition with *closed* arguments (free
--     variables all top-level) bakes those arguments into a memoized
--     clone named from the baked-argument canon ('originTag').
--     Self-calls in the clone still reference the
--     origin — the memo closes recursive loops. Memo keys are
--     alpha-canonical: arguments are renumbered from a disjoint unique
--     supply and printed. Baking is also how non-representable scalar
--     arguments (Integer, String) disappear: their call sites pass
--     closed values.
--
--   * 'purge': definitions unreachable from the device root (plus the
--     builtin signature carriers, which carry the builtins' type
--     assumptions to the Eidos-to-Hyle fold).
--
--   Join points are preserved (dead ones are dropped); datatypes are left
--   untouched.
module ReWire.Eidos.Simplify (simplify, purge, reduceProgram, reduceExp, dictTyCons, SimpT, runSimpT) where

import ReWire.Annotation (Annote, noAnn)
import ReWire.Builtins (builtins)
import ReWire.Error (AstError, MonadError, failAt)
import ReWire.Eidos.Naming (originTag)
import ReWire.Eidos.Pretty ()
import ReWire.Eidos.Subst (nextUniq, instantiateDefn, refreshDefn, refreshExp, substVars, substVarsRefreshing, occCounts, freeUniqs)
import ReWire.Eidos.Syntax
import ReWire.Eidos.Types (typeOf, flattenApp, flattenArrow, hasArrow, synthable)
import ReWire.Pretty (prettyPrint)

import Control.Monad ((>=>))
import Control.Monad.State.Strict (StateT, evalStateT, evalState, runStateT, get, put, gets, modify, lift)
import Data.HashMap.Strict (HashMap)
import Data.HashSet (HashSet)
import Data.List (sortOn)
import Data.Char (isDigit)
import Data.Maybe (fromMaybe, isNothing, mapMaybe)
import Data.Text (Text, intercalate)

import qualified Data.Text as T
import Numeric.Natural (Natural)

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

-- | Simplifier state: the unique supply, definitions minted this phase
--   (lifted functions and specializations, drained into the program), and
--   the persistent specialization memo.
data SimpSt = SimpSt
      { SimpSt -> Uniq
stSupply :: !Uniq
      , SimpSt -> [Defn]
stNew    :: ![Defn]
      , SimpSt -> HashMap (Uniq, Text) Defn
stMemo   :: !(HashMap (Uniq, Text) Defn)
      }

type SimpT m = StateT SimpSt m

-- | Run a simplifier action with a supply seeded above the program.
runSimpT :: Monad m => Program -> SimpT m a -> m a
runSimpT :: forall (m :: * -> *) a. Monad m => Program -> SimpT m a -> m a
runSimpT Program
p SimpT m a
m = SimpT m a -> SimpSt -> m a
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT SimpT m a
m (SimpSt -> m a) -> SimpSt -> m a
forall a b. (a -> b) -> a -> b
$ Uniq -> [Defn] -> HashMap (Uniq, Text) Defn -> SimpSt
SimpSt (Program -> Uniq
forall a. Data a => a -> Uniq
nextUniq Program
p) [] HashMap (Uniq, Text) Defn
forall a. Monoid a => a
mempty

-- | Run a Subst primitive on the simplifier's supply.
supplied :: Monad m => StateT Uniq m a -> SimpT m a
supplied :: forall (m :: * -> *) a. Monad m => StateT Uniq m a -> SimpT m a
supplied StateT Uniq m a
m = do
      st       <- StateT SimpSt m SimpSt
forall s (m :: * -> *). MonadState s m => m s
get
      (a, sup) <- lift $ runStateT m $ stSupply st
      put st { stSupply = sup }
      pure a

mint :: Monad m => Defn -> SimpT m ()
mint :: forall (m :: * -> *). Monad m => Defn -> SimpT m ()
mint Defn
d = (SimpSt -> SimpSt) -> StateT SimpSt m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((SimpSt -> SimpSt) -> StateT SimpSt m ())
-> (SimpSt -> SimpSt) -> StateT SimpSt m ()
forall a b. (a -> b) -> a -> b
$ \ SimpSt
st -> SimpSt
st { stNew = stNew st <> [d] }

drainNew :: Monad m => SimpT m [Defn]
drainNew :: forall (m :: * -> *). Monad m => SimpT m [Defn]
drainNew = do
      st <- StateT SimpSt m SimpSt
forall s (m :: * -> *). MonadState s m => m s
get
      put st { stNew = [] }
      pure $ stNew st

primNames :: HashSet Text
primNames :: HashSet Text
primNames = [Text] -> HashSet Text
forall a. (Eq a, Hashable a) => [a] -> HashSet a
Set.fromList ([Text] -> HashSet Text) -> [Text] -> HashSet Text
forall a b. (a -> b) -> a -> b
$ ((Text, Builtin) -> Text) -> [(Text, Builtin)] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text, Builtin) -> Text
forall a b. (a, b) -> a
fst [(Text, Builtin)]
builtins

---
--- The driver.
---

-- | Partial evaluation to the synthable, dictionary-free fixpoint,
--   bounded by the given depth.
simplify :: forall m. MonadError AstError m => Natural -> Program -> m Program
simplify :: forall (m :: * -> *).
MonadError AstError m =>
Natural -> Program -> m Program
simplify Natural
depth Program
p = Program -> SimpT m Program -> m Program
forall (m :: * -> *) a. Monad m => Program -> SimpT m a -> m a
runSimpT Program
p (SimpT m Program -> m Program) -> SimpT m Program -> m Program
forall a b. (a -> b) -> a -> b
$ Natural -> Program -> SimpT m Program
go (Natural -> Natural -> Natural
forall a. Ord a => a -> a -> a
max Natural
1 Natural
depth) Program
p
      where -- At least one round always runs (the retired loop's
            -- behavior): the done conditions are signature-level, so a
            -- program can arrive converged by them with unreduced bodies.
            go :: Natural -> Program -> SimpT m Program
            go :: Natural -> Program -> SimpT m Program
go Natural
n Program
pr = do
                  pr' <- Program -> SimpT m Program
step Program
pr
                  if | done pr'  -> pure pr'
                     | n <= 1    -> failAt noAnn
                           $ "Partial evaluation not terminating (mutually recursive definitions?). Not synthable: "
                           <> intercalate ", " (stuck pr')
                     | otherwise -> go (n - 1) pr'

            step :: Program -> SimpT m Program
            step :: Program -> SimpT m Program
step = Program -> SimpT m Program
forall (m :: * -> *).
MonadError AstError m =>
Program -> SimpT m Program
specialize (Program -> SimpT m Program)
-> (Program -> SimpT m Program) -> Program -> SimpT m Program
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> (Program -> SimpT m Program
forall a. a -> StateT SimpSt m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Program -> SimpT m Program)
-> (Program -> Program) -> Program -> SimpT m Program
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Program -> Program
purge) (Program -> SimpT m Program)
-> (Program -> SimpT m Program) -> Program -> SimpT m Program
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> Program -> SimpT m Program
forall (m :: * -> *).
MonadError AstError m =>
Program -> SimpT m Program
reduceProgram

            done :: Program -> Bool
            done :: Program -> Bool
done Program
pr = (Defn -> Bool) -> [Defn] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Defn -> Bool
synthableDefn (Program -> [Defn]
progDefns Program
pr) Bool -> Bool -> Bool
&& Program -> Bool
dictFree Program
pr

            -- Both ways a definition can hold up the fixpoint: a
            -- non-synthable signature, or a signature still mentioning a
            -- dictionary datatype (which alone would otherwise report an
            -- empty list).
            stuck :: Program -> [Text]
            stuck :: Program -> [Text]
stuck Program
pr = [ Id -> Text
forall a. Pretty a => a -> Text
prettyPrint (Defn -> Id
defnId Defn
d) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" :: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Ty -> Text
forall a. Pretty a => a -> Text
prettyPrint (Sig -> Ty
sigTy (Sig -> Ty) -> Sig -> Ty
forall a b. (a -> b) -> a -> b
$ Id -> Sig
idSig (Id -> Sig) -> Id -> Sig
forall a b. (a -> b) -> a -> b
$ Defn -> Id
defnId Defn
d)
                       | Defn
d <- Program -> [Defn]
progDefns Program
pr, Bool -> Bool
not (Defn -> Bool
synthableDefn Defn
d) Bool -> Bool -> Bool
|| Defn -> Bool
dictSig Defn
d ]
                  where bad :: HashSet Text
                        bad :: HashSet Text
bad = [DataDefn] -> HashSet Text
dictTyCons ([DataDefn] -> HashSet Text) -> [DataDefn] -> HashSet Text
forall a b. (a -> b) -> a -> b
$ Program -> [DataDefn]
progDatas Program
pr

                        dictSig :: Defn -> Bool
                        dictSig :: Defn -> Bool
dictSig Defn
d = (Text -> Bool) -> [Text] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Text -> HashSet Text -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
`Set.member` HashSet Text
bad) ([Text] -> Bool) -> [Text] -> Bool
forall a b. (a -> b) -> a -> b
$ Ty -> [Text]
tyConsOf (Ty -> [Text]) -> Ty -> [Text]
forall a b. (a -> b) -> a -> b
$ Sig -> Ty
sigTy (Sig -> Ty) -> Sig -> Ty
forall a b. (a -> b) -> a -> b
$ Id -> Sig
idSig (Id -> Sig) -> Id -> Sig
forall a b. (a -> b) -> a -> b
$ Defn -> Id
defnId Defn
d

-- | Prim-named definitions are signature carriers; everything else must
--   reach a representable type.
synthableDefn :: Defn -> Bool
synthableDefn :: Defn -> Bool
synthableDefn Defn
d = Text -> HashSet Text -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
Set.member (Id -> Text
idOcc (Id -> Text) -> Id -> Text
forall a b. (a -> b) -> a -> b
$ Defn -> Id
defnId Defn
d) HashSet Text
primNames
               Bool -> Bool -> Bool
|| Ty -> Bool
synthable (Sig -> Ty
sigTy (Sig -> Ty) -> Sig -> Ty
forall a b. (a -> b) -> a -> b
$ Id -> Sig
idSig (Id -> Sig) -> Id -> Sig
forall a b. (a -> b) -> a -> b
$ Defn -> Id
defnId Defn
d)

-- | No definition's signature mentions a datatype with a function-typed
--   constructor field (a class dictionary).
dictFree :: Program -> Bool
dictFree :: Program -> Bool
dictFree (Program [DataDefn]
datas [Defn]
defns Id
_) = HashSet Text -> Bool
forall a. HashSet a -> Bool
Set.null HashSet Text
bad Bool -> Bool -> Bool
|| (Defn -> Bool) -> [Defn] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Defn -> Bool
ok [Defn]
defns
      where bad :: HashSet Text
            bad :: HashSet Text
bad = [DataDefn] -> HashSet Text
dictTyCons [DataDefn]
datas

            ok :: Defn -> Bool
            ok :: Defn -> Bool
ok Defn
d = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (Text -> Bool) -> [Text] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Text -> HashSet Text -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
`Set.member` HashSet Text
bad) ([Text] -> Bool) -> [Text] -> Bool
forall a b. (a -> b) -> a -> b
$ Ty -> [Text]
tyConsOf (Ty -> [Text]) -> Ty -> [Text]
forall a b. (a -> b) -> a -> b
$ Sig -> Ty
sigTy (Sig -> Ty) -> Sig -> Ty
forall a b. (a -> b) -> a -> b
$ Id -> Sig
idSig (Id -> Sig) -> Id -> Sig
forall a b. (a -> b) -> a -> b
$ Defn -> Id
defnId Defn
d

-- | Names of the datatypes with a function-typed constructor field (class
--   dictionaries and their shapes): a value of such a type can never reach
--   Hyle, so bindings of these types must always reduce away.
dictTyCons :: [DataDefn] -> HashSet Text
dictTyCons :: [DataDefn] -> HashSet Text
dictTyCons [DataDefn]
datas = [Text] -> HashSet Text
forall a. (Eq a, Hashable a) => [a] -> HashSet a
Set.fromList [ DataDefn -> Text
dataName DataDefn
d | DataDefn
d <- [DataDefn]
datas, (DataCon -> Bool) -> [DataCon] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any DataCon -> Bool
arrowField ([DataCon] -> Bool) -> [DataCon] -> Bool
forall a b. (a -> b) -> a -> b
$ DataDefn -> [DataCon]
dataCons DataDefn
d ]
      where arrowField :: DataCon -> Bool
            arrowField :: DataCon -> Bool
arrowField (DataCon Annote
_ Text
_ (Sig [TyVar]
_ Ty
t)) = (Ty -> Bool) -> [Ty] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Ty -> Bool
hasArrow ([Ty] -> Bool) -> [Ty] -> Bool
forall a b. (a -> b) -> a -> b
$ ([Ty], Ty) -> [Ty]
forall a b. (a, b) -> a
fst (([Ty], Ty) -> [Ty]) -> ([Ty], Ty) -> [Ty]
forall a b. (a -> b) -> a -> b
$ Ty -> ([Ty], Ty)
flattenArrow Ty
t

-- | Every TyCon name a type mentions.
tyConsOf :: Ty -> [Text]
tyConsOf :: Ty -> [Text]
tyConsOf = \ case
      TyCon Annote
_ Text
c   -> [Text
c]
      TyApp Annote
_ Ty
a Ty
b -> Ty -> [Text]
tyConsOf Ty
a [Text] -> [Text] -> [Text]
forall a. Semigroup a => a -> a -> a
<> Ty -> [Text]
tyConsOf Ty
b
      Arrow Annote
_ Ty
a Ty
b -> Ty -> [Text]
tyConsOf Ty
a [Text] -> [Text] -> [Text]
forall a. Semigroup a => a -> a -> a
<> Ty -> [Text]
tyConsOf Ty
b
      Ty
_           -> []

---
--- Purge: definition-level DCE.
---

-- | Definitions transitively reachable from the device root and the
--   builtin signature carriers. Datatypes are left untouched.
purge :: Program -> Program
purge :: Program -> Program
purge (Program [DataDefn]
datas [Defn]
defns Id
top) = [DataDefn] -> [Defn] -> Id -> Program
Program [DataDefn]
datas [ Defn
d | Defn
d <- [Defn]
defns, Uniq -> IntSet -> Bool
IS.member (Id -> Uniq
idUniq (Id -> Uniq) -> Id -> Uniq
forall a b. (a -> b) -> a -> b
$ Defn -> Id
defnId Defn
d) IntSet
live ] Id
top
      where tops :: IS.IntSet
            tops :: IntSet
tops = [Uniq] -> IntSet
IS.fromList ([Uniq] -> IntSet) -> [Uniq] -> IntSet
forall a b. (a -> b) -> a -> b
$ (Defn -> Uniq) -> [Defn] -> [Uniq]
forall a b. (a -> b) -> [a] -> [b]
map (Id -> Uniq
idUniq (Id -> Uniq) -> (Defn -> Id) -> Defn -> Uniq
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Defn -> Id
defnId) [Defn]
defns

            refs :: IM.IntMap IS.IntSet
            refs :: IntMap IntSet
refs = [(Uniq, IntSet)] -> IntMap IntSet
forall a. [(Uniq, a)] -> IntMap a
IM.fromList [ (Id -> Uniq
idUniq (Id -> Uniq) -> Id -> Uniq
forall a b. (a -> b) -> a -> b
$ Defn -> Id
defnId Defn
d, Exp -> IntSet
freeUniqs (Defn -> Exp
defnBody Defn
d) IntSet -> IntSet -> IntSet
`IS.intersection` IntSet
tops) | Defn
d <- [Defn]
defns ]

            roots :: [Uniq]
            roots :: [Uniq]
roots = Id -> Uniq
idUniq Id
top
                  Uniq -> [Uniq] -> [Uniq]
forall a. a -> [a] -> [a]
: [ Id -> Uniq
idUniq (Id -> Uniq) -> Id -> Uniq
forall a b. (a -> b) -> a -> b
$ Defn -> Id
defnId Defn
d | Defn
d <- [Defn]
defns, Text -> HashSet Text -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
Set.member (Id -> Text
idOcc (Id -> Text) -> Id -> Text
forall a b. (a -> b) -> a -> b
$ Defn -> Id
defnId Defn
d) HashSet Text
primNames ]

            live :: IS.IntSet
            live :: IntSet
live = IntSet -> [Uniq] -> IntSet
go IntSet
forall a. Monoid a => a
mempty [Uniq]
roots
                  where go :: IS.IntSet -> [Uniq] -> IS.IntSet
                        go :: IntSet -> [Uniq] -> IntSet
go IntSet
seen []       = IntSet
seen
                        go IntSet
seen (Uniq
u : [Uniq]
us)
                              | Uniq -> IntSet -> Bool
IS.member Uniq
u IntSet
seen = IntSet -> [Uniq] -> IntSet
go IntSet
seen [Uniq]
us
                              | Bool
otherwise        = IntSet -> [Uniq] -> IntSet
go (Uniq -> IntSet -> IntSet
IS.insert Uniq
u IntSet
seen) ([Uniq] -> IntSet) -> [Uniq] -> IntSet
forall a b. (a -> b) -> a -> b
$ IntSet -> [Uniq]
IS.toList (IntSet -> Uniq -> IntMap IntSet -> IntSet
forall a. a -> Uniq -> IntMap a -> a
IM.findWithDefault IntSet
forall a. Monoid a => a
mempty Uniq
u IntMap IntSet
refs) [Uniq] -> [Uniq] -> [Uniq]
forall a. Semigroup a => a -> a -> a
<> [Uniq]
us

---
--- Reduce: let-preserving local reduction.
---

reduceProgram :: MonadError AstError m => Program -> SimpT m Program
reduceProgram :: forall (m :: * -> *).
MonadError AstError m =>
Program -> SimpT m Program
reduceProgram (Program [DataDefn]
datas [Defn]
defns Id
top) = do
      let tops :: IntSet
tops = [Uniq] -> IntSet
IS.fromList ([Uniq] -> IntSet) -> [Uniq] -> IntSet
forall a b. (a -> b) -> a -> b
$ (Defn -> Uniq) -> [Defn] -> [Uniq]
forall a b. (a -> b) -> [a] -> [b]
map (Id -> Uniq
idUniq (Id -> Uniq) -> (Defn -> Id) -> Defn -> Uniq
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Defn -> Id
defnId) [Defn]
defns
      defns' <- (Defn -> StateT SimpSt m Defn) -> [Defn] -> StateT SimpSt m [Defn]
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 (\ Defn
d -> (\ Exp
b -> Defn
d { defnBody = b }) (Exp -> Defn) -> StateT SimpSt m Exp -> StateT SimpSt m Defn
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IntSet -> HashSet Text -> Exp -> StateT SimpSt m Exp
forall (m :: * -> *).
MonadError AstError m =>
IntSet -> HashSet Text -> Exp -> SimpT m Exp
reduceExp IntSet
tops ([DataDefn] -> HashSet Text
dictTyCons [DataDefn]
datas) (Defn -> Exp
defnBody Defn
d)) [Defn]
defns
      new    <- drainNew
      pure $ Program datas (defns' <> new) top

-- | One bottom-up reduction pass over an expression: beta redexes become
--   lets; lets inline when dead, single-use, or atom-bound; multi-use
--   function-typed lets are lifted to top level (LiftNonRep); known-
--   constructor and known-literal cases select their alternative; lambdas
--   eta-reduce; dead join points drop. Top-level references are never
--   unfolded. The first argument is the set of top-level uniques (a
--   lift's captures are the free variables outside it); the second is the
--   dictionary-datatype names ('dictTyCons'): a dictionary-typed binding
--   is substituted at every occurrence rather than shared, because it can
--   never be represented -- multi-use superclass projections would
--   otherwise never see the constructor.
reduceExp :: forall m. MonadError AstError m => IS.IntSet -> HashSet Text -> Exp -> SimpT m Exp
reduceExp :: forall (m :: * -> *).
MonadError AstError m =>
IntSet -> HashSet Text -> Exp -> SimpT m Exp
reduceExp IntSet
tops HashSet Text
dicts = Exp -> SimpT m Exp
go
      where go :: Exp -> SimpT m Exp
            go :: Exp -> SimpT m Exp
go Exp
e = case Exp
e of
                  App Annote
an Exp
f Arg
a -> do
                        f' <- Exp -> SimpT m Exp
go Exp
f
                        a' <- goArg a
                        case (f', a') of
                              (Lam Annote
_ Id
x Exp
b, EArg Exp
rhs) -> Exp -> SimpT m Exp
go (Exp -> SimpT m Exp) -> Exp -> SimpT m Exp
forall a b. (a -> b) -> a -> b
$ Annote -> Bind -> Exp -> Exp
Let Annote
an (Id -> Exp -> Bind
NonRec Id
x Exp
rhs) Exp
b
                              (Exp, Arg)
_                     -> Exp -> SimpT m Exp
forall a. a -> StateT SimpSt m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> SimpT m Exp) -> Exp -> SimpT m Exp
forall a b. (a -> b) -> a -> b
$ Annote -> Exp -> Arg -> Exp
App Annote
an Exp
f' Arg
a'
                  Lam Annote
an Id
x Exp
b -> Annote -> Id -> Exp -> Exp
etaReduce Annote
an Id
x (Exp -> Exp) -> SimpT m Exp -> SimpT m Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> SimpT m Exp
go Exp
b
                  Let Annote
an (NonRec Id
x Exp
rhs) Exp
body -> do
                        rhs' <- Exp -> SimpT m Exp
go Exp
rhs
                        if atomic rhs'
                              then go $ substVars (IM.singleton (idUniq x) rhs') body
                              else do
                                    body' <- go body
                                    let occ = Uniq -> Uniq -> IntMap Uniq -> Uniq
forall a. a -> Uniq -> IntMap a -> a
IM.findWithDefault Uniq
0 (Id -> Uniq
idUniq Id
x) (IntMap Uniq -> Uniq) -> IntMap Uniq -> Uniq
forall a b. (a -> b) -> a -> b
$ Exp -> IntMap Uniq
occCounts Exp
body'
                                    if | occ == 0                   -> pure body'
                                       | occ == 1                   -> go $ substVars (IM.singleton (idUniq x) rhs') body'
                                       | dictBound x                -> go =<< supplied (substVarsRefreshing (IM.singleton (idUniq x) rhs') body')
                                       | hasArrow (sigTy $ idSig x) -> do
                                             ref <- liftNonRep an x rhs'
                                             pure $ substVars (IM.singleton (idUniq x) ref) body'
                                       | otherwise                  -> pure $ Let an (NonRec x rhs') body'
                  Let Annote
an (Rec [(Id, Exp)]
bs) Exp
body -> do
                        bs'   <- ((Id, Exp) -> StateT SimpSt m (Id, Exp))
-> [(Id, Exp)] -> StateT SimpSt 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)) -> SimpT m Exp -> StateT SimpSt m (Id, Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> SimpT m Exp
go Exp
rhs) [(Id, Exp)]
bs
                        body' <- go body
                        let occs = Exp -> IntMap Uniq
occCounts Exp
body'
                        if any (\ (Id
x, Exp
_) -> Uniq -> IntMap Uniq -> Bool
forall a. Uniq -> IntMap a -> Bool
IM.member (Id -> Uniq
idUniq Id
x) IntMap Uniq
occs) bs'
                              then pure $ Let an (Rec bs') body'
                              else pure body'
                  Let Annote
an (Join JoinId
j [Id]
ps Exp
b) Exp
body -> do
                        b'    <- Exp -> SimpT m Exp
go Exp
b
                        body' <- go body
                        if IM.member (idUniq $ jpId j) $ occCounts body'
                              then pure $ Let an (Join j ps b') body'
                              else pure body'
                  Jump Annote
an JoinId
j [Exp]
es -> Annote -> JoinId -> [Exp] -> Exp
Jump Annote
an JoinId
j ([Exp] -> Exp) -> StateT SimpSt m [Exp] -> SimpT m Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> SimpT m Exp) -> [Exp] -> StateT SimpSt 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 -> SimpT m Exp
go [Exp]
es
                  Case Annote
an Ty
t Exp
scrut Id
cb [Alt]
alts -> do
                        scrut' <- Exp -> SimpT m Exp
go Exp
scrut
                        case selectAlt scrut' alts of
                              Just (Alt Annote
aan AltCon
_ [Id]
xs Exp
rhs, [Exp]
args) -> Exp -> SimpT m Exp
go
                                    (Exp -> SimpT m Exp) -> Exp -> SimpT m Exp
forall a b. (a -> b) -> a -> b
$ Annote -> Bind -> Exp -> Exp
Let Annote
an (Id -> Exp -> Bind
NonRec Id
cb Exp
scrut')
                                    (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ ((Id, Exp) -> Exp -> Exp) -> Exp -> [(Id, Exp)] -> Exp
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\ (Id
x, Exp
a) Exp
acc -> Annote -> Bind -> Exp -> Exp
Let Annote
aan (Id -> Exp -> Bind
NonRec Id
x Exp
a) Exp
acc) Exp
rhs
                                    ([(Id, Exp)] -> Exp) -> [(Id, Exp)] -> Exp
forall a b. (a -> b) -> a -> b
$ [Id] -> [Exp] -> [(Id, Exp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Id]
xs [Exp]
args
                              Maybe (Alt, [Exp])
Nothing -> do
                                    alts' <- (Alt -> StateT SimpSt m Alt) -> [Alt] -> StateT SimpSt m [Alt]
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 (\ (Alt Annote
aan AltCon
c [Id]
xs Exp
b) -> Annote -> AltCon -> [Id] -> Exp -> Alt
Alt Annote
aan AltCon
c [Id]
xs (Exp -> Alt) -> SimpT m Exp -> StateT SimpSt m Alt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> SimpT m Exp
go Exp
b) [Alt]
alts
                                    pure $ Case an t scrut' cb alts'
                  LitList Annote
an Ty
t [Exp]
es -> Annote -> Ty -> [Exp] -> Exp
LitList Annote
an Ty
t ([Exp] -> Exp) -> StateT SimpSt m [Exp] -> SimpT m Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> SimpT m Exp) -> [Exp] -> StateT SimpSt 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 -> SimpT m Exp
go [Exp]
es
                  LitVec Annote
an Ty
t [Exp]
es  -> Annote -> Ty -> [Exp] -> Exp
LitVec Annote
an Ty
t ([Exp] -> Exp) -> StateT SimpSt m [Exp] -> SimpT m Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> SimpT m Exp) -> [Exp] -> StateT SimpSt 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 -> SimpT m Exp
go [Exp]
es
                  Var {}          -> Exp -> SimpT m Exp
forall a. a -> StateT SimpSt m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
                  Con {}          -> Exp -> SimpT m Exp
forall a. a -> StateT SimpSt m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
                  Prim {}         -> Exp -> SimpT m Exp
forall a. a -> StateT SimpSt m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
                  LitInt {}       -> Exp -> SimpT m Exp
forall a. a -> StateT SimpSt m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
                  LitStr {}       -> Exp -> SimpT m Exp
forall a. a -> StateT SimpSt m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e

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

            -- \ x -> f x  ==>  f   (x not free in f)
            etaReduce :: Annote -> Id -> Exp -> Exp
            etaReduce :: Annote -> Id -> Exp -> Exp
etaReduce Annote
an Id
x = \ case
                  App Annote
_ Exp
f (EArg (Var Annote
_ Id
x'))
                        | Id -> Uniq
idUniq Id
x' Uniq -> Uniq -> Bool
forall a. Eq a => a -> a -> Bool
== Id -> Uniq
idUniq Id
x
                        , Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Uniq -> IntMap Uniq -> Bool
forall a. Uniq -> IntMap a -> Bool
IM.member (Id -> Uniq
idUniq Id
x) (IntMap Uniq -> Bool) -> IntMap Uniq -> Bool
forall a b. (a -> b) -> a -> b
$ Exp -> IntMap Uniq
occCounts Exp
f -> Exp
f
                  Exp
b -> Annote -> Id -> Exp -> Exp
Lam Annote
an Id
x Exp
b

            -- A binder whose type mentions a dictionary datatype: never
            -- representable, so never worth sharing.
            dictBound :: Id -> Bool
            dictBound :: Id -> Bool
dictBound Id
x = (Text -> Bool) -> [Text] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Text -> HashSet Text -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
`Set.member` HashSet Text
dicts) ([Text] -> Bool) -> [Text] -> Bool
forall a b. (a -> b) -> a -> b
$ Ty -> [Text]
tyConsOf (Ty -> [Text]) -> Ty -> [Text]
forall a b. (a -> b) -> a -> b
$ Sig -> Ty
sigTy (Sig -> Ty) -> Sig -> Ty
forall a b. (a -> b) -> a -> b
$ Id -> Sig
idSig Id
x

            -- Substitutable everywhere without loss of sharing.
            atomic :: Exp -> Bool
            atomic :: Exp -> Bool
atomic = \ case
                  Var {}    -> Bool
True
                  LitInt {} -> Bool
True
                  LitStr {} -> Bool
True
                  Prim {}   -> Bool
True
                  Con Annote
_ Ty
t Text
_ -> [Ty] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([Ty] -> Bool) -> [Ty] -> Bool
forall a b. (a -> b) -> a -> b
$ ([Ty], Ty) -> [Ty]
forall a b. (a, b) -> a
fst (([Ty], Ty) -> [Ty]) -> ([Ty], Ty) -> [Ty]
forall a b. (a -> b) -> a -> b
$ Ty -> ([Ty], Ty)
flattenArrow Ty
t -- nullary constructor
                  Exp
_         -> Bool
False

            -- LiftNonRep: a multi-use function-typed binding becomes a
            -- top-level definition over its captured locals; occurrences
            -- become an application to those locals.
            liftNonRep :: Annote -> Id -> Exp -> SimpT m Exp
            liftNonRep :: Annote -> Id -> Exp -> SimpT m Exp
liftNonRep Annote
an Id
x Exp
rhs = do
                  let caps :: [Id]
caps = IntSet -> Exp -> [Id]
capturedIds (Exp -> IntSet
freeUniqs Exp
rhs IntSet -> IntSet -> IntSet
IS.\\ IntSet
tops) Exp
rhs
                      t :: Ty
t    = (Id -> Ty -> Ty) -> Ty -> [Id] -> Ty
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Annote -> Ty -> Ty -> Ty
Arrow Annote
an (Ty -> Ty -> Ty) -> (Id -> Ty) -> Id -> Ty -> Ty
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sig -> Ty
sigTy (Sig -> Ty) -> (Id -> Sig) -> Id -> Ty
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Id -> Sig
idSig) (Exp -> Ty
typeOf Exp
rhs) [Id]
caps
                  d <- StateT Uniq m Defn -> SimpT m Defn
forall (m :: * -> *) a. Monad m => StateT Uniq m a -> SimpT m a
supplied (StateT Uniq m Defn -> SimpT m Defn)
-> StateT Uniq m Defn -> SimpT m Defn
forall a b. (a -> b) -> a -> b
$ Defn -> StateT Uniq m Defn
forall (m :: * -> *). MonadState Uniq m => Defn -> m Defn
refreshDefn Defn
                        { defnAnnote :: Annote
defnAnnote = Annote
an
                        , defnId :: Id
defnId     = Id
x { idOcc = "$LL." <> idOcc x, idSig = Sig [] t }
                        , defnParams :: [Id]
defnParams = [Id]
caps
                        , defnBody :: Exp
defnBody   = Exp
rhs
                        , defnAttr :: Maybe DefnAttr
defnAttr   = Maybe DefnAttr
forall a. Maybe a
Nothing
                        , defnOrigin :: Maybe SpecOrigin
defnOrigin = Maybe SpecOrigin
forall a. Maybe a
Nothing
                        }
                  mint d
                  pure $ foldl (\ Exp
f Id
c -> Annote -> Exp -> Arg -> Exp
App Annote
an Exp
f (Arg -> Exp) -> Arg -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Arg
EArg (Exp -> Arg) -> Exp -> Arg
forall a b. (a -> b) -> a -> b
$ Annote -> Id -> Exp
Var Annote
an Id
c) (Var an $ defnId d) caps

            -- The captured locals' Ids: every occurrence carries its
            -- binder's signature, so the first occurrence serves.
            capturedIds :: IS.IntSet -> Exp -> [Id]
            capturedIds :: IntSet -> Exp -> [Id]
capturedIds IntSet
us Exp
rhs = (Uniq -> Maybe Id) -> [Uniq] -> [Id]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (Uniq -> IntMap Id -> Maybe Id
forall a. Uniq -> IntMap a -> Maybe a
`IM.lookup` IntMap Id
occIds) ([Uniq] -> [Id]) -> [Uniq] -> [Id]
forall a b. (a -> b) -> a -> b
$ IntSet -> [Uniq]
IS.toList IntSet
us
                  where occIds :: IM.IntMap Id
                        occIds :: IntMap Id
occIds = Exp -> IntMap Id
collect Exp
rhs

                        collect :: Exp -> IM.IntMap Id
                        collect :: Exp -> IntMap Id
collect = \ case
                              Var Annote
_ Id
v         -> Uniq -> Id -> IntMap Id
forall a. Uniq -> a -> IntMap a
IM.singleton (Id -> Uniq
idUniq Id
v) Id
v
                              App Annote
_ Exp
f Arg
a       -> Exp -> IntMap Id
collect Exp
f IntMap Id -> IntMap Id -> IntMap Id
forall a. Semigroup a => a -> a -> a
<> Arg -> IntMap Id
collectArg Arg
a
                              Lam Annote
_ Id
_ Exp
b       -> Exp -> IntMap Id
collect Exp
b
                              Let Annote
_ Bind
b Exp
body    -> Bind -> IntMap Id
collectBind Bind
b IntMap Id -> IntMap Id -> IntMap Id
forall a. Semigroup a => a -> a -> a
<> Exp -> IntMap Id
collect 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
collect [Exp]
es
                              Case Annote
_ Ty
_ Exp
s Id
_ [Alt]
as -> Exp -> IntMap Id
collect 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
collect 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
collect [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
collect [Exp]
es
                              Exp
_               -> IntMap Id
forall a. Monoid a => a
mempty

                        collectArg :: Arg -> IM.IntMap Id
                        collectArg :: Arg -> IntMap Id
collectArg = \ case
                              EArg Exp
e -> Exp -> IntMap Id
collect Exp
e
                              Arg
_      -> IntMap Id
forall a. Monoid a => a
mempty

                        collectBind :: Bind -> IM.IntMap Id
                        collectBind :: Bind -> IntMap Id
collectBind = \ case
                              NonRec Id
_ Exp
rhs' -> Exp -> IntMap Id
collect 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
collect (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
collect Exp
b

-- | Select the alternative for a known scrutinee: a constructor-spine
--   scrutinee picks its 'DataAlt' (fields as a parallel argument list) or
--   the default; a literal scrutinee picks its 'LitAlt' or the default.
--   Unknown scrutinees select nothing.
selectAlt :: Exp -> [Alt] -> Maybe (Alt, [Exp])
selectAlt :: Exp -> [Alt] -> Maybe (Alt, [Exp])
selectAlt Exp
scrut [Alt]
alts = case Exp -> (Exp, [Arg])
flattenApp Exp
scrut of
      (Con Annote
_ Ty
_ Text
c, [Arg]
args) ->
            let es :: [Exp]
es = [ Exp
a | EArg Exp
a <- [Arg]
args ]
            in case [ Alt
alt | alt :: Alt
alt@(Alt Annote
_ (DataAlt Text
c') [Id]
_ Exp
_) <- [Alt]
alts, Text
c' Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
c ] of
                  alt :: Alt
alt@(Alt Annote
_ AltCon
_ [Id]
xs Exp
_) : [Alt]
_
                        | [Id] -> Uniq
forall a. [a] -> Uniq
forall (t :: * -> *) a. Foldable t => t a -> Uniq
length [Id]
xs Uniq -> Uniq -> Bool
forall a. Eq a => a -> a -> Bool
== [Exp] -> Uniq
forall a. [a] -> Uniq
forall (t :: * -> *) a. Foldable t => t a -> Uniq
length [Exp]
es -> (Alt, [Exp]) -> Maybe (Alt, [Exp])
forall a. a -> Maybe a
Just (Alt
alt, [Exp]
es)
                        | Bool
otherwise              -> Maybe (Alt, [Exp])
forall a. Maybe a
Nothing -- ill-arity (unreachable on lint-clean input)
                  [] -> (, []) (Alt -> (Alt, [Exp])) -> Maybe Alt -> Maybe (Alt, [Exp])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Alt
defaultAlt
      (LitInt Annote
_ Ty
_ Integer
n, []) -> case [ Alt
alt | alt :: Alt
alt@(Alt Annote
_ (LitAlt Integer
n') [Id]
_ Exp
_) <- [Alt]
alts, Integer
n' Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
n ] of
                  Alt
alt : [Alt]
_ -> (Alt, [Exp]) -> Maybe (Alt, [Exp])
forall a. a -> Maybe a
Just (Alt
alt, [])
                  []      -> (, []) (Alt -> (Alt, [Exp])) -> Maybe Alt -> Maybe (Alt, [Exp])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Alt
defaultAlt
      (Exp, [Arg])
_ -> Maybe (Alt, [Exp])
forall a. Maybe a
Nothing
      where defaultAlt :: Maybe Alt
            defaultAlt :: Maybe Alt
defaultAlt = case [Alt]
alts of
                  alt :: Alt
alt@(Alt Annote
_ AltCon
DefaultAlt [Id]
_ Exp
_) : [Alt]
_ -> Alt -> Maybe Alt
forall a. a -> Maybe a
Just Alt
alt
                  [Alt]
_                              -> Maybe Alt
forall a. Maybe a
Nothing

---
--- Specialize on values (argument baking).
---

-- | Bake closed arguments of calls to top-level definitions into memoized
--   clones. An argument is bakeable when all its free variables are
--   top-level; a definition participates unless NOINLINE or prim-named.
--   Leading body lambdas are first promoted into the parameter telescope
--   (the bridge emits every definition with an empty telescope).
specialize :: forall m. MonadError AstError m => Program -> SimpT m Program
specialize :: forall (m :: * -> *).
MonadError AstError m =>
Program -> SimpT m Program
specialize (Program [DataDefn]
datas [Defn]
defns0 Id
top) = do
      defns' <- (Defn -> StateT SimpSt m Defn) -> [Defn] -> StateT SimpSt m [Defn]
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 Defn -> StateT SimpSt m Defn
specDefn [Defn]
defns
      new    <- drainNew
      -- Re-add memoized specializations dropped by an earlier purge but
      -- possibly referenced again (the retired specializer's behavior;
      -- the next purge removes unused ones). Sorted for determinism.
      let have = [Uniq] -> IntSet
IS.fromList ([Uniq] -> IntSet) -> [Uniq] -> IntSet
forall a b. (a -> b) -> a -> b
$ (Defn -> Uniq) -> [Defn] -> [Uniq]
forall a b. (a -> b) -> [a] -> [b]
map (Id -> Uniq
idUniq (Id -> Uniq) -> (Defn -> Id) -> Defn -> Uniq
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Defn -> Id
defnId) ([Defn] -> [Uniq]) -> [Defn] -> [Uniq]
forall a b. (a -> b) -> a -> b
$ [Defn]
defns' [Defn] -> [Defn] -> [Defn]
forall a. Semigroup a => a -> a -> a
<> [Defn]
new
      memoed <- gets $ sortOn (idUniq . defnId) . filter (\ Defn
d -> Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Uniq -> IntSet -> Bool
IS.member (Id -> Uniq
idUniq (Id -> Uniq) -> Id -> Uniq
forall a b. (a -> b) -> a -> b
$ Defn -> Id
defnId Defn
d) IntSet
have) . Map.elems . stMemo
      pure $ Program datas (defns' <> new <> memoed) top
      where defns :: [Defn]
            defns :: [Defn]
defns = (Defn -> Defn) -> [Defn] -> [Defn]
forall a b. (a -> b) -> [a] -> [b]
map Defn -> Defn
promote [Defn]
defns0

            -- Promote leading body lambdas into the parameter telescope,
            -- up to the signature's arity.
            promote :: Defn -> Defn
            promote :: Defn -> Defn
promote Defn
d = Defn
d { defnParams = defnParams d <> ps, defnBody = b }
                  where arity :: Uniq
arity     = [Ty] -> Uniq
forall a. [a] -> Uniq
forall (t :: * -> *) a. Foldable t => t a -> Uniq
length ([Ty] -> Uniq) -> [Ty] -> Uniq
forall a b. (a -> b) -> a -> b
$ ([Ty], Ty) -> [Ty]
forall a b. (a, b) -> a
fst (([Ty], Ty) -> [Ty]) -> ([Ty], Ty) -> [Ty]
forall a b. (a -> b) -> a -> b
$ Ty -> ([Ty], Ty)
flattenArrow (Ty -> ([Ty], Ty)) -> Ty -> ([Ty], Ty)
forall a b. (a -> b) -> a -> b
$ Sig -> Ty
sigTy (Sig -> Ty) -> Sig -> Ty
forall a b. (a -> b) -> a -> b
$ Id -> Sig
idSig (Id -> Sig) -> Id -> Sig
forall a b. (a -> b) -> a -> b
$ Defn -> Id
defnId Defn
d
                        ([Id]
ps, Exp
b)   = Uniq -> Exp -> ([Id], Exp)
peel (Uniq
arity Uniq -> Uniq -> Uniq
forall a. Num a => a -> a -> a
- [Id] -> Uniq
forall a. [a] -> Uniq
forall (t :: * -> *) a. Foldable t => t a -> Uniq
length (Defn -> [Id]
defnParams Defn
d)) (Exp -> ([Id], Exp)) -> Exp -> ([Id], Exp)
forall a b. (a -> b) -> a -> b
$ Defn -> Exp
defnBody Defn
d

                        peel :: Int -> Exp -> ([Id], Exp)
                        peel :: Uniq -> Exp -> ([Id], Exp)
peel Uniq
n (Lam Annote
_ Id
x Exp
e) | Uniq
n Uniq -> Uniq -> Bool
forall a. Ord a => a -> a -> Bool
> Uniq
0 = let ([Id]
xs, Exp
e') = Uniq -> Exp -> ([Id], Exp)
peel (Uniq
n Uniq -> Uniq -> Uniq
forall a. Num a => a -> a -> a
- Uniq
1) Exp
e in (Id
x Id -> [Id] -> [Id]
forall a. a -> [a] -> [a]
: [Id]
xs, Exp
e')
                        peel Uniq
_ Exp
e                   = ([], Exp
e)

            tops :: IS.IntSet
            tops :: IntSet
tops = [Uniq] -> IntSet
IS.fromList ([Uniq] -> IntSet) -> [Uniq] -> IntSet
forall a b. (a -> b) -> a -> b
$ (Defn -> Uniq) -> [Defn] -> [Uniq]
forall a b. (a -> b) -> [a] -> [b]
map (Id -> Uniq
idUniq (Id -> Uniq) -> (Defn -> Id) -> Defn -> Uniq
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Defn -> Id
defnId) [Defn]
defns

            gs :: IM.IntMap Defn
            gs :: IntMap Defn
gs = [(Uniq, Defn)] -> IntMap Defn
forall a. [(Uniq, a)] -> IntMap a
IM.fromList [ (Id -> Uniq
idUniq (Id -> Uniq) -> Id -> Uniq
forall a b. (a -> b) -> a -> b
$ Defn -> Id
defnId Defn
d, Defn
d) | Defn
d <- [Defn]
defns ]

            specDefn :: Defn -> SimpT m Defn
            specDefn :: Defn -> StateT SimpSt m Defn
specDefn Defn
d = do
                  b <- Exp -> SimpT m Exp
specExp (Exp -> SimpT m Exp) -> Exp -> SimpT m Exp
forall a b. (a -> b) -> a -> b
$ Defn -> Exp
defnBody Defn
d
                  pure d { defnBody = b }

            specExp :: Exp -> SimpT m Exp
            specExp :: Exp -> SimpT m Exp
specExp Exp
e = case Exp
e of
                  App Annote
an Exp
_ Arg
_ -> do
                        let (Exp
h, [Arg]
args) = Exp -> (Exp, [Arg])
flattenApp Exp
e
                        args' <- (Arg -> StateT SimpSt m Arg) -> [Arg] -> StateT SimpSt m [Arg]
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 Arg -> StateT SimpSt m Arg
specArg [Arg]
args
                        let eas = [ Exp
a | EArg Exp
a <- [Arg]
args' ]
                        case h of
                              Var Annote
_ Id
g | Just Defn
d <- Uniq -> IntMap Defn -> Maybe Defn
forall a. Uniq -> IntMap a -> Maybe a
IM.lookup (Id -> Uniq
idUniq Id
g) IntMap Defn
gs
                                        , Defn -> Bool
inlinable Defn
d
                                        , Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [Id] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([Id] -> Bool) -> [Id] -> Bool
forall a b. (a -> b) -> a -> b
$ Defn -> [Id]
defnParams Defn
d
                                        , [Maybe Exp]
bs <- (Id -> Exp -> Maybe Exp) -> [Id] -> [Exp] -> [Maybe Exp]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\ Id
_ Exp
a -> if Exp -> Bool
closed Exp
a then Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
a else Maybe Exp
forall a. Maybe a
Nothing) (Defn -> [Id]
defnParams Defn
d) [Exp]
eas
                                        , Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (Maybe Exp -> Bool) -> [Maybe Exp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Maybe Exp -> Bool
forall a. Maybe a -> Bool
isNothing [Maybe Exp]
bs -> do
                                    g' <- Defn -> [Maybe Exp] -> SimpT m Id
specialization Defn
d [Maybe Exp]
bs
                                    let kept = [ Exp -> Arg
EArg Exp
a | (Maybe Exp
b, Exp
a) <- [Maybe Exp] -> [Exp] -> [(Maybe Exp, Exp)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([Maybe Exp]
bs [Maybe Exp] -> [Maybe Exp] -> [Maybe Exp]
forall a. Semigroup a => a -> a -> a
<> Maybe Exp -> [Maybe Exp]
forall a. a -> [a]
repeat Maybe Exp
forall a. Maybe a
Nothing) [Exp]
eas, Maybe Exp -> Bool
forall a. Maybe a -> Bool
isNothing Maybe Exp
b ]
                                    pure $ foldl (App an) (Var an g') kept
                              Exp
_ -> do
                                    h' <- Exp -> SimpT m Exp
specExp Exp
h
                                    pure $ foldl (App an) h' args'
                  Lam Annote
an Id
x Exp
b      -> Annote -> Id -> Exp -> Exp
Lam Annote
an Id
x (Exp -> Exp) -> SimpT m Exp -> SimpT m Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> SimpT m Exp
specExp Exp
b
                  Let Annote
an Bind
b Exp
body   -> Annote -> Bind -> Exp -> Exp
Let Annote
an (Bind -> Exp -> Exp)
-> StateT SimpSt m Bind -> StateT SimpSt m (Exp -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bind -> StateT SimpSt m Bind
specBind Bind
b StateT SimpSt m (Exp -> Exp) -> SimpT m Exp -> SimpT m Exp
forall a b.
StateT SimpSt m (a -> b) -> StateT SimpSt m a -> StateT SimpSt m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> SimpT m Exp
specExp Exp
body
                  Jump Annote
an JoinId
j [Exp]
es    -> Annote -> JoinId -> [Exp] -> Exp
Jump Annote
an JoinId
j ([Exp] -> Exp) -> StateT SimpSt m [Exp] -> SimpT m Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> SimpT m Exp) -> [Exp] -> StateT SimpSt 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 -> SimpT m Exp
specExp [Exp]
es
                  Case Annote
an Ty
t Exp
s Id
cb [Alt]
alts -> do
                        s'    <- Exp -> SimpT m Exp
specExp Exp
s
                        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) -> SimpT m Exp -> StateT SimpSt m Alt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> SimpT m Exp
specExp Exp
b) alts
                        pure $ Case an t s' cb alts'
                  LitList Annote
an Ty
t [Exp]
es -> Annote -> Ty -> [Exp] -> Exp
LitList Annote
an Ty
t ([Exp] -> Exp) -> StateT SimpSt m [Exp] -> SimpT m Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> SimpT m Exp) -> [Exp] -> StateT SimpSt 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 -> SimpT m Exp
specExp [Exp]
es
                  LitVec Annote
an Ty
t [Exp]
es  -> Annote -> Ty -> [Exp] -> Exp
LitVec Annote
an Ty
t ([Exp] -> Exp) -> StateT SimpSt m [Exp] -> SimpT m Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> SimpT m Exp) -> [Exp] -> StateT SimpSt 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 -> SimpT m Exp
specExp [Exp]
es
                  Exp
_               -> Exp -> SimpT m Exp
forall a. a -> StateT SimpSt m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e

            specArg :: Arg -> SimpT m Arg
            specArg :: Arg -> StateT SimpSt m Arg
specArg = \ case
                  EArg Exp
a -> Exp -> Arg
EArg (Exp -> Arg) -> SimpT m Exp -> StateT SimpSt m Arg
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> SimpT m Exp
specExp Exp
a
                  Arg
t      -> Arg -> StateT SimpSt m Arg
forall a. a -> StateT SimpSt m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Arg
t

            specBind :: Bind -> SimpT m Bind
            specBind :: Bind -> StateT SimpSt m Bind
specBind = \ case
                  NonRec Id
x Exp
rhs -> Id -> Exp -> Bind
NonRec Id
x (Exp -> Bind) -> SimpT m Exp -> StateT SimpSt m Bind
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> SimpT m Exp
specExp Exp
rhs
                  Rec [(Id, Exp)]
bs       -> [(Id, Exp)] -> Bind
Rec ([(Id, Exp)] -> Bind)
-> StateT SimpSt m [(Id, Exp)] -> StateT SimpSt m Bind
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Id, Exp) -> StateT SimpSt m (Id, Exp))
-> [(Id, Exp)] -> StateT SimpSt 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)) -> SimpT m Exp -> StateT SimpSt m (Id, Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> SimpT m Exp
specExp Exp
rhs) [(Id, Exp)]
bs
                  Join JoinId
j [Id]
ps Exp
b  -> JoinId -> [Id] -> Exp -> Bind
Join JoinId
j [Id]
ps (Exp -> Bind) -> SimpT m Exp -> StateT SimpSt m Bind
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> SimpT m Exp
specExp Exp
b

            closed :: Exp -> Bool
            closed :: Exp -> Bool
closed Exp
a = Exp -> IntSet
freeUniqs Exp
a IntSet -> IntSet -> Bool
`IS.isSubsetOf` IntSet
tops

            inlinable :: Defn -> Bool
            inlinable :: Defn -> Bool
inlinable Defn
d = Defn -> Maybe DefnAttr
defnAttr Defn
d Maybe DefnAttr -> Maybe DefnAttr -> Bool
forall a. Eq a => a -> a -> Bool
/= DefnAttr -> Maybe DefnAttr
forall a. a -> Maybe a
Just DefnAttr
NoInline
                       Bool -> Bool -> Bool
&& Bool -> Bool
not (Text -> HashSet Text -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
Set.member (Id -> Text
idOcc (Id -> Text) -> Id -> Text
forall a b. (a -> b) -> a -> b
$ Defn -> Id
defnId Defn
d) HashSet Text
primNames)

            -- The memoized clone for a definition at the given argument
            -- baking (Just = bake).
            specialization :: Defn -> [Maybe Exp] -> SimpT m Id
            specialization :: Defn -> [Maybe Exp] -> SimpT m Id
specialization Defn
d [Maybe Exp]
bs = do
                  let key :: (Uniq, Text)
key = (Id -> Uniq
idUniq (Id -> Uniq) -> Id -> Uniq
forall a b. (a -> b) -> a -> b
$ Defn -> Id
defnId Defn
d, [Maybe Exp] -> Text
canonKey [Maybe Exp]
bs)
                  (SimpSt -> Maybe Defn) -> StateT SimpSt m (Maybe Defn)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((Uniq, Text) -> HashMap (Uniq, Text) Defn -> Maybe Defn
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup (Uniq, Text)
key (HashMap (Uniq, Text) Defn -> Maybe Defn)
-> (SimpSt -> HashMap (Uniq, Text) Defn) -> SimpSt -> Maybe Defn
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SimpSt -> HashMap (Uniq, Text) Defn
stMemo) StateT SimpSt m (Maybe Defn)
-> (Maybe Defn -> SimpT m Id) -> SimpT m Id
forall a b.
StateT SimpSt m a -> (a -> StateT SimpSt m b) -> StateT SimpSt m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ case
                        Just Defn
d' -> Id -> SimpT m Id
forall a. a -> StateT SimpSt m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Id -> SimpT m Id) -> Id -> SimpT m Id
forall a b. (a -> b) -> a -> b
$ Defn -> Id
defnId Defn
d'
                        Maybe Defn
Nothing -> do
                              d' <- Defn -> [Maybe Exp] -> StateT SimpSt m Defn
mkSpec Defn
d [Maybe Exp]
bs
                              modify $ \ SimpSt
st -> SimpSt
st { stMemo = Map.insert key d' $ stMemo st }
                              mint d'
                              pure $ defnId d'

            -- Bake: clone the definition (self-references still name the
            -- origin — the memo closes recursive loops on the next round),
            -- substitute the baked arguments for the corresponding cloned
            -- parameters, and drop them from the parameter list and the
            -- signature's arrow spine.
            mkSpec :: Defn -> [Maybe Exp] -> SimpT m Defn
            mkSpec :: Defn -> [Maybe Exp] -> StateT SimpSt m Defn
mkSpec Defn
d [Maybe Exp]
bs = do
                  -- The clone's display name carries a stable tag derived
                  -- from the baked-argument canon (usually its hash --
                  -- values render long), not a discovery counter: baking
                  -- the same arguments always yields the same name. The
                  -- canon's uniques are scrubbed first -- free (top-level)
                  -- references print as occ#uniq, and uniques shift on
                  -- unrelated edits (the memo key keeps them; only the
                  -- display name must not).
                  dr <- StateT Uniq m Defn -> StateT SimpSt m Defn
forall (m :: * -> *) a. Monad m => StateT Uniq m a -> SimpT m a
supplied (StateT Uniq m Defn -> StateT SimpSt m Defn)
-> StateT Uniq m Defn -> StateT SimpSt m Defn
forall a b. (a -> b) -> a -> b
$ Text -> [Ty] -> Defn -> StateT Uniq m Defn
forall (m :: * -> *).
MonadState Uniq m =>
Text -> [Ty] -> Defn -> m Defn
instantiateDefn (Id -> Text
idOcc (Defn -> Id
defnId Defn
d) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"$s" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Text] -> Text
originTag [Text -> Text
scrubUniqs (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ [Maybe Exp] -> Text
canonKey [Maybe Exp]
bs]) [] Defn
d
                  let an          = Defn -> Annote
defnAnnote Defn
dr
                      ps          = Defn -> [Id]
defnParams Defn
dr
                      bs'         = [Maybe Exp]
bs [Maybe Exp] -> [Maybe Exp] -> [Maybe Exp]
forall a. Semigroup a => a -> a -> a
<> Maybe Exp -> [Maybe Exp]
forall a. a -> [a]
repeat Maybe Exp
forall a. Maybe a
Nothing
                      baked       = [(Uniq, Exp)] -> IntMap Exp
forall a. [(Uniq, a)] -> IntMap a
IM.fromList [ (Id -> Uniq
idUniq Id
p, Exp
a) | (Id
p, Just Exp
a) <- [Id] -> [Maybe Exp] -> [(Id, Maybe Exp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Id]
ps [Maybe Exp]
bs ]
                      keptPs      = [ Id
p | (Id
p, Maybe Exp
b) <- [Id] -> [Maybe Exp] -> [(Id, Maybe Exp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Id]
ps [Maybe Exp]
bs', Maybe Exp -> Bool
forall a. Maybe a -> Bool
isNothing Maybe Exp
b ]
                      Sig _ t     = idSig $ defnId dr
                      (doms, res) = flattenArrow t
                      domsKept    = [ Ty
dom | (Ty
dom, Maybe Exp
b) <- [Ty] -> [Maybe Exp] -> [(Ty, Maybe Exp)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Uniq -> [Ty] -> [Ty]
forall a. Uniq -> [a] -> [a]
take ([Id] -> Uniq
forall a. [a] -> Uniq
forall (t :: * -> *) a. Foldable t => t a -> Uniq
length [Id]
ps) [Ty]
doms) [Maybe Exp]
bs', Maybe Exp -> Bool
forall a. Maybe a -> Bool
isNothing Maybe Exp
b ]
                                    [Ty] -> [Ty] -> [Ty]
forall a. Semigroup a => a -> a -> a
<> Uniq -> [Ty] -> [Ty]
forall a. Uniq -> [a] -> [a]
drop ([Id] -> Uniq
forall a. [a] -> Uniq
forall (t :: * -> *) a. Foldable t => t a -> Uniq
length [Id]
ps) [Ty]
doms
                      x'          = (Defn -> Id
defnId Defn
dr) { idSig = Sig [] $ foldr (Arrow an) res domsKept }
                  body' <- supplied $ substVarsRefreshing baked $ defnBody dr
                  pure $ Defn an x' keptPs body' (defnAttr d)
                       $ Just $ fromMaybe (BakeOrigin $ idOcc $ defnId d) $ defnOrigin d

            -- Alpha-canonical text of the baked-argument list: renumber
            -- binders from a disjoint (large negative) supply, then print.
            canonKey :: [Maybe Exp] -> Text
            canonKey :: [Maybe Exp] -> Text
canonKey = Text -> [Text] -> Text
intercalate Text
"|" ([Text] -> Text) -> ([Maybe Exp] -> [Text]) -> [Maybe Exp] -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Maybe Exp -> Text) -> [Maybe Exp] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text -> (Exp -> Text) -> Maybe Exp -> Text
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Text
"_" Exp -> Text
pp)
                  where pp :: Exp -> Text
                        pp :: Exp -> Text
pp Exp
a = Exp -> Text
forall a. Pretty a => a -> Text
prettyPrint (State Uniq Exp -> Uniq -> Exp
forall s a. State s a -> s -> a
evalState (Exp -> State Uniq Exp
forall (m :: * -> *). MonadState Uniq m => Exp -> m Exp
refreshExp Exp
a) Uniq
canonBase :: Exp)

canonBase :: Uniq
canonBase :: Uniq
canonBase = -Uniq
1000000000

-- | Strip the @#uniq@ suffixes from rendered names (@Main.f#3 7@ ->
--   @Main.f 7@): uniques are compile-relative and must not reach display
--   names.
scrubUniqs :: Text -> Text
scrubUniqs :: Text -> Text
scrubUniqs Text
t = case HasCallStack => Text -> Text -> (Text, Text)
Text -> Text -> (Text, Text)
T.breakOn Text
"#" Text
t of
      (Text
a, Text
"") -> Text
a
      (Text
a, Text
b)  -> Text
a Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text
scrubUniqs ((Char -> Bool) -> Text -> Text
T.dropWhile Char -> Bool
isDigit (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ (Char -> Bool) -> Text -> Text
T.dropWhile (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'-') (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ Uniq -> Text -> Text
T.drop Uniq
1 Text
b)