{-# LANGUAGE Safe #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
-- | Monomorphization by specialization (doc/eidos.md §4.1, §5, §8): a worklist
--   specializer driven by the type arguments the bridge keeps on
--   application spines. There is no inference and no unification: a
--   monomorphic definition's body mentions a polymorphic definition only
--   as a spine head saturated with closed type arguments (the linter's
--   poly-mode guarantee), so each such spine is a specialization request;
--   the clone is pure type substitution through the audited clone
--   primitive ('ReWire.Eidos.Subst.instantiateDefn'), memoized on the
--   normalized argument list. Clones carry their provenance ('SpecOrigin')
--   and are named @origin$tag@, where the tag renders the type arguments
--   ('ReWire.Eidos.Naming.originTag') — so an unrelated instantiation
--   elsewhere never renames a clone.
--
--   Polymorphic definitions are templates: they are dropped from the
--   output (their instantiations replace them), so the result is
--   monomorphic ("ReWire.Eidos.Lint" mono mode) — except the
--   builtin-named definitions (rwPrim*), whose polymorphic signatures
--   carry the builtins' type assumptions to the Eidos-to-Hyle fold (their
--   bodies are error stubs, and they are never referenced as variables —
--   references become 'Prim' occurrences at the bridge); they ride
--   through unchanged. The worklist runs in generations; an instantiation
--   chain deeper than the budget (the historical typechecker bound,
--   raised by @--depth@) is rejected with the retired specializer's
--   diagnostic.
module ReWire.Eidos.Spec (specialize) where

import ReWire.Annotation (noAnn, unAnn)
import ReWire.Builtins (builtins)
import ReWire.Error (AstError, MonadError, failAt)
import ReWire.Eidos.Naming (originTag)
import ReWire.Eidos.Pretty ()
import ReWire.Eidos.Subst (instantiateDefn, nextUniq)
import ReWire.Eidos.Syntax
import ReWire.Eidos.Types (natNorm, flattenApp, flattenArrow, flattenTyApp, evalNat)
import ReWire.Pretty (prettyPrint, showt)

import Control.Monad (foldM)
import Control.Monad.State.Strict (StateT, evalStateT)
import Data.HashMap.Strict (HashMap)
import Data.HashSet (HashSet)
import Data.List (partition)
import Data.Text (Text)
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

-- | A specialization request: a polymorphic definition (by unique) at a
--   closed, normalized type-argument list.
type Key = (Uniq, [Ty])

specialize :: forall m. MonadError AstError m => Natural -> Program -> m Program
specialize :: forall (m :: * -> *).
MonadError AstError m =>
Natural -> Program -> m Program
specialize Natural
bound p :: Program
p@(Program [DataDefn]
datas [Defn]
defns Id
top) = StateT Uniq m Program -> Uniq -> m Program
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT StateT Uniq m Program
go (Uniq -> m Program) -> Uniq -> m Program
forall a b. (a -> b) -> a -> b
$ Program -> Uniq
forall a. Data a => a -> Uniq
nextUniq Program
p
      where templates :: [Defn]
            kept      :: [Defn]
            ([Defn]
templates, [Defn]
kept) = (Defn -> Bool) -> [Defn] -> ([Defn], [Defn])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition Defn -> Bool
isTemplate [Defn]
defns

            -- Templates are the polymorphic definitions except the builtin
            -- signature carriers.
            isTemplate :: Defn -> Bool
            isTemplate :: Defn -> Bool
isTemplate Defn
d = Id -> Bool
isPoly (Defn -> Id
defnId Defn
d) Bool -> Bool -> Bool
&& Bool -> Bool
not (Text -> Bool
isPrimName (Text -> Bool) -> Text -> Bool
forall a b. (a -> b) -> a -> b
$ Id -> Text
idOcc (Id -> Text) -> Id -> Text
forall a b. (a -> b) -> a -> b
$ Defn -> Id
defnId Defn
d)

            polys :: IM.IntMap Defn
            polys :: IntMap Defn
