{-# LANGUAGE Safe #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
-- | A-normalization to purify's input form (doc/eidos.md §6): every
--   reactive definition body becomes a let chain over simple right-hand
--   sides ending in an atom (or a jump), by the spec's small ordered ruleset —
--   eta-expansion (definitions reach their signature arity, parameters in
--   the telescope), argument naming (the computed arguments of spines —
--   definition calls, constructor applications, cases — are let-bound),
--   subject naming (case scrutinees are atoms), head reduction (a
--   residual beta-redex becomes a let), let-flattening (no lets in
--   right-hand sides), and alternative flattening (alternative bodies are
--   themselves ANF). Primitive applications are transparent: a
--   primitive-headed argument is not named but normalized in place, so
--   the pure data path stays an expression tree of primitives over atoms
--   (the fold lowers it inline, and pattern-matches its idioms — bit
--   slices, finite literals — where they stand).
--
--   The *reactive* fragment is the deliberate exemption — it is what
--   purify consumes, and its structure must survive:
--
--   * a spine whose type mentions the reactive stack stays a spine; its
--     computed arguments are named, its lambda arguments (bind
--     continuations, like the higher-order primitives' function
--     arguments) stay in place with A-normalized bodies, and its
--     reactive arguments normalize recursively in place (a pure let may
--     wrap them);
--   * a case with a reactive result type stays in tail position (its
--     scrutinee is named, its alternatives normalize as tails) — purify
--     turns it into a terminator case; a pure-resulted case is let-bound
--     like any other computation.
--
--   Runs after the partial evaluator (ReWire.ModCache pass 6),
--   immediately before purify.
module ReWire.Eidos.ANF (normalize, hasJump, isAtom, isPrimExp) where

import ReWire.Annotation (Annote)
import ReWire.Error (AstError, MonadError, failAt)
import ReWire.Eidos.Subst (nextUniq)
import ReWire.Eidos.Syntax
import ReWire.Eidos.Types (typeOf, flattenApp, flattenArrow, hasArrow, reacOrStateT)

import Control.Monad.State.Strict (StateT, evalStateT, get, put)

-- | Only the reactive fragment normalizes: purify consumes the reactive
--   skeleton, and the Eidos-to-Hyle fold lowers pure expressions in any
--   shape, so naming every intermediate of the large,
--   partially-evaluated pure bodies would cost compile time for nothing.
normalize :: forall m. MonadError AstError m => Program -> m Program
normalize :: forall (m :: * -> *). MonadError AstError m => Program -> m Program
normalize 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 go :: StateT Uniq m Program
            go :: StateT Uniq m Program
go = do
                  defns' <- (Defn -> StateT Uniq m Defn) -> [Defn] -> StateT Uniq 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 -> if Defn -> Bool
reactiveDefn Defn
d then Defn -> StateT Uniq m Defn
forall (m :: * -> *). MonadError AstError m => Defn -> NM m Defn
normDefn Defn
d else Defn -> StateT Uniq m Defn
forall a. a -> StateT Uniq m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Defn
d) [Defn]
defns
                  pure $ Program datas defns' top

            reactiveDefn :: Defn -> Bool
            reactiveDefn :: Defn -> Bool
reactiveDefn Defn
d = Ty -> Bool
reacOrStateT (Ty -> Bool) -> Ty -> Bool
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

type NM m = StateT Uniq m

freshId :: Monad m => Ty -> NM m Id
freshId :: forall (m :: * -> *). Monad m => Ty -> NM m Id
freshId Ty
t = do
      u <- StateT Uniq m Uniq
forall s (m :: * -> *). MonadState s m => m s
get
      put $ u + 1
      pure $ Id { idOcc = "$a", idUniq = u, idSig = monoSig t }

