{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ViewPatterns #-}

-- Various optimizations that may be applied to ANF terms. Many of
-- these are able to be run after intermediate code generation and
-- before interpreter code generation, to improve the code without
-- affecting the communication protocol.
module Unison.Runtime.ANF.Optimize
  ( optimize,
    inline,
    Arities,
    InlineInfo (..),
    InlineClass (..),
    InlineInfos,
    OptInfos,
    optimizeHandler,
    buildOptInfos,
  )
where

import Control.Monad.State (get, modify, runState)
import Control.Monad.Writer (MonadWriter (..), Writer, WriterT (..), runWriter, tell)
import Data.Graph (SCC (..), stronglyConnComp)
import Data.Map qualified as Map
import Data.Monoid (All (..), Any (..))
import Data.Set qualified as Set
import Unison.ABT.Normalized qualified as ABTN
import Unison.Prelude
import Unison.Runtime.ANF
import Unison.Var (Var)
import Unison.Var qualified as Var

-- Characterizes the situations where it's acceptable to inline an
-- expression.
--
-- Anywhere means that something is safe to inline anywhere the
-- associated application occurs, because it won't change stack
-- contents.
--
-- Tail means that it's acceptable to inline into an application in
-- tail position, because any stack descrepancies wouldn't be noticed.
--
-- Don't means the expression shouldn't (normally) be inlined, and is
-- provided just for other analyses.
data InlineClass = AnywhereInl | TailInl | Don'tInl
  deriving (InlineClass -> InlineClass -> Bool
(InlineClass -> InlineClass -> Bool)
-> (InlineClass -> InlineClass -> Bool) -> Eq InlineClass
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: InlineClass -> InlineClass -> Bool
== :: InlineClass -> InlineClass -> Bool
$c/= :: InlineClass -> InlineClass -> Bool
/= :: InlineClass -> InlineClass -> Bool
Eq, Eq InlineClass
Eq InlineClass =>
(InlineClass -> InlineClass -> Ordering)
-> (InlineClass -> InlineClass -> Bool)
-> (InlineClass -> InlineClass -> Bool)
-> (InlineClass -> InlineClass -> Bool)
-> (InlineClass -> InlineClass -> Bool)
-> (InlineClass -> InlineClass -> InlineClass)
-> (InlineClass -> InlineClass -> InlineClass)
-> Ord InlineClass
InlineClass -> InlineClass -> Bool
InlineClass -> InlineClass -> Ordering
InlineClass -> InlineClass -> InlineClass
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: InlineClass -> InlineClass -> Ordering
compare :: InlineClass -> InlineClass -> Ordering
$c< :: InlineClass -> InlineClass -> Bool
< :: InlineClass -> InlineClass -> Bool
$c<= :: InlineClass -> InlineClass -> Bool
<= :: InlineClass -> InlineClass -> Bool
$c> :: InlineClass -> InlineClass -> Bool
> :: InlineClass -> InlineClass -> Bool
$c>= :: InlineClass -> InlineClass -> Bool
>= :: InlineClass -> InlineClass -> Bool
$cmax :: InlineClass -> InlineClass -> InlineClass
max :: InlineClass -> InlineClass -> InlineClass
$cmin :: InlineClass -> InlineClass -> InlineClass
min :: InlineClass -> InlineClass -> InlineClass
Ord, Int -> InlineClass -> ShowS
[InlineClass] -> ShowS
InlineClass -> String
(Int -> InlineClass -> ShowS)
-> (InlineClass -> String)
-> ([InlineClass] -> ShowS)
-> Show InlineClass
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> InlineClass -> ShowS
showsPrec :: Int -> InlineClass -> ShowS
$cshow :: InlineClass -> String
show :: InlineClass -> String
$cshowList :: [InlineClass] -> ShowS
showList :: [InlineClass] -> ShowS
Show)

instance Semigroup InlineClass where
  InlineClass
AnywhereInl <> :: InlineClass -> InlineClass -> InlineClass
<> InlineClass
c = InlineClass
c
  InlineClass
c <> InlineClass
AnywhereInl = InlineClass
c
  InlineClass
Don'tInl <> InlineClass
_ = InlineClass
Don'tInl
  InlineClass
_ <> InlineClass
Don'tInl = InlineClass
Don'tInl
  InlineClass
_ <> InlineClass
_ = InlineClass
TailInl

instance Monoid InlineClass where
  mempty :: InlineClass
mempty = InlineClass
AnywhereInl

data InlineInfo ref v = InlInfo
  { forall ref v. InlineInfo ref v -> InlineClass
_inlClass :: InlineClass,
    forall ref v. InlineInfo ref v -> ANormal ref v
inlExpr :: ANormal ref v
  }
  deriving (InlineInfo ref v -> InlineInfo ref v -> Bool
(InlineInfo ref v -> InlineInfo ref v -> Bool)
-> (InlineInfo ref v -> InlineInfo ref v -> Bool)
-> Eq (InlineInfo ref v)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall ref v.
(Var v, Eq ref) =>
InlineInfo ref v -> InlineInfo ref v -> Bool
$c== :: forall ref v.
(Var v, Eq ref) =>
InlineInfo ref v -> InlineInfo ref v -> Bool
== :: InlineInfo ref v -> InlineInfo ref v -> Bool
$c/= :: forall ref v.
(Var v, Eq ref) =>
InlineInfo ref v -> InlineInfo ref v -> Bool
/= :: InlineInfo ref v -> InlineInfo ref v -> Bool
Eq, Int -> InlineInfo ref v -> ShowS
[InlineInfo ref v] -> ShowS
InlineInfo ref v -> String
(Int -> InlineInfo ref v -> ShowS)
-> (InlineInfo ref v -> String)
-> ([InlineInfo ref v] -> ShowS)
-> Show (InlineInfo ref v)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall ref v.
(Show v, Show ref) =>
Int -> InlineInfo ref v -> ShowS
forall ref v. (Show v, Show ref) => [InlineInfo ref v] -> ShowS
forall ref v. (Show v, Show ref) => InlineInfo ref v -> String
$cshowsPrec :: forall ref v.
(Show v, Show ref) =>
Int -> InlineInfo ref v -> ShowS
showsPrec :: Int -> InlineInfo ref v -> ShowS
$cshow :: forall ref v. (Show v, Show ref) => InlineInfo ref v -> String
show :: InlineInfo ref v -> String
$cshowList :: forall ref v. (Show v, Show ref) => [InlineInfo ref v] -> ShowS
showList :: [InlineInfo ref v] -> ShowS
Show)

type Arities ref = Map ref Int

type InlineInfos ref v = Map ref (InlineInfo ref v)

type OptInfos ref v = (Arities ref, InlineInfos ref v)

-- Checks a SuperGroup makes it eligible to be inlined.
-- Unfortunately we need to be quite conservative about this.
--
-- The heuristic implemented below is as follows:
--
--   1. There are no local bindings, so only the 'entry point'
--      matters.
--   2. The entry point body is just a single expression, that is,
--      an application, variable or literal.
--
-- The first condition ensures that there isn't any need to jump
-- into a non-entrypoint from outside a group. These should be rare
-- anyway, because the local bindings are no longer used for
-- (unison-level) local function definitions (those are lifted
-- out). The second condition ensures that inlining the body should
-- have no effect on the runtime stack of of the function we're
-- inlining into, because the combinator is just a wrapper around
-- the simple expression.
--
-- Fortunately, it should be possible to make _most_ builtins have
-- this form, so that their instructions can be inlined directly
-- into the call sites when saturated.
--
-- The result of this function is the information necessary to
-- inline the combinator—an arity and the body expression with
-- bound variables. This should allow checking if the call is
-- saturated and make it possible to locally substitute for an
-- inlined expression.
--
-- The `Reference` argument allows us to check if the body is a
-- direct recursive call to the same function, which would result
-- in infinite inlining. This isn't the only such scenario, but
-- it's one we can opportunistically rule out.
inlineInfo ::
  (Var v) => Bool -> SuperGroup ref v -> Maybe (InlineInfo ref v)
inlineInfo :: forall v ref.
Var v =>
Bool -> SuperGroup ref v -> Maybe (InlineInfo ref v)
inlineInfo Bool
rec (Rec [] (Lambda [Mem]
_ body :: ANormal ref v
body@(ABTN.TAbss [v]
vs ANormal ref v
e)))
  | Just ANormal ref v
opt <- ANormal ref v -> Maybe (ANormal ref v)
forall v ref. Var v => ANormal ref v -> Maybe (ANormal ref v)
matchHandlerApp ANormal ref v
e =
      InlineInfo ref v -> Maybe (InlineInfo ref v)
forall a. a -> Maybe a
Just (InlineInfo ref v -> Maybe (InlineInfo ref v))
-> InlineInfo ref v -> Maybe (InlineInfo ref v)
forall a b. (a -> b) -> a -> b
$ InlineClass -> ANormal ref v -> InlineInfo ref v
forall ref v. InlineClass -> ANormal ref v -> InlineInfo ref v
InlInfo InlineClass
TailInl ([v] -> ANormal ref v -> ANormal ref v
forall v (f :: * -> * -> *). Var v => [v] -> Term f v -> Term f v
ABTN.TAbss [v]
vs ANormal ref v
opt)
  | Bool
otherwise =
      InlineInfo ref v -> Maybe (InlineInfo ref v)
forall a. a -> Maybe a
Just (InlineInfo ref v -> Maybe (InlineInfo ref v))
-> InlineInfo ref v -> Maybe (InlineInfo ref v)
forall a b. (a -> b) -> a -> b
$ InlineClass -> ANormal ref v -> InlineInfo ref v
forall ref v. InlineClass -> ANormal ref v -> InlineInfo ref v
InlInfo (Bool -> ANormal ref v -> InlineClass
forall v ref. Var v => Bool -> ANormal ref v -> InlineClass
classifyInline Bool
rec ANormal ref v
e) ANormal ref v
body
  where

inlineInfo Bool
_ SuperGroup ref v
_ = Maybe (InlineInfo ref v)
forall a. Maybe a
Nothing

-- Some special inline info that is relevant for optimizing recursive
-- groups, but should not be inlined in general (due to being
-- recursive).
recInlineInfo ::
  (Ord ref, Var v) =>
  Map ref (SuperGroup ref v) ->
  InlineInfos ref v
recInlineInfo :: forall ref v.
(Ord ref, Var v) =>
Map ref (SuperGroup ref v) -> InlineInfos ref v
recInlineInfo = (SuperGroup ref v -> Maybe (InlineInfo ref v))
-> Map ref (SuperGroup ref v) -> Map ref (InlineInfo ref v)
forall u v k. (u -> Maybe v) -> Map k u -> Map k v
mapMapMaybe SuperGroup ref v -> Maybe (InlineInfo ref v)
forall {v} {ref}.
Var v =>
SuperGroup ref v -> Maybe (InlineInfo ref v)
f
  where
    f :: SuperGroup ref v -> Maybe (InlineInfo ref v)
f (Rec [] (Lambda [Mem]
_ (ABTN.TAbss [v]
vs Term (ANormalF ref) v
e))) =
      InlineClass -> Term (ANormalF ref) v -> InlineInfo ref v
forall ref v. InlineClass -> ANormal ref v -> InlineInfo ref v
InlInfo InlineClass
TailInl (Term (ANormalF ref) v -> InlineInfo ref v)
-> (Term (ANormalF ref) v -> Term (ANormalF ref) v)
-> Term (ANormalF ref) v
-> InlineInfo ref v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [v] -> Term (ANormalF ref) v -> Term (ANormalF ref) v
forall v (f :: * -> * -> *). Var v => [v] -> Term f v -> Term f v
ABTN.TAbss [v]
vs (Term (ANormalF ref) v -> InlineInfo ref v)
-> Maybe (Term (ANormalF ref) v) -> Maybe (InlineInfo ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Term (ANormalF ref) v -> Maybe (Term (ANormalF ref) v)
forall v ref. Var v => ANormal ref v -> Maybe (ANormal ref v)
matchHandlerApp Term (ANormalF ref) v
e
    f SuperGroup ref v
_ = Maybe (InlineInfo ref v)
forall a. Maybe a
Nothing

arityInfo :: SuperGroup ref v -> Int
arityInfo :: forall ref v. SuperGroup ref v -> Int
arityInfo (Rec [(v, SuperNormal ref v)]
_ (Lambda [Mem]
ccs ANormal ref v
_)) = [Mem] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Mem]
ccs

-- This is a special inlining available for handlers, inlining the
-- entry point of the handler into the actual implementation. This
-- improves recursive handlers slightly, enables other optimizations,
-- and makes it easier to recognize affine handlers.
entryInfo :: (Var v) => SuperGroup ref v -> Maybe (InlineInfo ref v)
entryInfo :: forall {v} {ref}.
Var v =>
SuperGroup ref v -> Maybe (InlineInfo ref v)
entryInfo (Rec [(v
him, SuperNormal ref v
_)] (Lambda [Mem]
_ body :: ANormal ref v
body@(ABTN.TAbss [v]
vs ANormal ref v
e)))
  | v
req : [v]
_ <- [v] -> [v]
forall v. [v] -> [v]
shiftArgs [v]
vs,
    v -> v -> ANormal ref v -> Bool
forall v ref. Var v => v -> v -> ANormal ref v -> Bool
isHandlerEntry v
him v
req ANormal ref v
e =
      InlineInfo ref v -> Maybe (InlineInfo ref v)
forall a. a -> Maybe a
Just (InlineInfo ref v -> Maybe (InlineInfo ref v))
-> InlineInfo ref v -> Maybe (InlineInfo ref v)
forall a b. (a -> b) -> a -> b
$ InlineClass -> ANormal ref v -> InlineInfo ref v
forall ref v. InlineClass -> ANormal ref v -> InlineInfo ref v
InlInfo InlineClass
TailInl ANormal ref v
body
entryInfo SuperGroup ref v
_ = Maybe (InlineInfo ref v)
forall a. Maybe a
Nothing

-- Rewriting
-- ---------
--
-- These functions allow rewriting terms while keeping track of some
-- pertinent information. They use a technique to recognize whether
-- any rewriting has been applied on subterms, to e.g. just use the
-- original term structure if not (to preserve more sharing).

type Memo = MonadWriter Any

memo :: (Memo m) => a -> m a -> m a
memo :: forall (m :: * -> *) a. Memo m => a -> m a -> m a
memo a
orig m a
mod =
  m a -> m (a, Any)
forall a. m a -> m (a, Any)
forall w (m :: * -> *) a. MonadWriter w m => m a -> m (a, w)
listen m a
mod m (a, Any) -> ((a, Any) -> a) -> m a
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \(a
new, Any Bool
changed) ->
    if Bool
changed then a
new else a
orig

dirty :: (Memo m) => m ()
dirty :: forall (m :: * -> *). Memo m => m ()
dirty = Any -> m ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (Any -> m ()) -> Any -> m ()
forall a b. (a -> b) -> a -> b
$ Bool -> Any
Any Bool
True

runMemo :: Writer Any a -> a
runMemo :: forall a. Writer Any a -> a
runMemo = (a, Any) -> a
forall a b. (a, b) -> a
fst ((a, Any) -> a) -> (Writer Any a -> (a, Any)) -> Writer Any a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Writer Any a -> (a, Any)
forall w a. Writer w a -> (a, w)
runWriter

whenChanged :: (Memo m) => (a -> m a) -> m a -> m a
whenChanged :: forall (m :: * -> *) a. Memo m => (a -> m a) -> m a -> m a
whenChanged a -> m a
f m a
act = do
  (a
x, Any Bool
changed) <- m a -> m (a, Any)
forall a. m a -> m (a, Any)
forall w (m :: * -> *) a. MonadWriter w m => m a -> m (a, w)
listen m a
act
  if Bool
changed then a -> m a
f a
x else a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x

descend ::
  (Memo m, Var v) =>
  (Bool -> Set v -> ANormal ref v -> m (ANormal ref v)) ->
  Bool ->
  Set v ->
  ANormal ref v ->
  m (ANormal ref v)
descend :: forall (m :: * -> *) v ref.
(Memo m, Var v) =>
(Bool -> Set v -> ANormal ref v -> m (ANormal ref v))
-> Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
descend Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
rec Bool
tail Set v
bound ANormal ref v
tm = ANormal ref v -> m (ANormal ref v) -> m (ANormal ref v)
forall (m :: * -> *) a. Memo m => a -> m a -> m a
memo ANormal ref v
tm (m (ANormal ref v) -> m (ANormal ref v))
-> m (ANormal ref v) -> m (ANormal ref v)
forall a b. (a -> b) -> a -> b
$ case ANormal ref v
tm of
  TLets Direction Word16
d [v]
vs [Mem]
ccs ANormal ref v
bn ANormal ref v
bd ->
    Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
forall v ref.
Var v =>
Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
TLets Direction Word16
d [v]
vs [Mem]
ccs (ANormal ref v -> ANormal ref v -> ANormal ref v)
-> m (ANormal ref v) -> m (ANormal ref v -> ANormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
rec Bool
False Set v
bound ANormal ref v
bn m (ANormal ref v -> ANormal ref v)
-> m (ANormal ref v) -> m (ANormal ref v)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
rec Bool
tail Set v
bnd ANormal ref v
bd
    where
      bnd :: Set v
bnd = Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
Set.union ([v] -> Set v
forall a. Ord a => [a] -> Set a
Set.fromList [v]
vs) Set v
bound
  TName v
v Either ref v
f [v]
vs ANormal ref v
bd ->
    v -> Either ref v -> [v] -> ANormal ref v -> ANormal ref v
forall v ref.
Var v =>
v -> Either ref v -> [v] -> ANormal ref v -> ANormal ref v
TName v
v Either ref v
f [v]
vs (ANormal ref v -> ANormal ref v)
-> m (ANormal ref v) -> m (ANormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
rec Bool
tail (v -> Set v -> Set v
forall a. Ord a => a -> Set a -> Set a
Set.insert v
v Set v
bound) ANormal ref v
bd
  TMatch v
v Branched ref (ANormal ref v)
bs ->
    v -> Branched ref (ANormal ref v) -> ANormal ref v
forall v ref.
Var v =>
v -> Branched ref (ANormal ref v) -> ANormal ref v
TMatch v
v (Branched ref (ANormal ref v) -> ANormal ref v)
-> m (Branched ref (ANormal ref v)) -> m (ANormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ANormal ref v -> m (ANormal ref v))
-> Branched ref (ANormal ref v) -> m (Branched ref (ANormal ref v))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Branched ref a -> f (Branched ref b)
traverse (Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
rec Bool
tail Set v
bound) Branched ref (ANormal ref v)
bs
  TShift ref
r v
v ANormal ref v
bd ->
    ref -> v -> ANormal ref v -> ANormal ref v
forall v ref. Var v => ref -> v -> ANormal ref v -> ANormal ref v
TShift ref
r v
v (ANormal ref v -> ANormal ref v)
-> m (ANormal ref v) -> m (ANormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
rec Bool
tail (v -> Set v -> Set v
forall a. Ord a => a -> Set a -> Set a
Set.insert v
v Set v
bound) ANormal ref v
bd
  THnd [ref]
rs v
hn Maybe v
ha ANormal ref v
bd ->
    [ref] -> v -> Maybe v -> ANormal ref v -> ANormal ref v
forall v ref.
Var v =>
[ref] -> v -> Maybe v -> ANormal ref v -> ANormal ref v
THnd [ref]
rs v
hn Maybe v
ha (ANormal ref v -> ANormal ref v)
-> m (ANormal ref v) -> m (ANormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
rec Bool
tail Set v
bound ANormal ref v
bd
  TLocal v
v ANormal ref v
bd ->
    v -> ANormal ref v -> ANormal ref v
forall v ref. Var v => v -> ANormal ref v -> ANormal ref v
TLocal v
v (ANormal ref v -> ANormal ref v)
-> m (ANormal ref v) -> m (ANormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
rec Bool
tail Set v
bound ANormal ref v
bd
  ABTN.TAbs v
v (ABTN.TAbss [v]
vs ANormal ref v
bd) ->
    [v] -> ANormal ref v -> ANormal ref v
forall v (f :: * -> * -> *). Var v => [v] -> Term f v -> Term f v
ABTN.TAbss (v
v v -> [v] -> [v]
forall a. a -> [a] -> [a]
: [v]
vs) (ANormal ref v -> ANormal ref v)
-> m (ANormal ref v) -> m (ANormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
rec Bool
tail Set v
bnd ANormal ref v
bd
    where
      bnd :: Set v
bnd = Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
Set.union ([v] -> Set v
forall a. Ord a => [a] -> Set a
Set.fromList ([v] -> Set v) -> [v] -> Set v
forall a b. (a -> b) -> a -> b
$ v
v v -> [v] -> [v]
forall a. a -> [a] -> [a]
: [v]
vs) Set v
bound
  ANormal ref v
_ -> ANormal ref v -> m (ANormal ref v)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ANormal ref v
tm

-- Rewrites a term from the top down, first applying the step
-- transform given, then descending to children.
rewriteDown ::
  (Memo m, Var v) =>
  (Bool -> Set v -> ANormal ref v -> m (ANormal ref v)) ->
  ANormal ref v ->
  m (ANormal ref v)
rewriteDown :: forall (m :: * -> *) v ref.
(Memo m, Var v) =>
(Bool -> Set v -> ANormal ref v -> m (ANormal ref v))
-> ANormal ref v -> m (ANormal ref v)
rewriteDown Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
step = Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
go Bool
True Set v
forall a. Set a
Set.empty
  where
    go :: Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
go Bool
tail Set v
bound ANormal ref v
tm = Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
step Bool
tail Set v
bound ANormal ref v
tm m (ANormal ref v)
-> (ANormal ref v -> m (ANormal ref v)) -> m (ANormal ref v)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Bool -> Set v -> ANormal ref v -> m (ANormal ref v))
-> Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
forall (m :: * -> *) v ref.
(Memo m, Var v) =>
(Bool -> Set v -> ANormal ref v -> m (ANormal ref v))
-> Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
descend Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
go Bool
tail Set v
bound

rewriteUp ::
  (Memo m, Var v) =>
  (Bool -> Set v -> ANormal ref v -> m (ANormal ref v)) ->
  ANormal ref v ->
  m (ANormal ref v)
rewriteUp :: forall (m :: * -> *) v ref.
(Memo m, Var v) =>
(Bool -> Set v -> ANormal ref v -> m (ANormal ref v))
-> ANormal ref v -> m (ANormal ref v)
rewriteUp Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
step = Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
go Bool
True Set v
forall a. Set a
Set.empty
  where
    go :: Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
go Bool
tail Set v
bound ANormal ref v
tm =
      ANormal ref v -> m (ANormal ref v) -> m (ANormal ref v)
forall (m :: * -> *) a. Memo m => a -> m a -> m a
memo ANormal ref v
tm ((Bool -> Set v -> ANormal ref v -> m (ANormal ref v))
-> Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
forall (m :: * -> *) v ref.
(Memo m, Var v) =>
(Bool -> Set v -> ANormal ref v -> m (ANormal ref v))
-> Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
descend Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
go Bool
tail Set v
bound ANormal ref v
tm) m (ANormal ref v)
-> (ANormal ref v -> m (ANormal ref v)) -> m (ANormal ref v)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
step Bool
tail Set v
bound

-- Performs inlining on a `SuperGroup` using the inlining information
-- in the map. The map can be created from typical `SuperGroup` data
-- using the `buildInlineMap` function.
--
-- Inlining is capped at 30 iterations per site to avoid infinite
-- loops on recursive inlining situations that were not detected by
-- `builtInlineMap`.
inline ::
  (Ord ref, Memo m, Var v) =>
  Set v ->
  OptInfos ref v ->
  ANormal ref v ->
  m (ANormal ref v)
inline :: forall ref (m :: * -> *) v.
(Ord ref, Memo m, Var v) =>
Set v -> OptInfos ref v -> ANormal ref v -> m (ANormal ref v)
inline Set v
avoid (Arities ref
arities, InlineInfos ref v
inls) ANormal ref v
n0 = ANormal ref v -> m (ANormal ref v) -> m (ANormal ref v)
forall (m :: * -> *) a. Memo m => a -> m a -> m a
memo ANormal ref v
n0 (m (ANormal ref v) -> m (ANormal ref v))
-> m (ANormal ref v) -> m (ANormal ref v)
forall a b. (a -> b) -> a -> b
$ Int -> ANormal ref v -> m (ANormal ref v)
go (Int
30 :: Int) ANormal ref v
n0
  where
    go :: Int -> ANormal ref v -> m (ANormal ref v)
go Int
n ANormal ref v
tm
      | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = ANormal ref v -> m (ANormal ref v)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ANormal ref v
tm
      | Bool
otherwise = (Bool -> Set v -> ANormal ref v -> m (ANormal ref v))
-> ANormal ref v -> m (ANormal ref v)
forall (m :: * -> *) v ref.
(Memo m, Var v) =>
(Bool -> Set v -> ANormal ref v -> m (ANormal ref v))
-> ANormal ref v -> m (ANormal ref v)
rewriteUp (Int -> Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
step Int
n) ANormal ref v
tm

    step :: Int -> Bool -> Set v -> ANormal ref v -> m (ANormal ref v)
step Int
n Bool
tail Set v
bound (TApp (FComb ref
r) [v]
args)
      | Just ANormal ref v
new <- Bool -> Set v -> ref -> [v] -> Maybe (ANormal ref v)
findInline Bool
tail Set v
bound ref
r [v]
args =
          m ()
forall (m :: * -> *). Memo m => m ()
dirty m () -> m (ANormal ref v) -> m (ANormal ref v)
forall a b. m a -> m b -> m b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Int -> ANormal ref v -> m (ANormal ref v)
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ANormal ref v
new
    step Int
_ Bool
_tail Set v
_bound ANormal ref v
tm = ANormal ref v -> m (ANormal ref v)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ANormal ref v
tm

    findInline :: Bool -> Set v -> ref -> [v] -> Maybe (ANormal ref v)
findInline Bool
tail Set v
bound ref
r [v]
args = do
      InlineInfo ref v
info <- ref -> InlineInfos ref v -> Maybe (InlineInfo ref v)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup ref
r InlineInfos ref v
inls
      Int
arity <- ref -> Arities ref -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup ref
r Arities ref
arities
      Bool
-> Set v -> [v] -> Int -> InlineInfo ref v -> Maybe (ANormal ref v)
tweak Bool
tail Set v
bound [v]
args Int
arity InlineInfo ref v
info

    don'tInline :: InlineClass -> Bool -> Bool
don'tInline InlineClass
Don'tInl Bool
_ = Bool
True
    don'tInline InlineClass
TailInl Bool
isTail = Bool -> Bool
not Bool
isTail
    don'tInline InlineClass
AnywhereInl Bool
_ = Bool
False

    -- Note: we use `renameAvoiding` because by adding `entryInfo` to
    -- the inlining map, we may be inlining terms with free variables
    -- referring to the floated handler code. This happens over
    -- multiple inlining steps, so we freshen anything else we inline
    -- to not be capable of capturing the variables from the entry
    -- code.
    tweak :: Bool
-> Set v -> [v] -> Int -> InlineInfo ref v -> Maybe (ANormal ref v)
tweak Bool
isTail Set v
bound [v]
args Int
arity (InlInfo InlineClass
clazz (ABTN.TAbss [v]
vs ANormal ref v
body))
      | InlineClass -> Bool -> Bool
don'tInline InlineClass
clazz Bool
isTail = Maybe (ANormal ref v)
forall a. Maybe a
Nothing
      -- exactly saturated
      | [v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
args Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
arity,
        Map v v
rn <- [(v, v)] -> Map v v
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([v] -> [v] -> [(v, v)]
forall a b. [a] -> [b] -> [(a, b)]
zip [v]
vs [v]
args) =
          ANormal ref v -> Maybe (ANormal ref v)
forall a. a -> Maybe a
Just (ANormal ref v -> Maybe (ANormal ref v))
-> ANormal ref v -> Maybe (ANormal ref v)
forall a b. (a -> b) -> a -> b
$ Set v -> Map v v -> ANormal ref v -> ANormal ref v
forall v (f :: * -> * -> *).
(Var v, Bifunctor f, Bifoldable f) =>
Set v -> Map v v -> Term f v -> Term f v
ABTN.renamesAvoiding (Set v
avoid Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
`Set.union` Set v
bound) Map v v
rn ANormal ref v
body
      -- oversaturated, only makes sense if body is a call
      | [v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
args Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
arity,
        ([v]
pre, [v]
post) <- Int -> [v] -> ([v], [v])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
arity [v]
args,
        Map v v
rn <- [(v, v)] -> Map v v
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([v] -> [v] -> [(v, v)]
forall a b. [a] -> [b] -> [(a, b)]
zip [v]
vs [v]
pre),
        TApp Func ref v
f [v]
pre <- Set v -> Map v v -> ANormal ref v -> ANormal ref v
forall v (f :: * -> * -> *).
(Var v, Bifunctor f, Bifoldable f) =>
Set v -> Map v v -> Term f v -> Term f v
ABTN.renamesAvoiding (Set v
avoid Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
`Set.union` Set v
bound) Map v v
rn ANormal ref v
body =
          ANormal ref v -> Maybe (ANormal ref v)
forall a. a -> Maybe a
Just (ANormal ref v -> Maybe (ANormal ref v))
-> ANormal ref v -> Maybe (ANormal ref v)
forall a b. (a -> b) -> a -> b
$ Func ref v -> [v] -> ANormal ref v
forall v ref. Var v => Func ref v -> [v] -> ANormal ref v
TApp Func ref v
f ([v]
pre [v] -> [v] -> [v]
forall a. [a] -> [a] -> [a]
++ [v]
post)
      | Bool
otherwise = Maybe (ANormal ref v)
forall a. Maybe a
Nothing

-- Performs peephole optimizations on a `SuperGroup`. The term is
-- traversed and certain patterns are looked for, and replaced with
-- more optimized patterns. Some patterns rely on inlining to have
-- already happened to work, so the former should be called first, and
-- it may be beneficial to repeat the process after peephole
-- optimization runs in case more inlining became available.
--
-- The `affine` flag specifies whether the optimizations are for code
-- that will be used in an affine handler context. Such code will
-- never have its stack captured, so more radical optimizations can be
-- performed.
peephole ::
  (Ord ref, Memo m, Var v) =>
  Arities ref ->
  Bool ->
  ANormal ref v ->
  m (ANormal ref v)
peephole :: forall ref (m :: * -> *) v.
(Ord ref, Memo m, Var v) =>
Arities ref -> Bool -> ANormal ref v -> m (ANormal ref v)
peephole Arities ref
arities Bool
affine ANormal ref v
n0 = ANormal ref v -> m (ANormal ref v) -> m (ANormal ref v)
forall (m :: * -> *) a. Memo m => a -> m a -> m a
memo ANormal ref v
n0 (m (ANormal ref v) -> m (ANormal ref v))
-> m (ANormal ref v) -> m (ANormal ref v)
forall a b. (a -> b) -> a -> b
$ Int -> ANormal ref v -> m (ANormal ref v)
go (Int
30 :: Int) ANormal ref v
n0
  where
    go :: Int -> ANormal ref v -> m (ANormal ref v)
go Int
0 = ANormal ref v -> m (ANormal ref v)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    go Int
n =
      (ANormal ref v -> m (ANormal ref v))
-> m (ANormal ref v) -> m (ANormal ref v)
forall (m :: * -> *) a. Memo m => (a -> m a) -> m a -> m a
whenChanged (Int -> ANormal ref v -> m (ANormal ref v)
go (Int -> ANormal ref v -> m (ANormal ref v))
-> Int -> ANormal ref v -> m (ANormal ref v)
forall a b. (a -> b) -> a -> b
$ Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (m (ANormal ref v) -> m (ANormal ref v))
-> (ANormal ref v -> m (ANormal ref v))
-> ANormal ref v
-> m (ANormal ref v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Bool -> Set v -> ANormal ref v -> m (ANormal ref v))
-> ANormal ref v -> m (ANormal ref v)
forall (m :: * -> *) v ref.
(Memo m, Var v) =>
(Bool -> Set v -> ANormal ref v -> m (ANormal ref v))
-> ANormal ref v -> m (ANormal ref v)
rewriteDown \Bool
tail Set v
_bound -> \case
        -- eliminate `v = u` bindings in affine contexts
        TLet Direction Word16
_ v
v Mem
_ (TVar v
u) ANormal ref v
bd
          | Bool
affine -> v -> v -> ANormal ref v -> ANormal ref v
forall v (f :: * -> * -> *).
(Var v, Ord v, Bifunctor f, Bifoldable f) =>
v -> v -> Term f v -> Term f v
ABTN.rename v
v v
u ANormal ref v
bd ANormal ref v -> m () -> m (ANormal ref v)
forall a b. a -> m b -> m a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ m ()
forall (m :: * -> *). Memo m => m ()
dirty
        -- turn `let v = ex in v` into just `ex` when in tail position
        TLet Direction Word16
_ v
v Mem
_ ANormal ref v
bn (TVar v
u)
          | v
v v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
u, Bool
tail -> ANormal ref v
bn ANormal ref v -> m () -> m (ANormal ref v)
forall a b. a -> m b -> m a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ m ()
forall (m :: * -> *). Memo m => m ()
dirty
        -- promote 'indirect' bindings to 'direct' bindings if possible
        TLets (Indirect Word16
_) [v]
vs [Mem]
ccs ANormal ref v
bn ANormal ref v
bd
          | ANormal ref v -> Bool
forall v ref. Var v => ANormal ref v -> Bool
directAllowed ANormal ref v
bn ->
              Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
forall v ref.
Var v =>
Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
TLets Direction Word16
forall a. Direction a
Direct [v]
vs [Mem]
ccs ANormal ref v
bn ANormal ref v
bd ANormal ref v -> m () -> m (ANormal ref v)
forall a b. a -> m b -> m a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ m ()
forall (m :: * -> *). Memo m => m ()
dirty
        -- elminate unused bindings if they have no effects and no one
        -- will be able to see the stack differences.
        TLets Direction Word16
_ [v]
vs [Mem]
_ ANormal ref v
bn ANormal ref v
bd
          | (v -> Bool) -> [v] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.notMember` ANormal ref v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
ABTN.freeVars ANormal ref v
bd) [v]
vs,
            ANormal ref v -> Bool
forall v ref. Var v => ANormal ref v -> Bool
effectless ANormal ref v
bn,
            Bool
affine Bool -> Bool -> Bool
|| ANormal ref v -> Bool
forall v ref. Var v => ANormal ref v -> Bool
stackInsensitive ANormal ref v
bd ->
              ANormal ref v
bd ANormal ref v -> m () -> m (ANormal ref v)
forall a b. a -> m b -> m a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ m ()
forall (m :: * -> *). Memo m => m ()
dirty
        -- turn the code generated by handler applications into
        -- a slightly nicer form. This is always in the tail position.
        HandlerApp ANormal ref v
rw -> ANormal ref v
rw ANormal ref v -> m () -> m (ANormal ref v)
forall a b. a -> m b -> m a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ m ()
forall (m :: * -> *). Memo m => m ()
dirty
        -- Rewrite a handler resumption into a more direct form that
        -- avoids allocating and immediately forcing a lazy value.
        -- This is always a tail occurrence.
        HandlerResume v
lz Either ref v
f [v]
as v
lh Either ref v
h [v]
bs [ref]
rs
          | (v -> Bool) -> Either ref v -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (v -> v -> Bool
forall a. Eq a => a -> a -> Bool
/= v
lz) Either ref v
h,
            (v -> Bool) -> [v] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (v -> v -> Bool
forall a. Eq a => a -> a -> Bool
/= v
lz) [v]
bs -> do
              m ()
forall (m :: * -> *). Memo m => m ()
dirty
              ANormal ref v -> m (ANormal ref v)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ANormal ref v -> m (ANormal ref v))
-> (ANormal ref v -> ANormal ref v)
-> ANormal ref v
-> m (ANormal ref v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v -> Either ref v -> [v] -> ANormal ref v -> ANormal ref v
forall v ref.
Var v =>
v -> Either ref v -> [v] -> ANormal ref v -> ANormal ref v
TName v
lh Either ref v
h [v]
bs (ANormal ref v -> ANormal ref v)
-> (ANormal ref v -> ANormal ref v)
-> ANormal ref v
-> ANormal ref v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ref] -> v -> Maybe v -> ANormal ref v -> ANormal ref v
forall v ref.
Var v =>
[ref] -> v -> Maybe v -> ANormal ref v -> ANormal ref v
THnd [ref]
rs v
lh Maybe v
forall a. Maybe a
Nothing (ANormal ref v -> m (ANormal ref v))
-> ANormal ref v -> m (ANormal ref v)
forall a b. (a -> b) -> a -> b
$ Func ref v -> [v] -> ANormal ref v
forall v ref. Var v => Func ref v -> [v] -> ANormal ref v
TApp (Either ref v -> Func ref v
forall {ref} {v}. Either ref v -> Func ref v
Nameable Either ref v
f) [v]
as
        -- Eliminate the allocation of a delayed thunk if it is
        -- subsequently forced exactly once in a handler, and no one
        -- will be able to see the stack differences.
        HandledThunk ref
r Int
n Bool
safe ANormal ref v
expr
          | Bool
safe Bool -> Bool -> Bool
|| Bool
affine,
            Just Int
arity <- ref -> Arities ref -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup ref
r Arities ref
arities,
            Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
arity ->
              ANormal ref v
expr ANormal ref v -> m () -> m (ANormal ref v)
forall a b. a -> m b -> m a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ m ()
forall (m :: * -> *). Memo m => m ()
dirty
        ANormal ref v
tm -> ANormal ref v -> m (ANormal ref v)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ANormal ref v
tm

-- Alternates inlining and peephole optimizations until neither makes
-- any changes to the term.
optNormal ::
  (Ord ref, Memo m, Var v) =>
  OptInfos ref v ->
  Set v ->
  Bool ->
  ANormal ref v ->
  m (ANormal ref v)
optNormal :: forall ref (m :: * -> *) v.
(Ord ref, Memo m, Var v) =>
OptInfos ref v
-> Set v -> Bool -> ANormal ref v -> m (ANormal ref v)
optNormal opts :: OptInfos ref v
opts@(Arities ref
arities, InlineInfos ref v
_) Set v
avoid Bool
affine ANormal ref v
n0 =
  ANormal ref v -> m (ANormal ref v)
peep (ANormal ref v -> m (ANormal ref v))
-> m (ANormal ref v) -> m (ANormal ref v)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Set v -> OptInfos ref v -> ANormal ref v -> m (ANormal ref v)
forall ref (m :: * -> *) v.
(Ord ref, Memo m, Var v) =>
Set v -> OptInfos ref v -> ANormal ref v -> m (ANormal ref v)
inline Set v
avoid OptInfos ref v
opts ANormal ref v
n0
  where
    inl :: ANormal ref v -> m (ANormal ref v)
inl ANormal ref v
n = (ANormal ref v -> m (ANormal ref v))
-> m (ANormal ref v) -> m (ANormal ref v)
forall (m :: * -> *) a. Memo m => (a -> m a) -> m a -> m a
whenChanged ANormal ref v -> m (ANormal ref v)
peep (m (ANormal ref v) -> m (ANormal ref v))
-> m (ANormal ref v) -> m (ANormal ref v)
forall a b. (a -> b) -> a -> b
$ Set v -> OptInfos ref v -> ANormal ref v -> m (ANormal ref v)
forall ref (m :: * -> *) v.
(Ord ref, Memo m, Var v) =>
Set v -> OptInfos ref v -> ANormal ref v -> m (ANormal ref v)
inline Set v
avoid OptInfos ref v
opts ANormal ref v
n
    peep :: ANormal ref v -> m (ANormal ref v)
peep ANormal ref v
n = (ANormal ref v -> m (ANormal ref v))
-> m (ANormal ref v) -> m (ANormal ref v)
forall (m :: * -> *) a. Memo m => (a -> m a) -> m a -> m a
whenChanged ANormal ref v -> m (ANormal ref v)
inl (m (ANormal ref v) -> m (ANormal ref v))
-> m (ANormal ref v) -> m (ANormal ref v)
forall a b. (a -> b) -> a -> b
$ Arities ref -> Bool -> ANormal ref v -> m (ANormal ref v)
forall ref (m :: * -> *) v.
(Ord ref, Memo m, Var v) =>
Arities ref -> Bool -> ANormal ref v -> m (ANormal ref v)
peephole Arities ref
arities Bool
affine ANormal ref v
n

-- Optimizes a single definition
optSuper ::
  (Ord ref, Memo m, Var v) =>
  OptInfos ref v ->
  Set v ->
  Bool ->
  SuperNormal ref v ->
  m (SuperNormal ref v)
optSuper :: forall ref (m :: * -> *) v.
(Ord ref, Memo m, Var v) =>
OptInfos ref v
-> Set v -> Bool -> SuperNormal ref v -> m (SuperNormal ref v)
optSuper OptInfos ref v
opts Set v
avoid0 Bool
affine sn :: SuperNormal ref v
sn@(Lambda [Mem]
ccs (ABTN.TAbss [v]
vs Term (ANormalF ref) v
bd)) =
  SuperNormal ref v -> m (SuperNormal ref v) -> m (SuperNormal ref v)
forall (m :: * -> *) a. Memo m => a -> m a -> m a
memo SuperNormal ref v
sn (m (SuperNormal ref v) -> m (SuperNormal ref v))
-> m (SuperNormal ref v) -> m (SuperNormal ref v)
forall a b. (a -> b) -> a -> b
$
    [Mem] -> Term (ANormalF ref) v -> SuperNormal ref v
forall ref v. [Mem] -> ANormal ref v -> SuperNormal ref v
Lambda [Mem]
ccs (Term (ANormalF ref) v -> SuperNormal ref v)
-> (Term (ANormalF ref) v -> Term (ANormalF ref) v)
-> Term (ANormalF ref) v
-> SuperNormal ref v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [v] -> Term (ANormalF ref) v -> Term (ANormalF ref) v
forall v (f :: * -> * -> *). Var v => [v] -> Term f v -> Term f v
ABTN.TAbss [v]
vs (Term (ANormalF ref) v -> SuperNormal ref v)
-> m (Term (ANormalF ref) v) -> m (SuperNormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OptInfos ref v
-> Set v
-> Bool
-> Term (ANormalF ref) v
-> m (Term (ANormalF ref) v)
forall ref (m :: * -> *) v.
(Ord ref, Memo m, Var v) =>
OptInfos ref v
-> Set v -> Bool -> ANormal ref v -> m (ANormal ref v)
optNormal OptInfos ref v
opts Set v
avoid Bool
affine Term (ANormalF ref) v
bd
  where
    avoid :: Set v
avoid = Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
Set.union ([v] -> Set v
forall a. Ord a => [a] -> Set a
Set.fromList [v]
vs) Set v
avoid0

-- Optimizes a single group
optGroup ::
  (Ord ref, Var v) =>
  OptInfos ref v ->
  ref ->
  SuperGroup ref v ->
  SuperGroup ref v
optGroup :: forall ref v.
(Ord ref, Var v) =>
OptInfos ref v -> ref -> SuperGroup ref v -> SuperGroup ref v
optGroup (Arities ref
arities, InlineInfos ref v
inls0) ref
self grp :: SuperGroup ref v
grp@(Rec [(v, SuperNormal ref v)]
bs SuperNormal ref v
en) =
  Writer Any (SuperGroup ref v) -> SuperGroup ref v
forall a. Writer Any a -> a
runMemo (Writer Any (SuperGroup ref v) -> SuperGroup ref v)
-> (Writer Any (SuperGroup ref v) -> Writer Any (SuperGroup ref v))
-> Writer Any (SuperGroup ref v)
-> SuperGroup ref v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SuperGroup ref v
-> Writer Any (SuperGroup ref v) -> Writer Any (SuperGroup ref v)
forall (m :: * -> *) a. Memo m => a -> m a -> m a
memo SuperGroup ref v
grp (Writer Any (SuperGroup ref v) -> SuperGroup ref v)
-> Writer Any (SuperGroup ref v) -> SuperGroup ref v
forall a b. (a -> b) -> a -> b
$
    [(v, SuperNormal ref v)] -> SuperNormal ref v -> SuperGroup ref v
forall ref v.
[(v, SuperNormal ref v)] -> SuperNormal ref v -> SuperGroup ref v
Rec
      ([(v, SuperNormal ref v)] -> SuperNormal ref v -> SuperGroup ref v)
-> WriterT Any Identity [(v, SuperNormal ref v)]
-> WriterT Any Identity (SuperNormal ref v -> SuperGroup ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (((v, SuperNormal ref v)
 -> WriterT Any Identity (v, SuperNormal ref v))
-> [(v, SuperNormal ref v)]
-> WriterT Any Identity [(v, SuperNormal ref v)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (((v, SuperNormal ref v)
  -> WriterT Any Identity (v, SuperNormal ref v))
 -> [(v, SuperNormal ref v)]
 -> WriterT Any Identity [(v, SuperNormal ref v)])
-> ((SuperNormal ref v -> WriterT Any Identity (SuperNormal ref v))
    -> (v, SuperNormal ref v)
    -> WriterT Any Identity (v, SuperNormal ref v))
-> (SuperNormal ref v -> WriterT Any Identity (SuperNormal ref v))
-> [(v, SuperNormal ref v)]
-> WriterT Any Identity [(v, SuperNormal ref v)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SuperNormal ref v -> WriterT Any Identity (SuperNormal ref v))
-> (v, SuperNormal ref v)
-> WriterT Any Identity (v, SuperNormal ref v)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> (v, a) -> f (v, b)
traverse) ((Arities ref, InlineInfos ref v)
-> Set v
-> Bool
-> SuperNormal ref v
-> WriterT Any Identity (SuperNormal ref v)
forall ref (m :: * -> *) v.
(Ord ref, Memo m, Var v) =>
OptInfos ref v
-> Set v -> Bool -> SuperNormal ref v -> m (SuperNormal ref v)
optSuper (Arities ref, InlineInfos ref v)
opts Set v
avoid Bool
False) [(v, SuperNormal ref v)]
bs
      WriterT Any Identity (SuperNormal ref v -> SuperGroup ref v)
-> WriterT Any Identity (SuperNormal ref v)
-> Writer Any (SuperGroup ref v)
forall a b.
WriterT Any Identity (a -> b)
-> WriterT Any Identity a -> WriterT Any Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Arities ref, InlineInfos ref v)
-> Set v
-> Bool
-> SuperNormal ref v
-> WriterT Any Identity (SuperNormal ref v)
forall ref (m :: * -> *) v.
(Ord ref, Memo m, Var v) =>
OptInfos ref v
-> Set v -> Bool -> SuperNormal ref v -> m (SuperNormal ref v)
optSuper (Arities ref, InlineInfos ref v)
opts Set v
avoid Bool
False SuperNormal ref v
en
  where
    avoid :: Set v
avoid = [v] -> Set v
forall a. Ord a => [a] -> Set a
Set.fromList ([v] -> Set v) -> [v] -> Set v
forall a b. (a -> b) -> a -> b
$ (v, SuperNormal ref v) -> v
forall a b. (a, b) -> a
fst ((v, SuperNormal ref v) -> v) -> [(v, SuperNormal ref v)] -> [v]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(v, SuperNormal ref v)]
bs

    opts :: (Arities ref, InlineInfos ref v)
opts = (Arities ref
arities, InlineInfos ref v
inls)
    inls :: InlineInfos ref v
inls = (InlineInfos ref v -> InlineInfos ref v)
-> (InlineInfo ref v -> InlineInfos ref v -> InlineInfos ref v)
-> Maybe (InlineInfo ref v)
-> InlineInfos ref v
-> InlineInfos ref v
forall b a. b -> (a -> b) -> Maybe a -> b
maybe InlineInfos ref v -> InlineInfos ref v
forall a. a -> a
id (ref -> InlineInfo ref v -> InlineInfos ref v -> InlineInfos ref v
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert ref
self) (SuperGroup ref v -> Maybe (InlineInfo ref v)
forall {v} {ref}.
Var v =>
SuperGroup ref v -> Maybe (InlineInfo ref v)
entryInfo SuperGroup ref v
grp) InlineInfos ref v
inls0

optimize ::
  (Ord ref, Var v) =>
  Map ref (SuperGroup ref v) ->
  OptInfos ref v ->
  (Map ref (SuperGroup ref v), OptInfos ref v)
optimize :: forall ref v.
(Ord ref, Var v) =>
Map ref (SuperGroup ref v)
-> OptInfos ref v -> (Map ref (SuperGroup ref v), OptInfos ref v)
optimize Map ref (SuperGroup ref v)
gs = State (OptInfos ref v) (Map ref (SuperGroup ref v))
-> OptInfos ref v -> (Map ref (SuperGroup ref v), OptInfos ref v)
forall s a. State s a -> s -> (a, s)
runState do
  -- add new arities
  (OptInfos ref v -> OptInfos ref v)
-> StateT (OptInfos ref v) Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((OptInfos ref v -> OptInfos ref v)
 -> StateT (OptInfos ref v) Identity ())
-> (OptInfos ref v -> OptInfos ref v)
-> StateT (OptInfos ref v) Identity ()
forall a b. (a -> b) -> a -> b
$ (Map ref Int -> Map ref Int) -> OptInfos ref v -> OptInfos ref v
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (Map ref Int -> Map ref Int -> Map ref Int
forall k a. Ord k => Map k a -> Map k a -> Map k a
Map.union (Map ref Int -> Map ref Int -> Map ref Int)
-> Map ref Int -> Map ref Int -> Map ref Int
forall a b. (a -> b) -> a -> b
$ SuperGroup ref v -> Int
forall ref v. SuperGroup ref v -> Int
arityInfo (SuperGroup ref v -> Int)
-> Map ref (SuperGroup ref v) -> Map ref Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map ref (SuperGroup ref v)
gs)
  let doOpt :: OptInfos ref v
-> (ref, SuperGroup ref v) -> (ref, SuperGroup ref v)
doOpt OptInfos ref v
opts (ref
r, SuperGroup ref v
sg) = (ref
r, OptInfos ref v -> ref -> SuperGroup ref v -> SuperGroup ref v
forall ref v.
(Ord ref, Var v) =>
OptInfos ref v -> ref -> SuperGroup ref v -> SuperGroup ref v
optGroup OptInfos ref v
opts ref
r SuperGroup ref v
sg)
  -- Note: inlining info is never available to self-inline
  [[(ref, SuperGroup ref v)]]
ngs <- [SCC (ref, SuperGroup ref v)]
-> (SCC (ref, SuperGroup ref v)
    -> StateT (OptInfos ref v) Identity [(ref, SuperGroup ref v)])
-> StateT (OptInfos ref v) Identity [[(ref, SuperGroup ref v)]]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [SCC (ref, SuperGroup ref v)]
sccs \case
    AcyclicSCC (ref, SuperGroup ref v)
p -> do
      OptInfos ref v
opts <- StateT (OptInfos ref v) Identity (OptInfos ref v)
forall s (m :: * -> *). MonadState s m => m s
get
      let rgs :: [(ref, SuperGroup ref v)]
rgs = [OptInfos ref v
-> (ref, SuperGroup ref v) -> (ref, SuperGroup ref v)
forall {v} {ref}.
(Var v, Ord ref) =>
OptInfos ref v
-> (ref, SuperGroup ref v) -> (ref, SuperGroup ref v)
doOpt OptInfos ref v
opts (ref, SuperGroup ref v)
p]
      [(ref, SuperGroup ref v)]
rgs [(ref, SuperGroup ref v)]
-> StateT (OptInfos ref v) Identity ()
-> StateT (OptInfos ref v) Identity [(ref, SuperGroup ref v)]
forall a b.
a
-> StateT (OptInfos ref v) Identity b
-> StateT (OptInfos ref v) Identity a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Bool
-> [(ref, SuperGroup ref v)] -> StateT (OptInfos ref v) Identity ()
forall {t :: * -> *} {f :: * -> *} {p :: * -> * -> *} {a} {k} {ref}
       {v}.
(Foldable t, MonadState (p a (Map k (InlineInfo ref v))) f,
 Bifunctor p, Var v, Ord k) =>
Bool -> t (k, SuperGroup ref v) -> f ()
addInls Bool
False [(ref, SuperGroup ref v)]
rgs
    CyclicSCC [(ref, SuperGroup ref v)]
rgs0 -> do
      OptInfos ref v
opts <- [(ref, SuperGroup ref v)] -> OptInfos ref v -> OptInfos ref v
forall {v} {k} {a}.
(Var v, Ord k) =>
[(k, SuperGroup k v)]
-> (a, Map k (InlineInfo k v)) -> (a, Map k (InlineInfo k v))
augment [(ref, SuperGroup ref v)]
rgs0 (OptInfos ref v -> OptInfos ref v)
-> StateT (OptInfos ref v) Identity (OptInfos ref v)
-> StateT (OptInfos ref v) Identity (OptInfos ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT (OptInfos ref v) Identity (OptInfos ref v)
forall s (m :: * -> *). MonadState s m => m s
get
      let rgs :: [(ref, SuperGroup ref v)]
rgs = OptInfos ref v
-> (ref, SuperGroup ref v) -> (ref, SuperGroup ref v)
forall {v} {ref}.
(Var v, Ord ref) =>
OptInfos ref v
-> (ref, SuperGroup ref v) -> (ref, SuperGroup ref v)
doOpt OptInfos ref v
opts ((ref, SuperGroup ref v) -> (ref, SuperGroup ref v))
-> [(ref, SuperGroup ref v)] -> [(ref, SuperGroup ref v)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(ref, SuperGroup ref v)]
rgs0
      [(ref, SuperGroup ref v)]
rgs [(ref, SuperGroup ref v)]
-> StateT (OptInfos ref v) Identity ()
-> StateT (OptInfos ref v) Identity [(ref, SuperGroup ref v)]
forall a b.
a
-> StateT (OptInfos ref v) Identity b
-> StateT (OptInfos ref v) Identity a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Bool
-> [(ref, SuperGroup ref v)] -> StateT (OptInfos ref v) Identity ()
forall {t :: * -> *} {f :: * -> *} {p :: * -> * -> *} {a} {k} {ref}
       {v}.
(Foldable t, MonadState (p a (Map k (InlineInfo ref v))) f,
 Bifunctor p, Var v, Ord k) =>
Bool -> t (k, SuperGroup ref v) -> f ()
addInls Bool
True [(ref, SuperGroup ref v)]
rgs
  Map ref (SuperGroup ref v)
-> State (OptInfos ref v) (Map ref (SuperGroup ref v))
forall a. a -> StateT (OptInfos ref v) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Map ref (SuperGroup ref v)
 -> State (OptInfos ref v) (Map ref (SuperGroup ref v)))
-> ([(ref, SuperGroup ref v)] -> Map ref (SuperGroup ref v))
-> [(ref, SuperGroup ref v)]
-> State (OptInfos ref v) (Map ref (SuperGroup ref v))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(ref, SuperGroup ref v)] -> Map ref (SuperGroup ref v)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(ref, SuperGroup ref v)]
 -> State (OptInfos ref v) (Map ref (SuperGroup ref v)))
-> [(ref, SuperGroup ref v)]
-> State (OptInfos ref v) (Map ref (SuperGroup ref v))
forall a b. (a -> b) -> a -> b
$ [[(ref, SuperGroup ref v)]] -> [(ref, SuperGroup ref v)]
forall m. Monoid m => [m] -> m
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold [[(ref, SuperGroup ref v)]]
ngs
  where
    f :: (b, SuperGroup ref v) -> ((b, SuperGroup ref v), b, [ref])
f p :: (b, SuperGroup ref v)
p@(b
r, SuperGroup ref v
sg) = ((b, SuperGroup ref v)
p, b
r, SuperGroup ref v -> [ref]
forall ref v. (Ord ref, Var v) => SuperGroup ref v -> [ref]
groupTermLinks SuperGroup ref v
sg)

    -- Process code in dependency order.
    sccs :: [SCC (ref, SuperGroup ref v)]
sccs = [((ref, SuperGroup ref v), ref, [ref])]
-> [SCC (ref, SuperGroup ref v)]
forall key node. Ord key => [(node, key, [key])] -> [SCC node]
stronglyConnComp ([((ref, SuperGroup ref v), ref, [ref])]
 -> [SCC (ref, SuperGroup ref v)])
-> ([(ref, SuperGroup ref v)]
    -> [((ref, SuperGroup ref v), ref, [ref])])
-> [(ref, SuperGroup ref v)]
-> [SCC (ref, SuperGroup ref v)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((ref, SuperGroup ref v) -> ((ref, SuperGroup ref v), ref, [ref]))
-> [(ref, SuperGroup ref v)]
-> [((ref, SuperGroup ref v), ref, [ref])]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ref, SuperGroup ref v) -> ((ref, SuperGroup ref v), ref, [ref])
forall {v} {ref} {b}.
(Var v, Ord ref) =>
(b, SuperGroup ref v) -> ((b, SuperGroup ref v), b, [ref])
f ([(ref, SuperGroup ref v)] -> [SCC (ref, SuperGroup ref v)])
-> [(ref, SuperGroup ref v)] -> [SCC (ref, SuperGroup ref v)]
forall a b. (a -> b) -> a -> b
$ Map ref (SuperGroup ref v) -> [(ref, SuperGroup ref v)]
forall k a. Map k a -> [(k, a)]
Map.toList Map ref (SuperGroup ref v)
gs

    addInls :: Bool -> t (k, SuperGroup ref v) -> f ()
addInls Bool
rec t (k, SuperGroup ref v)
rgs =
      t (k, SuperGroup ref v) -> ((k, SuperGroup ref v) -> f ()) -> f ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ t (k, SuperGroup ref v)
rgs \(k
r, SuperGroup ref v
sg) -> do
        (p a (Map k (InlineInfo ref v)) -> p a (Map k (InlineInfo ref v)))
-> f ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((p a (Map k (InlineInfo ref v)) -> p a (Map k (InlineInfo ref v)))
 -> f ())
-> (Maybe (InlineInfo ref v)
    -> p a (Map k (InlineInfo ref v))
    -> p a (Map k (InlineInfo ref v)))
-> Maybe (InlineInfo ref v)
-> f ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map k (InlineInfo ref v) -> Map k (InlineInfo ref v))
-> p a (Map k (InlineInfo ref v)) -> p a (Map k (InlineInfo ref v))
forall b c a. (b -> c) -> p a b -> p a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ((Map k (InlineInfo ref v) -> Map k (InlineInfo ref v))
 -> p a (Map k (InlineInfo ref v))
 -> p a (Map k (InlineInfo ref v)))
-> (Maybe (InlineInfo ref v)
    -> Map k (InlineInfo ref v) -> Map k (InlineInfo ref v))
-> Maybe (InlineInfo ref v)
-> p a (Map k (InlineInfo ref v))
-> p a (Map k (InlineInfo ref v))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map k (InlineInfo ref v) -> Map k (InlineInfo ref v))
-> (InlineInfo ref v
    -> Map k (InlineInfo ref v) -> Map k (InlineInfo ref v))
-> Maybe (InlineInfo ref v)
-> Map k (InlineInfo ref v)
-> Map k (InlineInfo ref v)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Map k (InlineInfo ref v) -> Map k (InlineInfo ref v)
forall a. a -> a
id (k
-> InlineInfo ref v
-> Map k (InlineInfo ref v)
-> Map k (InlineInfo ref v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
r) (Maybe (InlineInfo ref v) -> f ())
-> Maybe (InlineInfo ref v) -> f ()
forall a b. (a -> b) -> a -> b
$ Bool -> SuperGroup ref v -> Maybe (InlineInfo ref v)
forall v ref.
Var v =>
Bool -> SuperGroup ref v -> Maybe (InlineInfo ref v)
inlineInfo Bool
rec SuperGroup ref v
sg

    augment :: [(k, SuperGroup k v)]
-> (a, Map k (InlineInfo k v)) -> (a, Map k (InlineInfo k v))
augment [(k, SuperGroup k v)]
rgs (a
ars, Map k (InlineInfo k v)
inls) =
      (a
ars, Map k (SuperGroup k v) -> Map k (InlineInfo k v)
forall ref v.
(Ord ref, Var v) =>
Map ref (SuperGroup ref v) -> InlineInfos ref v
recInlineInfo ([(k, SuperGroup k v)] -> Map k (SuperGroup k v)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(k, SuperGroup k v)]
rgs) Map k (InlineInfo k v)
-> Map k (InlineInfo k v) -> Map k (InlineInfo k v)
forall k a. Ord k => Map k a -> Map k a -> Map k a
`Map.union` Map k (InlineInfo k v)
inls)

-- Does further optimizations on the matching portion of a handler in
-- preparation for recognizing affine handlers. Some optimizations are
-- not safe to do in general, but are in an affine context, and make
-- it possible to recognize more affine handlers (as well as making
-- them perform better). So, we apply these optimizations before doing
-- the affine translation.
affinePreOptimize ::
  (Ord ref, Var v) =>
  OptInfos ref v ->
  Set v ->
  SuperNormal ref v ->
  SuperNormal ref v
affinePreOptimize :: forall ref v.
(Ord ref, Var v) =>
OptInfos ref v -> Set v -> SuperNormal ref v -> SuperNormal ref v
affinePreOptimize OptInfos ref v
opts Set v
avoid =
  Writer Any (SuperNormal ref v) -> SuperNormal ref v
forall a. Writer Any a -> a
runMemo (Writer Any (SuperNormal ref v) -> SuperNormal ref v)
-> (SuperNormal ref v -> Writer Any (SuperNormal ref v))
-> SuperNormal ref v
-> SuperNormal ref v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OptInfos ref v
-> Set v
-> Bool
-> SuperNormal ref v
-> Writer Any (SuperNormal ref v)
forall ref (m :: * -> *) v.
(Ord ref, Memo m, Var v) =>
OptInfos ref v
-> Set v -> Bool -> SuperNormal ref v -> m (SuperNormal ref v)
optSuper (((Map ref (InlineInfo ref v) -> Map ref (InlineInfo ref v))
-> OptInfos ref v -> OptInfos ref v
forall a b. (a -> b) -> (Arities ref, a) -> (Arities ref, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Map ref (InlineInfo ref v) -> Map ref (InlineInfo ref v))
 -> OptInfos ref v -> OptInfos ref v)
-> ((InlineInfo ref v -> InlineInfo ref v)
    -> Map ref (InlineInfo ref v) -> Map ref (InlineInfo ref v))
-> (InlineInfo ref v -> InlineInfo ref v)
-> OptInfos ref v
-> OptInfos ref v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (InlineInfo ref v -> InlineInfo ref v)
-> Map ref (InlineInfo ref v) -> Map ref (InlineInfo ref v)
forall a b. (a -> b) -> Map ref a -> Map ref b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) InlineInfo ref v -> InlineInfo ref v
forall {v} {ref}. Var v => InlineInfo ref v -> InlineInfo ref v
reclassify OptInfos ref v
opts) Set v
avoid Bool
True
  where
    -- reclassifies certain inlinings that would be unsafe in general,
    -- but are fine in an affine context.
    reclassify :: InlineInfo ref v -> InlineInfo ref v
reclassify (InlInfo InlineClass
Don'tInl ANormal ref v
bd)
      | ANormal ref v -> Bool
forall v ref. Var v => ANormal ref v -> Bool
isExtendedHandlerApp ANormal ref v
bd = InlineClass -> ANormal ref v -> InlineInfo ref v
forall ref v. InlineClass -> ANormal ref v -> InlineInfo ref v
InlInfo InlineClass
TailInl ANormal ref v
bd
    reclassify InlineInfo ref v
info = InlineInfo ref v
info

-- Applies optimizations one more time to the rewritten affine
-- handler, since some of the transformations may introduce some new
-- opportunities.
affinePostOptimize ::
  (Ord ref, Var v) =>
  OptInfos ref v ->
  Set v ->
  SuperNormal ref v ->
  SuperNormal ref v
affinePostOptimize :: forall ref v.
(Ord ref, Var v) =>
OptInfos ref v -> Set v -> SuperNormal ref v -> SuperNormal ref v
affinePostOptimize OptInfos ref v
opts Set v
avoid = Writer Any (SuperNormal ref v) -> SuperNormal ref v
forall a. Writer Any a -> a
runMemo (Writer Any (SuperNormal ref v) -> SuperNormal ref v)
-> (SuperNormal ref v -> Writer Any (SuperNormal ref v))
-> SuperNormal ref v
-> SuperNormal ref v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OptInfos ref v
-> Set v
-> Bool
-> SuperNormal ref v
-> Writer Any (SuperNormal ref v)
forall ref (m :: * -> *) v.
(Ord ref, Memo m, Var v) =>
OptInfos ref v
-> Set v -> Bool -> SuperNormal ref v -> m (SuperNormal ref v)
optSuper OptInfos ref v
opts Set v
avoid Bool
True

effectless :: (Var v) => ANormal ref v -> Bool
effectless :: forall v ref. Var v => ANormal ref v -> Bool
effectless (TCon {}) = Bool
True
effectless (TLit {}) = Bool
True
effectless (TBLit {}) = Bool
True
effectless (TVar {}) = Bool
True
effectless ANormal ref v
_ = Bool
False

-- Determines whether a term cannot possibly depend on the stack
-- contents. Useful if you want to decide whether you can eliminate an
-- unused binding, for example (because if you can only do so if it
-- not being on the stack is okay).
stackInsensitive :: (Var v) => ANormal ref v -> Bool
stackInsensitive :: forall v ref. Var v => ANormal ref v -> Bool
stackInsensitive = Bool -> ANormal ref v -> Bool
forall {v} {ref}. Var v => Bool -> ANormal ref v -> Bool
go Bool
True
  where
    go :: Bool -> ANormal ref v -> Bool
go Bool
tail = \case
      TCon {} -> Bool
True
      TLit {} -> Bool
True
      TBLit {} -> Bool
True
      TVar {} -> Bool
True
      TPrm {} -> Bool
True
      TName v
_ Either ref v
_ [v]
_ ANormal ref v
bd -> Bool -> ANormal ref v -> Bool
go Bool
tail ANormal ref v
bd
      TLets Direction Word16
_ [v]
_ [Mem]
_ ANormal ref v
bn ANormal ref v
bd -> Bool -> ANormal ref v -> Bool
go Bool
False ANormal ref v
bn Bool -> Bool -> Bool
&& Bool -> ANormal ref v -> Bool
go Bool
tail ANormal ref v
bd
      TApv {} -> Bool
tail
      TCom {} -> Bool
tail
      TMatch v
_ Branched ref (ANormal ref v)
bs -> (ANormal ref v -> Bool) -> Branched ref (ANormal ref v) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Bool -> ANormal ref v -> Bool
go Bool
tail) Branched ref (ANormal ref v)
bs
      ABTN.TAbs v
_ (ABTN.TAbss [v]
_ ANormal ref v
bd) -> Bool -> ANormal ref v -> Bool
go Bool
tail ANormal ref v
bd
      ANormal ref v
_ -> Bool
False

-- Recognize Func that can be used in a `TName`
nameable :: Func ref v -> Maybe (Either ref v)
nameable :: forall ref v. Func ref v -> Maybe (Either ref v)
nameable (FVar v
v) = Either ref v -> Maybe (Either ref v)
forall a. a -> Maybe a
Just (Either ref v -> Maybe (Either ref v))
-> Either ref v -> Maybe (Either ref v)
forall a b. (a -> b) -> a -> b
$ v -> Either ref v
forall a b. b -> Either a b
Right v
v
nameable (FComb ref
r) = Either ref v -> Maybe (Either ref v)
forall a. a -> Maybe a
Just (Either ref v -> Maybe (Either ref v))
-> Either ref v -> Maybe (Either ref v)
forall a b. (a -> b) -> a -> b
$ ref -> Either ref v
forall a b. a -> Either a b
Left ref
r
nameable Func ref v
_ = Maybe (Either ref v)
forall a. Maybe a
Nothing

pattern $mNameable :: forall {r} {ref} {v}.
Func ref v -> (Either ref v -> r) -> ((# #) -> r) -> r
$bNameable :: forall {ref} {v}. Either ref v -> Func ref v
Nameable e <- (nameable -> Just e)
  where
    Nameable Either ref v
e = (ref -> Func ref v)
-> (v -> Func ref v) -> Either ref v -> Func ref v
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either ref -> Func ref v
forall ref v. ref -> Func ref v
FComb v -> Func ref v
forall ref v. v -> Func ref v
FVar Either ref v
e

-- Recognize a handler call with a given delayed value and handler for
-- it, yielding the specified handled references
handlerResumption ::
  (Var v) =>
  v ->
  v ->
  ANormal ref v ->
  Maybe [ref]
handlerResumption :: forall v ref. Var v => v -> v -> ANormal ref v -> Maybe [ref]
handlerResumption v
lz0 v
lh0 (THnd [ref]
rs v
lh1 Maybe v
Nothing (TFrc v
lz1)) =
  [ref]
rs [ref] -> Maybe () -> Maybe [ref]
forall a b. a -> Maybe b -> Maybe a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (v
lz0 v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
lz1 Bool -> Bool -> Bool
&& v
lh0 v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
lh1)
handlerResumption v
_ v
_ ANormal ref v
_ = Maybe [ref]
forall a. Maybe a
Nothing

-- Tail of a handler implementation resuming itself. This is the
-- unoptimized version from initial code generation.
--
--   lazy lz := f as
--   lazy lh := h bs
--   handle{rs} !lz with lh
pattern $mHandlerResume :: forall {r} {v} {ref}.
Var v =>
ANormal ref v
-> (v
    -> Either ref v -> [v] -> v -> Either ref v -> [v] -> [ref] -> r)
-> ((# #) -> r)
-> r
HandlerResume lz f as lh h bs rs <-
  TName lz f as (TName lh h bs (handlerResumption lz lh -> Just rs))

-- Recognize a possibly underapplied combinator that is used at the
-- end of a handler block. The result is the combinator reference, the
-- length of the variables it's applied to, and a rewritten expression
-- that inlines the application to the single use site, with in
-- indication of whether it's safe to use the rewritten expression in
-- non-affine contexts.
matchHandledThunk ::
  (Ord ref, Var v) =>
  ANormal ref v ->
  Maybe (ref, Int, Bool, ANormal ref v)
matchHandledThunk :: forall ref v.
(Ord ref, Var v) =>
ANormal ref v -> Maybe (ref, Int, Bool, ANormal ref v)
matchHandledThunk (TLet Direction Word16
_ v
th Mem
_ (TCom ref
r [v]
vs) ANormal ref v
bd) =
  (ANormal ref v, All) -> (ref, Int, Bool, ANormal ref v)
final ((ANormal ref v, All) -> (ref, Int, Bool, ANormal ref v))
-> Maybe (ANormal ref v, All)
-> Maybe (ref, Int, Bool, ANormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> WriterT All Maybe (ANormal ref v) -> Maybe (ANormal ref v, All)
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (Renaming v -> ANormal ref v -> WriterT All Maybe (ANormal ref v)
prefix (Set v -> Renaming v
forall v. Set v -> Renaming v
ABTN.avoiding (Set v -> Renaming v) -> ([v] -> Set v) -> [v] -> Renaming v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [v] -> Set v
forall a. Ord a => [a] -> Set a
Set.fromList ([v] -> Renaming v) -> [v] -> Renaming v
forall a b. (a -> b) -> a -> b
$ v
th v -> [v] -> [v]
forall a. a -> [a] -> [a]
: [v]
vs) ANormal ref v
bd)
  where
    none :: WriterT w Maybe a
none = Maybe (a, w) -> WriterT w Maybe a
forall w (m :: * -> *) a. m (a, w) -> WriterT w m a
WriterT Maybe (a, w)
forall a. Maybe a
Nothing

    final :: (ANormal ref v, All) -> (ref, Int, Bool, ANormal ref v)
final (ANormal ref v
bd, All Bool
lazy) = (ref
r, [v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
vs, Bool
lazy, ANormal ref v
bd)

    prefix :: Renaming v -> ANormal ref v -> WriterT All Maybe (ANormal ref v)
prefix Renaming v
rn (TName v
v Either ref v
g [v]
bs ANormal ref v
bd)
      | v
v v -> v -> Bool
forall a. Eq a => a -> a -> Bool
/= v
th,
        (v -> Bool) -> Either ref v -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (v -> v -> Bool
forall a. Eq a => a -> a -> Bool
/= v
th) Either ref v
g,
        (v -> Bool) -> [v] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (v -> v -> Bool
forall a. Eq a => a -> a -> Bool
/= v
th) [v]
bs,
        Either ref v
g <- Renaming v -> v -> v
forall v. Var v => Renaming v -> v -> v
ABTN.renameVar Renaming v
rn (v -> v) -> Either ref v -> Either ref v
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Either ref v
g,
        [v]
bs <- Renaming v -> v -> v
forall v. Var v => Renaming v -> v -> v
ABTN.renameVar Renaming v
rn (v -> v) -> [v] -> [v]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [v]
bs,
        (Renaming v
rn, v
v) <- Set v -> Renaming v -> v -> (Renaming v, v)
forall v. Var v => Set v -> Renaming v -> v -> (Renaming v, v)
ABTN.freshenBinder (ANormal ref v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
ABTN.freeVars ANormal ref v
bd) Renaming v
rn v
v =
          v -> Either ref v -> [v] -> ANormal ref v -> ANormal ref v
forall v ref.
Var v =>
v -> Either ref v -> [v] -> ANormal ref v -> ANormal ref v
TName v
v Either ref v
g [v]
bs (ANormal ref v -> ANormal ref v)
-> WriterT All Maybe (ANormal ref v)
-> WriterT All Maybe (ANormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Renaming v -> ANormal ref v -> WriterT All Maybe (ANormal ref v)
prefix Renaming v
rn ANormal ref v
bd
    prefix Renaming v
rn (TLets Direction Word16
d [v]
vs [Mem]
ccs ANormal ref v
bn ANormal ref v
bd)
      | (v -> Bool) -> [v] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (v -> v -> Bool
forall a. Eq a => a -> a -> Bool
/= v
th) [v]
vs,
        v
th v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.notMember` ANormal ref v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
ABTN.freeVars ANormal ref v
bn,
        ANormal ref v
bn <- Renaming v -> ANormal ref v -> ANormal ref v
forall v (f :: * -> * -> *).
(Var v, Bifunctor f, Bifoldable f) =>
Renaming v -> Term f v -> Term f v
ABTN.renamesAndFreshen0 Renaming v
rn ANormal ref v
bn,
        (Renaming v
rn, [v]
vs) <- Set v -> Renaming v -> [v] -> (Renaming v, [v])
forall v. Var v => Set v -> Renaming v -> [v] -> (Renaming v, [v])
ABTN.freshenBinders (ANormal ref v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
ABTN.freeVars ANormal ref v
bd) Renaming v
rn [v]
vs =
          Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
forall v ref.
Var v =>
Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
TLets Direction Word16
d [v]
vs [Mem]
ccs ANormal ref v
bn (ANormal ref v -> ANormal ref v)
-> WriterT All Maybe (ANormal ref v)
-> WriterT All Maybe (ANormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Renaming v -> ANormal ref v -> WriterT All Maybe (ANormal ref v)
prefix Renaming v
rn ANormal ref v
bd WriterT All Maybe (ANormal ref v)
-> WriterT All Maybe () -> WriterT All Maybe (ANormal ref v)
forall a b.
WriterT All Maybe a -> WriterT All Maybe b -> WriterT All Maybe a
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* All -> WriterT All Maybe ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (Bool -> All
All Bool
False)
    prefix Renaming v
rn (TMatch v
sc Branched ref (ANormal ref v)
bs) =
      v -> Branched ref (ANormal ref v) -> ANormal ref v
forall v ref.
Var v =>
v -> Branched ref (ANormal ref v) -> ANormal ref v
TMatch (Renaming v -> v -> v
forall v. Var v => Renaming v -> v -> v
ABTN.renameVar Renaming v
rn v
sc) (Branched ref (ANormal ref v) -> ANormal ref v)
-> WriterT All Maybe (Branched ref (ANormal ref v))
-> WriterT All Maybe (ANormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ANormal ref v -> WriterT All Maybe (ANormal ref v))
-> Branched ref (ANormal ref v)
-> WriterT All Maybe (Branched ref (ANormal ref v))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Branched ref a -> f (Branched ref b)
traverse ANormal ref v -> WriterT All Maybe (ANormal ref v)
under Branched ref (ANormal ref v)
bs
      where
        under :: ANormal ref v -> WriterT All Maybe (ANormal ref v)
under (ABTN.TAbss [v]
vs ANormal ref v
bd)
          | (v -> Bool) -> [v] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (v -> v -> Bool
forall a. Eq a => a -> a -> Bool
/= v
th) [v]
vs,
            (Renaming v
rn, [v]
vs) <- Set v -> Renaming v -> [v] -> (Renaming v, [v])
forall v. Var v => Set v -> Renaming v -> [v] -> (Renaming v, [v])
ABTN.freshenBinders (ANormal ref v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
ABTN.freeVars ANormal ref v
bd) Renaming v
rn [v]
vs =
              [v] -> ANormal ref v -> ANormal ref v
forall v (f :: * -> * -> *). Var v => [v] -> Term f v -> Term f v
ABTN.TAbss [v]
vs (ANormal ref v -> ANormal ref v)
-> WriterT All Maybe (ANormal ref v)
-> WriterT All Maybe (ANormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Renaming v -> ANormal ref v -> WriterT All Maybe (ANormal ref v)
prefix Renaming v
rn ANormal ref v
bd
        under ANormal ref v
_ = WriterT All Maybe (ANormal ref v)
forall {w} {a}. WriterT w Maybe a
none
    prefix Renaming v
rn (THnd [ref]
rs v
nh Maybe v
ah ANormal ref v
bd)
      | v
nh v -> v -> Bool
forall a. Eq a => a -> a -> Bool
/= v
th,
        (v -> Bool) -> Maybe v -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (v -> v -> Bool
forall a. Eq a => a -> a -> Bool
/= v
th) Maybe v
ah,
        Maybe v
ah <- Renaming v -> v -> v
forall v. Var v => Renaming v -> v -> v
ABTN.renameVar Renaming v
rn (v -> v) -> Maybe v -> Maybe v
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe v
ah,
        v
nh <- Renaming v -> v -> v
forall v. Var v => Renaming v -> v -> v
ABTN.renameVar Renaming v
rn v
nh =
          [ref] -> v -> Maybe v -> ANormal ref v -> ANormal ref v
forall v ref.
Var v =>
[ref] -> v -> Maybe v -> ANormal ref v -> ANormal ref v
THnd [ref]
rs v
nh Maybe v
ah (ANormal ref v -> ANormal ref v)
-> WriterT All Maybe (ANormal ref v)
-> WriterT All Maybe (ANormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Renaming v -> ANormal ref v -> WriterT All Maybe (ANormal ref v)
suffix Renaming v
rn ANormal ref v
bd
    prefix Renaming v
_ ANormal ref v
_ = WriterT All Maybe (ANormal ref v)
forall {w} {a}. WriterT w Maybe a
none

    -- Some values may be bound before the thunk call as long as
    -- they're 'direct' calls that can't capture stacks and reveal
    -- that we've changed the convention.
    suffix :: Renaming v -> ANormal ref v -> WriterT All Maybe (ANormal ref v)
suffix Renaming v
rn (TLets Direction Word16
d [v]
vs [Mem]
ccs ANormal ref v
bn ANormal ref v
bd)
      | (v -> Bool) -> [v] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (v -> v -> Bool
forall a. Eq a => a -> a -> Bool
/= v
th) [v]
vs,
        v
th v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.notMember` ANormal ref v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
ABTN.freeVars ANormal ref v
bn,
        ANormal ref v
bn <- Renaming v -> ANormal ref v -> ANormal ref v
forall v (f :: * -> * -> *).
(Var v, Bifunctor f, Bifoldable f) =>
Renaming v -> Term f v -> Term f v
ABTN.renamesAndFreshen0 Renaming v
rn ANormal ref v
bn,
        (Renaming v
rn, [v]
vs) <- Set v -> Renaming v -> [v] -> (Renaming v, [v])
forall v. Var v => Set v -> Renaming v -> [v] -> (Renaming v, [v])
ABTN.freshenBinders (ANormal ref v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
ABTN.freeVars ANormal ref v
bd) Renaming v
rn [v]
vs =
          Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
forall v ref.
Var v =>
Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
TLets Direction Word16
d [v]
vs [Mem]
ccs ANormal ref v
bn (ANormal ref v -> ANormal ref v)
-> WriterT All Maybe (ANormal ref v)
-> WriterT All Maybe (ANormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Renaming v -> ANormal ref v -> WriterT All Maybe (ANormal ref v)
suffix Renaming v
rn ANormal ref v
bd WriterT All Maybe (ANormal ref v)
-> WriterT All Maybe () -> WriterT All Maybe (ANormal ref v)
forall a b.
WriterT All Maybe a -> WriterT All Maybe b -> WriterT All Maybe a
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* All -> WriterT All Maybe ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (Bool -> All
All (Bool -> All) -> Bool -> All
forall a b. (a -> b) -> a -> b
$ Direction Word16
d Direction Word16 -> Direction Word16 -> Bool
forall a. Eq a => a -> a -> Bool
== Direction Word16
forall a. Direction a
Direct)
    -- final expression in handle body is a call to the thunk.
    suffix Renaming v
rn (TApv v
h [v]
us)
      | v
h v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
th,
        (v -> Bool) -> [v] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (v -> v -> Bool
forall a. Eq a => a -> a -> Bool
/= v
th) [v]
us =
          ANormal ref v -> WriterT All Maybe (ANormal ref v)
forall a. a -> WriterT All Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ANormal ref v -> WriterT All Maybe (ANormal ref v))
-> ([v] -> ANormal ref v)
-> [v]
-> WriterT All Maybe (ANormal ref v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ref -> [v] -> ANormal ref v
forall v ref. Var v => ref -> [v] -> ANormal ref v
TCom ref
r ([v] -> WriterT All Maybe (ANormal ref v))
-> [v] -> WriterT All Maybe (ANormal ref v)
forall a b. (a -> b) -> a -> b
$ [v]
vs [v] -> [v] -> [v]
forall a. [a] -> [a] -> [a]
++ (v -> v) -> [v] -> [v]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Renaming v -> v -> v
forall v. Var v => Renaming v -> v -> v
ABTN.renameVar Renaming v
rn) [v]
us
    suffix Renaming v
_ ANormal ref v
_ = WriterT All Maybe (ANormal ref v)
forall {w} {a}. WriterT w Maybe a
none
matchHandledThunk ANormal ref v
_ = Maybe (ref, Int, Bool, ANormal ref v)
forall a. Maybe a
Nothing

--  th = f <vs> -- undersaturated
--  v = ...
--  ...
--  handle
--    w = Con ...
--    x = Lit
--    ...
--    th w x ...
--  with ...
--
--  ==>
--
--  v = ...
--  ...
--  handle
--    ...
--    f <vs> w x ...
--  with ...
--
-- `safe` indicates whether the rewritten expression is safe to use in
-- arbitrary contexts, or whether it should only be applied in affine
-- contexts where nothing in it might observe odd stack behavior.
pattern $mHandledThunk :: forall {r} {v} {ref}.
(Var v, Ord ref) =>
ANormal ref v
-> (ref -> Int -> Bool -> ANormal ref v -> r) -> ((# #) -> r) -> r
HandledThunk ref ar safe expr <-
  (matchHandledThunk -> Just (ref, ar, safe, expr))

-- Builds a basic optimization map. Assumes the code in question is
-- not recursive, and makes no effort to optimize the code, so it
-- should be used only for something like builtins.
buildOptInfos ::
  (Var v) =>
  Map ref (SuperGroup ref v) ->
  OptInfos ref v
buildOptInfos :: forall v ref. Var v => Map ref (SuperGroup ref v) -> OptInfos ref v
buildOptInfos Map ref (SuperGroup ref v)
sgs =
  (SuperGroup ref v -> Int
forall ref v. SuperGroup ref v -> Int
arityInfo (SuperGroup ref v -> Int)
-> Map ref (SuperGroup ref v) -> Map ref Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map ref (SuperGroup ref v)
sgs, (SuperGroup ref v -> Maybe (InlineInfo ref v))
-> Map ref (SuperGroup ref v) -> Map ref (InlineInfo ref v)
forall u v k. (u -> Maybe v) -> Map k u -> Map k v
mapMapMaybe (Bool -> SuperGroup ref v -> Maybe (InlineInfo ref v)
forall v ref.
Var v =>
Bool -> SuperGroup ref v -> Maybe (InlineInfo ref v)
inlineInfo Bool
False) Map ref (SuperGroup ref v)
sgs)

mapMapMaybe :: (u -> Maybe v) -> Map k u -> Map k v
mapMapMaybe :: forall u v k. (u -> Maybe v) -> Map k u -> Map k v
mapMapMaybe u -> Maybe v
f = Identity (Map k v) -> Map k v
forall a. Identity a -> a
runIdentity (Identity (Map k v) -> Map k v)
-> (Map k u -> Identity (Map k v)) -> Map k u -> Map k v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (k -> u -> Identity (Maybe v)) -> Map k u -> Identity (Map k v)
forall (f :: * -> *) k a b.
Applicative f =>
(k -> a -> f (Maybe b)) -> Map k a -> f (Map k b)
Map.traverseMaybeWithKey (\k
_ -> Maybe v -> Identity (Maybe v)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe v -> Identity (Maybe v))
-> (u -> Maybe v) -> u -> Identity (Maybe v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. u -> Maybe v
f)

-- Classifies an expression with regard to inlining. Generally:
--
--   - Constants, variables and applications can be inlined anywhere,
--     because they're just replacing one call with another, or a
--     non-call
--   - More complex expressions that match or allocate extra values
--     before doing something can be tail inlined, because the
--     differences in the stack will be erased by the final jump
--   - Expressions that call multiple complex functions can't be
--     inlined easily because they add to the return-points of the
--     combinator they're inlined to.
classifyInline :: (Var v) => Bool -> ANormal ref v -> InlineClass
classifyInline :: forall v ref. Var v => Bool -> ANormal ref v -> InlineClass
classifyInline Bool
rec = \case
  -- Don't inline rec functions
  TCom ref
_ [v]
_ -> if Bool
rec then InlineClass
Don'tInl else InlineClass
AnywhereInl
  TApp {} -> InlineClass
AnywhereInl
  TBLit {} -> InlineClass
AnywhereInl
  TLit {} -> InlineClass
AnywhereInl
  TVar {} -> InlineClass
AnywhereInl
  TName v
_ Either ref v
_ [v]
_ ANormal ref v
bd -> InlineClass
TailInl InlineClass -> InlineClass -> InlineClass
forall a. Semigroup a => a -> a -> a
<> Bool -> ANormal ref v -> InlineClass
forall v ref. Var v => Bool -> ANormal ref v -> InlineClass
classifyInline Bool
rec ANormal ref v
bd
  TLets Direction Word16
Direct [v]
_ [Mem]
_ ANormal ref v
bn ANormal ref v
bd ->
    InlineClass
TailInl InlineClass -> InlineClass -> InlineClass
forall a. Semigroup a => a -> a -> a
<> Bool -> ANormal ref v -> InlineClass
forall v ref. Var v => Bool -> ANormal ref v -> InlineClass
classifyInline Bool
rec ANormal ref v
bn InlineClass -> InlineClass -> InlineClass
forall a. Semigroup a => a -> a -> a
<> Bool -> ANormal ref v -> InlineClass
forall v ref. Var v => Bool -> ANormal ref v -> InlineClass
classifyInline Bool
rec ANormal ref v
bd
  TLets {} -> InlineClass
Don'tInl
  TMatch v
_ Branched ref (ANormal ref v)
bs ->
    InlineClass
TailInl InlineClass -> InlineClass -> InlineClass
forall a. Semigroup a => a -> a -> a
<> (ANormal ref v -> InlineClass)
-> Branched ref (ANormal ref v) -> InlineClass
forall m a. Monoid m => (a -> m) -> Branched ref a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\(ABTN.TAbss [v]
_ ANormal ref v
bd) -> Bool -> ANormal ref v -> InlineClass
forall v ref. Var v => Bool -> ANormal ref v -> InlineClass
classifyInline Bool
rec ANormal ref v
bd) Branched ref (ANormal ref v)
bs
  TShift {} -> InlineClass
Don'tInl
  THnd {} -> InlineClass
Don'tInl
  TFrc {} -> InlineClass
Don'tInl
  TDiscard {} -> InlineClass
Don'tInl
  TLocal {} -> InlineClass
Don'tInl
  TUpdate {} -> InlineClass
Don'tInl
  ABTN.TAbs v
_ ANormal ref v
bd -> Bool -> ANormal ref v -> InlineClass
forall v ref. Var v => Bool -> ANormal ref v -> InlineClass
classifyInline Bool
rec ANormal ref v
bd

-- Recognizes the form resulting from certain `handle` calls in the
-- surface syntax. The recognized pattern is:
--
--   h = hh <us>
--   lazy lz = th <vs>
--   h lz
--
-- We take the opportunity to rearrange the call a bit, making it:
--
--   lazy lz = th <vs>
--   hh <us> lz
--
-- This can be used both for inlining and improving the tail of a
-- handler.
matchHandlerApp :: (Var v) => ANormal ref v -> Maybe (ANormal ref v)
matchHandlerApp :: forall v ref. Var v => ANormal ref v -> Maybe (ANormal ref v)
matchHandlerApp ANormal ref v
tm
  | TLet Direction Word16
_ v
h0 Mem
_ (TCom ref
r [v]
us) ANormal ref v
bd <- ANormal ref v
tm,
    TName v
lz0 Either ref v
th [v]
vs ANormal ref v
bd <- ANormal ref v
bd,
    TApv v
h1 [v
lz1] <- ANormal ref v
bd,
    v
h0 v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
h1,
    v
lz0 v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
lz1,
    -- binding/shadowing
    (v -> Bool) -> Either ref v -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (v -> v -> Bool
forall a. Eq a => a -> a -> Bool
/= v
h0) Either ref v
th,
    (v -> Bool) -> [v] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (v -> v -> Bool
forall a. Eq a => a -> a -> Bool
/= v
h0) [v]
vs,
    (v -> Bool) -> [v] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (v -> v -> Bool
forall a. Eq a => a -> a -> Bool
/= v
lz0) [v]
us =
      ANormal ref v -> Maybe (ANormal ref v)
forall a. a -> Maybe a
Just (ANormal ref v -> Maybe (ANormal ref v))
-> (ANormal ref v -> ANormal ref v)
-> ANormal ref v
-> Maybe (ANormal ref v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v -> Either ref v -> [v] -> ANormal ref v -> ANormal ref v
forall v ref.
Var v =>
v -> Either ref v -> [v] -> ANormal ref v -> ANormal ref v
TName v
lz0 Either ref v
th [v]
vs (ANormal ref v -> Maybe (ANormal ref v))
-> ANormal ref v -> Maybe (ANormal ref v)
forall a b. (a -> b) -> a -> b
$ ref -> [v] -> ANormal ref v
forall v ref. Var v => ref -> [v] -> ANormal ref v
TCom ref
r ([v]
us [v] -> [v] -> [v]
forall a. [a] -> [a] -> [a]
++ [v
lz1])
  | Bool
otherwise = Maybe (ANormal ref v)
forall a. Maybe a
Nothing

pattern $mHandlerApp :: forall {r} {v} {ref}.
Var v =>
ANormal ref v -> (ANormal ref v -> r) -> ((# #) -> r) -> r
HandlerApp rw <- (matchHandlerApp -> Just rw)

-- Recognizes a rewritten handler app preceded by some additional
-- miscellaneous code. Some forms of recursive handlers do not just
-- have the handler app as the entire body of a function. In affine
-- scenarios, it is acceptable to inline these, and is essential to
-- optimizing them.
isExtendedHandlerApp :: (Var v) => ANormal ref v -> Bool
isExtendedHandlerApp :: forall v ref. Var v => ANormal ref v -> Bool
isExtendedHandlerApp = \case
  -- the actual handler application
  TName v
lz0 Either ref v
_ [v]
_ (TCom ref
_ [v]
as)
    | v
lz1 : [v]
us <- [v] -> [v]
forall v. [v] -> [v]
shiftArgs [v]
as ->
        v
lz0 v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
lz1 Bool -> Bool -> Bool
&& (v -> Bool) -> [v] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (v -> v -> Bool
forall a. Eq a => a -> a -> Bool
/= v
lz0) [v]
us
  TLets Direction Word16
_ [v]
_ [Mem]
_ ANormal ref v
_ ANormal ref v
bd -> ANormal ref v -> Bool
forall v ref. Var v => ANormal ref v -> Bool
isExtendedHandlerApp ANormal ref v
bd
  TName v
_ Either ref v
_ [v]
_ ANormal ref v
bd -> ANormal ref v -> Bool
forall v ref. Var v => ANormal ref v -> Bool
isExtendedHandlerApp ANormal ref v
bd
  TMatch v
_ Branched ref (ANormal ref v)
bs -> (ANormal ref v -> Bool) -> Branched ref (ANormal ref v) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ANormal ref v -> Bool
forall v ref. Var v => ANormal ref v -> Bool
isExtendedHandlerApp Branched ref (ANormal ref v)
bs
  ABTN.TAbs v
_ (ABTN.TAbss [v]
_ ANormal ref v
bd) -> ANormal ref v -> Bool
forall v ref. Var v => ANormal ref v -> Bool
isExtendedHandlerApp ANormal ref v
bd
  ANormal ref v
_ -> Bool
False

directAllowed :: (Var v) => ANormal ref v -> Bool
directAllowed :: forall v ref. Var v => ANormal ref v -> Bool
directAllowed TLit {} = Bool
True
directAllowed TBLit {} = Bool
True
directAllowed TPrm {} = Bool
True
directAllowed TFOp {} = Bool
True
directAllowed TCon {} = Bool
True
directAllowed ANormal ref v
_ = Bool
False

-- Recognizes the entry point of a handler, for inlining into the
-- actual handler if applicable.
isHandlerEntry :: (Var v) => v -> v -> ANormal ref v -> Bool
isHandlerEntry :: forall v ref. Var v => v -> v -> ANormal ref v -> Bool
isHandlerEntry v
him0 v
req0 ANormal ref v
tm
  | TName v
lzh0 (Right v
him1) [v]
_ ANormal ref v
tm <- ANormal ref v
tm,
    THnd [ref]
_ v
lzh1 Maybe v
Nothing (TFrc v
req1) <- ANormal ref v
tm =
      v
lzh0 v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
lzh1 Bool -> Bool -> Bool
&& v
him0 v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
him1 Bool -> Bool -> Bool
&& v
req0 v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
req1
  | Bool
otherwise = Bool
False

-- If the provided SuperGroup is recognized as a handler, applies
-- optimizations to improve it, like adding better code for affine
-- handlers.
optimizeHandler ::
  (Ord ref, Var v) =>
  (Text -> ref) ->
  OptInfos ref v ->
  ref ->
  SuperGroup ref v ->
  SuperGroup ref v
optimizeHandler :: forall ref v.
(Ord ref, Var v) =>
(Text -> ref)
-> OptInfos ref v -> ref -> SuperGroup ref v -> SuperGroup ref v
optimizeHandler Text -> ref
bi OptInfos ref v
opts ref
self SuperGroup ref v
group =
  SuperGroup ref v -> Maybe (SuperGroup ref v) -> SuperGroup ref v
forall a. a -> Maybe a -> a
fromMaybe SuperGroup ref v
group (Maybe (SuperGroup ref v) -> SuperGroup ref v)
-> Maybe (SuperGroup ref v) -> SuperGroup ref v
forall a b. (a -> b) -> a -> b
$ (Text -> ref)
-> OptInfos ref v
-> ref
-> SuperGroup ref v
-> Maybe (SuperGroup ref v)
forall ref v.
(Ord ref, Var v) =>
(Text -> ref)
-> OptInfos ref v
-> ref
-> SuperGroup ref v
-> Maybe (SuperGroup ref v)
augmentHandler Text -> ref
bi OptInfos ref v
opts ref
self SuperGroup ref v
group

-- moves the last value of a list to the start, for easier matching
shiftArgs :: [v] -> [v]
shiftArgs :: forall v. [v] -> [v]
shiftArgs [v]
vs = case [v] -> [v]
forall v. [v] -> [v]
reverse [v]
vs of
  v
v : [v]
vs -> v
v v -> [v] -> [v]
forall a. a -> [a] -> [a]
: [v] -> [v]
forall v. [v] -> [v]
reverse [v]
vs
  [] -> []

-- Checks if the group represents a handler, and if so, tries to add
-- optimized affine code.
augmentHandler ::
  (Ord ref, Var v) =>
  (Text -> ref) ->
  OptInfos ref v ->
  ref ->
  SuperGroup ref v ->
  Maybe (SuperGroup ref v)
augmentHandler :: forall ref v.
(Ord ref, Var v) =>
(Text -> ref)
-> OptInfos ref v
-> ref
-> SuperGroup ref v
-> Maybe (SuperGroup ref v)
augmentHandler Text -> ref
bi (Arities ref
arities, InlineInfos ref v
inls0) ref
self SuperGroup ref v
group
  | Rec [(v
mv0, SuperNormal ref v
matcher)] SuperNormal ref v
entry <- SuperGroup ref v
group,
    Lambda [Mem]
ccs (ABTN.TAbss [v]
args Term (ANormalF ref) v
body) <- SuperNormal ref v
entry,
    v
thunk : [v]
_ <- [v] -> [v]
forall v. [v] -> [v]
shiftArgs [v]
args,
    Just Term (ANormalF ref) v
body <- v
-> v -> v -> Term (ANormalF ref) v -> Maybe (Term (ANormalF ref) v)
forall v ref.
Var v =>
v -> v -> v -> ANormal ref v -> Maybe (ANormal ref v)
augmentHandlerEntry v
thunk v
mv0 v
ah Term (ANormalF ref) v
body,
    SuperNormal ref v
omatcher <- (Arities ref, InlineInfos ref v)
-> Set v -> SuperNormal ref v -> SuperNormal ref v
forall ref v.
(Ord ref, Var v) =>
OptInfos ref v -> Set v -> SuperNormal ref v -> SuperNormal ref v
affinePreOptimize (Arities ref, InlineInfos ref v)
opts (v -> Set v
forall a. a -> Set a
Set.singleton v
mv0) SuperNormal ref v
matcher,
    Just SuperNormal ref v
amatcher <- (Text -> ref)
-> (Arities ref, InlineInfos ref v)
-> v
-> v
-> SuperNormal ref v
-> Maybe (SuperNormal ref v)
forall ref v.
(Ord ref, Var v) =>
(Text -> ref)
-> OptInfos ref v
-> v
-> v
-> SuperNormal ref v
-> Maybe (SuperNormal ref v)
translateHandlerMatch Text -> ref
bi (Arities ref, InlineInfos ref v)
opts v
mv0 v
ah SuperNormal ref v
omatcher,
    SuperNormal ref v
amatcher <-
      (Arities ref, InlineInfos ref v)
-> Set v -> SuperNormal ref v -> SuperNormal ref v
forall ref v.
(Ord ref, Var v) =>
OptInfos ref v -> Set v -> SuperNormal ref v -> SuperNormal ref v
affinePostOptimize (Arities ref, InlineInfos ref v)
opts ([v] -> Set v
forall a. Ord a => [a] -> Set a
Set.fromList [v
mv0, v
ah]) SuperNormal ref v
amatcher =
      SuperGroup ref v -> Maybe (SuperGroup ref v)
forall a. a -> Maybe a
Just
        (SuperGroup ref v -> Maybe (SuperGroup ref v))
-> (Term (ANormalF ref) v -> SuperGroup ref v)
-> Term (ANormalF ref) v
-> Maybe (SuperGroup ref v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(v, SuperNormal ref v)] -> SuperNormal ref v -> SuperGroup ref v
forall ref v.
[(v, SuperNormal ref v)] -> SuperNormal ref v -> SuperGroup ref v
Rec [(v
mv0, SuperNormal ref v
matcher), (v
ah, SuperNormal ref v
amatcher)]
        (SuperNormal ref v -> SuperGroup ref v)
-> (Term (ANormalF ref) v -> SuperNormal ref v)
-> Term (ANormalF ref) v
-> SuperGroup ref v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Mem] -> Term (ANormalF ref) v -> SuperNormal ref v
forall ref v. [Mem] -> ANormal ref v -> SuperNormal ref v
Lambda [Mem]
ccs
        (Term (ANormalF ref) v -> Maybe (SuperGroup ref v))
-> Term (ANormalF ref) v -> Maybe (SuperGroup ref v)
forall a b. (a -> b) -> a -> b
$ [v] -> Term (ANormalF ref) v -> Term (ANormalF ref) v
forall v (f :: * -> * -> *). Var v => [v] -> Term f v -> Term f v
ABTN.TAbss [v]
args Term (ANormalF ref) v
body
  | Bool
otherwise = Maybe (SuperGroup ref v)
forall a. Maybe a
Nothing
  where
    ah :: v
ah = Word64 -> v
forall v. Var v => Word64 -> v
freshAff Word64
0
    opts :: (Arities ref, InlineInfos ref v)
opts = (Arities ref
arities, InlineInfos ref v
inls)
    inls :: InlineInfos ref v
inls = (InlineInfos ref v -> InlineInfos ref v)
-> (InlineInfo ref v -> InlineInfos ref v -> InlineInfos ref v)
-> Maybe (InlineInfo ref v)
-> InlineInfos ref v
-> InlineInfos ref v
forall b a. b -> (a -> b) -> Maybe a -> b
maybe InlineInfos ref v -> InlineInfos ref v
forall a. a -> a
id (ref -> InlineInfo ref v -> InlineInfos ref v -> InlineInfos ref v
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert ref
self) (SuperGroup ref v -> Maybe (InlineInfo ref v)
forall {v} {ref}.
Var v =>
SuperGroup ref v -> Maybe (InlineInfo ref v)
entryInfo SuperGroup ref v
group) InlineInfos ref v
inls0

-- Recognizes the matching portion of a handler, and produces an
-- optimized affine version if possible.
translateHandlerMatch ::
  (Ord ref, Var v) =>
  (Text -> ref) ->
  OptInfos ref v ->
  v ->
  v ->
  SuperNormal ref v ->
  Maybe (SuperNormal ref v)
translateHandlerMatch :: forall ref v.
(Ord ref, Var v) =>
(Text -> ref)
-> OptInfos ref v
-> v
-> v
-> SuperNormal ref v
-> Maybe (SuperNormal ref v)
translateHandlerMatch Text -> ref
bi OptInfos ref v
opts v
self v
ah (Lambda [Mem]
ccs (ABTN.TAbss [v]
args Term (ANormalF ref) v
body))
  | v
v : [v]
vs <- [v] -> [v]
forall v. [v] -> [v]
shiftArgs [v]
args,
    Set v
bound <- [v] -> Set v
forall a. Ord a => [a] -> Set a
Set.fromList (v
self v -> [v] -> [v]
forall a. a -> [a] -> [a]
: [v]
args),
    TMatch v
u Branched ref (Term (ANormalF ref) v)
branches <- Term (ANormalF ref) v
body,
    v
u v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
v,
    MatchRequest [(ref, EnumMap CTag ([Mem], Term (ANormalF ref) v))]
cs Term (ANormalF ref) v
df <- Branched ref (Term (ANormalF ref) v)
branches,
    [v]
args <- [v]
vs [v] -> [v] -> [v]
forall a. [a] -> [a] -> [a]
++ [v
ar, v
v],
    [Mem]
ccs <- [Mem]
ccs [Mem] -> [Mem] -> [Mem]
forall a. [a] -> [a] -> [a]
++ [Mem
BX] =
      [Mem] -> Term (ANormalF ref) v -> SuperNormal ref v
forall ref v. [Mem] -> ANormal ref v -> SuperNormal ref v
Lambda [Mem]
ccs
        (Term (ANormalF ref) v -> SuperNormal ref v)
-> ([(ref, EnumMap CTag ([Mem], Term (ANormalF ref) v))]
    -> Term (ANormalF ref) v)
-> [(ref, EnumMap CTag ([Mem], Term (ANormalF ref) v))]
-> SuperNormal ref v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [v] -> Term (ANormalF ref) v -> Term (ANormalF ref) v
forall v (f :: * -> * -> *). Var v => [v] -> Term f v -> Term f v
ABTN.TAbss [v]
args
        (Term (ANormalF ref) v -> Term (ANormalF ref) v)
-> ([(ref, EnumMap CTag ([Mem], Term (ANormalF ref) v))]
    -> Term (ANormalF ref) v)
-> [(ref, EnumMap CTag ([Mem], Term (ANormalF ref) v))]
-> Term (ANormalF ref) v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v -> Branched ref (Term (ANormalF ref) v) -> Term (ANormalF ref) v
forall v ref.
Var v =>
v -> Branched ref (ANormal ref v) -> ANormal ref v
TMatch v
u
        (Branched ref (Term (ANormalF ref) v) -> Term (ANormalF ref) v)
-> ([(ref, EnumMap CTag ([Mem], Term (ANormalF ref) v))]
    -> Branched ref (Term (ANormalF ref) v))
-> [(ref, EnumMap CTag ([Mem], Term (ANormalF ref) v))]
-> Term (ANormalF ref) v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([(ref, EnumMap CTag ([Mem], Term (ANormalF ref) v))]
 -> Term (ANormalF ref) v -> Branched ref (Term (ANormalF ref) v))
-> Term (ANormalF ref) v
-> [(ref, EnumMap CTag ([Mem], Term (ANormalF ref) v))]
-> Branched ref (Term (ANormalF ref) v)
forall a b c. (a -> b -> c) -> b -> a -> c
flip [(ref, EnumMap CTag ([Mem], Term (ANormalF ref) v))]
-> Term (ANormalF ref) v -> Branched ref (Term (ANormalF ref) v)
forall ref e.
[(ref, EnumMap CTag ([Mem], e))] -> e -> Branched ref e
MatchRequest Term (ANormalF ref) v
df
        ([(ref, EnumMap CTag ([Mem], Term (ANormalF ref) v))]
 -> SuperNormal ref v)
-> Maybe [(ref, EnumMap CTag ([Mem], Term (ANormalF ref) v))]
-> Maybe (SuperNormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Term (ANormalF ref) v -> Maybe (Term (ANormalF ref) v))
-> [(ref, EnumMap CTag ([Mem], Term (ANormalF ref) v))]
-> Maybe [(ref, EnumMap CTag ([Mem], Term (ANormalF ref) v))]
forall {a} {b}.
(a -> Maybe b)
-> [(ref, EnumMap CTag ([Mem], a))]
-> Maybe [(ref, EnumMap CTag ([Mem], b))]
traverse4 ((Text -> ref)
-> OptInfos ref v
-> v
-> Set v
-> [v]
-> v
-> Term (ANormalF ref) v
-> Maybe (Term (ANormalF ref) v)
forall ref v.
(Ord ref, Var v) =>
(Text -> ref)
-> OptInfos ref v
-> v
-> Set v
-> [v]
-> v
-> ANormal ref v
-> Maybe (ANormal ref v)
affineHandlerCase Text -> ref
bi OptInfos ref v
opts v
self Set v
bound [v]
vs v
ah) [(ref, EnumMap CTag ([Mem], Term (ANormalF ref) v))]
cs
  | Bool
otherwise = Maybe (SuperNormal ref v)
forall a. Maybe a
Nothing
  where
    ar :: v
ar = Word64 -> v
forall v. Var v => Word64 -> v
freshAff Word64
2
    traverse4 :: (a -> Maybe b)
-> [(ref, EnumMap CTag ([Mem], a))]
-> Maybe [(ref, EnumMap CTag ([Mem], b))]
traverse4 = ((ref, EnumMap CTag ([Mem], a))
 -> Maybe (ref, EnumMap CTag ([Mem], b)))
-> [(ref, EnumMap CTag ([Mem], a))]
-> Maybe [(ref, EnumMap CTag ([Mem], b))]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (((ref, EnumMap CTag ([Mem], a))
  -> Maybe (ref, EnumMap CTag ([Mem], b)))
 -> [(ref, EnumMap CTag ([Mem], a))]
 -> Maybe [(ref, EnumMap CTag ([Mem], b))])
-> ((a -> Maybe b)
    -> (ref, EnumMap CTag ([Mem], a))
    -> Maybe (ref, EnumMap CTag ([Mem], b)))
-> (a -> Maybe b)
-> [(ref, EnumMap CTag ([Mem], a))]
-> Maybe [(ref, EnumMap CTag ([Mem], b))]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (EnumMap CTag ([Mem], a) -> Maybe (EnumMap CTag ([Mem], b)))
-> (ref, EnumMap CTag ([Mem], a))
-> Maybe (ref, EnumMap CTag ([Mem], b))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> (ref, a) -> f (ref, b)
traverse ((EnumMap CTag ([Mem], a) -> Maybe (EnumMap CTag ([Mem], b)))
 -> (ref, EnumMap CTag ([Mem], a))
 -> Maybe (ref, EnumMap CTag ([Mem], b)))
-> ((a -> Maybe b)
    -> EnumMap CTag ([Mem], a) -> Maybe (EnumMap CTag ([Mem], b)))
-> (a -> Maybe b)
-> (ref, EnumMap CTag ([Mem], a))
-> Maybe (ref, EnumMap CTag ([Mem], b))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (([Mem], a) -> Maybe ([Mem], b))
-> EnumMap CTag ([Mem], a) -> Maybe (EnumMap CTag ([Mem], b))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> EnumMap CTag a -> f (EnumMap CTag b)
traverse ((([Mem], a) -> Maybe ([Mem], b))
 -> EnumMap CTag ([Mem], a) -> Maybe (EnumMap CTag ([Mem], b)))
-> ((a -> Maybe b) -> ([Mem], a) -> Maybe ([Mem], b))
-> (a -> Maybe b)
-> EnumMap CTag ([Mem], a)
-> Maybe (EnumMap CTag ([Mem], b))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Maybe b) -> ([Mem], a) -> Maybe ([Mem], b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> ([Mem], a) -> f ([Mem], b)
traverse

-- Recognizes the entry combinator of a compiled handler. If it is
-- one, then the result is a modified version with an affine handler
-- filled in.
augmentHandlerEntry ::
  (Var v) => v -> v -> v -> ANormal ref v -> Maybe (ANormal ref v)
augmentHandlerEntry :: forall v ref.
Var v =>
v -> v -> v -> ANormal ref v -> Maybe (ANormal ref v)
augmentHandlerEntry v
thunk0 v
mv0 v
ah ANormal ref v
body
  | TName v
hv (Right v
mv1) [v]
us ANormal ref v
body <- ANormal ref v
body,
    THnd [ref]
rs v
nh Maybe v
Nothing (TFrc v
thunk1) <- ANormal ref v
body,
    v
mv0 v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
mv1,
    v
nh v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
hv,
    v
thunk0 v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
thunk1 =
      ANormal ref v -> Maybe (ANormal ref v)
forall a. a -> Maybe a
Just
        (ANormal ref v -> Maybe (ANormal ref v))
-> (ANormal ref v -> ANormal ref v)
-> ANormal ref v
-> Maybe (ANormal ref v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v -> Either ref v -> [v] -> ANormal ref v -> ANormal ref v
forall v ref.
Var v =>
v -> Either ref v -> [v] -> ANormal ref v -> ANormal ref v
TName v
hv (v -> Either ref v
forall a b. b -> Either a b
Right v
mv1) [v]
us
        (ANormal ref v -> ANormal ref v)
-> (ANormal ref v -> ANormal ref v)
-> ANormal ref v
-> ANormal ref v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v -> Either ref v -> [v] -> ANormal ref v -> ANormal ref v
forall v ref.
Var v =>
v -> Either ref v -> [v] -> ANormal ref v -> ANormal ref v
TName v
ahp (v -> Either ref v
forall a b. b -> Either a b
Right v
ah) [v]
us
        (ANormal ref v -> Maybe (ANormal ref v))
-> ANormal ref v -> Maybe (ANormal ref v)
forall a b. (a -> b) -> a -> b
$ [ref] -> v -> Maybe v -> ANormal ref v -> ANormal ref v
forall v ref.
Var v =>
[ref] -> v -> Maybe v -> ANormal ref v -> ANormal ref v
THnd [ref]
rs v
nh (v -> Maybe v
forall a. a -> Maybe a
Just v
ahp) (v -> ANormal ref v
forall v ref. Var v => v -> ANormal ref v
TFrc v
thunk1)
  | Bool
otherwise = Maybe (ANormal ref v)
forall a. Maybe a
Nothing
  where
    ahp :: v
ahp = Word64 -> v
forall v. Var v => Word64 -> v
freshAff Word64
1

-- Recognizes an affine handler case, yielding a translated efficient
-- version if it is one.
affineHandlerCase ::
  (Ord ref, Var v) =>
  (Text -> ref) ->
  OptInfos ref v ->
  v ->
  Set v ->
  [v] ->
  v ->
  ANormal ref v ->
  Maybe (ANormal ref v)
affineHandlerCase :: forall ref v.
(Ord ref, Var v) =>
(Text -> ref)
-> OptInfos ref v
-> v
-> Set v
-> [v]
-> v
-> ANormal ref v
-> Maybe (ANormal ref v)
affineHandlerCase Text -> ref
builtin opts :: OptInfos ref v
opts@(Arities ref
_, InlineInfos ref v
_) v
self Set v
bound [v]
vs v
rec ANormal ref v
br
  | ABTN.TAbss [v]
us ANormal ref v
body <- ANormal ref v
br,
    TShift ref
_ v
kf0 ANormal ref v
body <- ANormal ref v
body,
    TName v
kf (Left ref
jc) [v
kf1] ANormal ref v
body <- ANormal ref v
body,
    ref
jc ref -> ref -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> ref
builtin Text
"jumpCont",
    Set v
bound <- Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
Set.union Set v
bound ([v] -> Set v
forall a. Ord a => [a] -> Set a
Set.fromList (v
kf0 v -> [v] -> [v]
forall a. a -> [a] -> [a]
: v
kf v -> [v] -> [v]
forall a. a -> [a] -> [a]
: [v]
us)),
    v
kf0 v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
kf1 =
      [v] -> ANormal ref v -> ANormal ref v
forall v (f :: * -> * -> *). Var v => [v] -> Term f v -> Term f v
ABTN.TAbss [v]
us
        (ANormal ref v -> ANormal ref v)
-> Maybe (ANormal ref v) -> Maybe (ANormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OptInfos ref v
-> v
-> Set v
-> [v]
-> v
-> v
-> v
-> ANormal ref v
-> Maybe (ANormal ref v)
forall ref v.
(Ord ref, Var v) =>
OptInfos ref v
-> v
-> Set v
-> [v]
-> v
-> v
-> v
-> ANormal ref v
-> Maybe (ANormal ref v)
affinePreBranch OptInfos ref v
opts v
self Set v
bound [v]
vs v
rec v
ar v
kf ANormal ref v
body
  | Bool
otherwise = Maybe (ANormal ref v)
forall a. Maybe a
Nothing
  where
    ar :: v
ar = Word64 -> v
forall v. Var v => Word64 -> v
freshAff Word64
2

-- Allows for having multiple branches that differ in the exact type
-- of affine handler recognized.
--
-- If the entire term doesn't use the continuation, then an irrelevant
-- handler is generated.
--
-- If the immediate term is a match, then we delay the choice of which
-- type of handler to generate into each branch.
--
-- If neither of the above cases hold, then we look for a linear case.
affinePreBranch ::
  (Ord ref, Var v) =>
  OptInfos ref v ->
  v ->
  Set v ->
  [v] ->
  v ->
  v ->
  v ->
  ANormal ref v ->
  Maybe (ANormal ref v)
affinePreBranch :: forall ref v.
(Ord ref, Var v) =>
OptInfos ref v
-> v
-> Set v
-> [v]
-> v
-> v
-> v
-> ANormal ref v
-> Maybe (ANormal ref v)
affinePreBranch OptInfos ref v
opts v
self Set v
bound [v]
vs v
rec v
ar v
kf ANormal ref v
bd
  | Just ANormal ref v
it <- v -> v -> ANormal ref v -> Maybe (ANormal ref v)
forall v ref.
Var v =>
v -> v -> ANormal ref v -> Maybe (ANormal ref v)
irrelevantTail v
ar v
kf ANormal ref v
bd = ANormal ref v -> Maybe (ANormal ref v)
forall a. a -> Maybe a
Just ANormal ref v
it
  | TMatch v
v Branched ref (ANormal ref v)
bs <- ANormal ref v
bd =
      v -> Branched ref (ANormal ref v) -> ANormal ref v
forall v ref.
Var v =>
v -> Branched ref (ANormal ref v) -> ANormal ref v
TMatch v
v
        (Branched ref (ANormal ref v) -> ANormal ref v)
-> Maybe (Branched ref (ANormal ref v)) -> Maybe (ANormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Branched ref (ANormal ref v)
-> (ANormal ref v -> Maybe (ANormal ref v))
-> Maybe (Branched ref (ANormal ref v))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for Branched ref (ANormal ref v)
bs \case
          ABTN.TAbss [v]
us ANormal ref v
bd ->
            [v] -> ANormal ref v -> ANormal ref v
forall v (f :: * -> * -> *). Var v => [v] -> Term f v -> Term f v
ABTN.TAbss [v]
us
              (ANormal ref v -> ANormal ref v)
-> Maybe (ANormal ref v) -> Maybe (ANormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OptInfos ref v
-> v
-> Set v
-> [v]
-> v
-> v
-> v
-> ANormal ref v
-> Maybe (ANormal ref v)
forall ref v.
(Ord ref, Var v) =>
OptInfos ref v
-> v
-> Set v
-> [v]
-> v
-> v
-> v
-> ANormal ref v
-> Maybe (ANormal ref v)
affinePreBranch OptInfos ref v
opts v
self Set v
bound' [v]
vs v
rec v
ar v
kf ANormal ref v
bd
            where
              bound' :: Set v
bound' = Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
Set.union ([v] -> Set v
forall a. Ord a => [a] -> Set a
Set.fromList [v]
us) Set v
bound
  | Bool
otherwise =
      (ANormal ref v, Any) -> ANormal ref v
localize
        ((ANormal ref v, Any) -> ANormal ref v)
-> Maybe (ANormal ref v, Any) -> Maybe (ANormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> WriterT Any Maybe (ANormal ref v) -> Maybe (ANormal ref v, Any)
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (OptInfos ref v
-> v
-> Set v
-> [v]
-> v
-> v
-> v
-> ANormal ref v
-> WriterT Any Maybe (ANormal ref v)
forall ref v.
(Ord ref, Var v) =>
OptInfos ref v
-> v
-> Set v
-> [v]
-> v
-> v
-> v
-> ANormal ref v
-> WriterT Any Maybe (ANormal ref v)
translateLinear OptInfos ref v
opts v
self Set v
bound [v]
vs v
rec v
ar v
kf ANormal ref v
bd)
  where
    localize :: (ANormal ref v, Any) -> ANormal ref v
localize (ANormal ref v
tm, Any Bool
True) = v -> ANormal ref v -> ANormal ref v
forall v ref. Var v => v -> ANormal ref v -> ANormal ref v
TLocal v
ar ANormal ref v
tm
    localize (ANormal ref v
tm, Any Bool
False) = ANormal ref v
tm

translateLinear ::
  (Ord ref, Var v) =>
  OptInfos ref v ->
  v ->
  Set v ->
  [v] ->
  v ->
  v ->
  v ->
  ANormal ref v ->
  WriterT Any Maybe (ANormal ref v)
translateLinear :: forall ref v.
(Ord ref, Var v) =>
OptInfos ref v
-> v
-> Set v
-> [v]
-> v
-> v
-> v
-> ANormal ref v
-> WriterT Any Maybe (ANormal ref v)
translateLinear OptInfos ref v
opts v
self Set v
bound0 [v]
vs v
rec v
ar v
kf = Set v -> ANormal ref v -> WriterT Any Maybe (ANormal ref v)
go Set v
bound0
  where
    go :: Set v -> ANormal ref v -> WriterT Any Maybe (ANormal ref v)
go Set v
bound ANormal ref v
body
      | Just ANormal ref v
lt <- OptInfos ref v
-> v
-> [v]
-> Set v
-> v
-> v
-> v
-> ANormal ref v
-> Maybe (ANormal ref v)
forall ref v.
(Ord ref, Var v) =>
OptInfos ref v
-> v
-> [v]
-> Set v
-> v
-> v
-> v
-> ANormal ref v
-> Maybe (ANormal ref v)
linearTail OptInfos ref v
opts v
self [v]
vs Set v
bound v
rec v
ar v
kf ANormal ref v
body =
          ANormal ref v
lt ANormal ref v
-> WriterT Any Maybe () -> WriterT Any Maybe (ANormal ref v)
forall a b. a -> WriterT Any Maybe b -> WriterT Any Maybe a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Any -> WriterT Any Maybe ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (Bool -> Any
Any Bool
True)
      | Just ANormal ref v
it <- v -> v -> ANormal ref v -> Maybe (ANormal ref v)
forall v ref.
Var v =>
v -> v -> ANormal ref v -> Maybe (ANormal ref v)
irrelevantTail v
ar v
kf ANormal ref v
body = ANormal ref v -> WriterT Any Maybe (ANormal ref v)
forall a. a -> WriterT Any Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ANormal ref v
it
      | TLet Direction Word16
d v
v Mem
cc ANormal ref v
e ANormal ref v
body <- ANormal ref v
body,
        v
kf v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.notMember` ANormal ref v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
ABTN.freeVars ANormal ref v
e =
          Direction Word16
-> v -> Mem -> ANormal ref v -> ANormal ref v -> ANormal ref v
forall v ref.
Var v =>
Direction Word16
-> v -> Mem -> ANormal ref v -> ANormal ref v -> ANormal ref v
TLet Direction Word16
d v
v Mem
cc ANormal ref v
e (ANormal ref v -> ANormal ref v)
-> WriterT Any Maybe (ANormal ref v)
-> WriterT Any Maybe (ANormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set v -> ANormal ref v -> WriterT Any Maybe (ANormal ref v)
go (v -> Set v -> Set v
forall a. Ord a => a -> Set a -> Set a
Set.insert v
v Set v
bound) ANormal ref v
body
      | TName v
v Either ref v
f [v]
us ANormal ref v
body <- ANormal ref v
body,
        (v -> Bool) -> Either ref v -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (v
kf v -> v -> Bool
forall a. Eq a => a -> a -> Bool
/=) Either ref v
f,
        (v -> Bool) -> [v] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (v
kf v -> v -> Bool
forall a. Eq a => a -> a -> Bool
/=) [v]
us =
          v -> Either ref v -> [v] -> ANormal ref v -> ANormal ref v
forall v ref.
Var v =>
v -> Either ref v -> [v] -> ANormal ref v -> ANormal ref v
TName v
v Either ref v
f [v]
us (ANormal ref v -> ANormal ref v)
-> WriterT Any Maybe (ANormal ref v)
-> WriterT Any Maybe (ANormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set v -> ANormal ref v -> WriterT Any Maybe (ANormal ref v)
go (v -> Set v -> Set v
forall a. Ord a => a -> Set a -> Set a
Set.insert v
v Set v
bound) ANormal ref v
body
      | TMatch v
v Branched ref (ANormal ref v)
bs <- ANormal ref v
body =
          v -> Branched ref (ANormal ref v) -> ANormal ref v
forall v ref.
Var v =>
v -> Branched ref (ANormal ref v) -> ANormal ref v
TMatch v
v
            (Branched ref (ANormal ref v) -> ANormal ref v)
-> WriterT Any Maybe (Branched ref (ANormal ref v))
-> WriterT Any Maybe (ANormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Branched ref (ANormal ref v)
-> (ANormal ref v -> WriterT Any Maybe (ANormal ref v))
-> WriterT Any Maybe (Branched ref (ANormal ref v))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for Branched ref (ANormal ref v)
bs \case
              ABTN.TAbss [v]
us ANormal ref v
bd ->
                [v] -> ANormal ref v -> ANormal ref v
forall v (f :: * -> * -> *). Var v => [v] -> Term f v -> Term f v
ABTN.TAbss [v]
us
                  (ANormal ref v -> ANormal ref v)
-> WriterT Any Maybe (ANormal ref v)
-> WriterT Any Maybe (ANormal ref v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set v -> ANormal ref v -> WriterT Any Maybe (ANormal ref v)
go ([v] -> Set v
forall a. Ord a => [a] -> Set a
Set.fromList [v]
us Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
`Set.union` Set v
bound) ANormal ref v
bd
      | Bool
otherwise = WriterT Any Maybe (ANormal ref v)
forall a. WriterT Any Maybe a
forall (m :: * -> *) a. MonadPlus m => m a
mzero

-- Recognizes the tail of a linear handler case, where the
-- continuation is called once in tail position. Returns a transformed
-- version if a match is found.
--
-- Arguments:
--   self: Reference to handler combinator
--   bound: arguments bound since header
--   vs: arguments to handler combinator
--   rec: local variable for affine handler
--   ar: argument variable for affine handler info
--   kf0: continuation variable
--   tm: term to transform
--
-- Note: this relies on inlining into the thunked continuation call to
-- avoid see exactly what the `k result` call is, rather than it
-- having multiple forms depending on the variable order.
linearTail ::
  (Ord ref, Var v) =>
  OptInfos ref v ->
  v ->
  [v] ->
  Set v ->
  v ->
  v ->
  v ->
  ANormal ref v ->
  Maybe (ANormal ref v)
linearTail :: forall ref v.
(Ord ref, Var v) =>
OptInfos ref v
-> v
-> [v]
-> Set v
-> v
-> v
-> v
-> ANormal ref v
-> Maybe (ANormal ref v)
linearTail OptInfos ref v
opts v
self [v]
vs Set v
bound v
rec v
ar v
kf0 ANormal ref v
tm
  | TName v
rh (Right v
f) [v]
as ANormal ref v
tm <- ANormal ref v
tm,
    v
f v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
self, -- recursive handler call
    (v -> Bool) -> [v] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (v -> v -> Bool
forall a. Eq a => a -> a -> Bool
/= v
kf0) [v]
as,
    v
rh v -> v -> Bool
forall a. Eq a => a -> a -> Bool
/= v
kf0, -- no shadowing or non-linearity
    THnd [ref]
_rs v
hh Maybe v
Nothing ANormal ref v
bd <- ANormal ref v
tm,
    v
rh v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
hh, -- handle recursively
    Set v
avoid <- v -> Set v -> Set v
forall a. Ord a => a -> Set a -> Set a
Set.insert v
rh Set v
bound,
    ANormal ref v
bd <- OptInfos ref v -> Set v -> ANormal ref v -> ANormal ref v
forall ref v.
(Ord ref, Var v) =>
OptInfos ref v -> Set v -> ANormal ref v -> ANormal ref v
replaceLinearBody OptInfos ref v
opts Set v
avoid ANormal ref v
bd,
    SimpleBody ANormal ref v -> ANormal ref v
pre Bool
ind Set v
shad Set v
free v
kf1 v
result <- ANormal ref v
bd, -- simple enough body
    v
kf0 v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.notMember` Set v
shad, -- kf is not shadowed in body
    v
kf0 v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.notMember` Set v
free, -- kf is not free in `pre`
    v
kf0 v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
kf1 -- continuation is called in tail position
    =
      ANormal ref v -> Maybe (ANormal ref v)
forall a. a -> Maybe a
Just (ANormal ref v -> Maybe (ANormal ref v))
-> (ANormal ref v -> ANormal ref v)
-> ANormal ref v
-> Maybe (ANormal ref v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> v -> [v] -> ANormal ref v -> ANormal ref v
update Bool
ind v
rh [v]
as (ANormal ref v -> ANormal ref v)
-> (ANormal ref v -> ANormal ref v)
-> ANormal ref v
-> ANormal ref v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ANormal ref v -> ANormal ref v
pre (ANormal ref v -> Maybe (ANormal ref v))
-> ANormal ref v -> Maybe (ANormal ref v)
forall a b. (a -> b) -> a -> b
$ v -> ANormal ref v
forall v ref. Var v => v -> ANormal ref v
TVar v
result
  | Bool
otherwise = Maybe (ANormal ref v)
forall a. Maybe a
Nothing
  where
    update :: Bool -> v -> [v] -> ANormal ref v -> ANormal ref v
update Bool
ind v
huv [v]
us
      -- recursive call with identical, non-shadowed variables and no
      -- indirect body calls; no need to update
      | Bool -> Bool
not Bool
ind,
        [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
Prelude.and ((v -> v -> Bool) -> [v] -> [v] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith v -> v -> Bool
forall a. Eq a => a -> a -> Bool
(==) [v]
us [v]
vs),
        (v -> Bool) -> [v] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.notMember` Set v
bound) [v]
us =
          ANormal ref v -> ANormal ref v
forall a. a -> a
id
      -- repurpose hr0 variable for update call
      | Bool
otherwise =
          v -> Either ref v -> [v] -> ANormal ref v -> ANormal ref v
forall v ref.
Var v =>
v -> Either ref v -> [v] -> ANormal ref v -> ANormal ref v
TName v
huv (v -> Either ref v
forall a b. b -> Either a b
Right v
rec) ([v]
us [v] -> [v] -> [v]
forall a. [a] -> [a] -> [a]
++ [v
ar])
            (ANormal ref v -> ANormal ref v)
-> (ANormal ref v -> ANormal ref v)
-> ANormal ref v
-> ANormal ref v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
forall v ref.
Var v =>
Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
TLets Direction Word16
forall a. Direction a
Direct [] [] (Bool -> v -> v -> ANormal ref v
forall v ref. Var v => Bool -> v -> v -> Term (ANormalF ref) v
TUpdate Bool
ind v
ar v
huv)

-- Does one level of inlining to revert any floating of the handler
-- body. This does not obey the normal inlining classification, so it
-- may inline a function that makes indirect calls, and needs to
-- account for that. This allows us to see more linear situations.
--
-- Note: this does _not_ fix the `Direction` numbering in the result.
-- The numbering is intended to allow returning to specific points
-- within a function; each number mapping to a code section that is
-- the rest of the function from that point.
--
-- However, this is not used during normal code execution. `Let` in
-- code directly stores its body. The numbering is used to reconstruct
-- continuations from the interchange format, since they only send the
-- numbers. However, _affine_ handlers are only use in contexts where
-- continuations aren't captured, so we don't actually need a correct
-- numbering. If this is ever changed, then the numbering here must be
-- adjusted.
replaceLinearBody ::
  (Ord ref, Var v) =>
  OptInfos ref v ->
  Set v ->
  ANormal ref v ->
  ANormal ref v
replaceLinearBody :: forall ref v.
(Ord ref, Var v) =>
OptInfos ref v -> Set v -> ANormal ref v -> ANormal ref v
replaceLinearBody opts :: OptInfos ref v
opts@(Arities ref
arities, InlineInfos ref v
inls) Set v
avoid ANormal ref v
bd
  | TLetD v
v Mem
cc ANormal ref v
bn ANormal ref v
bd <- ANormal ref v
bd =
      v -> Mem -> ANormal ref v -> ANormal ref v -> ANormal ref v
forall v ref.
Var v =>
v -> Mem -> ANormal ref v -> ANormal ref v -> ANormal ref v
TLetD v
v Mem
cc ANormal ref v
bn (ANormal ref v -> ANormal ref v) -> ANormal ref v -> ANormal ref v
forall a b. (a -> b) -> a -> b
$ OptInfos ref v -> Set v -> ANormal ref v -> ANormal ref v
forall ref v.
(Ord ref, Var v) =>
OptInfos ref v -> Set v -> ANormal ref v -> ANormal ref v
replaceLinearBody OptInfos ref v
opts (v -> Set v -> Set v
forall a. Ord a => a -> Set a -> Set a
Set.insert v
v Set v
avoid) ANormal ref v
bd
  | TCom ref
r [v]
vs <- ANormal ref v
bd,
    Just Int
n <- ref -> Arities ref -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup ref
r Arities ref
arities,
    [v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
vs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n,
    Just (InlInfo InlineClass
_ (ABTN.TAbss [v]
us ANormal ref v
expr)) <- ref -> InlineInfos ref v -> Maybe (InlineInfo ref v)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup ref
r InlineInfos ref v
inls,
    Map v v
rn <- [(v, v)] -> Map v v
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([v] -> [v] -> [(v, v)]
forall a b. [a] -> [b] -> [(a, b)]
zip [v]
us [v]
vs) =
      Set v -> Map v v -> ANormal ref v -> ANormal ref v
forall v (f :: * -> * -> *).
(Var v, Bifunctor f, Bifoldable f) =>
Set v -> Map v v -> Term f v -> Term f v
ABTN.renamesAvoiding Set v
avoid Map v v
rn ANormal ref v
expr
replaceLinearBody OptInfos ref v
_ Set v
_ ANormal ref v
bd = ANormal ref v
bd

parseSimpleHandlerBody ::
  (Var v) =>
  ANormal ref v ->
  Maybe (ANormal ref v -> ANormal ref v, Bool, Set v, Set v, v, v)
parseSimpleHandlerBody :: forall v ref.
Var v =>
ANormal ref v
-> Maybe (ANormal ref v -> ANormal ref v, Bool, Set v, Set v, v, v)
parseSimpleHandlerBody = \case
  TLet Direction Word16
d v
u Mem
cc ANormal ref v
bn ANormal ref v
bd ->
    v
-> Direction Word16
-> Set v
-> (ANormal ref v -> ANormal ref v)
-> (ANormal ref v -> ANormal ref v, Bool, Set v, Set v, v, v)
-> (ANormal ref v -> ANormal ref v, Bool, Set v, Set v, v, v)
forall {a} {a} {b} {c} {a} {e} {f}.
Ord a =>
a
-> Direction a
-> Set a
-> (b -> c)
-> (a -> b, Bool, Set a, Set a, e, f)
-> (a -> c, Bool, Set a, Set a, e, f)
tweak v
u Direction Word16
d (ANormal ref v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
ABTN.freeVars ANormal ref v
bn) (Direction Word16
-> v -> Mem -> ANormal ref v -> ANormal ref v -> ANormal ref v
forall v ref.
Var v =>
Direction Word16
-> v -> Mem -> ANormal ref v -> ANormal ref v -> ANormal ref v
TLet Direction Word16
d v
u Mem
cc ANormal ref v
bn)
      ((ANormal ref v -> ANormal ref v, Bool, Set v, Set v, v, v)
 -> (ANormal ref v -> ANormal ref v, Bool, Set v, Set v, v, v))
-> Maybe (ANormal ref v -> ANormal ref v, Bool, Set v, Set v, v, v)
-> Maybe (ANormal ref v -> ANormal ref v, Bool, Set v, Set v, v, v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ANormal ref v
-> Maybe (ANormal ref v -> ANormal ref v, Bool, Set v, Set v, v, v)
forall v ref.
Var v =>
ANormal ref v
-> Maybe (ANormal ref v -> ANormal ref v, Bool, Set v, Set v, v, v)
parseSimpleHandlerBody ANormal ref v
bd
  TApv v
u [v
result] ->
    (ANormal ref v -> ANormal ref v, Bool, Set v, Set v, v, v)
-> Maybe (ANormal ref v -> ANormal ref v, Bool, Set v, Set v, v, v)
forall a. a -> Maybe a
Just (ANormal ref v -> ANormal ref v
forall a. a -> a
id, Bool
False, Set v
forall a. Monoid a => a
mempty, Set v
forall a. Monoid a => a
mempty, v
u, v
result)
  ANormal ref v
_ -> Maybe (ANormal ref v -> ANormal ref v, Bool, Set v, Set v, v, v)
forall a. Maybe a
Nothing
  where
    tweak :: a
-> Direction a
-> Set a
-> (b -> c)
-> (a -> b, Bool, Set a, Set a, e, f)
-> (a -> c, Bool, Set a, Set a, e, f)
tweak a
w Direction a
d Set a
fvs1 b -> c
f (a -> b
g, Bool
ind, Set a
sh, Set a
fvs0, e
u, f
v) =
      (b -> c
f (b -> c) -> (a -> b) -> a -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
g, Bool
ind Bool -> Bool -> Bool
|| Direction a -> Bool
forall {a}. Direction a -> Bool
isIndirect Direction a
d, a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
Set.insert a
w Set a
sh, Set a
fvs, e
u, f
v)
      where
        fvs :: Set a
fvs = Set a
fvs1 Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
`Set.union` a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
Set.delete a
w Set a
fvs0

        isIndirect :: Direction a -> Bool
isIndirect Direction a
Direct = Bool
False
        isIndirect Direction a
_ = Bool
True

pattern $mSimpleBody :: forall {r} {v} {ref}.
Var v =>
ANormal ref v
-> ((ANormal ref v -> ANormal ref v)
    -> Bool -> Set v -> Set v -> v -> v -> r)
-> ((# #) -> r)
-> r
SimpleBody head ind shad free kf result <-
  (parseSimpleHandlerBody -> Just (head, ind, shad, free, kf, result))

irrelevantTail :: (Var v) => v -> v -> ANormal ref v -> Maybe (ANormal ref v)
irrelevantTail :: forall v ref.
Var v =>
v -> v -> ANormal ref v -> Maybe (ANormal ref v)
irrelevantTail v
ar v
kf ANormal ref v
tm
  | v
kf v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.notMember` ANormal ref v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
ABTN.freeVars ANormal ref v
tm =
      ANormal ref v -> Maybe (ANormal ref v)
forall a. a -> Maybe a
Just (ANormal ref v -> Maybe (ANormal ref v))
-> ANormal ref v -> Maybe (ANormal ref v)
forall a b. (a -> b) -> a -> b
$ Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
forall v ref.
Var v =>
Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
TLets Direction Word16
forall a. Direction a
Direct [] [] (v -> ANormal ref v
forall v ref. Var v => v -> ANormal ref v
TDiscard v
ar) ANormal ref v
tm
  | Bool
otherwise = Maybe (ANormal ref v)
forall a. Maybe a
Nothing

freshAff :: (Var v) => Word64 -> v
freshAff :: forall v. Var v => Word64 -> v
freshAff Word64
fr = Word64 -> v -> v
forall v. Var v => Word64 -> v -> v
Var.freshenId Word64
fr (v -> v) -> v -> v
forall a b. (a -> b) -> a -> b
$ Type -> v
forall v. Var v => Type -> v
Var.typed Type
Var.AffBlank