polys = [(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]
templates ]

            isPoly :: Id -> Bool
            isPoly :: Id -> Bool
isPoly = Bool -> Bool
not (Bool -> Bool) -> (Id -> Bool) -> Id -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TyVar] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([TyVar] -> Bool) -> (Id -> [TyVar]) -> Id -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sig -> [TyVar]
sigTVs (Sig -> [TyVar]) -> (Id -> Sig) -> Id -> [TyVar]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Id -> Sig
idSig

            isPrimName :: Text -> Bool
            isPrimName :: Text -> Bool
isPrimName = (Text -> HashSet Text -> Bool) -> HashSet Text -> Text -> Bool
forall a b c. (a -> b -> c) -> b -> a -> c
flip Text -> HashSet Text -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
Set.member HashSet Text
primNames

            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

            go :: StateT Uniq m Program
            go :: StateT Uniq m Program
go = do
                  (table, clones) <- Natural
-> HashMap (Uniq, [Ty]) Id
-> [Defn]
-> [(Uniq, [Ty])]
-> StateT Uniq m (HashMap (Uniq, [Ty]) Id, [Defn])
rounds Natural
0 HashMap (Uniq, [Ty]) Id
forall a. Monoid a => a
mempty [] ([(Uniq, [Ty])] -> StateT Uniq m (HashMap (Uniq, [Ty]) Id, [Defn]))
-> [(Uniq, [Ty])]
-> StateT Uniq m (HashMap (Uniq, [Ty]) Id, [Defn])
forall a b. (a -> b) -> a -> b
$ (Defn -> [(Uniq, [Ty])]) -> [Defn] -> [(Uniq, [Ty])]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Exp -> [(Uniq, [Ty])]
requests (Exp -> [(Uniq, [Ty])]) -> (Defn -> Exp) -> Defn -> [(Uniq, [Ty])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Defn -> Exp
defnBody) [Defn]
kept
                  pure $ Program datas (map (rewrite table) $ kept <> clones) top

            -- One worklist generation per round.
            rounds :: Natural -> HashMap Key Id -> [Defn] -> [Key] -> StateT Uniq m (HashMap Key Id, [Defn])
            rounds :: Natural
-> HashMap (Uniq, [Ty]) Id
-> [Defn]
-> [(Uniq, [Ty])]
-> StateT Uniq m (HashMap (Uniq, [Ty]) Id, [Defn])
rounds Natural
n HashMap (Uniq, [Ty]) Id
table [Defn]
clones [(Uniq, [Ty])]
reqs
                  | [(Uniq, [Ty])] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(Uniq, [Ty])]
new   = (HashMap (Uniq, [Ty]) Id, [Defn])
-> StateT Uniq m (HashMap (Uniq, [Ty]) Id, [Defn])
forall a. a -> StateT Uniq m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (HashMap (Uniq, [Ty]) Id
table, [Defn]
clones)
                  | Natural
n Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
>= Natural
bound = Annote -> Text -> StateT Uniq m (HashMap (Uniq, [Ty]) Id, [Defn])
forall (m :: * -> *) an a.
(MonadError AstError m, Annotation an) =>
an -> Text -> m a
failAt Annote
noAnn Text
"polymorphic function instantiation not terminating (mutually recursive definitions?)."
                  | Bool
otherwise  = do
                        (table', batch) <- ((HashMap (Uniq, [Ty]) Id, [Defn])
 -> (Uniq, [Ty]) -> StateT Uniq m (HashMap (Uniq, [Ty]) Id, [Defn]))
-> (HashMap (Uniq, [Ty]) Id, [Defn])
-> [(Uniq, [Ty])]
-> StateT Uniq m (HashMap (Uniq, [Ty]) Id, [Defn])
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (HashMap (Uniq, [Ty]) Id, [Defn])
-> (Uniq, [Ty]) -> StateT Uniq m (HashMap (Uniq, [Ty]) Id, [Defn])
step (HashMap (Uniq, [Ty]) Id
table, []) [(Uniq, [Ty])]
new
                        let batch' = [Defn] -> [Defn]