-- | Eta-expand to signature arity (leading body lambdas promote into the
--   parameter telescope; missing parameters are minted and applied), then
--   normalize the body as a tail.
normDefn :: MonadError AstError m => Defn -> NM m Defn
normDefn :: forall (m :: * -> *). MonadError AstError m => Defn -> NM m Defn
normDefn Defn
d = do
      let ([Ty]
doms, Ty
_)  = 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
body) = Uniq -> Exp -> ([Id], Exp)
peel ([Ty] -> Uniq
forall a. [a] -> Uniq
forall (t :: * -> *) a. Foldable t => t a -> Uniq
length [Ty]
doms 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
          params :: [Id]
params     = Defn -> [Id]
defnParams Defn
d [Id] -> [Id] -> [Id]
forall a. Semigroup a => a -> a -> a
<> [Id]
ps
          an :: Annote
an         = Defn -> Annote
defnAnnote Defn
d
      etaPs <- (Ty -> StateT Uniq m Id) -> [Ty] -> StateT Uniq m [Id]
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 Ty -> StateT Uniq m Id
forall (m :: * -> *). Monad m => Ty -> NM m Id
freshId ([Ty] -> StateT Uniq m [Id]) -> [Ty] -> StateT Uniq m [Id]
forall a b. (a -> b) -> a -> b
$ 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]
params) [Ty]
doms
      let body' = (Exp -> Id -> Exp) -> Exp -> [Id] -> Exp
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\ Exp
e Id
x -> Annote -> Exp -> Arg -> Exp
App Annote
an Exp
e (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
x) Exp
body [Id]
etaPs
      nb <- normTail body'
      pure d { defnParams = params <> etaPs, defnBody = nb }
      where 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)

-- | Hoisted bindings, innermost last.
type Hoist = [(Annote, Id, Exp)]

wrapLets :: Hoist -> Exp -> Exp
wrapLets :: Hoist -> Exp -> Exp
wrapLets Hoist
bs Exp
e = ((Annote, Id, Exp) -> Exp -> Exp) -> Exp -> Hoist -> Exp
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\ (Annote
an, Id
x, Exp
r) Exp
acc -> Annote -> Bind -> Exp -> Exp
Let Annote
an (Id -> Exp -> Bind
NonRec Id
x Exp
r) Exp
acc) Exp
e Hoist
bs

-- | Is the expression already an atom (§6)? Variables and literals;
--   nullary constructor and bare primitive occurrences; list and vector
--   literals count as literals once their elements are atoms.
isAtom :: Exp -> Bool
isAtom :: Exp -> Bool
isAtom = \ case
      Var {}          -> Bool
True
      LitInt {}       -> Bool
True
      LitStr {}       -> 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
      Prim Annote
_ Ty
t Builtin
_      -> [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
      LitList Annote
_ Ty
_ [Exp]
es  -> (Exp -> Bool) -> [Exp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Exp -> Bool
isAtom [Exp]
es
      LitVec Annote
_ Ty
_ [Exp]
es   -> (Exp -> Bool) -> [Exp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Exp -> Bool
isAtom [Exp]
es
      Exp
_               -> Bool
False

-- | A primitive expression: a primitive applied to arguments. Transparent
--   to naming (its arguments normalize in place), so it may nest: the
--   pure data path is a tree of primitives over atoms.
isPrimExp :: Exp -> Bool
isPrimExp :: Exp -> Bool
isPrimExp Exp
e = case Exp -> (Exp, [Arg])
flattenApp Exp
e of
      (Prim {}, Arg
_ : [Arg]
_) -> Bool
True
      (Exp, [Arg])
_                -> Bool
False

reactive :: Exp -> Bool
reactive :: Exp -> Bool
reactive = Ty -> Bool
reacOrStateT (Ty -> Bool) -> (Exp -> Ty) -> Exp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> Ty
typeOf

-- | Commute an application over a case or let head: a function-typed
--   case (or a let/join wrapping one — GHC's pattern-match-failure join
--   shape) applied to arguments pushes the arguments inside — into every
--   non-jump alternative, through let bodies, and into join bodies (the
--   join's signature and its jumps' carried ids are rebuilt to the
--   applied type). Scope-safe under the uniqueness discipline: the
--   arguments predate the binders. Commuting here is what puts the
--   reactive case or spine back in tail position for purify.
commuteCaseApp :: Exp -> Maybe Exp
commuteCaseApp :: Exp -> Maybe Exp
commuteCaseApp Exp
e = case Exp -> (Exp, [Arg])
flattenApp Exp
e of
      (h :: Exp
h@Case {}, args :: [Arg]
args@(Arg
_ : [Arg]
_)) -> Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Exp -> Maybe Exp) -> Exp -> Maybe Exp
forall a b. (a -> b) -> a -> b
$ Exp -> [Arg] -> Exp
pushE Exp
h [Arg]
args
      (h :: Exp
h@Let {},  args :: [Arg]
args@(Arg
_ : [Arg]
_)) -> Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Exp -> Maybe Exp) -> Exp -> Maybe Exp
forall a b. (a -> b) -> a -> b
$ Exp -> [Arg] -> Exp
pushE Exp
h [Arg]
args
      (Exp, [Arg])
_                         -> Maybe Exp
forall a. Maybe a
Nothing
      where pushE :: Exp -> [Arg] -> Exp
            pushE :: Exp -> [Arg] -> Exp
pushE Exp
h [Arg]
args = case Exp
h of
                  Let Annote
an (Join JoinId
j [Id]
ps Exp
b) Exp
body ->
                        let b' :: Exp
b'  = Exp -> [Arg] -> Exp
pushE Exp
b [Arg]
args
                            tvs :: [Ty]
tvs = (Id -> Ty) -> [Id] -> [Ty]
forall a b. (a -> b) -> [a] -> [b]
map (Sig -> Ty
sigTy (Sig -> Ty) -> (Id -> Sig) -> Id -> Ty
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Id -> Sig
idSig) [Id]
ps
                            j' :: JoinId
j'  = JoinId
j { jpId = (jpId j) { idSig = monoSig $ foldr (Arrow an) (typeOf b') tvs } }
                        in Annote -> Bind -> Exp -> Exp
Let Annote
an (JoinId -> [Id] -> Exp -> Bind
Join JoinId
j' [Id]
ps Exp
b') (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ JoinId -> JoinId -> Exp -> Exp
rejump JoinId
j JoinId
j' (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> [Arg] -> Exp
pushE Exp
body [Arg]
args
                  Let Annote
an Bind
bnd Exp
body -> Annote -> Bind -> Exp -> Exp
Let Annote
an Bind
bnd (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> [Arg] -> Exp
pushE Exp
body [Arg]
args
                  Case Annote
an Ty
t Exp
s Id
cb [Alt]
alts ->
                        Annote -> Ty -> Exp -> Id -> [Alt] -> Exp
Case Annote
an (Uniq -> Ty -> Ty
peelArrows ([()] -> Uniq
forall a. [a] -> Uniq
forall (t :: * -> *) a. Foldable t => t a -> Uniq
length [ () | EArg Exp
_ <- [Arg]
args ]) Ty
t) Exp
s Id
cb
                              [ Annote -> AltCon -> [Id] -> Exp -> Alt
Alt Annote
aan AltCon
c [Id]
xs (if Exp -> Bool
isJumpTail Exp
b then Exp
b else Exp -> [Arg] -> Exp
pushE Exp
b [Arg]
args) | Alt Annote
aan AltCon
c [Id]
xs Exp
b <- [Alt]
alts ]
                  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 -> Exp -> Arg -> Exp) -> Annote -> Exp -> Arg -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Annote
ann' Exp
h) Exp
h [Arg]
args

            isJumpTail :: Exp -> Bool
            isJumpTail :: Exp -> Bool
isJumpTail = \ case
                  Jump {}      -> Bool
True
                  Let Annote
_ Bind
_ Exp
body -> Exp -> Bool
isJumpTail Exp
body
                  Exp
_            -> Bool
False

            peelArrows :: Int -> Ty -> Ty
            peelArrows :: Uniq -> Ty -> Ty
peelArrows Uniq
0 Ty
t             = Ty
t
            peelArrows Uniq
n (Arrow Annote
_ Ty
_ Ty
u) = Uniq -> Ty -> Ty
peelArrows (Uniq
n Uniq -> Uniq -> Uniq
forall a. Num a => a -> a -> a
- Uniq
1) Ty
u
            peelArrows Uniq
_ Ty
t             = Ty
t

            -- Rewrite jump occurrences of the old join id to the rebuilt one.
            rejump :: JoinId -> JoinId -> Exp -> Exp
            rejump :: JoinId -> JoinId -> Exp -> Exp
rejump JoinId
j JoinId
j' = Exp -> Exp
go
                  where go :: Exp -> Exp
                        go :: Exp -> Exp
go Exp
ex = case Exp
ex of
                              Jump Annote
an JoinId
jj [Exp]
es
                                    | Id -> Uniq
idUniq (JoinId -> Id
jpId JoinId
jj) Uniq -> Uniq -> Bool
forall a. Eq a => a -> a -> Bool
== Id -> Uniq
idUniq (JoinId -> Id
jpId JoinId
j) -> 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
                                    | Bool
otherwise                           -> Annote -> JoinId -> [Exp] -> Exp
Jump Annote
an JoinId
jj ([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
goA 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
bnd Exp
body  -> Annote -> Bind -> Exp -> Exp
Let Annote
an (Bind -> Bind
goB Bind
bnd) (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
go Exp
body
                              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 ]
                              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
                              Exp
_                -> Exp
ex

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

                        goB :: Bind -> Bind
                        goB :: Bind -> Bind
goB = \ 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
jj [Id]
ps Exp
b -> JoinId -> [Id] -> Exp -> Bind
Join JoinId
jj [Id]
ps (Exp -> Bind) -> Exp -> Bind
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
go Exp
b

-- | A residual beta-redex — a lambda head applied to arguments (the
--   simplifier's single round can leave one) — is a let, whose right-hand
--   side and body then normalize like any other.
betaHead :: Exp -> Maybe Exp
betaHead :: Exp -> Maybe Exp
betaHead Exp
e = case Exp -> (Exp, [Arg])
flattenApp Exp
e of
      (Lam Annote
an Id
x Exp
b, EArg Exp
a : [Arg]
rest) -> Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Exp -> Maybe Exp) -> Exp -> Maybe Exp
forall a b. (a -> b) -> a -> b
$ Annote -> Bind -> Exp -> Exp
Let Annote
an (Id -> Exp -> Bind
NonRec Id
x Exp
a) (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ (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
b [Arg]
rest
      (Exp, [Arg])
_                           -> Maybe Exp
forall a. Maybe a
Nothing

-- | Does the expression contain a jump (anywhere)? Jumps are tail-only
--   on lint-clean input, so a jump-containing case is a join point's
--   scope and must stay in tail position.
hasJump :: Exp -> Bool
hasJump :: Exp -> Bool
hasJump = \ case
      Jump {}         -> Bool
True
      App Annote
_ Exp
f Arg
a       -> Exp -> Bool
hasJump Exp
f Bool -> Bool -> Bool
|| (case Arg
a of { EArg Exp
x -> Exp -> Bool
hasJump Exp
x; Arg
_ -> Bool
False })
      Lam Annote
_ Id
_ Exp
b       -> Exp -> Bool
hasJump Exp
b
      Let Annote
_ Bind
b Exp
body    -> Bind -> Bool
bindJump Bind
b Bool -> Bool -> Bool
|| Exp -> Bool
hasJump Exp
body
      Case Annote
_ Ty
_ Exp
s Id
_ [Alt]
as -> Exp -> Bool
hasJump Exp
s Bool -> Bool -> Bool
|| [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or [ Exp -> Bool
hasJump Exp
b | Alt Annote
_ AltCon
_ [Id]
_ Exp
b <- [Alt]
as ]
      LitList Annote
_ Ty
_ [Exp]
es  -> (Exp -> Bool) -> [Exp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Exp -> Bool
hasJump [Exp]
es
      LitVec Annote
_ Ty
_ [Exp]
es   -> (Exp -> Bool) -> [Exp] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Exp -> Bool
hasJump [Exp]
es
      Exp
_               -> Bool
False
      where bindJump :: Bind -> Bool
            bindJump :: Bind -> Bool
bindJump = \ case
                  NonRec Id
_ Exp
rhs -> Exp -> Bool
hasJump Exp
rhs
                  Rec [(Id, Exp)]
bs       -> ((Id, Exp) -> Bool) -> [(Id, Exp)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Exp -> Bool
hasJump (Exp -> Bool) -> ((Id, Exp) -> Exp) -> (Id, Exp) -> Bool
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 -> Bool
hasJump Exp
b

-- | Normalize a tail position: a let chain ending in an atom, a jump, a
--   reactive spine, or a reactive case.
normTail :: forall m. MonadError AstError m => Exp -> NM m Exp
normTail :: forall (m :: * -> *). MonadError AstError m => Exp -> NM m Exp
normTail Exp
e = case Exp
e of
      Exp
_ | Exp -> Bool
isAtom Exp
e   -> do
            (bs, a) <- Exp -> NM m (Hoist, Exp)
forall (m :: * -> *).
MonadError AstError m =>
Exp -> NM m (Hoist, Exp)
atomize Exp
e -- names non-atom literal elements
            pure $ wrapLets bs a
      Lam Annote
an Id
x Exp
b     -> Annote -> Id -> Exp -> Exp
Lam Annote
an Id
x (Exp -> Exp) -> NM m Exp -> NM m Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> NM m Exp
forall (m :: * -> *). MonadError AstError m => Exp -> NM m Exp
normTail Exp
b -- residual (reactive continuation) lambda
      Let Annote
an (NonRec Id
x Exp
rhs) Exp
body -> do
            (bs, r) <- Exp -> NM m (Hoist, Exp)
forall (m :: * -> *).
MonadError AstError m =>
Exp -> NM m (Hoist, Exp)
normR Exp
rhs
            body'   <- normTail body
            pure $ wrapLets bs $ Let an (NonRec x r) body'
      Let Annote
an (Join JoinId
j [Id]
ps Exp
b) Exp
body -> do
            b'    <- Exp -> NM m Exp
forall (m :: * -> *). MonadError AstError m => Exp -> NM m Exp
normTail Exp
b
            body' <- normTail body
            pure $ Let an (Join j ps b') body'
      Let Annote
an (Rec [(Id, Exp)]
bs) Exp
body -> do
            -- Local recursion never survives the simplifier on the
            -- pipeline; normalize structurally for fixture robustness.
            bs'   <- ((Id, Exp) -> StateT Uniq m (Id, Exp))
-> [(Id, Exp)] -> StateT Uniq 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)) -> NM m Exp -> StateT Uniq m (Id, Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> NM m Exp
forall (m :: * -> *). MonadError AstError m => Exp -> NM m Exp
normTail Exp
rhs) [(Id, Exp)]
bs
            body' <- normTail body
            pure $ Let an (Rec bs') body'
      Jump Annote
an JoinId
j [Exp]
es   -> do
            (bs, as) <- [Exp] -> NM m (Hoist, [Exp])
forall (m :: * -> *).
MonadError AstError m =>
[Exp] -> NM m (Hoist, [Exp])
atomizeMany [Exp]
es
            pure $ wrapLets bs $ Jump an j as
      -- A reactive-resulted case (a terminator case, after purify) and a
      -- case whose alternatives jump (the scope of a join point — jumps
      -- are tail-only, so the case cannot be let-bound) stay in tail
      -- position.
      Case Annote
an Ty
t Exp
s Id
cb [Alt]
alts
            | Ty -> Bool
reacOrStateT Ty
t Bool -> Bool -> Bool
|| Exp -> Bool
hasJump Exp
e -> do
                  (bs, sa) <- Exp -> NM m (Hoist, Exp)
forall (m :: * -> *).
MonadError AstError m =>
Exp -> NM m (Hoist, Exp)
atomize 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) -> NM m Exp -> StateT Uniq m Alt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> NM m Exp
forall (m :: * -> *). MonadError AstError m => Exp -> NM m Exp
normTail Exp
b) alts
                  pure $ wrapLets bs $ Case an t sa cb alts'
            | Bool
otherwise -> NM m Exp
named
      App {}
            | Just Exp
e' <- Exp -> Maybe Exp
commuteCaseApp Exp
e       -> Exp -> NM m Exp
forall (m :: * -> *). MonadError AstError m => Exp -> NM m Exp
normTail Exp
e'
            | Just Exp
e' <- Exp -> Maybe Exp
betaHead Exp
e             -> Exp -> NM m Exp
forall (m :: * -> *). MonadError AstError m => Exp -> NM m Exp
normTail Exp
e'
            | Exp -> Bool
reactive Exp
e Bool -> Bool -> Bool
|| Ty -> Bool
hasArrow (Exp -> Ty
typeOf Exp
e) -> do
                  (bs, e') <- Exp -> NM m (Hoist, Exp)
forall (m :: * -> *).
MonadError AstError m =>
Exp -> NM m (Hoist, Exp)
normSpine Exp
e
                  pure $ wrapLets bs e'
            | Bool
otherwise  -> NM m Exp
named
      Exp
_ -> NM m Exp
named
      where -- A pure computation in tail position is named: let a = r in a.
            named :: NM m Exp
            named :: NM m Exp
named = do
                  (bs, r) <- Exp -> NM m (Hoist, Exp)
forall (m :: * -> *).
MonadError AstError m =>
Exp -> NM m (Hoist, Exp)
normR Exp
e
                  if isAtom r then pure $ wrapLets bs r else do
                        x <- freshId $ typeOf r
                        pure $ wrapLets bs $ Let (ann' r) (NonRec x r) $ Var (ann' r) x

-- | Normalize to a right-hand side (§6 r-forms), hoisting bindings.
normR :: forall m. MonadError AstError m => Exp -> NM m (Hoist, Exp)
normR :: forall (m :: * -> *).
MonadError AstError m =>
Exp -> NM m (Hoist, Exp)
normR Exp
e = case Exp
e of
      Exp
_ | Exp -> Bool
isAtom Exp
e -> Exp -> NM m (Hoist, Exp)
forall (m :: * -> *).
MonadError AstError m =>
Exp -> NM m (Hoist, Exp)
atomize Exp
e
      Lam {}       -> (Hoist, Exp) -> NM m (Hoist, Exp)
forall a. a -> StateT Uniq m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], Exp
e) -- residual lambda (never let-bound on the pipeline)
      Let Annote
_ (NonRec Id
x Exp
rhs) Exp
body -> do
            -- Let-flattening: the binding hoists out of the right-hand side.
            (bs, r)    <- Exp -> NM m (Hoist, Exp)
forall (m :: * -> *).
MonadError AstError m =>
Exp -> NM m (Hoist, Exp)
normR Exp
rhs
            (bs', r'') <- normR body
            pure (bs <> [(ann' rhs, x, r)] <> bs', r'')
      Let {}       -> ([], ) (Exp -> (Hoist, Exp)) -> StateT Uniq m Exp -> NM m (Hoist, Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> StateT Uniq m Exp
forall (m :: * -> *). MonadError AstError m => Exp -> NM m Exp
normTail Exp
e -- join/rec in rhs position: structural
      Jump {}      -> ([], ) (Exp -> (Hoist, Exp)) -> StateT Uniq m Exp -> NM m (Hoist, Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> StateT Uniq m Exp
forall (m :: * -> *). MonadError AstError m => Exp -> NM m Exp
normTail Exp
e
      Case Annote
an Ty
t Exp
s Id
cb [Alt]
alts -> do
            (bs, sa) <- Exp -> NM m (Hoist, Exp)
forall (m :: * -> *).
MonadError AstError m =>
Exp -> NM m (Hoist, Exp)
atomize 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) -> StateT Uniq m Exp -> StateT Uniq m Alt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> StateT Uniq m Exp
forall (m :: * -> *). MonadError AstError m => Exp -> NM m Exp
normTail Exp
b) alts
            pure (bs, Case an t sa cb alts')
      App {}
            | Just Exp
e' <- Exp -> Maybe Exp
commuteCaseApp Exp
e -> Exp -> NM m (Hoist, Exp)
forall (m :: * -> *).
MonadError AstError m =>
Exp -> NM m (Hoist, Exp)
normR Exp
e'
            | Just Exp
e' <- Exp -> Maybe Exp
betaHead Exp
e       -> Exp -> NM m (Hoist, Exp)
forall (m :: * -> *).
MonadError AstError m =>
Exp -> NM m (Hoist, Exp)
normR Exp
e'
            | Bool
otherwise                   -> Exp -> NM m (Hoist, Exp)
forall (m :: * -> *).
MonadError AstError m =>
Exp -> NM m (Hoist, Exp)
normSpine Exp
e
      LitList {}   -> Exp -> NM m (Hoist, Exp)
forall (m :: * -> *).
MonadError AstError m =>
Exp -> NM m (Hoist, Exp)
atomize Exp
e
      LitVec {}    -> Exp -> NM m (Hoist, Exp)
forall (m :: * -> *).
MonadError AstError m =>
Exp -> NM m (Hoist, Exp)
atomize Exp
e
      Exp
_            -> (Hoist, Exp) -> NM m (Hoist, Exp)
forall a. a -> StateT Uniq m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], Exp
e)

-- | A spine normalized in place (the head — a variable, constructor, or
--   primitive once case, let, and lambda heads are commuted or reduced —
--   is kept): computed arguments are named and hoisted (argument naming);
--   primitive-headed arguments normalize in place (primitive expressions
--   are transparent); lambda arguments (continuations, the higher-order
--   primitives' function arguments) keep their place with normalized
--   bodies; function-typed arguments (partial applications, definition,
--   constructor, and operator references) keep their place with their
--   own arguments normalized — any other function-typed form (a function
--   chosen by a case) is rejected, since no consumer can lower it;
--   reactive arguments normalize recursively in place.
normSpine :: forall m. MonadError AstError m => Exp -> NM m (Hoist, Exp)
normSpine :: forall (m :: * -> *).
MonadError AstError m =>
Exp -> NM m (Hoist, Exp)
normSpine Exp
e = do
      let (Exp
h, [Arg]
args) = Exp -> (Exp, [Arg])
flattenApp Exp
e
      (bs, args') <- [Arg] -> NM m (Hoist, [Arg])
go [Arg]
args
      pure (bs, foldl (App $ ann' e) h args')
      where go :: [Arg] -> NM m (Hoist, [Arg])
            go :: [Arg] -> NM m (Hoist, [Arg])
go []             = (Hoist, [Arg]) -> NM m (Hoist, [Arg])
forall a. a -> StateT Uniq m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], [])
            go (TArg Ty
t : [Arg]
as)  = ([Arg] -> [Arg]) -> (Hoist, [Arg]) -> (Hoist, [Arg])
forall a b. (a -> b) -> (Hoist, a) -> (Hoist, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Ty -> Arg
TArg Ty
t Arg -> [Arg] -> [Arg]
forall a. a -> [a] -> [a]
:) ((Hoist, [Arg]) -> (Hoist, [Arg]))
-> NM m (Hoist, [Arg]) -> NM m (Hoist, [Arg])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Arg] -> NM m (Hoist, [Arg])
go [Arg]
as
            go (EArg Exp
a : [Arg]
as)  = do
                  (bs, a')   <- Exp -> StateT Uniq m (Hoist, Exp)
one Exp
a
                  (bs', as') <- go as
                  pure (bs <> bs', EArg a' : as')

            one :: Exp -> NM m (Hoist, Exp)
            one :: Exp -> StateT Uniq m (Hoist, Exp)
one Exp
a | Exp -> Bool
isAtom Exp
a               = Exp -> StateT Uniq m (Hoist, Exp)
forall (m :: * -> *).
MonadError AstError m =>
Exp -> NM m (Hoist, Exp)
atomize Exp
a
                  | App {} <- Exp
a
                  , Just Exp
a' <- Exp -> Maybe Exp
commuteCaseApp Exp
a = Exp -> StateT Uniq m (Hoist, Exp)
one Exp
a'
                  | App {} <- Exp
a
                  , Just Exp
a' <- Exp -> Maybe Exp
betaHead Exp
a  = Exp -> StateT Uniq m (Hoist, Exp)
one Exp
a'
                  | Lam {} <- Exp
a            = ([], ) (Exp -> (Hoist, Exp))
-> StateT Uniq m Exp -> StateT Uniq m (Hoist, Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> StateT Uniq m Exp
forall (m :: * -> *). MonadError AstError m => Exp -> NM m Exp
normTail Exp
a
                  | Exp -> Bool
isPrimExp Exp
a            = Exp -> StateT Uniq m (Hoist, Exp)
forall (m :: * -> *).
MonadError AstError m =>
Exp -> NM m (Hoist, Exp)
normSpine Exp
a
                  | App {} <- Exp
a
                  , Ty -> Bool
hasArrow (Ty -> Bool) -> Ty -> Bool
forall a b. (a -> b) -> a -> b
$ Exp -> Ty
typeOf Exp
a    = Exp -> StateT Uniq m (Hoist, Exp)
forall (m :: * -> *).
MonadError AstError m =>
Exp -> NM m (Hoist, Exp)
normSpine Exp
a
                  | Con {} <- Exp
a
                  , Ty -> Bool
hasArrow (Ty -> Bool) -> Ty -> Bool
forall a b. (a -> b) -> a -> b
$ Exp -> Ty
typeOf Exp
a    = (Hoist, Exp) -> StateT Uniq m (Hoist, Exp)
forall a. a -> StateT Uniq m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], Exp
a) -- a bare constructor reference
                  | Prim {} <- Exp
a
                  , Ty -> Bool
hasArrow (Ty -> Bool) -> Ty -> Bool
forall a b. (a -> b) -> a -> b
$ Exp -> Ty
typeOf Exp
a    = (Hoist, Exp) -> StateT Uniq m (Hoist, Exp)
forall a. a -> StateT Uniq m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], Exp
a) -- a bare operator primitive
                  | Ty -> Bool
hasArrow (Ty -> Bool) -> Ty -> Bool
forall a b. (a -> b) -> a -> b
$ Exp -> Ty
typeOf Exp
a    = Annote -> Text -> StateT Uniq m (Hoist, Exp)
forall (m :: * -> *) an a.
(MonadError AstError m, Annotation an) =>
an -> Text -> m a
failAt (Exp -> Annote
ann' Exp
a)
                        Text
"unsupported function-typed argument: a function chosen by a case or let cannot be lowered (use a lambda, a partial application, or a named definition)"
                  | Exp -> Bool
reactive Exp
a             = ([], ) (Exp -> (Hoist, Exp))
-> StateT Uniq m Exp -> StateT Uniq m (Hoist, Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> StateT Uniq m Exp
forall (m :: * -> *). MonadError AstError m => Exp -> NM m Exp
normTail Exp
a
                  | Bool
otherwise              = Exp -> StateT Uniq m (Hoist, Exp)
forall (m :: * -> *).
MonadError AstError m =>
Exp -> NM m (Hoist, Exp)
atomize Exp
a

-- | Normalize to an atom, hoisting a binding when the result is not
--   already one.
atomize :: forall m. MonadError AstError m => Exp -> NM m (Hoist, Exp)
atomize :: forall (m :: * -> *).
MonadError AstError m =>
Exp -> NM m (Hoist, Exp)
atomize Exp
e = case Exp
e of
      Var {}         -> (Hoist, Exp) -> NM m (Hoist, Exp)
forall a. a -> StateT Uniq m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], Exp
e)
      LitInt {}      -> (Hoist, Exp) -> NM m (Hoist, Exp)
forall a. a -> StateT Uniq m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], Exp
e)
      LitStr {}      -> (Hoist, Exp) -> NM m (Hoist, Exp)
forall a. a -> StateT Uniq m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], Exp
e)
      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 -> (Hoist, Exp) -> NM m (Hoist, Exp)
forall a. a -> StateT Uniq m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], Exp
e)
      Prim Annote
_ Ty
t Builtin
_ | [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 -> (Hoist, Exp) -> NM m (Hoist, Exp)
forall a. a -> StateT Uniq m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], Exp
e)
      LitList Annote
an Ty
t [Exp]
es -> do
            (bs, as) <- [Exp] -> NM m (Hoist, [Exp])
forall (m :: * -> *).
MonadError AstError m =>
[Exp] -> NM m (Hoist, [Exp])
atomizeMany [Exp]
es
            pure (bs, LitList an t as)
      LitVec Annote
an Ty
t [Exp]
es  -> do
            (bs, as) <- [Exp] -> NM m (Hoist, [Exp])
forall (m :: * -> *).
MonadError AstError m =>
[Exp] -> NM m (Hoist, [Exp])
atomizeMany [Exp]
es
            pure (bs, LitVec an t as)
      Exp
_              -> do
            (bs, r) <- Exp -> NM m (Hoist, Exp)
forall (m :: * -> *).
MonadError AstError m =>
Exp -> NM m (Hoist, Exp)
normR Exp
e
            if isAtom r then pure (bs, r) else do
                  x <- freshId $ typeOf r
                  pure (bs <> [(ann' r, x, r)], Var (ann' r) x)

atomizeMany :: MonadError AstError m => [Exp] -> NM m (Hoist, [Exp])
atomizeMany :: forall (m :: * -> *).
MonadError AstError m =>
[Exp] -> NM m (Hoist, [Exp])
atomizeMany []       = (Hoist, [Exp]) -> StateT Uniq m (Hoist, [Exp])
forall a. a -> StateT Uniq m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], [])
atomizeMany (Exp
e : [Exp]
es) = do
      (bs, a)   <- Exp -> NM m (Hoist, Exp)
forall (m :: * -> *).
MonadError AstError m =>
Exp -> NM m (Hoist, Exp)
atomize Exp
e
      (bs', as) <- atomizeMany es
      pure (bs <> bs', a : as)

ann' :: Exp -> Annote
ann' :: Exp -> Annote
ann' = \ case
      Var Annote
an Id
_        -> Annote
an
      Con Annote
an Ty
_ Text
_      -> Annote
an
      Prim Annote
an Ty
_ Builtin
_     -> Annote
an
      LitInt Annote
an Ty
_ Integer
_   -> Annote
an
      LitStr Annote
an Text
_     -> Annote
an
      LitList Annote
an Ty
_ [Exp]
_  -> Annote
an
      LitVec Annote
an Ty
_ [Exp]
_   -> Annote
an
      App Annote
an Exp
_ Arg
_      -> Annote
an
      Lam Annote
an Id
_ Exp
_      -> Annote
an
      Let Annote
an Bind
_ Exp
_      -> Annote
an
      Jump Annote
an JoinId
_ [Exp]
_     -> Annote
an
      Case Annote
an Ty
_ Exp
_ Id
_ [Alt]
_ -> Annote
an