forall a. [a] -> [a]
reverse [Defn]
batch
                        rounds (n + 1) table' (clones <> batch') $ concatMap (requests . defnBody) batch'
                  where new :: [(Uniq, [Ty])]
new = HashMap (Uniq, [Ty]) Id -> [(Uniq, [Ty])] -> [(Uniq, [Ty])]
dedupe HashMap (Uniq, [Ty]) Id
table [(Uniq, [Ty])]
reqs

            step :: (HashMap Key Id, [Defn]) -> Key -> StateT Uniq m (HashMap Key Id, [Defn])
            step :: (HashMap (Uniq, [Ty]) Id, [Defn])
-> (Uniq, [Ty]) -> StateT Uniq m (HashMap (Uniq, [Ty]) Id, [Defn])
step (HashMap (Uniq, [Ty]) Id
table, [Defn]
acc) k :: (Uniq, [Ty])
k@(Uniq
u, [Ty]
ts) = case Uniq -> IntMap Defn -> Maybe Defn
forall a. Uniq -> IntMap a -> Maybe a
IM.lookup Uniq
u IntMap Defn
polys of
                  Maybe Defn
Nothing -> (HashMap (Uniq, [Ty]) Id, [Defn])
-> StateT Uniq m (HashMap (Uniq, [Ty]) Id, [Defn])
forall a. a -> StateT Uniq m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((HashMap (Uniq, [Ty]) Id
table), [Defn]
acc) -- unreachable on lint-clean input
                  Just Defn
d  -> do
                        -- The clone's display name carries its
                        -- instantiation ('grev$Vec_8_Bool'), not a
                        -- discovery counter: adding an unrelated
                        -- instantiation elsewhere must not rename this
                        -- one. Renderings that collide after sanitizing
                        -- are disambiguated at the machine fold.
                        let occ :: Text
occ = Id -> Text
idOcc (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
<> [Text] -> Text
originTag ((Ty -> Text) -> [Ty] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map Ty -> Text
renderTy [Ty]
ts)
                        d' <- Text -> [Ty] -> Defn -> StateT Uniq m Defn
forall (m :: * -> *).
MonadState Uniq m =>
Text -> [Ty] -> Defn -> m Defn
instantiateDefn Text
occ [Ty]
ts Defn
d
                        let d'' = Defn
d' { defnOrigin = Just $ SpecOrigin (idOcc $ defnId d) ts }
                        pure (Map.insert k (defnId d'') table, d'' : acc)

            -- New keys, in discovery order.
            dedupe :: HashMap Key Id -> [Key] -> [Key]
            dedupe :: HashMap (Uniq, [Ty]) Id -> [(Uniq, [Ty])] -> [(Uniq, [Ty])]
dedupe HashMap (Uniq, [Ty]) Id
table = HashMap (Uniq, [Ty]) () -> [(Uniq, [Ty])] -> [(Uniq, [Ty])]
go' HashMap (Uniq, [Ty]) ()
forall a. Monoid a => a
mempty
                  where go' :: HashMap Key () -> [Key] -> [Key]
                        go' :: HashMap (Uniq, [Ty]) () -> [(Uniq, [Ty])] -> [(Uniq, [Ty])]
go' HashMap (Uniq, [Ty]) ()
_ [] = []
                        go' HashMap (Uniq, [Ty]) ()
seen ((Uniq, [Ty])
k : [(Uniq, [Ty])]
ks)
                              | (Uniq, [Ty]) -> HashMap (Uniq, [Ty]) Id -> Bool
forall k a. (Eq k, Hashable k) => k -> HashMap k a -> Bool
Map.member (Uniq, [Ty])
k HashMap (Uniq, [Ty]) Id
table Bool -> Bool -> Bool
|| (Uniq, [Ty]) -> HashMap (Uniq, [Ty]) () -> Bool
forall k a. (Eq k, Hashable k) => k -> HashMap k a -> Bool
Map.member (Uniq, [Ty])
k HashMap (Uniq, [Ty]) ()
seen = HashMap (Uniq, [Ty]) () -> [(Uniq, [Ty])] -> [(Uniq, [Ty])]
go' HashMap (Uniq, [Ty]) ()
seen [(Uniq, [Ty])]
ks
                              | Bool
otherwise                               = (Uniq, [Ty])
k (Uniq, [Ty]) -> [(Uniq, [Ty])] -> [(Uniq, [Ty])]
forall a. a -> [a] -> [a]
: HashMap (Uniq, [Ty]) () -> [(Uniq, [Ty])] -> [(Uniq, [Ty])]
go' ((Uniq, [Ty])
-> () -> HashMap (Uniq, [Ty]) () -> HashMap (Uniq, [Ty]) ()
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
Map.insert (Uniq, [Ty])
k () HashMap (Uniq, [Ty]) ()
seen) [(Uniq, [Ty])]
ks

            -- Specialization requests in an expression, in traversal order:
            -- every spine headed by a saturated polymorphic top-level
            -- reference.
            requests :: Exp -> [Key]
            requests :: Exp -> [(Uniq, [Ty])]
requests Exp
e = case Exp
e of
                  App {}          -> Exp -> [(Uniq, [Ty])]
spine Exp
e
                  Var {}          -> []
                  Con {}          -> []
                  Prim {}         -> []
                  LitInt {}       -> []
                  LitStr {}       -> []
                  LitList Annote
_ Ty
_ [Exp]
es  -> (Exp -> [(Uniq, [Ty])]) -> [Exp] -> [(Uniq, [Ty])]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Exp -> [(Uniq, [Ty])]
requests [Exp]
es
                  LitVec Annote
_ Ty
_ [Exp]
es   -> (Exp -> [(Uniq, [Ty])]) -> [Exp] -> [(Uniq, [Ty])]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Exp -> [(Uniq, [Ty])]
requests [Exp]
es
                  Lam Annote
_ Id
_ Exp
b       -> Exp -> [(Uniq, [Ty])]
requests Exp
b
                  Let Annote
_ Bind
b Exp
body    -> Bind -> [(Uniq, [Ty])]
bindReqs Bind
b [(Uniq, [Ty])] -> [(Uniq, [Ty])] -> [(Uniq, [Ty])]
forall a. Semigroup a => a -> a -> a
<> Exp -> [(Uniq, [Ty])]
requests Exp
body
                  Jump Annote
_ JoinId
_ [Exp]
es     -> (Exp -> [(Uniq, [Ty])]) -> [Exp] -> [(Uniq, [Ty])]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Exp -> [(Uniq, [Ty])]
requests [Exp]
es
                  Case Annote
_ Ty
_ Exp
s Id
_ [Alt]
as -> Exp -> [(Uniq, [Ty])]
requests Exp
s [(Uniq, [Ty])] -> [(Uniq, [Ty])] -> [(Uniq, [Ty])]
forall a. Semigroup a => a -> a -> a
<> [[(Uniq, [Ty])]] -> [(Uniq, [Ty])]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [ Exp -> [(Uniq, [Ty])]
requests Exp
b | Alt Annote
_ AltCon
_ [Id]
_ Exp
b <- [Alt]
as ]
                  where spine :: Exp -> [Key]
                        spine :: Exp -> [(Uniq, [Ty])]
spine Exp
se = [(Uniq, [Ty])]
hk [(Uniq, [Ty])] -> [(Uniq, [Ty])] -> [(Uniq, [Ty])]
forall a. Semigroup a => a -> a -> a
<> (Exp -> [(Uniq, [Ty])]) -> [Exp] -> [(Uniq, [Ty])]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Exp -> [(Uniq, [Ty])]
requests [Exp]
eas
                              where (Exp
h, [Arg]
args) = Exp -> (Exp, [Arg])
flattenApp Exp
se
                                    tys :: [Ty]
tys       = (Ty -> Ty) -> [Ty] -> [Ty]
forall a b. (a -> b) -> [a] -> [b]
map Ty -> Ty
natNorm [ Ty
t | TArg Ty
t <- (Arg -> Bool) -> [Arg] -> [Arg]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile Arg -> Bool
isTArg [Arg]
args ]
                                    eas :: [Exp]
eas       = [ Exp
a | EArg Exp
a <- [Arg]
args ]
                                    hk :: [(Uniq, [Ty])]
hk        = case Exp
h of
                                          Var Annote
_ Id
x | Bool -> Bool
not ([Ty] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Ty]
tys), Uniq -> IntMap Defn -> Bool
forall a. Uniq -> IntMap a -> Bool
IM.member (Id -> Uniq
idUniq Id
x) IntMap Defn
polys -> [(Id -> Uniq
idUniq Id
x, [Ty]
tys)]
                                          Var {}                                               -> []
                                          Exp
_                                                    -> Exp -> [(Uniq, [Ty])]
requests Exp
h

                        bindReqs :: Bind -> [Key]
                        bindReqs :: Bind -> [(Uniq, [Ty])]
bindReqs = \ case
                              NonRec Id
_ Exp
rhs -> Exp -> [(Uniq, [Ty])]
requests Exp
rhs
                              Rec [(Id, Exp)]
bs       -> ((Id, Exp) -> [(Uniq, [Ty])]) -> [(Id, Exp)] -> [(Uniq, [Ty])]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Exp -> [(Uniq, [Ty])]
requests (Exp -> [(Uniq, [Ty])])
-> ((Id, Exp) -> Exp) -> (Id, Exp) -> [(Uniq, [Ty])]
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 -> [(Uniq, [Ty])]
requests Exp
b

            -- Rewrite every specialized spine: the head becomes the clone
            -- and the type arguments are erased.
            rewrite :: HashMap Key Id -> Defn -> Defn
            rewrite :: HashMap (Uniq, [Ty]) Id -> Defn -> Defn
rewrite HashMap (Uniq, [Ty]) Id
table Defn
d = Defn
d { defnBody = rw $ defnBody d }
                  where rw :: Exp -> Exp
                        rw :: Exp -> Exp
rw Exp
e = case Exp
e of
                              App Annote
an Exp
_ Arg
_ ->
                                    let (Exp
h, [Arg]
args)  = Exp -> (Exp, [Arg])
flattenApp Exp
e
                                        ([Arg]
tas, [Arg]
eas) = (Arg -> Bool) -> [Arg] -> ([Arg], [Arg])
forall a. (a -> Bool) -> [a] -> ([a], [a])
span Arg -> Bool
isTArg [Arg]
args
                                        tys :: [Ty]
tys        = (Ty -> Ty) -> [Ty] -> [Ty]
forall a b. (a -> b) -> [a] -> [b]
map Ty -> Ty
natNorm [ Ty
t | TArg Ty
t <- [Arg]
tas ]
                                    in case Exp
h of
                                          Var Annote
van Id
x | Bool -> Bool
not ([Ty] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Ty]
tys), Just Id
x' <- (Uniq, [Ty]) -> HashMap (Uniq, [Ty]) Id -> Maybe Id
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup (Id -> Uniq
idUniq Id
x, [Ty]
tys) HashMap (Uniq, [Ty]) Id
table ->
                                                (Exp -> Arg -> Exp) -> Exp -> [Arg] -> Exp
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (Annote -> Exp -> Arg -> Exp
App Annote
an) (Annote -> Id -> Exp
Var Annote
van Id
x') [ Exp -> Arg
EArg (Exp -> Arg) -> Exp -> Arg
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
rw Exp
a | EArg Exp
a <- [Arg]
eas ]
                                          Exp
_ -> (Exp -> Arg -> Exp) -> Exp -> [Arg] -> Exp
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (Annote -> Exp -> Arg -> Exp
App Annote
an) (Exp -> Exp
rw Exp
h) ([Arg] -> Exp) -> [Arg] -> Exp
forall a b. (a -> b) -> a -> b
$ (Arg -> Arg) -> [Arg] -> [Arg]
forall a b. (a -> b) -> [a] -> [b]
map Arg -> Arg
rwArg [Arg]
args
                              Var {}          -> Exp
e
                              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
rw [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
rw [Exp]
es
                              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
rw Exp
b
                              Let Annote
an Bind
b Exp
body   -> Annote -> Bind -> Exp -> Exp
Let Annote
an (Bind -> Bind
rwBind Bind
b) (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
rw 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
rw [Exp]
es
                              Case Annote
an Ty
t Exp
s Id
cb [Alt]
alts -> Annote -> Ty -> Exp -> Id -> [Alt] -> Exp
Case Annote
an Ty
t (Exp -> Exp
rw Exp
s) Id
cb [ Annote -> AltCon -> [Id] -> Exp -> Alt
Alt Annote
aan AltCon
c [Id]
xs (Exp -> Exp
rw Exp
b) | Alt Annote
aan AltCon
c [Id]
xs Exp
b <- [Alt]
alts ]

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

                        rwBind :: Bind -> Bind
                        rwBind :: Bind -> Bind
rwBind = \ 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
rw Exp
rhs
                              Rec [(Id, Exp)]
bs       -> [(Id, Exp)] -> Bind
Rec [ (Id
x, Exp -> Exp
rw 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
rw Exp
b

            isTArg :: Arg -> Bool
            isTArg :: Arg -> Bool
isTArg = \ case
                  TArg Ty
_ -> Bool
True
                  Arg
_      -> Bool
False

-- | A type argument's display text for instance names: compact forms for
--   the common cases (@Vec 8 Bool@ -> @W8@, @Vec 4 (Vec 8 Bool)@ ->
--   @V4W8@, arrows @_@-joined), the pretty-printed type otherwise
--   (annotations scrubbed so positions can't leak into names; 'originTag'
--   sanitizes and caps the result).
renderTy :: Ty -> Text
renderTy :: Ty -> Text
renderTy Ty
t
      | Just Natural
k <- Ty -> Maybe Natural
evalNat Ty
t                = Natural -> Text
forall a. TextShow a => a -> Text
showt Natural
k
      | (ds :: [Ty]
ds@(Ty
_ : [Ty]
_), Ty
cd) <- Ty -> ([Ty], Ty)
flattenArrow Ty
t = Text -> [Text] -> Text
T.intercalate Text
"_" ([Text] -> Text) -> [Text] -> Text
forall a b. (a -> b) -> a -> b
$ (Ty -> Text) -> [Ty] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map Ty -> Text
renderTy ([Ty] -> [Text]) -> [Ty] -> [Text]
forall a b. (a -> b) -> a -> b
$ [Ty]
ds [Ty] -> [Ty] -> [Ty]
forall a. Semigroup a => a -> a -> a
<> [Ty
cd]
      | Bool
otherwise = case Ty -> (Ty, [Ty])
flattenTyApp Ty
t of
            (TyCon Annote
_ Text
"Vec", [Ty
n, TyCon Annote
_ Text
"Bool"]) | Just Natural
k <- Ty -> Maybe Natural
evalNat Ty
n -> Text
"W" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Natural -> Text
forall a. TextShow a => a -> Text
showt Natural
k
            (TyCon Annote
_ Text
"Vec", [Ty
n, Ty
el])             | Just Natural
k <- Ty -> Maybe Natural
evalNat Ty
n -> Text
"V" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Natural -> Text
forall a. TextShow a => a -> Text
showt Natural
k Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Ty -> Text
renderTy Ty
el
            (TyCon Annote
_ Text
c, [Ty]
args)                                          -> (Char -> Bool) -> Text -> Text
T.takeWhileEnd (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
'.') Text
c Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Text] -> Text
forall a. Monoid a => [a] -> a
mconcat ((Ty -> Text) -> [Ty] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map Ty -> Text
renderTy [Ty]
args)
            (Ty, [Ty])
_                                                          -> Ty -> Text
forall a. Pretty a => a -> Text
prettyPrint (Ty -> Ty
forall d. Data d => d -> d
unAnn Ty
t :: Ty)