{-# LANGUAGE PatternSynonyms #-}

-- | Loop state exploded into parameters: call-pattern specialisation.
--
-- A function of a batch that some saturated call in the batch passes a
-- product or constructor built for the call (a let of a constructor
-- application or a record literal, @(acc, i)@ in @go (acc, i)@, @Left x@
-- in @go (Left x)@, or a record update of such a value) gets a
-- specialisation for that call pattern: a worker, a combinator of the
-- group at an index of its own ('specIndex'), whose parameters are the
-- pattern's fields and whose body is the lambda's with each match on the
-- argument replaced by the arm for the pattern's constructor and each
-- update of it by a literal of the fields. Call sites that build the
-- pattern call the worker with the field variables, which are in scope
-- where the value was built, and the let that built it is dropped when
-- nothing else uses it ('dropDead'); so a loop whose state is a tuple or
-- a record allocates nothing per round, and a loop whose state is a sum
-- runs as one worker per constructor, mutually recursive through the
-- sites, with no tag anywhere. The lambda's own index keeps the original
-- body, for callers the batch cannot see (a later batch, a function
-- value, a partial application) and for sites that pass a value they did
-- not build. Fields that are products built for the call explode in
-- turn, which is what a tuple's nesting needs. This is GHC's SpecConstr
-- (Peyton Jones, "Call-pattern specialisation for Haskell programs",
-- 2007), done for the call sites the loader can see ('explosions', over
-- each batch 'Unison.Runtime.Machine.cacheAdd' loads) and without any
-- change to how values are represented: a worker's parameters are
-- ordinary slots, so nothing here depends on the wide convention
-- (Unison.Runtime.Wide). A continuation captured in rewritten code is
-- written out as the unoptimized code's frames, which is what 'Sources'
-- and 'Prepared' are for ('Unison.Runtime.MCode.FrameLayout', design.md
-- "Continuations across processes").
--
-- A pattern is taken apart only as far as the body consumes it: an
-- argument the body passes on, stores or returns on some path stays a
-- plain parameter ('usable'), so nothing is ever rebuilt. The patterns of
-- a function are a closure: the ones its callers' bodies build, then the
-- ones each worker's own body builds, up to a cap per function
-- ('maxSpecs'). A sum built in the branches of a match and passed at once,
-- @go (if p then Left x else Right y)@, is first rewritten so that each
-- branch makes the call ('sinkCalls'), which is what gives each branch its
-- worker.
--
-- The return side is the same idea for results ('returns'): a function
-- whose every return builds a product, through its own tail calls, gets a
-- /product-return worker/, a combinator of its group at 'retIndex' whose
-- tails yield the product's leaves as several values instead of packing
-- them ('returnWorker'), while the entry keeps the original body. A caller
-- in the batch that takes the result apart at once calls the worker with
-- a let of the leaves ('useReturns'), so a loop and its step function
-- exchange a tuple with nothing allocated in either direction; a result
-- that escapes keeps the entry, and one that is only passed on to an
-- exploded position of another worker is rebuilt in the caller and
-- dropped like any other product built for such a call. A function whose
-- own recursion embeds its result, a tree insert rebuilding the path, gets
-- no worker ('selfResultEscapes'): the worker's recursion would yield the
-- leaves and pack them again at every level.
--
-- Everything here is a rewrite of the ANF, and the result says what it
-- means: a call of a worker is @TApp (FWorker r i) args@, the combinator
-- at index @i@ of the group of @r@ with exactly the arguments it takes,
-- and a tail that returns several values is @TApp FYield vs@. The emitter
-- (Unison.Runtime.MCode) turns each into one instruction and checks the
-- counts ('Unison.Runtime.MCode.resultCount'); it decides nothing. The
-- two forms exist only in a batch's code after it is loaded, never in
-- what the cache stores or a continuation names, so the serializers
-- refuse them.
--
-- What it does not change: the groups the cache stores and hashes (the
-- rewrite is a function of a batch's optimized code, like inlining, run
-- again at every load), a compiled program's file (its MCode has the
-- workers in it), and the serialized form of a continuation (a worker's
-- frame holds plain slots; its index is a combinator of its group like any
-- other). A continuation captured inside a worker names the worker's
-- index, so a process that loads the same code with the rewrite off
-- cannot resume it, as with any difference in the optimizer between two
-- versions.
--
-- Two rewrites of a body on its own come first ('knownCons',
-- 'splitJoins'; design.md "Known constructors and case of case"): a let
-- binding a constructor application or a record literal has every match
-- on it replaced by the arm, the arm's variables renamed to the fields
-- (the argument side's 'rewrite'), so that @match (a, b) with ..@, which
-- builds the tuple and takes it apart at once, allocates nothing and
-- tests no tag; and a let whose binding is a match (or anything whose
-- every tail builds a constructor of the matched type) and whose body is
-- a match on it (the let's body, or the binding of the next let) has
-- the outer match moved into those tails, where the constructor is
-- known, an outer arm that several tails reach becoming a /join point/,
-- a combinator of the group at 'joinIndex' that the tails call with the
-- arm's fields and the variables the arm uses, so that an arm's code is
-- not copied into every tail (an arm that is a single expression, or
-- that one tail reaches, is copied instead), and so does the rest of the
-- body when the match's value was let-bound. The join points of a lambda are made last,
-- after the other rewrites, with their parameters taken from the
-- unoptimized arm ('caseSites'), so that every variable the unoptimized
-- arm reads is in the join point's frame and a continuation captured
-- inside it can be written out ('Join').
--
-- @UNISON_OPTIMIZATIONS=-explode_args@ and @-explode_returns@ switch the
-- two sides off, and @-known_constructor@ the two rewrites above
-- ('explodeArgsOn', 'explodeReturnsOn', 'knownConstructorOn';
-- Unison.Runtime.Optimizations), for measuring and for a process that has
-- to agree with one that ran without them.
module Unison.Runtime.MCode.Explode
  ( explodeArgsOn,
    explodeReturnsOn,
    knownConstructorOn,
    rewritesOn,
    specIndex,
    retIndex,
    joinIndex,
    isJoinIndex,
    Shape (..),
    Pattern,
    shapeLeaves,
    Batch (..),
    Explosions,
    explosions,
    Returns,
    returns,
    prepare,
    specialise,
    returnWorker,
    Prepared (..),
    Join (..),
    JoinSetting (..),
    Sources,
    Built (..),
    builtFields,
    conLets,
    uniqueBody,
    universe,
  )
where

import Control.Monad (foldM)
import Control.Monad.State.Strict (State, evalState, get, modify, put, runState, state)
import Data.Bifunctor (first)
import Data.Bits (shiftL, shiftR, (.|.))
import Data.Containers.ListUtils (nubOrd)
import Data.Foldable (toList)
import Data.List (elemIndex)
import Data.Map.Strict qualified as M
import Data.Maybe (fromMaybe, isJust, listToMaybe)
import Data.Monoid (Sum (..))
import Data.Set qualified as Set
import Data.Word (Word16, Word64)
import Unison.ABT.Normalized (pattern TAbss)
import Unison.ABT.Normalized qualified as ABTN
import Unison.Reference (Reference)
import Unison.Runtime.ANF
import Unison.Runtime.ANF qualified as ANF
import Unison.Runtime.InternalError (internalBug)
import Unison.Runtime.Optimizations (Optimization (..), optimizationSetting)
import Unison.Util.EnumContainers qualified as EC
import Unison.Var (Var)
import Unison.Var qualified as Var

-- | Whether loop state is exploded in this process (the argument side:
-- 'explosions', 'specialise'): on unless @UNISON_OPTIMIZATIONS@ says
-- @-explode_args@.
explodeArgsOn :: Bool
explodeArgsOn :: Bool
explodeArgsOn = Bool -> Maybe Bool -> Bool
forall a. a -> Maybe a -> a
fromMaybe Bool
True (Optimization -> Maybe Bool
optimizationSetting Optimization
ExplodeArgs)

-- | Whether product results are exploded (the return side: 'returns',
-- 'returnWorker'): on unless @UNISON_OPTIMIZATIONS@ says @-explode_returns@.
explodeReturnsOn :: Bool
explodeReturnsOn :: Bool
explodeReturnsOn = Bool -> Maybe Bool -> Bool
forall a. a -> Maybe a -> a
fromMaybe Bool
True (Optimization -> Maybe Bool
optimizationSetting Optimization
ExplodeReturns)

-- | Whether a product or sum built and matched in the same body is
-- matched where it is built, and a match on the value of a match moved
-- into its tails ('knownCons', 'splitJoins'): on unless
-- @UNISON_OPTIMIZATIONS@ says @-known_constructor@.
knownConstructorOn :: Bool
knownConstructorOn :: Bool
knownConstructorOn = Bool -> Maybe Bool -> Bool
forall a. a -> Maybe a -> a
fromMaybe Bool
True (Optimization -> Maybe Bool
optimizationSetting Optimization
KnownConstructor)

-- | Whether any rewrite of the ANF is on: with none, every lambda is
-- emitted as its unoptimized form and no section gets a deoptimization
-- table (the emitter checks this).
rewritesOn :: Bool
rewritesOn :: Bool
rewritesOn = Bool
explodeArgsOn Bool -> Bool -> Bool
|| Bool
explodeReturnsOn Bool -> Bool -> Bool
|| Bool
knownConstructorOn

-- | The combinator index of a group's k-th specialisation (k from 1): in
-- the high bits like a local function's, with a bit no let index reaches
-- (the lets of a combinator are numbered from 1 in its low sixteen bits)
-- and no local function has. The JIT treats a worker like the entry point
-- it stands for (Compile's @isEntry@).
specIndex :: Int -> Word64
specIndex :: Int -> Word64
specIndex Int
k = (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftL` Int
16) Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Word64
0x8000

-- | The combinator index of a group's product-return worker ('returns'):
-- the specialisation slot no call pattern takes, since those are numbered
-- from 1.
retIndex :: Word64
retIndex :: Word64
retIndex = Int -> Word64
specIndex Int
0

-- | The combinator index of the k-th join point (k from 1) of the lambda
-- at an index of the group ('splitJoins'): the lambda's index with k in
-- the bits above any a local function's or a specialisation's number
-- uses, so that the low fifteen bits are the owner's and the JIT compiles
-- it as it does a worker (Compile's @isEntry@), and the lets of the join
-- point are numbered under it like any combinator's.
joinIndex :: Word64 -> Int -> Word64
joinIndex :: Word64 -> Int -> Word64
joinIndex Word64
owner Int
k = Word64
owner Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftL` Int
48)

isJoinIndex :: Word64 -> Bool
isJoinIndex :: Word64 -> Bool
isJoinIndex Word64
i = Word64
i Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
48 Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word64
0

-- | The most specialisations a function gets.
maxSpecs :: Int
maxSpecs :: Int
maxSpecs = Int
8

-- | The most join points a lambda gets (what 'joinIndex' has bits for,
-- with room to spare).
maxJoins :: Int
maxJoins :: Int
maxJoins = Int
15

-- | The most leaves a returned product may have.
maxLeaves :: Int
maxLeaves :: Int
maxLeaves = Int
8

-- | How an argument position is taken apart: left as one slot, or a
-- product, a data constructor or a record (@Fields@, its fields in the
-- schema's order, which is the field names' order), whose fields are
-- taken apart in turn.
data Shape = Plain | Con !Reference !CTag [Shape] | Fields !RecordSchema [Shape]
  deriving (Shape -> Shape -> Bool
(Shape -> Shape -> Bool) -> (Shape -> Shape -> Bool) -> Eq Shape
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Shape -> Shape -> Bool
== :: Shape -> Shape -> Bool
$c/= :: Shape -> Shape -> Bool
/= :: Shape -> Shape -> Bool
Eq, Eq Shape
Eq Shape =>
(Shape -> Shape -> Ordering)
-> (Shape -> Shape -> Bool)
-> (Shape -> Shape -> Bool)
-> (Shape -> Shape -> Bool)
-> (Shape -> Shape -> Bool)
-> (Shape -> Shape -> Shape)
-> (Shape -> Shape -> Shape)
-> Ord Shape
Shape -> Shape -> Bool
Shape -> Shape -> Ordering
Shape -> Shape -> Shape
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 :: Shape -> Shape -> Ordering
compare :: Shape -> Shape -> Ordering
$c< :: Shape -> Shape -> Bool
< :: Shape -> Shape -> Bool
$c<= :: Shape -> Shape -> Bool
<= :: Shape -> Shape -> Bool
$c> :: Shape -> Shape -> Bool
> :: Shape -> Shape -> Bool
$c>= :: Shape -> Shape -> Bool
>= :: Shape -> Shape -> Bool
$cmax :: Shape -> Shape -> Shape
max :: Shape -> Shape -> Shape
$cmin :: Shape -> Shape -> Shape
min :: Shape -> Shape -> Shape
Ord, Int -> Shape -> ShowS
[Shape] -> ShowS
Shape -> String
(Int -> Shape -> ShowS)
-> (Shape -> String) -> ([Shape] -> ShowS) -> Show Shape
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Shape -> ShowS
showsPrec :: Int -> Shape -> ShowS
$cshow :: Shape -> String
show :: Shape -> String
$cshowList :: [Shape] -> ShowS
showList :: [Shape] -> ShowS
Show)

-- | A call pattern: a shape per argument.
type Pattern = [Shape]

isPlain :: Shape -> Bool
isPlain :: Shape -> Bool
isPlain Shape
Plain = Bool
True
isPlain Shape
_ = Bool
False

subShapes :: Shape -> [Shape]
subShapes :: Shape -> [Shape]
subShapes Shape
Plain = []
subShapes (Con Reference
_ CTag
_ [Shape]
ss) = [Shape]
ss
subShapes (Fields RecordSchema
_ [Shape]
ss) = [Shape]
ss

withSubShapes :: Shape -> [Shape] -> Shape
withSubShapes :: Shape -> [Shape] -> Shape
withSubShapes Shape
Plain [Shape]
_ = Shape
Plain
withSubShapes (Con Reference
r CTag
t [Shape]
_) [Shape]
ss = Reference -> CTag -> [Shape] -> Shape
Con Reference
r CTag
t [Shape]
ss
withSubShapes (Fields RecordSchema
rs [Shape]
_) [Shape]
ss = RecordSchema -> [Shape] -> Shape
Fields RecordSchema
rs [Shape]
ss

-- | Whether two shapes are the same product at the top.
sameProduct :: Shape -> Shape -> Bool
sameProduct :: Shape -> Shape -> Bool
sameProduct (Con Reference
r CTag
t [Shape]
as) (Con Reference
r' CTag
t' [Shape]
bs) = Reference
r Reference -> Reference -> Bool
forall a. Eq a => a -> a -> Bool
== Reference
r' Bool -> Bool -> Bool
&& CTag
t CTag -> CTag -> Bool
forall a. Eq a => a -> a -> Bool
== CTag
t' Bool -> Bool -> Bool
&& [Shape] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Shape]
as Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Shape] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Shape]
bs
sameProduct (Fields RecordSchema
rs [Shape]
as) (Fields RecordSchema
rs' [Shape]
bs) = RecordSchema
rs RecordSchema -> RecordSchema -> Bool
forall a. Eq a => a -> a -> Bool
== RecordSchema
rs' Bool -> Bool -> Bool
&& [Shape] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Shape]
as Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Shape] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Shape]
bs
sameProduct Shape
_ Shape
_ = Bool
False

-- | Whether a site's shape offers everything a pattern's shape takes
-- apart.
refines :: Shape -> Shape -> Bool
refines :: Shape -> Shape -> Bool
refines Shape
_ Shape
Plain = Bool
True
refines Shape
s Shape
p = Shape -> Shape -> Bool
sameProduct Shape
s Shape
p Bool -> Bool -> Bool
&& [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Shape -> Shape -> Bool) -> [Shape] -> [Shape] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Shape -> Shape -> Bool
refines (Shape -> [Shape]
subShapes Shape
s) (Shape -> [Shape]
subShapes Shape
p))

-- | The values a shape is taken apart into: one for a plain position, the
-- leaves of its fields for a product.
shapeLeaves :: Shape -> Int
shapeLeaves :: Shape -> Int
shapeLeaves Shape
Plain = Int
1
shapeLeaves Shape
s = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Shape -> Int) -> [Shape] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Shape -> Int
shapeLeaves (Shape -> [Shape]
subShapes Shape
s))

-- | What two shapes agree on: the product where both have the same one,
-- with the meet of their fields, and plain anywhere else.
meetShape :: Shape -> Shape -> Shape
meetShape :: Shape -> Shape -> Shape
meetShape Shape
a Shape
b
  | Shape -> Shape -> Bool
sameProduct Shape
a Shape
b = Shape -> [Shape] -> Shape
withSubShapes Shape
a ((Shape -> Shape -> Shape) -> [Shape] -> [Shape] -> [Shape]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Shape -> Shape -> Shape
meetShape (Shape -> [Shape]
subShapes Shape
a) (Shape -> [Shape]
subShapes Shape
b))
  | Bool
otherwise = Shape
Plain

-- | What the rewrites know about the batch, each by a function's
-- reference and Nothing for a function outside the batch or without one:
-- the call patterns it has a specialisation for, in worker order
-- ('explosions'), the product it returns on every path ('returns'), and
-- its arity. The analyses run with the parts not yet known empty.
data Batch = Batch
  { Batch -> Reference -> Maybe [[Shape]]
batchSpecs :: Reference -> Maybe [Pattern],
    Batch -> Reference -> Maybe Shape
batchRets :: Reference -> Maybe Shape,
    Batch -> Reference -> Maybe Int
batchArity :: Reference -> Maybe Int
  }

-- | The specialisations of each function of a batch, by its reference:
-- the patterns, in the order their workers are numbered ('specIndex').
type Explosions = M.Map Reference [Pattern]

-- | A product a let built: which one, and its field variables, in the
-- shape's field order.
data Built v = BuiltCon Reference CTag [v] | BuiltRec RecordSchema [v]
  deriving (Int -> Built v -> ShowS
[Built v] -> ShowS
Built v -> String
(Int -> Built v -> ShowS)
-> (Built v -> String) -> ([Built v] -> ShowS) -> Show (Built v)
forall v. Show v => Int -> Built v -> ShowS
forall v. Show v => [Built v] -> ShowS
forall v. Show v => Built v -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall v. Show v => Int -> Built v -> ShowS
showsPrec :: Int -> Built v -> ShowS
$cshow :: forall v. Show v => Built v -> String
show :: Built v -> String
$cshowList :: forall v. Show v => [Built v] -> ShowS
showList :: [Built v] -> ShowS
Show)

-- | Where a variable of the unoptimized body went, for the frames of a
-- continuation captured in the optimized code to be written out as the
-- unoptimized code's (design.md, "Continuations across processes";
-- 'Unison.Runtime.MCode.FrameLayout'): renamed to another variable (a
-- match arm's field variable to the leaf or product that stands for it, a
-- field read's let to the parameter it reads), or replaced by a product
-- the optimized body no longer names, given as the product with the
-- names of its fields in the optimized body (an exploded parameter, a
-- result the leaves replace, a product of either that a match took
-- apart in turn). A variable with no entry has the same name in both.
type Sources v = M.Map v (Either v (Built v))

-- | A lambda as the emitter sees it, with what its rewrite changed: the
-- sources of the unoptimized body's variables, the binders of each let
-- the rewrite changed, mapped to the binders of the unoptimized let it
-- stands for (a let of a result's leaves, to the let of the result), and
-- the variables the rewrite found dead wherever a frame could hold them
-- (the let of a site 'splitJoins' rewrote, used only by the match that
-- is gone), and its join points, the k-th at 'joinIndex' k of the
-- lambda's index.
data Prepared v = Prepared
  { forall v. Prepared v -> SuperNormal Reference v
preparedLambda :: SuperNormal Reference v,
    forall v. Prepared v -> Sources v
preparedSources :: Sources v,
    forall v. Prepared v -> Map [v] [v]
preparedLets :: M.Map [v] [v],
    forall v. Prepared v -> Set v
preparedDead :: Set.Set v,
    forall v. Prepared v -> [Join v]
preparedJoins :: [Join v]
  }

-- | A join point of a lambda ('splitJoins'): an arm of a match on the
-- value of a match, as a combinator of the group whose parameters are
-- the arm's variables and then the variables the arm uses; the binders
-- of the unoptimized let whose binding the match was in, if any, which is
-- where the unoptimized frames of a continuation captured inside the
-- join point stop (the emitter's 'deoptOf'); and the source of the
-- match's variable, the constructor this arm is for over the arm's
-- variables (none for the default arm, whose frames never read it).
data Join v = Join
  { forall v. Join v -> SuperNormal Reference v
joinLambda :: SuperNormal Reference v,
    forall v. Join v -> Maybe [v]
joinCut :: Maybe [v],
    forall v. Join v -> Sources v
joinSources :: Sources v
  }

-- | Where a lambda's join points go: the lambda's index in its group,
-- the group's reference, and the names of the group's local functions,
-- which a body refers to as variables but which no frame holds.
data JoinSetting v = JoinSetting
  { forall v. JoinSetting v -> Word64
jsOwner :: Word64,
    forall v. JoinSetting v -> Reference
jsGroup :: Reference,
    forall v. JoinSetting v -> Set v
jsLocals :: Set.Set v
  }

-- | What 'splitJoins' knows of the unoptimized body: the setting, its
-- sites ('caseSites'), the products it built ('conLets'), and the
-- sources the rewrites before it recorded.
data Joining v = Joining (JoinSetting v) (Sites v) (ConLets v) (Sources v)

builtFields :: Built v -> [v]
builtFields :: forall v. Built v -> [v]
builtFields (BuiltCon Reference
_ CTag
_ [v]
fs) = [v]
fs
builtFields (BuiltRec RecordSchema
_ [v]
fs) = [v]
fs

builtShape :: Built v -> [Shape] -> Shape
builtShape :: forall v. Built v -> [Shape] -> Shape
builtShape (BuiltCon Reference
r CTag
t [v]
_) = Reference -> CTag -> [Shape] -> Shape
Con Reference
r CTag
t
builtShape (BuiltRec RecordSchema
rs [v]
_) = RecordSchema -> [Shape] -> Shape
Fields RecordSchema
rs

matchesBuilt :: Shape -> Built v -> Bool
matchesBuilt :: forall v. Shape -> Built v -> Bool
matchesBuilt (Con Reference
r CTag
t [Shape]
ss) (BuiltCon Reference
r' CTag
t' [v]
fs) = Reference
r Reference -> Reference -> Bool
forall a. Eq a => a -> a -> Bool
== Reference
r' Bool -> Bool -> Bool
&& CTag
t CTag -> CTag -> Bool
forall a. Eq a => a -> a -> Bool
== CTag
t' Bool -> Bool -> Bool
&& [Shape] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Shape]
ss Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
fs
matchesBuilt (Fields RecordSchema
rs [Shape]
ss) (BuiltRec RecordSchema
rs' [v]
fs) = RecordSchema
rs RecordSchema -> RecordSchema -> Bool
forall a. Eq a => a -> a -> Bool
== RecordSchema
rs' Bool -> Bool -> Bool
&& [Shape] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Shape]
ss Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
fs
matchesBuilt Shape
_ Built v
_ = Bool
False

-- | The product lets of a lambda: for each variable bound to a constructor
-- application, a record literal, or a record update of a variable whose
-- fields are known, the product and its field variables. Binders are
-- unique ('prepare').
type ConLets v = M.Map v (Built v)

conLets :: (Var v) => ANormal Reference v -> ConLets v
conLets :: forall v. Var v => ANormal Reference v -> ConLets v
conLets = Map v (Built v) -> ANormal Reference v -> Map v (Built v)
forall {k}.
Var k =>
Map k (Built k) -> ANormal Reference k -> Map k (Built k)
go Map v (Built v)
forall k a. Map k a
M.empty
  where
    go :: Map k (Built k) -> ANormal Reference k -> Map k (Built k)
go Map k (Built k)
acc ANormal Reference k
tm = case ANormal Reference k
tm of
      TLets Direction Word16
_ [k
u] [Mem]
_ ANormal Reference k
bn ANormal Reference k
bo
        | TApp (FCon Reference
r CTag
t) [k]
args <- ANormal Reference k
bn -> Map k (Built k) -> ANormal Reference k -> Map k (Built k)
go (k -> Built k -> Map k (Built k) -> Map k (Built k)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
u (Reference -> CTag -> [k] -> Built k
forall v. Reference -> CTag -> [v] -> Built v
BuiltCon Reference
r CTag
t [k]
args) Map k (Built k)
acc) ANormal Reference k
bo
        | TApp (FRec RecordSchema
rs) [k]
args <- ANormal Reference k
bn -> Map k (Built k) -> ANormal Reference k -> Map k (Built k)
go (k -> Built k -> Map k (Built k) -> Map k (Built k)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
u (RecordSchema -> [k] -> Built k
forall v. RecordSchema -> [v] -> Built v
BuiltRec RecordSchema
rs [k]
args) Map k (Built k)
acc) ANormal Reference k
bo
        | TApp (FRecSet Text
f) [k
base, k
x] <- ANormal Reference k
bn,
          Just (BuiltRec RecordSchema
rs [k]
fs) <- k -> Map k (Built k) -> Maybe (Built k)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
base Map k (Built k)
acc,
          Just Int
j <- Text -> [Text] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
elemIndex Text
f (RecordSchema -> [Text]
schemaFields RecordSchema
rs) ->
            Map k (Built k) -> ANormal Reference k -> Map k (Built k)
go (k -> Built k -> Map k (Built k) -> Map k (Built k)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
u (RecordSchema -> [k] -> Built k
forall v. RecordSchema -> [v] -> Built v
BuiltRec RecordSchema
rs (Int -> [k] -> [k]
forall a. Int -> [a] -> [a]
take Int
j [k]
fs [k] -> [k] -> [k]
forall a. [a] -> [a] -> [a]
++ k
x k -> [k] -> [k]
forall a. a -> [a] -> [a]
: Int -> [k] -> [k]
forall a. Int -> [a] -> [a]
drop (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [k]
fs)) Map k (Built k)
acc) ANormal Reference k
bo
      ANormal Reference k
_ -> (Map k (Built k) -> ANormal Reference k -> Map k (Built k))
-> Map k (Built k) -> [ANormal Reference k] -> Map k (Built k)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Map k (Built k) -> ANormal Reference k -> Map k (Built k)
go Map k (Built k)
acc (ANormal Reference k -> [ANormal Reference k]
forall v. ANormal Reference v -> [ANormal Reference v]
children ANormal Reference k
tm)

-- | Whether a let binds a product ('conLets' records it): a binding with
-- no effect, so that its let can be left out when nothing uses it
-- ('dropDead').
isBuiltLet :: (Var v) => ANormal Reference v -> Bool
isBuiltLet :: forall v. Var v => ANormal Reference v -> Bool
isBuiltLet = \case
  TApp (FCon Reference
_ CTag
_) [v]
_ -> Bool
True
  TApp (FRec RecordSchema
_) [v]
_ -> Bool
True
  TApp (FRecSet Text
_) [v]
_ -> Bool
True
  ANormal Reference v
_ -> Bool
False

schemaFields :: RecordSchema -> [FieldName]
schemaFields :: RecordSchema -> [Text]
schemaFields (RecordSchema Set Text
fields) = Set Text -> [Text]
forall a. Set a -> [a]
Set.toAscList Set Text
fields

-- | The shape a call site offers at an argument: the product the variable
-- was built from, its fields' shapes below, to a depth that bounds the
-- worker's arity.
siteShape :: (Var v) => ConLets v -> Int -> v -> Shape
siteShape :: forall v. Var v => ConLets v -> Int -> v -> Shape
siteShape ConLets v
cons Int
depth v
a
  | Int
depth Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
maxDepth, Just Built v
b <- v -> ConLets v -> Maybe (Built v)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup v
a ConLets v
cons = Built v -> [Shape] -> Shape
forall v. Built v -> [Shape] -> Shape
builtShape Built v
b ((v -> Shape) -> [v] -> [Shape]
forall a b. (a -> b) -> [a] -> [b]
map (ConLets v -> Int -> v -> Shape
forall v. Var v => ConLets v -> Int -> v -> Shape
siteShape ConLets v
cons (Int
depth Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) (Built v -> [v]
forall v. Built v -> [v]
builtFields Built v
b))
  | Bool
otherwise = Shape
Plain
  where
    maxDepth :: Int
maxDepth = Int
8

-- | The saturated calls of a body to functions of the batch, each with the
-- pattern its arguments offer, where some argument offers a product. A
-- call already rewritten to a worker's is not one.
sites :: (Var v) => Batch -> ANormal Reference v -> [(Reference, Pattern)]
sites :: forall v.
Var v =>
Batch -> ANormal Reference v -> [(Reference, [Shape])]
sites Batch
b ANormal Reference v
bd =
  [ (Reference
r, [Shape]
pat)
  | TApp (FComb Reference
r) [v]
args <- ANormal Reference v -> [ANormal Reference v]
forall v. ANormal Reference v -> [ANormal Reference v]
universe ANormal Reference v
bd,
    Just Int
n <- [Batch -> Reference -> Maybe Int
batchArity Batch
b Reference
r],
    Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
args,
    let pat :: [Shape]
pat = (v -> Shape) -> [v] -> [Shape]
forall a b. (a -> b) -> [a] -> [b]
map (ConLets v -> Int -> v -> Shape
forall v. Var v => ConLets v -> Int -> v -> Shape
siteShape ConLets v
cons Int
0) [v]
args,
    Bool -> Bool
not ((Shape -> Bool) -> [Shape] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Shape -> Bool
isPlain [Shape]
pat)
  ]
  where
    cons :: ConLets v
cons = ANormal Reference v -> ConLets v
forall v. Var v => ANormal Reference v -> ConLets v
conLets ANormal Reference v
bd

-- | The specialisations of a batch of groups: for each group's entry
-- point, the patterns its callers in the batch build, each narrowed to
-- what its body consumes, and then the patterns each worker's own body
-- builds, until there are no new ones.
explosions :: (Var v) => Returns -> M.Map Reference (SuperGroup Reference v) -> Explosions
explosions :: forall v.
Var v =>
Returns -> Map Reference (SuperGroup Reference v) -> Explosions
explosions Returns
rets Map Reference (SuperGroup Reference v)
groups
  | Bool -> Bool
not Bool
explodeArgsOn = Explosions
forall k a. Map k a
M.empty
  | Bool
otherwise = Explosions -> [(Reference, [Shape])] -> Explosions
close Explosions
forall k a. Map k a
M.empty [(Reference, [Shape])]
initial
  where
    b :: Batch
b = (Reference -> Maybe [[Shape]])
-> (Reference -> Maybe Shape) -> (Reference -> Maybe Int) -> Batch
Batch (Maybe [[Shape]] -> Reference -> Maybe [[Shape]]
forall a b. a -> b -> a
const Maybe [[Shape]]
forall a. Maybe a
Nothing) (Reference -> Returns -> Maybe Shape
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Returns
rets) (Map Reference (SuperGroup Reference v) -> Reference -> Maybe Int
forall v.
Map Reference (SuperGroup Reference v) -> Reference -> Maybe Int
groupArity Map Reference (SuperGroup Reference v)
groups)
    prep :: ANormal Reference v -> ANormal Reference v
prep ANormal Reference v
bd = let (ANormal Reference v
bd', Sources v
_, Map [v] [v]
_, Set v
_, [Join v]
_) = Batch
-> Maybe Shape
-> Maybe (Joining v)
-> ANormal Reference v
-> (ANormal Reference v, Sources v, Map [v] [v], Set v, [Join v])
forall v.
Var v =>
Batch
-> Maybe Shape
-> Maybe (Joining v)
-> ANormal Reference v
-> (ANormal Reference v, Sources v, Map [v] [v], Set v, [Join v])
prepareBody Batch
b Maybe Shape
forall a. Maybe a
Nothing Maybe (Joining v)
forall a. Maybe a
Nothing ANormal Reference v
bd in ANormal Reference v
bd'
    initial :: [(Reference, [Shape])]
initial = [[(Reference, [Shape])]] -> [(Reference, [Shape])]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [Batch -> ANormal Reference v -> [(Reference, [Shape])]
forall v.
Var v =>
Batch -> ANormal Reference v -> [(Reference, [Shape])]
sites Batch
b (ANormal Reference v -> ANormal Reference v
prep (SuperNormal Reference v -> ANormal Reference v
forall v. Var v => SuperNormal Reference v -> ANormal Reference v
lambdaBody SuperNormal Reference v
lam)) | SuperGroup Reference v
g <- Map Reference (SuperGroup Reference v) -> [SuperGroup Reference v]
forall k a. Map k a -> [a]
M.elems Map Reference (SuperGroup Reference v)
groups, SuperNormal Reference v
lam <- SuperGroup Reference v -> [SuperNormal Reference v]
forall v. SuperGroup Reference v -> [SuperNormal Reference v]
lambdas SuperGroup Reference v
g]
    close :: Explosions -> [(Reference, [Shape])] -> Explosions
close Explosions
specs [] = Explosions
specs
    close Explosions
specs ((Reference
r, [Shape]
site) : [(Reference, [Shape])]
rest)
      | Just (Rec [(v, SuperNormal Reference v)]
_ ent :: SuperNormal Reference v
ent@(Lambda [Mem]
_ (TAbss [v]
vs ANormal Reference v
_))) <- Reference
-> Map Reference (SuperGroup Reference v)
-> Maybe (SuperGroup Reference v)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Reference
r Map Reference (SuperGroup Reference v)
groups,
        [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
== [Shape] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Shape]
site,
        let pat :: [Shape]
pat = (Shape -> v -> Shape) -> [Shape] -> [v] -> [Shape]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Shape
s v
v -> Shape -> v -> ANormal Reference v -> Shape
forall v. Var v => Shape -> v -> ANormal Reference v -> Shape
usable Shape
s v
v (ANormal Reference v -> ANormal Reference v
prep (SuperNormal Reference v -> ANormal Reference v
forall v. Var v => SuperNormal Reference v -> ANormal Reference v
lambdaBody SuperNormal Reference v
ent))) [Shape]
site [v]
vs,
        Bool -> Bool
not ((Shape -> Bool) -> [Shape] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Shape -> Bool
isPlain [Shape]
pat),
        let have :: [[Shape]]
have = [[Shape]] -> Reference -> Explosions -> [[Shape]]
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault [] Reference
r Explosions
specs,
        [Shape]
pat [Shape] -> [[Shape]] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [[Shape]]
have,
        [[Shape]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Shape]]
have Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
maxSpecs =
          Explosions -> [(Reference, [Shape])] -> Explosions
close (Reference -> [[Shape]] -> Explosions -> Explosions
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Reference
r ([[Shape]]
have [[Shape]] -> [[Shape]] -> [[Shape]]
forall a. [a] -> [a] -> [a]
++ [[Shape]
pat]) Explosions
specs) ([(Reference, [Shape])]
rest [(Reference, [Shape])]
-> [(Reference, [Shape])] -> [(Reference, [Shape])]
forall a. [a] -> [a] -> [a]
++ Batch -> ANormal Reference v -> [(Reference, [Shape])]
forall v.
Var v =>
Batch -> ANormal Reference v -> [(Reference, [Shape])]
sites Batch
b (SuperNormal Reference v -> ANormal Reference v
forall v. Var v => SuperNormal Reference v -> ANormal Reference v
lambdaBody (Prepared v -> SuperNormal Reference v
forall v. Prepared v -> SuperNormal Reference v
preparedLambda (Batch
-> Maybe (JoinSetting v)
-> [Shape]
-> SuperNormal Reference v
-> Prepared v
forall v.
Var v =>
Batch
-> Maybe (JoinSetting v)
-> [Shape]
-> SuperNormal Reference v
-> Prepared v
specialise Batch
b Maybe (JoinSetting v)
forall a. Maybe a
Nothing [Shape]
pat SuperNormal Reference v
ent))))
      | Bool
otherwise = Explosions -> [(Reference, [Shape])] -> Explosions
close Explosions
specs [(Reference, [Shape])]
rest

-- | The arity of a group's entry point.
groupArity :: M.Map Reference (SuperGroup Reference v) -> Reference -> Maybe Int
groupArity :: forall v.
Map Reference (SuperGroup Reference v) -> Reference -> Maybe Int
groupArity Map Reference (SuperGroup Reference v)
groups Reference
r = case SuperGroup Reference v -> [Int]
forall ref v. SuperGroup ref v -> [Int]
ANF.arities (SuperGroup Reference v -> [Int])
-> Maybe (SuperGroup Reference v) -> Maybe [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Reference
-> Map Reference (SuperGroup Reference v)
-> Maybe (SuperGroup Reference v)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Reference
r Map Reference (SuperGroup Reference v)
groups of
  Just (Int
n : [Int]
_) -> Int -> Maybe Int
forall a. a -> Maybe a
Just Int
n
  Maybe [Int]
_ -> Maybe Int
forall a. Maybe a
Nothing

-- | The product each function of a batch returns on every path, by its
-- reference: the shape every tail of its body builds (a constructor
-- application or a record literal, or a variable bound to one, its fields
-- taken apart in turn where they are built for it), met over the tails,
-- where the other tails are saturated calls of functions of the batch that
-- return the same product. A closure from the shapes the tails build: a
-- callee's shape narrows its callers' until nothing changes, and a
-- function with a tail that returns anything else, or calls a function
-- that does, is dropped. Products of one leaf are left to the wide
-- convention and products of none are nothing to return; a function of no
-- arguments is a value evaluated once when it loads, so it keeps its entry.
type Returns = M.Map Reference Shape

returns :: (Var v) => M.Map Reference (SuperGroup Reference v) -> Returns
returns :: forall v.
Var v =>
Map Reference (SuperGroup Reference v) -> Returns
returns Map Reference (SuperGroup Reference v)
groups
  | Bool -> Bool
not Bool
explodeReturnsOn = Returns
forall k a. Map k a
M.empty
  | Bool
otherwise = ((Shape, [Reference]) -> Shape)
-> Map Reference (Shape, [Reference]) -> Returns
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (Shape, [Reference]) -> Shape
forall a b. (a, b) -> a
fst (Map Reference (Shape, [Reference])
-> Map Reference (Shape, [Reference])
forall {t :: * -> *} {k}.
(Traversable t, Ord k, Eq (t k)) =>
Map k (Shape, t k) -> Map k (Shape, t k)
settle ((Reference -> SuperGroup Reference v -> Maybe (Shape, [Reference]))
-> Map Reference (SuperGroup Reference v)
-> Map Reference (Shape, [Reference])
forall k a b. (k -> a -> Maybe b) -> Map k a -> Map k b
M.mapMaybeWithKey Reference -> SuperGroup Reference v -> Maybe (Shape, [Reference])
candidate Map Reference (SuperGroup Reference v)
groups))
  where
    b :: Batch
b = (Reference -> Maybe [[Shape]])
-> (Reference -> Maybe Shape) -> (Reference -> Maybe Int) -> Batch
Batch (Maybe [[Shape]] -> Reference -> Maybe [[Shape]]
forall a b. a -> b -> a
const Maybe [[Shape]]
forall a. Maybe a
Nothing) (Maybe Shape -> Reference -> Maybe Shape
forall a b. a -> b -> a
const Maybe Shape
forall a. Maybe a
Nothing) (Map Reference (SuperGroup Reference v) -> Reference -> Maybe Int
forall v.
Map Reference (SuperGroup Reference v) -> Reference -> Maybe Int
groupArity Map Reference (SuperGroup Reference v)
groups)
    candidate :: Reference -> SuperGroup Reference v -> Maybe (Shape, [Reference])
candidate Reference
r (Rec [(v, SuperNormal Reference v)]
_ (Lambda [Mem]
_ (TAbss [v]
vs Term (ANormalF Reference) v
bd0)))
      | [v] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [v]
vs = Maybe (Shape, [Reference])
forall a. Maybe a
Nothing
      | Just Int
n <- Batch -> Reference -> Maybe Int
batchArity Batch
b Reference
r, Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
vs = Maybe (Shape, [Reference])
forall a. Maybe a
Nothing
      | Reference -> Term (ANormalF Reference) v -> Bool
selfResultEscapes Reference
r Term (ANormalF Reference) v
bd = Maybe (Shape, [Reference])
forall a. Maybe a
Nothing
      | Just [Either Shape Reference]
ts <- [Maybe (Either Shape Reference)] -> Maybe [Either Shape Reference]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence (Batch
-> ConLets v
-> Term (ANormalF Reference) v
-> [Maybe (Either Shape Reference)]
forall v.
Var v =>
Batch
-> ConLets v
-> ANormal Reference v
-> [Maybe (Either Shape Reference)]
tails Batch
b ConLets v
cons Term (ANormalF Reference) v
bd),
        built :: [Shape]
built@(Shape
s0 : [Shape]
_) <- [Shape
s | Left Shape
s <- [Either Shape Reference]
ts],
        let s :: Shape
s = (Shape -> Shape -> Shape) -> Shape -> [Shape] -> Shape
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Shape -> Shape -> Shape
meetShape Shape
s0 [Shape]
built,
        Shape -> Bool
returnable Shape
s =
          (Shape, [Reference]) -> Maybe (Shape, [Reference])
forall a. a -> Maybe a
Just (Shape
s, [Reference
r' | Right Reference
r' <- [Either Shape Reference]
ts])
      | Bool
otherwise = Maybe (Shape, [Reference])
forall a. Maybe a
Nothing
      where
        (Term (ANormalF Reference) v
bd, Sources v
_, Map [v] [v]
_, Set v
_, [Join v]
_) = Batch
-> Maybe Shape
-> Maybe (Joining v)
-> Term (ANormalF Reference) v
-> (Term (ANormalF Reference) v, Sources v, Map [v] [v], Set v,
    [Join v])
forall v.
Var v =>
Batch
-> Maybe Shape
-> Maybe (Joining v)
-> ANormal Reference v
-> (ANormal Reference v, Sources v, Map [v] [v], Set v, [Join v])
prepareBody Batch
b Maybe Shape
forall a. Maybe a
Nothing Maybe (Joining v)
forall a. Maybe a
Nothing Term (ANormalF Reference) v
bd0
        cons :: ConLets v
cons = Term (ANormalF Reference) v -> ConLets v
forall v. Var v => ANormal Reference v -> ConLets v
conLets Term (ANormalF Reference) v
bd
    -- A function that lets the result of a saturated call of itself and
    -- uses it whole (a field of what a tail builds, as in a tree insert
    -- whose recursive call rebuilds the path) gets no worker: its own
    -- recursion would yield the leaves and pack them again at every
    -- level, and a result that escapes at once is better boxed by the
    -- callee. A result taken apart (matched, or passed to a function of
    -- the batch, where an exploded position may take it) is what the
    -- worker is for.
    selfResultEscapes :: Reference -> Term (ANormalF Reference) v -> Bool
selfResultEscapes Reference
r = Term (ANormalF Reference) v -> Bool
go
      where
        go :: Term (ANormalF Reference) v -> Bool
go Term (ANormalF Reference) v
tm = case Term (ANormalF Reference) v
tm of
          TLets Direction Word16
_ [v
u] [Mem]
_ (TApp (FComb Reference
r') [v]
args) Term (ANormalF Reference) v
bo
            | Reference
r' Reference -> Reference -> Bool
forall a. Eq a => a -> a -> Bool
== Reference
r,
              Batch -> Reference -> Maybe Int
batchArity Batch
b Reference
r Maybe Int -> Maybe Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int -> Maybe Int
forall a. a -> Maybe a
Just ([v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
args),
              v -> Term (ANormalF Reference) v -> Int
forall v. Var v => v -> ANormal Reference v -> Int
occurrences v
u Term (ANormalF Reference) v
bo Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= v -> Term (ANormalF Reference) v -> Int
forall {a} {c}. (Var a, Num c) => a -> ANormal Reference a -> c
matchUses v
u Term (ANormalF Reference) v
bo Int -> Int -> Int
forall a. Num a => a -> a -> a
+ v -> Term (ANormalF Reference) v -> Int
callUses v
u Term (ANormalF Reference) v
bo ->
                Bool
True
          TLets Direction Word16
_ [v]
_ [Mem]
_ Term (ANormalF Reference) v
bn Term (ANormalF Reference) v
bo -> Term (ANormalF Reference) v -> Bool
go Term (ANormalF Reference) v
bn Bool -> Bool -> Bool
|| Term (ANormalF Reference) v -> Bool
go Term (ANormalF Reference) v
bo
          ABTN.Term Set v
_ (ABTN.Abs v
_ Term (ANormalF Reference) v
e) -> Term (ANormalF Reference) v -> Bool
go Term (ANormalF Reference) v
e
          ABTN.Term Set v
_ (ABTN.Tm ANormalF Reference v (Term (ANormalF Reference) v)
f) -> (Term (ANormalF Reference) v -> Bool)
-> ANormalF Reference v (Term (ANormalF Reference) v) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Term (ANormalF Reference) v -> Bool
go ANormalF Reference v (Term (ANormalF Reference) v)
f
        matchUses :: a -> ANormal Reference a -> c
matchUses a
u = Sum c -> c
forall a. Sum a -> a
getSum (Sum c -> c)
-> (ANormal Reference a -> Sum c) -> ANormal Reference a -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a
-> (ANormalF Reference a (ANormal Reference a) -> Sum c)
-> ANormal Reference a
-> Sum c
forall v m.
(Var v, Monoid m) =>
v
-> (ANormalF Reference v (ANormal Reference v) -> m)
-> ANormal Reference v
-> m
foldMentions a
u \case AMatch a
v Branched Reference (ANormal Reference a)
_ | a
v a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
u -> c -> Sum c
forall a. a -> Sum a
Sum c
1; ANormalF Reference a (ANormal Reference a)
_ -> c -> Sum c
forall a. a -> Sum a
Sum c
0
        callUses :: v -> Term (ANormalF Reference) v -> Int
callUses v
u =
          Sum Int -> Int
forall a. Sum a -> a
getSum (Sum Int -> Int)
-> (Term (ANormalF Reference) v -> Sum Int)
-> Term (ANormalF Reference) v
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v
-> (ANormalF Reference v (Term (ANormalF Reference) v) -> Sum Int)
-> Term (ANormalF Reference) v
-> Sum Int
forall v m.
(Var v, Monoid m) =>
v
-> (ANormalF Reference v (ANormal Reference v) -> m)
-> ANormal Reference v
-> m
foldMentions v
u \case
            AApp (FComb Reference
r') [v]
args | Batch -> Reference -> Maybe Int
batchArity Batch
b Reference
r' Maybe Int -> Maybe Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int -> Maybe Int
forall a. a -> Maybe a
Just ([v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
args) -> Int -> Sum Int
forall a. a -> Sum a
Sum ([v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ((v -> Bool) -> [v] -> [v]
forall a. (a -> Bool) -> [a] -> [a]
filter (v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
u) [v]
args))
            ANormalF Reference v (Term (ANormalF Reference) v)
_ -> Int -> Sum Int
forall a. a -> Sum a
Sum Int
0
    settle :: Map k (Shape, t k) -> Map k (Shape, t k)
settle Map k (Shape, t k)
cs
      | Map k (Shape, t k)
cs' Map k (Shape, t k) -> Map k (Shape, t k) -> Bool
forall a. Eq a => a -> a -> Bool
== Map k (Shape, t k)
cs = Map k (Shape, t k)
cs
      | Bool
otherwise = Map k (Shape, t k) -> Map k (Shape, t k)
settle Map k (Shape, t k)
cs'
      where
        cs' :: Map k (Shape, t k)
cs' = ((Shape, t k) -> Maybe (Shape, t k))
-> Map k (Shape, t k) -> Map k (Shape, t k)
forall a b k. (a -> Maybe b) -> Map k a -> Map k b
M.mapMaybe (Shape, t k) -> Maybe (Shape, t k)
step Map k (Shape, t k)
cs
        step :: (Shape, t k) -> Maybe (Shape, t k)
step (Shape
s, t k
calls) = do
          t Shape
ss <- (k -> Maybe Shape) -> t k -> Maybe (t Shape)
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) -> t a -> f (t b)
traverse (\k
r -> (Shape, t k) -> Shape
forall a b. (a, b) -> a
fst ((Shape, t k) -> Shape) -> Maybe (Shape, t k) -> Maybe Shape
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> k -> Map k (Shape, t k) -> Maybe (Shape, t k)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
r Map k (Shape, t k)
cs) t k
calls
          let s' :: Shape
s' = (Shape -> Shape -> Shape) -> Shape -> t Shape -> Shape
forall b a. (b -> a -> b) -> b -> t a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Shape -> Shape -> Shape
meetShape Shape
s t Shape
ss
          if Shape -> Bool
returnable Shape
s' then (Shape, t k) -> Maybe (Shape, t k)
forall a. a -> Maybe a
Just (Shape
s', t k
calls) else Maybe (Shape, t k)
forall a. Maybe a
Nothing
    returnable :: Shape -> Bool
returnable Shape
s = Bool -> Bool
not (Shape -> Bool
isPlain Shape
s) Bool -> Bool -> Bool
&& Shape -> Int
shapeLeaves Shape
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2 Bool -> Bool -> Bool
&& Shape -> Int
shapeLeaves Shape
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxLeaves

-- | The tails of a body: the shape each builds (Left), the function of the
-- batch each calls saturated (Right), or Nothing for any other tail.
-- 'yieldLeaves' rewrites exactly these.
tails :: (Var v) => Batch -> ConLets v -> ANormal Reference v -> [Maybe (Either Shape Reference)]
tails :: forall v.
Var v =>
Batch
-> ConLets v
-> ANormal Reference v
-> [Maybe (Either Shape Reference)]
tails Batch
b ConLets v
cons = ANormal Reference v -> [Maybe (Either Shape Reference)]
go
  where
    go :: ANormal Reference v -> [Maybe (Either Shape Reference)]
go ANormal Reference v
tm = case ANormal Reference v
tm of
      TLets Direction Word16
_ [v]
_ [Mem]
_ ANormal Reference v
_ ANormal Reference v
bo -> ANormal Reference v -> [Maybe (Either Shape Reference)]
go ANormal Reference v
bo
      TName v
_ Either Reference v
_ [v]
_ ANormal Reference v
bo -> ANormal Reference v -> [Maybe (Either Shape Reference)]
go ANormal Reference v
bo
      TLocal v
_ ANormal Reference v
bo -> ANormal Reference v -> [Maybe (Either Shape Reference)]
go ANormal Reference v
bo
      TMatch v
_ Branched Reference (ANormal Reference v)
bs -> (ANormal Reference v -> [Maybe (Either Shape Reference)])
-> [ANormal Reference v] -> [Maybe (Either Shape Reference)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(TAbss [v]
_ ANormal Reference v
bo) -> ANormal Reference v -> [Maybe (Either Shape Reference)]
go ANormal Reference v
bo) (Branched Reference (ANormal Reference v) -> [ANormal Reference v]
forall a. Branched Reference a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Branched Reference (ANormal Reference v)
bs)
      TVar v
u | Just Built v
_ <- v -> ConLets v -> Maybe (Built v)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup v
u ConLets v
cons -> [Either Shape Reference -> Maybe (Either Shape Reference)
forall a. a -> Maybe a
Just (Shape -> Either Shape Reference
forall a b. a -> Either a b
Left (ConLets v -> Int -> v -> Shape
forall v. Var v => ConLets v -> Int -> v -> Shape
siteShape ConLets v
cons Int
0 v
u))]
      TApp (FCon Reference
r CTag
t) [v]
args -> [Either Shape Reference -> Maybe (Either Shape Reference)
forall a. a -> Maybe a
Just (Shape -> Either Shape Reference
forall a b. a -> Either a b
Left (Reference -> CTag -> [Shape] -> Shape
Con Reference
r CTag
t ((v -> Shape) -> [v] -> [Shape]
forall a b. (a -> b) -> [a] -> [b]
map (ConLets v -> Int -> v -> Shape
forall v. Var v => ConLets v -> Int -> v -> Shape
siteShape ConLets v
cons Int
1) [v]
args)))]
      TApp (FRec RecordSchema
rs) [v]
args -> [Either Shape Reference -> Maybe (Either Shape Reference)
forall a. a -> Maybe a
Just (Shape -> Either Shape Reference
forall a b. a -> Either a b
Left (RecordSchema -> [Shape] -> Shape
Fields RecordSchema
rs ((v -> Shape) -> [v] -> [Shape]
forall a b. (a -> b) -> [a] -> [b]
map (ConLets v -> Int -> v -> Shape
forall v. Var v => ConLets v -> Int -> v -> Shape
siteShape ConLets v
cons Int
1) [v]
args)))]
      TApp (FComb Reference
r) [v]
args | Batch -> Reference -> Maybe Int
batchArity Batch
b Reference
r Maybe Int -> Maybe Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int -> Maybe Int
forall a. a -> Maybe a
Just ([v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
args) -> [Either Shape Reference -> Maybe (Either Shape Reference)
forall a. a -> Maybe a
Just (Reference -> Either Shape Reference
forall a b. b -> Either a b
Right Reference
r)]
      ANormal Reference v
_ -> [Maybe (Either Shape Reference)
forall a. Maybe a
Nothing]

-- | The values a product-return worker yields for a tail that builds a
-- product: the leaves of the shape, a field built for the tail taken apart
-- in turn, any other field as itself.
returnLeaves :: (Var v) => ConLets v -> Shape -> Built v -> [v]
returnLeaves :: forall v. Var v => ConLets v -> Shape -> Built v -> [v]
returnLeaves ConLets v
cons Shape
s Built v
b = [[v]] -> [v]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ((Shape -> v -> [v]) -> [Shape] -> [v] -> [[v]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Shape -> v -> [v]
leaf (Shape -> [Shape]
subShapes Shape
s) (Built v -> [v]
forall v. Built v -> [v]
builtFields Built v
b))
  where
    leaf :: Shape -> v -> [v]
leaf Shape
Plain v
a = [v
a]
    leaf Shape
sub v
a
      | Just Built v
b' <- v -> ConLets v -> Maybe (Built v)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup v
a ConLets v
cons, Shape -> Built v -> Bool
forall v. Shape -> Built v -> Bool
matchesBuilt Shape
sub Built v
b' = ConLets v -> Shape -> Built v -> [v]
forall v. Var v => ConLets v -> Shape -> Built v -> [v]
returnLeaves ConLets v
cons Shape
sub Built v
b'
      | Bool
otherwise = [v
a]

-- | The body of a product-return worker for a shape: each tail that builds
-- the product (the constructor application or record literal itself, or a
-- variable bound to one) yields its leaves as that many values, a field
-- built for the tail taken apart in turn ('returnLeaves'), and each tail
-- call of a function with a worker for the same product calls that
-- worker. Only the tails change: a product built in a let's binding, or
-- under a handler or a shift, is packed as before. The analysis
-- ('returns', through 'tails') chose the shape from these tails, so a tail
-- of any other form is a bug.
yieldLeaves :: (Var v) => Batch -> Shape -> ANormal Reference v -> ANormal Reference v
yieldLeaves :: forall v.
Var v =>
Batch -> Shape -> ANormal Reference v -> ANormal Reference v
yieldLeaves Batch
b Shape
s ANormal Reference v
bd = ANormal Reference v -> ANormal Reference v
go ANormal Reference v
bd
  where
    cons :: ConLets v
cons = ANormal Reference v -> ConLets v
forall v. Var v => ANormal Reference v -> ConLets v
conLets ANormal Reference v
bd
    go :: ANormal Reference v -> ANormal Reference v
go ANormal Reference v
tm = case ANormal Reference v
tm of
      TLets Direction Word16
d [v]
us [Mem]
ms ANormal Reference v
bn ANormal Reference v
bo -> Direction Word16
-> [v]
-> [Mem]
-> ANormal Reference v
-> ANormal Reference v
-> ANormal Reference v
forall v ref.
Var v =>
Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
TLets Direction Word16
d [v]
us [Mem]
ms ANormal Reference v
bn (ANormal Reference v -> ANormal Reference v
go ANormal Reference v
bo)
      TName v
v Either Reference v
f [v]
as ANormal Reference v
bo -> v
-> Either Reference v
-> [v]
-> ANormal Reference v
-> ANormal Reference v
forall v ref.
Var v =>
v -> Either ref v -> [v] -> ANormal ref v -> ANormal ref v
TName v
v Either Reference v
f [v]
as (ANormal Reference v -> ANormal Reference v
go ANormal Reference v
bo)
      TLocal v
v ANormal Reference v
bo -> v -> ANormal Reference v -> ANormal Reference v
forall v ref. Var v => v -> ANormal ref v -> ANormal ref v
TLocal v
v (ANormal Reference v -> ANormal Reference v
go ANormal Reference v
bo)
      TMatch v
v Branched Reference (ANormal Reference v)
bs -> v
-> Branched Reference (ANormal Reference v) -> ANormal Reference v
forall v ref.
Var v =>
v -> Branched ref (ANormal ref v) -> ANormal ref v
TMatch v
v ((ANormal Reference v -> ANormal Reference v)
-> Branched Reference (ANormal Reference v)
-> Branched Reference (ANormal Reference v)
forall a b.
(a -> b) -> Branched Reference a -> Branched Reference b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ANormal Reference v -> ANormal Reference v
go Branched Reference (ANormal Reference v)
bs)
      ABTN.Term Set v
_ (ABTN.Abs v
u ANormal Reference v
e) -> v -> ANormal Reference v -> ANormal Reference v
forall v (f :: * -> * -> *). Var v => v -> Term f v -> Term f v
ABTN.TAbs v
u (ANormal Reference v -> ANormal Reference v
go ANormal Reference v
e)
      TVar v
u | Just Built v
bt <- v -> ConLets v -> Maybe (Built v)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup v
u ConLets v
cons, Shape -> Built v -> Bool
forall v. Shape -> Built v -> Bool
matchesBuilt Shape
s Built v
bt -> [v] -> ANormal Reference v
yield (ConLets v -> Shape -> Built v -> [v]
forall v. Var v => ConLets v -> Shape -> Built v -> [v]
returnLeaves ConLets v
cons Shape
s Built v
bt)
      TApp (FCon Reference
r CTag
t) [v]
args | Shape -> Built v -> Bool
forall v. Shape -> Built v -> Bool
matchesBuilt Shape
s (Reference -> CTag -> [v] -> Built v
forall v. Reference -> CTag -> [v] -> Built v
BuiltCon Reference
r CTag
t [v]
args) -> [v] -> ANormal Reference v
yield (ConLets v -> Shape -> Built v -> [v]
forall v. Var v => ConLets v -> Shape -> Built v -> [v]
returnLeaves ConLets v
cons Shape
s (Reference -> CTag -> [v] -> Built v
forall v. Reference -> CTag -> [v] -> Built v
BuiltCon Reference
r CTag
t [v]
args))
      TApp (FRec RecordSchema
rs) [v]
args | Shape -> Built v -> Bool
forall v. Shape -> Built v -> Bool
matchesBuilt Shape
s (RecordSchema -> [v] -> Built v
forall v. RecordSchema -> [v] -> Built v
BuiltRec RecordSchema
rs [v]
args) -> [v] -> ANormal Reference v
yield (ConLets v -> Shape -> Built v -> [v]
forall v. Var v => ConLets v -> Shape -> Built v -> [v]
returnLeaves ConLets v
cons Shape
s (RecordSchema -> [v] -> Built v
forall v. RecordSchema -> [v] -> Built v
BuiltRec RecordSchema
rs [v]
args))
      TApp (FComb Reference
r) [v]
args | Just Shape
_ <- Batch -> Reference -> Maybe Shape
batchRets Batch
b Reference
r, Batch -> Reference -> Maybe Int
batchArity Batch
b Reference
r Maybe Int -> Maybe Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int -> Maybe Int
forall a. a -> Maybe a
Just ([v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
args) -> Func Reference v -> [v] -> ANormal Reference v
forall v ref. Var v => Func ref v -> [v] -> ANormal ref v
TApp (Reference -> Word64 -> Func Reference v
forall ref v. ref -> Word64 -> Func ref v
FWorker Reference
r Word64
retIndex) [v]
args
      ANormal Reference v
_ -> [Word] -> String -> ANormal Reference v
forall a. HasCallStack => [Word] -> String -> a
internalBug [] (String -> ANormal Reference v) -> String -> ANormal Reference v
forall a b. (a -> b) -> a -> b
$ String
"yieldLeaves: a tail of a product-return worker that does not build " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Shape -> String
forall a. Show a => a -> String
show Shape
s String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
": " String -> ShowS
forall a. [a] -> [a] -> [a]
++ ANormal Reference v -> String
forall a. Show a => a -> String
show ANormal Reference v
tm
    yield :: [v] -> ANormal Reference v
yield [v]
leaves
      | [v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
leaves Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Shape -> Int
shapeLeaves Shape
s = Func Reference v -> [v] -> ANormal Reference v
forall v ref. Var v => Func ref v -> [v] -> ANormal ref v
TApp Func Reference v
forall ref v. Func ref v
FYield [v]
leaves
      | Bool
otherwise = [Word] -> String -> ANormal Reference v
forall a. HasCallStack => [Word] -> String -> a
internalBug [] (String -> ANormal Reference v) -> String -> ANormal Reference v
forall a b. (a -> b) -> a -> b
$ String
"yieldLeaves: a tail building " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show ([v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
leaves) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" leaves for a product of " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (Shape -> Int
shapeLeaves Shape
s)

-- | A body with each let of a saturated call of a function that has a
-- product-return worker rewritten to a let of the result's leaves over a
-- call of the worker, when the body takes the result apart (every match
-- on it, and on the fields the shape takes apart, replaced by its arm:
-- 'rewrite') and what is left of it is only passed to saturated calls of
-- functions of the batch, where an exploded position may take it; the
-- product, and any field the body still uses as a whole, is rebuilt from
-- the leaves before the body ('rebuild'). A result used any other way
-- keeps the call of the entry.
useReturns :: (Var v) => Batch -> ANormal Reference v -> (ANormal Reference v, Sources v, M.Map [v] [v])
useReturns :: forall v.
Var v =>
Batch
-> ANormal Reference v
-> (ANormal Reference v, Sources v, Map [v] [v])
useReturns Batch
b ANormal Reference v
bd = (ANormal Reference v
bd', Map v (Either v (Built v))
srcs, Map [v] [v]
lets)
  where
    (ANormal Reference v
bd', (Set v
_, Map v (Either v (Built v))
srcs, Map [v] [v]
lets)) = State
  (Set v, Map v (Either v (Built v)), Map [v] [v])
  (ANormal Reference v)
-> (Set v, Map v (Either v (Built v)), Map [v] [v])
-> (ANormal Reference v,
    (Set v, Map v (Either v (Built v)), Map [v] [v]))
forall s a. State s a -> s -> (a, s)
runState (ANormal Reference v
-> State
     (Set v, Map v (Either v (Built v)), Map [v] [v])
     (ANormal Reference v)
go ANormal Reference v
bd) ([Set v] -> Set v
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions (ANormal Reference v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
ABTN.freeVars ANormal Reference v
bd Set v -> [Set v] -> [Set v]
forall a. a -> [a] -> [a]
: (ANormal Reference v -> Set v) -> [ANormal Reference v] -> [Set v]
forall a b. (a -> b) -> [a] -> [b]
map ANormal Reference v -> Set v
forall v. Ord v => ANormal Reference v -> Set v
boundVars (ANormal Reference v -> [ANormal Reference v]
forall v. ANormal Reference v -> [ANormal Reference v]
universe ANormal Reference v
bd)), Map v (Either v (Built v))
forall k a. Map k a
M.empty, Map [v] [v]
forall k a. Map k a
M.empty)
    go :: ANormal Reference v
-> State
     (Set v, Map v (Either v (Built v)), Map [v] [v])
     (ANormal Reference v)
go ANormal Reference v
tm = case ANormal Reference v
tm of
      TLets Direction Word16
d [v
u] [Mem
BX] bn :: ANormal Reference v
bn@(TApp (FComb Reference
r) [v]
args) ANormal Reference v
bo
        | Just Shape
s <- Batch -> Reference -> Maybe Shape
batchRets Batch
b Reference
r,
          Batch -> Reference -> Maybe Int
batchArity Batch
b Reference
r Maybe Int -> Maybe Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int -> Maybe Int
forall a. a -> Maybe a
Just ([v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
args) -> do
            (Set v
avoid, Map v (Either v (Built v))
srcs0, Map [v] [v]
lets0) <- StateT
  (Set v, Map v (Either v (Built v)), Map [v] [v])
  Identity
  (Set v, Map v (Either v (Built v)), Map [v] [v])
forall s (m :: * -> *). MonadState s m => m s
get
            let (Tree v
tree0, Set v
avoid') = Set v -> Shape -> (Tree v, Set v)
forall v. Var v => Set v -> Shape -> (Tree v, Set v)
treeOf Set v
avoid Shape
s
                tree :: Tree v
tree = v -> Tree v -> Tree v
forall {v}. v -> Tree v -> Tree v
renameRoot v
u Tree v
tree0
                (ANormal Reference v
rewritten, Map v (Either v (Built v))
armSrcs) = State (Map v (Either v (Built v))) (ANormal Reference v)
-> Map v (Either v (Built v))
-> (ANormal Reference v, Map v (Either v (Built v)))
forall s a. State s a -> s -> (a, s)
runState (v
-> Tree v
-> ANormal Reference v
-> State (Map v (Either v (Built v))) (ANormal Reference v)
forall v.
Var v =>
v
-> Tree v
-> ANormal Reference v
-> State (Sources v) (ANormal Reference v)
rewrite v
u Tree v
tree ANormal Reference v
bo) Map v (Either v (Built v))
forall k a. Map k a
M.empty
                bo' :: ANormal Reference v
bo' = Tree v -> ANormal Reference v -> ANormal Reference v
forall v.
Var v =>
Tree v -> ANormal Reference v -> ANormal Reference v
rebuild Tree v
tree ANormal Reference v
rewritten
            -- the uses are counted before 'rebuild', which binds the
            -- variable again; counted after it, there are none
            if v -> ANormal Reference v -> Int
forall v. Var v => v -> ANormal Reference v -> Int
occurrences v
u ANormal Reference v
rewritten Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== v -> ANormal Reference v -> Int
callUses v
u ANormal Reference v
rewritten
              then do
                (Set v, Map v (Either v (Built v)), Map [v] [v])
-> StateT
     (Set v, Map v (Either v (Built v)), Map [v] [v]) Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Set v
avoid', [Map v (Either v (Built v))] -> Map v (Either v (Built v))
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
M.unions [Map v (Either v (Built v))
srcs0, Map v (Either v (Built v))
armSrcs, v -> Tree v -> Map v (Either v (Built v))
forall v. Ord v => v -> Tree v -> Sources v
treeSources v
u Tree v
tree], [v] -> [v] -> Map [v] [v] -> Map [v] [v]
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (Tree v -> [v]
forall v. Tree v -> [v]
leafNames Tree v
tree) [v
u] Map [v] [v]
lets0)
                Direction Word16
-> [v]
-> [Mem]
-> ANormal Reference v
-> ANormal Reference v
-> ANormal Reference v
forall v ref.
Var v =>
Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
TLets Direction Word16
d (Tree v -> [v]
forall v. Tree v -> [v]
leafNames Tree v
tree) ((v -> Mem) -> [v] -> [Mem]
forall a b. (a -> b) -> [a] -> [b]
map (Mem -> v -> Mem
forall a b. a -> b -> a
const Mem
BX) (Tree v -> [v]
forall v. Tree v -> [v]
leafNames Tree v
tree)) (Func Reference v -> [v] -> ANormal Reference v
forall v ref. Var v => Func ref v -> [v] -> ANormal ref v
TApp (Reference -> Word64 -> Func Reference v
forall ref v. ref -> Word64 -> Func ref v
FWorker Reference
r Word64
retIndex) [v]
args) (ANormal Reference v -> ANormal Reference v)
-> State
     (Set v, Map v (Either v (Built v)), Map [v] [v])
     (ANormal Reference v)
-> State
     (Set v, Map v (Either v (Built v)), Map [v] [v])
     (ANormal Reference v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ANormal Reference v
-> State
     (Set v, Map v (Either v (Built v)), Map [v] [v])
     (ANormal Reference v)
go ANormal Reference v
bo'
              else do
                (Set v, Map v (Either v (Built v)), Map [v] [v])
-> StateT
     (Set v, Map v (Either v (Built v)), Map [v] [v]) Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Set v
avoid', Map v (Either v (Built v))
srcs0, Map [v] [v]
lets0)
                Direction Word16
-> [v]
-> [Mem]
-> ANormal Reference v
-> ANormal Reference v
-> ANormal Reference v
forall v ref.
Var v =>
Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
TLets Direction Word16
d [v
u] [Mem
BX] ANormal Reference v
bn (ANormal Reference v -> ANormal Reference v)
-> State
     (Set v, Map v (Either v (Built v)), Map [v] [v])
     (ANormal Reference v)
-> State
     (Set v, Map v (Either v (Built v)), Map [v] [v])
     (ANormal Reference v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ANormal Reference v
-> State
     (Set v, Map v (Either v (Built v)), Map [v] [v])
     (ANormal Reference v)
go ANormal Reference v
bo
      ABTN.Term Set v
_ (ABTN.Abs v
v ANormal Reference v
e) -> v -> ANormal Reference v -> ANormal Reference v
forall v (f :: * -> * -> *). Var v => v -> Term f v -> Term f v
ABTN.TAbs v
v (ANormal Reference v -> ANormal Reference v)
-> State
     (Set v, Map v (Either v (Built v)), Map [v] [v])
     (ANormal Reference v)
-> State
     (Set v, Map v (Either v (Built v)), Map [v] [v])
     (ANormal Reference v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ANormal Reference v
-> State
     (Set v, Map v (Either v (Built v)), Map [v] [v])
     (ANormal Reference v)
go ANormal Reference v
e
      ABTN.Term Set v
_ (ABTN.Tm ANormalF Reference v (ANormal Reference v)
f) -> ANormalF Reference v (ANormal Reference v) -> ANormal Reference v
forall v (f :: * -> * -> *).
(Var v, Bifoldable f) =>
f v (Term f v) -> Term f v
ABTN.TTm (ANormalF Reference v (ANormal Reference v) -> ANormal Reference v)
-> StateT
     (Set v, Map v (Either v (Built v)), Map [v] [v])
     Identity
     (ANormalF Reference v (ANormal Reference v))
-> State
     (Set v, Map v (Either v (Built v)), Map [v] [v])
     (ANormal Reference v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ANormal Reference v
 -> State
      (Set v, Map v (Either v (Built v)), Map [v] [v])
      (ANormal Reference v))
-> ANormalF Reference v (ANormal Reference v)
-> StateT
     (Set v, Map v (Either v (Built v)), Map [v] [v])
     Identity
     (ANormalF Reference v (ANormal Reference 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) -> ANormalF Reference v a -> f (ANormalF Reference v b)
traverse ANormal Reference v
-> State
     (Set v, Map v (Either v (Built v)), Map [v] [v])
     (ANormal Reference v)
go ANormalF Reference v (ANormal Reference v)
f
    callUses :: v -> ANormal Reference v -> Int
callUses v
u =
      Sum Int -> Int
forall a. Sum a -> a
getSum (Sum Int -> Int)
-> (ANormal Reference v -> Sum Int) -> ANormal Reference v -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v
-> (ANormalF Reference v (ANormal Reference v) -> Sum Int)
-> ANormal Reference v
-> Sum Int
forall v m.
(Var v, Monoid m) =>
v
-> (ANormalF Reference v (ANormal Reference v) -> m)
-> ANormal Reference v
-> m
foldMentions v
u \case
        AApp (FComb Reference
r) [v]
args | Batch -> Reference -> Maybe Int
batchArity Batch
b Reference
r Maybe Int -> Maybe Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int -> Maybe Int
forall a. a -> Maybe a
Just ([v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
args) -> Int -> Sum Int
forall a. a -> Sum a
Sum ([v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ((v -> Bool) -> [v] -> [v]
forall a. (a -> Bool) -> [a] -> [a]
filter (v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
u) [v]
args))
        ANormalF Reference v (ANormal Reference v)
_ -> Int -> Sum Int
forall a. a -> Sum a
Sum Int
0
    renameRoot :: v -> Tree v -> Tree v
renameRoot v
u = \case
      Node v
_ Shape
sh [Tree v]
ts -> v -> Shape -> [Tree v] -> Tree v
forall v. v -> Shape -> [Tree v] -> Tree v
Node v
u Shape
sh [Tree v]
ts
      Tree v
t -> Tree v
t

-- | A body with a let rebuilding each product of the tree that the body
-- still names, from its fields, an inner product before the one that
-- holds it.
rebuild :: (Var v) => Tree v -> ANormal Reference v -> ANormal Reference v
rebuild :: forall v.
Var v =>
Tree v -> ANormal Reference v -> ANormal Reference v
rebuild (Leaf v
_) ANormal Reference v
e = ANormal Reference v
e
rebuild (Node v
v Shape
sh [Tree v]
ts) ANormal Reference v
e = (Tree v -> ANormal Reference v -> ANormal Reference v)
-> ANormal Reference v -> [Tree v] -> ANormal Reference v
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Tree v -> ANormal Reference v -> ANormal Reference v
forall v.
Var v =>
Tree v -> ANormal Reference v -> ANormal Reference v
rebuild ANormal Reference v
withNode [Tree v]
ts
  where
    withNode :: ANormal Reference v
withNode
      | v -> ANormal Reference v -> Bool
forall v. Var v => v -> ANormal Reference v -> Bool
mentions v
v ANormal Reference v
e = Direction Word16
-> [v]
-> [Mem]
-> ANormal Reference v
-> ANormal Reference v
-> ANormal Reference 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
v] [Mem
BX] (Shape -> [v] -> ANormal Reference v
packOf Shape
sh ((Tree v -> v) -> [Tree v] -> [v]
forall a b. (a -> b) -> [a] -> [b]
map Tree v -> v
forall v. Tree v -> v
treeName [Tree v]
ts)) ANormal Reference v
e
      | Bool
otherwise = ANormal Reference v
e
    packOf :: Shape -> [v] -> ANormal Reference v
packOf (Con Reference
r CTag
t [Shape]
_) [v]
fs = Func Reference v -> [v] -> ANormal Reference v
forall v ref. Var v => Func ref v -> [v] -> ANormal ref v
TApp (Reference -> CTag -> Func Reference v
forall ref v. ref -> CTag -> Func ref v
FCon Reference
r CTag
t) [v]
fs
    packOf (Fields RecordSchema
rs [Shape]
_) [v]
fs = Func Reference v -> [v] -> ANormal Reference v
forall v ref. Var v => Func ref v -> [v] -> ANormal ref v
TApp (RecordSchema -> Func Reference v
forall ref v. RecordSchema -> Func ref v
FRec RecordSchema
rs) [v]
fs
    packOf Shape
Plain [v]
_ = v -> ANormal Reference v
forall v ref. Var v => v -> ANormal ref v
TVar v
v

-- | A body with each saturated call of a function of the batch that builds
-- one of its call patterns rewritten to a call of that specialisation with
-- the fields of what it built ('siteSpec', 'explodedArgs').
useSpecs :: (Var v) => Batch -> ANormal Reference v -> ANormal Reference v
useSpecs :: forall v.
Var v =>
Batch -> ANormal Reference v -> ANormal Reference v
useSpecs Batch
b ANormal Reference v
bd = ANormal Reference v -> ANormal Reference v
go ANormal Reference v
bd
  where
    cons :: ConLets v
cons = ANormal Reference v -> ConLets v
forall v. Var v => ANormal Reference v -> ConLets v
conLets ANormal Reference v
bd
    go :: ANormal Reference v -> ANormal Reference v
go ANormal Reference v
tm = case ANormal Reference v
tm of
      TApp (FComb Reference
r) [v]
args
        | Just [[Shape]]
pats <- Batch -> Reference -> Maybe [[Shape]]
batchSpecs Batch
b Reference
r,
          Just (Int
k, [Shape]
pat) <- ConLets v -> [[Shape]] -> [v] -> Maybe (Int, [Shape])
forall v.
Var v =>
ConLets v -> [[Shape]] -> [v] -> Maybe (Int, [Shape])
siteSpec ConLets v
cons [[Shape]]
pats [v]
args,
          Just [v]
fields <- ConLets v -> [Shape] -> [v] -> Maybe [v]
forall v. Var v => ConLets v -> [Shape] -> [v] -> Maybe [v]
explodedArgs ConLets v
cons [Shape]
pat [v]
args ->
            Func Reference v -> [v] -> ANormal Reference v
forall v ref. Var v => Func ref v -> [v] -> ANormal ref v
TApp (Reference -> Word64 -> Func Reference v
forall ref v. ref -> Word64 -> Func ref v
FWorker Reference
r (Int -> Word64
specIndex Int
k)) [v]
fields
      ABTN.Term Set v
_ (ABTN.Abs v
u ANormal Reference v
e) -> v -> ANormal Reference v -> ANormal Reference v
forall v (f :: * -> * -> *). Var v => v -> Term f v -> Term f v
ABTN.TAbs v
u (ANormal Reference v -> ANormal Reference v
go ANormal Reference v
e)
      ABTN.Term Set v
_ (ABTN.Tm ANormalF Reference v (ANormal Reference v)
f) -> ANormalF Reference v (ANormal Reference v) -> ANormal Reference v
forall v (f :: * -> * -> *).
(Var v, Bifoldable f) =>
f v (Term f v) -> Term f v
ABTN.TTm ((ANormal Reference v -> ANormal Reference v)
-> ANormalF Reference v (ANormal Reference v)
-> ANormalF Reference v (ANormal Reference v)
forall a b.
(a -> b) -> ANormalF Reference v a -> ANormalF Reference v b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ANormal Reference v -> ANormal Reference v
go ANormalF Reference v (ANormal Reference v)
f)

-- | The body without the product lets nothing uses: what was built for a
-- call that takes its fields instead, or for a tail that yields its
-- leaves, and, once such a let is gone, a field only it held, or the
-- record only an update of it updated. A product let has no effect, so
-- one with no use is left out whatever it was built for. The body is
-- walked from the inside out, so that a let is judged once the lets in
-- its body are settled; binders are unique ('prepare'), so a use is a
-- use by name.
dropDead :: (Var v) => ANormal Reference v -> ANormal Reference v
dropDead :: forall v. Var v => ANormal Reference v -> ANormal Reference v
dropDead ANormal Reference v
tm = case ANormal Reference v
tm of
  TLets Direction Word16
d [v]
us [Mem]
ms ANormal Reference v
bn ANormal Reference v
bo
    | [v
u] <- [v]
us, ANormal Reference v -> Bool
forall v. Var v => ANormal Reference v -> Bool
isBuiltLet ANormal Reference v
bn, Bool -> Bool
not (v -> ANormal Reference v -> Bool
forall v. Var v => v -> ANormal Reference v -> Bool
mentions v
u ANormal Reference v
bo') -> ANormal Reference v
bo'
    | Bool
otherwise -> Direction Word16
-> [v]
-> [Mem]
-> ANormal Reference v
-> ANormal Reference v
-> ANormal Reference v
forall v ref.
Var v =>
Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
TLets Direction Word16
d [v]
us [Mem]
ms (ANormal Reference v -> ANormal Reference v
forall v. Var v => ANormal Reference v -> ANormal Reference v
dropDead ANormal Reference v
bn) ANormal Reference v
bo'
    where
      bo' :: ANormal Reference v
bo' = ANormal Reference v -> ANormal Reference v
forall v. Var v => ANormal Reference v -> ANormal Reference v
dropDead ANormal Reference v
bo
  ABTN.Term Set v
_ (ABTN.Abs v
u ANormal Reference v
e) -> v -> ANormal Reference v -> ANormal Reference v
forall v (f :: * -> * -> *). Var v => v -> Term f v -> Term f v
ABTN.TAbs v
u (ANormal Reference v -> ANormal Reference v
forall v. Var v => ANormal Reference v -> ANormal Reference v
dropDead ANormal Reference v
e)
  ABTN.Term Set v
_ (ABTN.Tm ANormalF Reference v (ANormal Reference v)
f) -> ANormalF Reference v (ANormal Reference v) -> ANormal Reference v
forall v (f :: * -> * -> *).
(Var v, Bifoldable f) =>
f v (Term f v) -> Term f v
ABTN.TTm ((ANormal Reference v -> ANormal Reference v)
-> ANormalF Reference v (ANormal Reference v)
-> ANormalF Reference v (ANormal Reference v)
forall a b.
(a -> b) -> ANormalF Reference v a -> ANormalF Reference v b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ANormal Reference v -> ANormal Reference v
forall v. Var v => ANormal Reference v -> ANormal Reference v
dropDead ANormalF Reference v (ANormal Reference v)
f)

-- | The part of a shape the body consumes: a product node stays when every
-- occurrence of the variable is a data match with an arm for the
-- constructor, a record match (of any of its fields), or a record update,
-- and its fields are what the arms' bodies consume of them; an update
-- passes every field on into the new record, so it leaves the fields
-- plain. An argument used any other way, or matched without an arm for
-- the constructor it is built from, is left plain.
usable :: (Var v) => Shape -> v -> ANormal Reference v -> Shape
usable :: forall v. Var v => Shape -> v -> ANormal Reference v -> Shape
usable Shape
Plain v
_ ANormal Reference v
_ = Shape
Plain
usable Shape
s v
v ANormal Reference v
bd
  | v -> ANormal Reference v -> Int
forall v. Var v => v -> ANormal Reference v -> Int
occurrences v
v ANormal Reference v
bd Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Branched Reference (ANormal Reference v)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Branched Reference (ANormal Reference v)]
matches Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [()] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [()]
updates = Shape
Plain
  | ([([(Int, v)], ANormal Reference v)] -> Bool)
-> [[([(Int, v)], ANormal Reference v)]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any [([(Int, v)], ANormal Reference v)] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[([(Int, v)], ANormal Reference v)]]
arms = Shape
Plain
  | Bool -> Bool
not ([()] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [()]
updates) = Shape -> [Shape] -> Shape
withSubShapes Shape
s ((Shape -> Shape) -> [Shape] -> [Shape]
forall a b. (a -> b) -> [a] -> [b]
map (Shape -> Shape -> Shape
forall a b. a -> b -> a
const Shape
Plain) [Shape]
subs)
  | Bool
otherwise = Shape -> [Shape] -> Shape
withSubShapes Shape
s [(Shape -> Shape -> Shape) -> Shape -> [Shape] -> Shape
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Shape -> Shape -> Shape
meet Shape
sub [Shape -> v -> ANormal Reference v -> Shape
forall v. Var v => Shape -> v -> ANormal Reference v -> Shape
usable Shape
sub v
f ANormal Reference v
bo | (v
f, ANormal Reference v
bo) <- Int -> [(v, ANormal Reference v)]
fieldArms Int
j] | (Int
j, Shape
sub) <- [Int] -> [Shape] -> [(Int, Shape)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [Shape]
subs]
  where
    subs :: [Shape]
subs = Shape -> [Shape]
subShapes Shape
s
    matches :: [Branched Reference (ANormal Reference v)]
matches = v
-> (ANormalF Reference v (ANormal Reference v)
    -> [Branched Reference (ANormal Reference v)])
-> ANormal Reference v
-> [Branched Reference (ANormal Reference v)]
forall v m.
(Var v, Monoid m) =>
v
-> (ANormalF Reference v (ANormal Reference v) -> m)
-> ANormal Reference v
-> m
foldMentions v
v (\case AMatch v
u Branched Reference (ANormal Reference v)
bs | v
u v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
v -> [Branched Reference (ANormal Reference v)
bs]; ANormalF Reference v (ANormal Reference v)
_ -> []) ANormal Reference v
bd
    updates :: [()]
updates = v
-> (ANormalF Reference v (ANormal Reference v) -> [()])
-> ANormal Reference v
-> [()]
forall v m.
(Var v, Monoid m) =>
v
-> (ANormalF Reference v (ANormal Reference v) -> m)
-> ANormal Reference v
-> m
foldMentions v
v (\case ALet Direction Word16
_ [Mem
_] (TApp (FRecSet Text
f) [v
base, v
_]) ANormal Reference v
_ | v
base v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
v, Fields RecordSchema
rs [Shape]
_ <- Shape
s, Text
f Text -> [Text] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` RecordSchema -> [Text]
schemaFields RecordSchema
rs -> [()]; ANormalF Reference v (ANormal Reference v)
_ -> []) ANormal Reference v
bd
    -- an arm as the fields it binds, each with its position in the product
    arms :: [[([(Int, v)], ANormal Reference v)]]
arms = (Branched Reference (ANormal Reference v)
 -> [([(Int, v)], ANormal Reference v)])
-> [Branched Reference (ANormal Reference v)]
-> [[([(Int, v)], ANormal Reference v)]]
forall a b. (a -> b) -> [a] -> [b]
map Branched Reference (ANormal Reference v)
-> [([(Int, v)], ANormal Reference v)]
armFor [Branched Reference (ANormal Reference v)]
matches
    armFor :: Branched Reference (ANormal Reference v)
-> [([(Int, v)], ANormal Reference v)]
armFor Branched Reference (ANormal Reference v)
bs = case (Shape
s, Branched Reference (ANormal Reference v)
bs) of
      (Con Reference
r CTag
t [Shape]
_, MatchData Reference
r' EnumMap CTag ([Mem], ANormal Reference v)
cs Maybe (ANormal Reference v)
_)
        | Reference
r' Reference -> Reference -> Bool
forall a. Eq a => a -> a -> Bool
== Reference
r,
          Just ([Mem]
_, TAbss [v]
fs ANormal Reference v
bo) <- CTag
-> EnumMap CTag ([Mem], ANormal Reference v)
-> Maybe ([Mem], ANormal Reference v)
forall k a. EnumKey k => k -> EnumMap k a -> Maybe a
EC.lookup CTag
t EnumMap CTag ([Mem], ANormal Reference v)
cs,
          [v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
fs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Shape] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Shape]
subs ->
            [([Int] -> [v] -> [(Int, v)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [v]
fs, ANormal Reference v
bo)]
      (Fields RecordSchema
rs [Shape]
_, MatchRec RecordSchema
sch (TAbss [v]
fs ANormal Reference v
bo))
        | Just [Int]
js <- (Text -> Maybe Int) -> [Text] -> Maybe [Int]
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 (Text -> [Text] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
`elemIndex` RecordSchema -> [Text]
schemaFields RecordSchema
rs) (RecordSchema -> [Text]
schemaFields RecordSchema
sch),
          [v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
fs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
js ->
            [([Int] -> [v] -> [(Int, v)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
js [v]
fs, ANormal Reference v
bo)]
      (Shape, Branched Reference (ANormal Reference v))
_ -> []
    fieldArms :: Int -> [(v, ANormal Reference v)]
fieldArms Int
j = [(v
f, ANormal Reference v
bo) | ([(Int, v)]
jfs, ANormal Reference v
bo) <- [[([(Int, v)], ANormal Reference v)]]
-> [([(Int, v)], ANormal Reference v)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[([(Int, v)], ANormal Reference v)]]
arms, (Int
j', v
f) <- [(Int, v)]
jfs, Int
j' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j]
    meet :: Shape -> Shape -> Shape
meet Shape
a Shape
b
      | Shape -> Shape -> Bool
sameProduct Shape
a Shape
b = Shape -> [Shape] -> Shape
withSubShapes Shape
a ((Shape -> Shape -> Shape) -> [Shape] -> [Shape] -> [Shape]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Shape -> Shape -> Shape
meet (Shape -> [Shape]
subShapes Shape
a) (Shape -> [Shape]
subShapes Shape
b))
      | Bool
otherwise = Shape
Plain

-- | The worker of a lambda for a pattern: its parameters are the plain
-- arguments and the fields of the exploded ones, each a fresh name, and
-- its body is the lambda's with the exploded arguments rewritten to their
-- fields, prepared again afterwards, since a record update the rewrite
-- turned into a literal is a product built for whatever it is passed to.
-- The join points are made by the second pass, from the sites of the
-- unoptimized body.
specialise :: (Var v) => Batch -> Maybe (JoinSetting v) -> Pattern -> SuperNormal Reference v -> Prepared v
specialise :: forall v.
Var v =>
Batch
-> Maybe (JoinSetting v)
-> [Shape]
-> SuperNormal Reference v
-> Prepared v
specialise Batch
b Maybe (JoinSetting v)
js [Shape]
pat (Lambda [Mem]
ccs (TAbss [v]
vs Term (ANormalF Reference) v
bd0)) =
  SuperNormal Reference v
-> Sources v -> Map [v] [v] -> Set v -> [Join v] -> Prepared v
forall v.
SuperNormal Reference v
-> Sources v -> Map [v] [v] -> Set v -> [Join v] -> Prepared v
Prepared
    ([Mem] -> Term (ANormalF Reference) v -> SuperNormal Reference v
forall ref v. [Mem] -> ANormal ref v -> SuperNormal ref v
Lambda (((v, Mem) -> Mem) -> [(v, Mem)] -> [Mem]
forall a b. (a -> b) -> [a] -> [b]
map (v, Mem) -> Mem
forall a b. (a, b) -> b
snd [(v, Mem)]
params) ([v] -> Term (ANormalF Reference) v -> Term (ANormalF Reference) v
forall v (f :: * -> * -> *). Var v => [v] -> Term f v -> Term f v
TAbss (((v, Mem) -> v) -> [(v, Mem)] -> [v]
forall a b. (a -> b) -> [a] -> [b]
map (v, Mem) -> v
forall a b. (a, b) -> a
fst [(v, Mem)]
params) Term (ANormalF Reference) v
body))
    ([Sources v] -> Sources v
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
M.unions [Sources v
srcs0, Sources v
treeSrcs, Sources v
armSrcs, Sources v
srcs1])
    (Map [v] [v]
lets0 Map [v] [v] -> Map [v] [v] -> Map [v] [v]
forall a. Semigroup a => a -> a -> a
<> Map [v] [v]
lets1)
    Set v
dead
    [Join v]
joins
  where
    (Term (ANormalF Reference) v
bd, Sources v
srcs0, Map [v] [v]
lets0, Set v
_, [Join v]
_) = Batch
-> Maybe Shape
-> Maybe (Joining v)
-> Term (ANormalF Reference) v
-> (Term (ANormalF Reference) v, Sources v, Map [v] [v], Set v,
    [Join v])
forall v.
Var v =>
Batch
-> Maybe Shape
-> Maybe (Joining v)
-> ANormal Reference v
-> (ANormal Reference v, Sources v, Map [v] [v], Set v, [Join v])
prepareBody Batch
b Maybe Shape
forall a. Maybe a
Nothing Maybe (Joining v)
forall a. Maybe a
Nothing Term (ANormalF Reference) v
bd0
    treeSrcs :: Sources v
treeSrcs = [Sources v] -> Sources v
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
M.unions [v -> Tree v -> Sources v
forall v. Ord v => v -> Tree v -> Sources v
treeSources v
v Tree v
tree | (v
v, Mem
_, Tree v
tree) <- [(v, Mem, Tree v)]
trees]
    avoid0 :: Set v
avoid0 = [Set v] -> Set v
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions (Term (ANormalF Reference) v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
ABTN.freeVars Term (ANormalF Reference) v
bd Set v -> [Set v] -> [Set v]
forall a. a -> [a] -> [a]
: [v] -> Set v
forall a. Ord a => [a] -> Set a
Set.fromList [v]
vs Set v -> [Set v] -> [Set v]
forall a. a -> [a] -> [a]
: (Term (ANormalF Reference) v -> Set v)
-> [Term (ANormalF Reference) v] -> [Set v]
forall a b. (a -> b) -> [a] -> [b]
map Term (ANormalF Reference) v -> Set v
forall v. Ord v => ANormal Reference v -> Set v
boundVars (Term (ANormalF Reference) v -> [Term (ANormalF Reference) v]
forall v. ANormal Reference v -> [ANormal Reference v]
universe Term (ANormalF Reference) v
bd))
    (Set v
_, [(v, Mem, Tree v)]
trees) = ((Set v, [(v, Mem, Tree v)])
 -> (v, Mem, Shape) -> (Set v, [(v, Mem, Tree v)]))
-> (Set v, [(v, Mem, Tree v)])
-> [(v, Mem, Shape)]
-> (Set v, [(v, Mem, Tree v)])
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (Set v, [(v, Mem, Tree v)])
-> (v, Mem, Shape) -> (Set v, [(v, Mem, Tree v)])
forall {v} {b}.
Var v =>
(Set v, [(v, b, Tree v)])
-> (v, b, Shape) -> (Set v, [(v, b, Tree v)])
step (Set v
avoid0, []) ([v] -> [Mem] -> [Shape] -> [(v, Mem, Shape)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [v]
vs [Mem]
ccs [Shape]
pat)
    step :: (Set v, [(v, b, Tree v)])
-> (v, b, Shape) -> (Set v, [(v, b, Tree v)])
step (Set v
avoid, [(v, b, Tree v)]
acc) (v
v, b
cc, Shape
s) = case Shape
s of
      Shape
Plain -> (Set v
avoid, [(v, b, Tree v)]
acc [(v, b, Tree v)] -> [(v, b, Tree v)] -> [(v, b, Tree v)]
forall a. [a] -> [a] -> [a]
++ [(v
v, b
cc, v -> Tree v
forall v. v -> Tree v
Leaf v
v)])
      Shape
_ ->
        let (Tree v
tree, Set v
avoid') = Set v -> Shape -> (Tree v, Set v)
forall v. Var v => Set v -> Shape -> (Tree v, Set v)
treeOf Set v
avoid Shape
s
         in (Set v
avoid', [(v, b, Tree v)]
acc [(v, b, Tree v)] -> [(v, b, Tree v)] -> [(v, b, Tree v)]
forall a. [a] -> [a] -> [a]
++ [(v
v, b
cc, Tree v
tree)])
    params :: [(v, Mem)]
params = [[(v, Mem)]] -> [(v, Mem)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [Tree v -> [(v, Mem)]
forall v. Tree v -> [(v, Mem)]
leaves Tree v
tree | (v
_, Mem
_, Tree v
tree) <- [(v, Mem, Tree v)]
trees]
    (Term (ANormalF Reference) v
rewritten, Sources v
armSrcs) = State (Sources v) (Term (ANormalF Reference) v)
-> Sources v -> (Term (ANormalF Reference) v, Sources v)
forall s a. State s a -> s -> (a, s)
runState ((Term (ANormalF Reference) v
 -> (v, Mem, Tree v)
 -> State (Sources v) (Term (ANormalF Reference) v))
-> Term (ANormalF Reference) v
-> [(v, Mem, Tree v)]
-> State (Sources v) (Term (ANormalF Reference) v)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\Term (ANormalF Reference) v
e (v
v, Mem
_, Tree v
tree) -> v
-> Tree v
-> Term (ANormalF Reference) v
-> State (Sources v) (Term (ANormalF Reference) v)
forall v.
Var v =>
v
-> Tree v
-> ANormal Reference v
-> State (Sources v) (ANormal Reference v)
rewrite v
v Tree v
tree Term (ANormalF Reference) v
e) Term (ANormalF Reference) v
bd [(v, Mem, Tree v)]
trees) Sources v
forall k a. Map k a
M.empty
    -- the second pass rewrites only lets the first left (a let of a
    -- call of a worker is not a let of a call), so its sources and lets
    -- add to the first's
    (Term (ANormalF Reference) v
body, Sources v
srcs1, Map [v] [v]
lets1, Set v
dead, [Join v]
joins) = Batch
-> Maybe Shape
-> Maybe (Joining v)
-> Term (ANormalF Reference) v
-> (Term (ANormalF Reference) v, Sources v, Map [v] [v], Set v,
    [Join v])
forall v.
Var v =>
Batch
-> Maybe Shape
-> Maybe (Joining v)
-> ANormal Reference v
-> (ANormal Reference v, Sources v, Map [v] [v], Set v, [Join v])
prepareBody Batch
b Maybe Shape
forall a. Maybe a
Nothing (Maybe (JoinSetting v)
-> Term (ANormalF Reference) v -> Sources v -> Maybe (Joining v)
forall v.
Var v =>
Maybe (JoinSetting v)
-> ANormal Reference v -> Sources v -> Maybe (Joining v)
joiningOf Maybe (JoinSetting v)
js Term (ANormalF Reference) v
bd0 ([Sources v] -> Sources v
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
M.unions [Sources v
srcs0, Sources v
treeSrcs, Sources v
armSrcs])) Term (ANormalF Reference) v
rewritten

-- | The product-return worker of a lambda that returns a product of the
-- shape on every path: the lambda with its tails yielding the leaves. It
-- gets no join points: its tails are what yield the leaves.
returnWorker :: (Var v) => Batch -> Shape -> SuperNormal Reference v -> Prepared v
returnWorker :: forall v.
Var v =>
Batch -> Shape -> SuperNormal Reference v -> Prepared v
returnWorker Batch
b Shape
s (Lambda [Mem]
ccs (TAbss [v]
vs Term (ANormalF Reference) v
bd)) = SuperNormal Reference v
-> Sources v -> Map [v] [v] -> Set v -> [Join v] -> Prepared v
forall v.
SuperNormal Reference v
-> Sources v -> Map [v] [v] -> Set v -> [Join v] -> Prepared v
Prepared ([Mem] -> Term (ANormalF Reference) v -> SuperNormal Reference v
forall ref v. [Mem] -> ANormal ref v -> SuperNormal ref v
Lambda [Mem]
ccs ([v] -> Term (ANormalF Reference) v -> Term (ANormalF Reference) v
forall v (f :: * -> * -> *). Var v => [v] -> Term f v -> Term f v
TAbss [v]
vs Term (ANormalF Reference) v
bd')) Sources v
srcs Map [v] [v]
lets Set v
forall a. Set a
Set.empty []
  where
    (Term (ANormalF Reference) v
bd', Sources v
srcs, Map [v] [v]
lets, Set v
_, [Join v]
_) = Batch
-> Maybe Shape
-> Maybe (Joining v)
-> Term (ANormalF Reference) v
-> (Term (ANormalF Reference) v, Sources v, Map [v] [v], Set v,
    [Join v])
forall v.
Var v =>
Batch
-> Maybe Shape
-> Maybe (Joining v)
-> ANormal Reference v
-> (ANormal Reference v, Sources v, Map [v] [v], Set v, [Join v])
prepareBody Batch
b (Shape -> Maybe Shape
forall a. a -> Maybe a
Just Shape
s) Maybe (Joining v)
forall a. Maybe a
Nothing Term (ANormalF Reference) v
bd

-- | A tree of fresh names for an exploded argument or result: a leaf is a
-- worker parameter (or a yielded value), a node a product whose fields are
-- below it, named so that a body can refer to the product as a whole
-- ('rebuild').
data Tree v = Leaf v | Node v Shape [Tree v]

treeOf :: (Var v) => Set.Set v -> Shape -> (Tree v, Set.Set v)
treeOf :: forall v. Var v => Set v -> Shape -> (Tree v, Set v)
treeOf Set v
avoid Shape
Plain = (v -> Tree v) -> (v, Set v) -> (Tree v, Set 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 v -> Tree v
forall v. v -> Tree v
Leaf (Set v -> (v, Set v)
forall v. Var v => Set v -> (v, Set v)
freshVar Set v
avoid)
treeOf Set v
avoid Shape
s = let (v
nv, Set v
avoid') = Set v -> (v, Set v)
forall v. Var v => Set v -> (v, Set v)
freshVar Set v
avoid in ([Tree v] -> Tree v) -> ([Tree v], Set v) -> (Tree v, Set 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 (v -> Shape -> [Tree v] -> Tree v
forall v. v -> Shape -> [Tree v] -> Tree v
Node v
nv Shape
s) (Set v -> [Shape] -> ([Tree v], Set v)
forall {v}. Var v => Set v -> [Shape] -> ([Tree v], Set v)
go Set v
avoid' (Shape -> [Shape]
subShapes Shape
s))
  where
    go :: Set v -> [Shape] -> ([Tree v], Set v)
go Set v
av [] = ([], Set v
av)
    go Set v
av (Shape
sub : [Shape]
ss) = let (Tree v
tr, Set v
av') = Set v -> Shape -> (Tree v, Set v)
forall v. Var v => Set v -> Shape -> (Tree v, Set v)
treeOf Set v
av Shape
sub; ([Tree v]
trs, Set v
av'') = Set v -> [Shape] -> ([Tree v], Set v)
go Set v
av' [Shape]
ss in (Tree v
tr Tree v -> [Tree v] -> [Tree v]
forall a. a -> [a] -> [a]
: [Tree v]
trs, Set v
av'')

treeName :: Tree v -> v
treeName :: forall v. Tree v -> v
treeName (Leaf v
v) = v
v
treeName (Node v
v Shape
_ [Tree v]
_) = v
v

leaves :: Tree v -> [(v, Mem)]
leaves :: forall v. Tree v -> [(v, Mem)]
leaves (Leaf v
v) = [(v
v, Mem
BX)]
leaves (Node v
_ Shape
_ [Tree v]
ts) = (Tree v -> [(v, Mem)]) -> [Tree v] -> [(v, Mem)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Tree v -> [(v, Mem)]
forall v. Tree v -> [(v, Mem)]
leaves [Tree v]
ts

leafNames :: Tree v -> [v]
leafNames :: forall v. Tree v -> [v]
leafNames = ((v, Mem) -> v) -> [(v, Mem)] -> [v]
forall a b. (a -> b) -> [a] -> [b]
map (v, Mem) -> v
forall a b. (a, b) -> a
fst ([(v, Mem)] -> [v]) -> (Tree v -> [(v, Mem)]) -> Tree v -> [v]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree v -> [(v, Mem)]
forall v. Tree v -> [(v, Mem)]
leaves

-- | The sources ('Sources') of a variable a tree replaces and of the
-- tree's inner products: each is the product of its children's names.
-- A leaf is a name of its own and needs none.
treeSources :: (Ord v) => v -> Tree v -> Sources v
treeSources :: forall v. Ord v => v -> Tree v -> Sources v
treeSources v
_ (Leaf v
_) = Map v (Either v (Built v))
forall k a. Map k a
M.empty
treeSources v
v (Node v
_ Shape
sh [Tree v]
ts) =
  v
-> Either v (Built v)
-> Map v (Either v (Built v))
-> Map v (Either v (Built v))
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert v
v (Built v -> Either v (Built v)
forall a b. b -> Either a b
Right (Shape -> [v] -> Built v
forall {v}. Shape -> [v] -> Built v
packOf Shape
sh ((Tree v -> v) -> [Tree v] -> [v]
forall a b. (a -> b) -> [a] -> [b]
map Tree v -> v
forall v. Tree v -> v
treeName [Tree v]
ts))) ([Map v (Either v (Built v))] -> Map v (Either v (Built v))
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
M.unions [v -> Tree v -> Map v (Either v (Built v))
forall v. Ord v => v -> Tree v -> Sources v
treeSources (Tree v -> v
forall v. Tree v -> v
treeName Tree v
t) Tree v
t | Tree v
t <- [Tree v]
ts])
  where
    packOf :: Shape -> [v] -> Built v
packOf (Con Reference
r CTag
t [Shape]
_) [v]
fs = Reference -> CTag -> [v] -> Built v
forall v. Reference -> CTag -> [v] -> Built v
BuiltCon Reference
r CTag
t [v]
fs
    packOf (Fields RecordSchema
rs [Shape]
_) [v]
fs = RecordSchema -> [v] -> Built v
forall v. RecordSchema -> [v] -> Built v
BuiltRec RecordSchema
rs [v]
fs
    packOf Shape
Plain [v]
_ = [Word] -> String -> Built v
forall a. HasCallStack => [Word] -> String -> a
internalBug [] String
"treeSources: a product node of plain shape"

-- | The body with every use of the variable rewritten to its fields: a
-- match becomes the arm for its product, a leaf field's arm variable
-- renamed to the parameter and a nested field's arm variable rewritten in
-- turn; a record update becomes a record literal of the parameters with
-- one replaced (the fields of an updated record are leaves: 'usable').
rewrite :: (Var v) => v -> Tree v -> ANormal Reference v -> State (Sources v) (ANormal Reference v)
rewrite :: forall v.
Var v =>
v
-> Tree v
-> ANormal Reference v
-> State (Sources v) (ANormal Reference v)
rewrite v
_ (Leaf v
_) ANormal Reference v
e = ANormal Reference v
-> StateT (Sources v) Identity (ANormal Reference v)
forall a. a -> StateT (Sources v) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ANormal Reference v
e
rewrite v
v (Node v
_ Shape
s [Tree v]
ts) ANormal Reference v
e = ANormal Reference v
-> StateT (Sources v) Identity (ANormal Reference v)
go ANormal Reference v
e
  where
    fields :: [Text]
fields = case Shape
s of
      Fields RecordSchema
rs [Shape]
_ -> RecordSchema -> [Text]
schemaFields RecordSchema
rs
      Shape
_ -> []
    -- a subterm that does not use the variable has no match on it
    go :: ANormal Reference v
-> StateT (Sources v) Identity (ANormal Reference v)
go ANormal Reference v
tm | Bool -> Bool
not (v -> ANormal Reference v -> Bool
forall v. Var v => v -> ANormal Reference v -> Bool
mentions v
v ANormal Reference v
tm) = ANormal Reference v
-> StateT (Sources v) Identity (ANormal Reference v)
forall a. a -> StateT (Sources v) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ANormal Reference v
tm
    go ANormal Reference v
tm = case ANormal Reference v
tm of
      TMatch v
u (MatchData Reference
r' EnumMap CTag ([Mem], ANormal Reference v)
cs Maybe (ANormal Reference v)
_)
        | v
u v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
v,
          Con Reference
r CTag
t [Shape]
_ <- Shape
s,
          Reference
r' Reference -> Reference -> Bool
forall a. Eq a => a -> a -> Bool
== Reference
r,
          Just ([Mem]
_, TAbss [v]
fs ANormal Reference v
bo) <- CTag
-> EnumMap CTag ([Mem], ANormal Reference v)
-> Maybe ([Mem], ANormal Reference v)
forall k a. EnumKey k => k -> EnumMap k a -> Maybe a
EC.lookup CTag
t EnumMap CTag ([Mem], ANormal Reference v)
cs,
          [v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
fs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Tree v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Tree v]
ts ->
            [(v, Tree v)]
-> ANormal Reference v
-> StateT (Sources v) Identity (ANormal Reference v)
arm ([v] -> [Tree v] -> [(v, Tree v)]
forall a b. [a] -> [b] -> [(a, b)]
zip [v]
fs [Tree v]
ts) ANormal Reference v
bo
      TMatch v
u (MatchRec RecordSchema
sch (TAbss [v]
fs ANormal Reference v
bo))
        | v
u v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
v,
          Fields RecordSchema
_ [Shape]
_ <- Shape
s,
          Just [Int]
js <- (Text -> Maybe Int) -> [Text] -> Maybe [Int]
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 (Text -> [Text] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
`elemIndex` [Text]
fields) (RecordSchema -> [Text]
schemaFields RecordSchema
sch),
          [v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
fs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
js ->
            [(v, Tree v)]
-> ANormal Reference v
-> StateT (Sources v) Identity (ANormal Reference v)
arm [(v
f, [Tree v]
ts [Tree v] -> Int -> Tree v
forall a. HasCallStack => [a] -> Int -> a
!! Int
j) | (v
f, Int
j) <- [v] -> [Int] -> [(v, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [v]
fs [Int]
js] ANormal Reference v
bo
      TLets Direction Word16
d [v
u] [Mem]
ms (TApp (FRecSet Text
f) [v
base, v
x]) ANormal Reference v
bo
        | v
base v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
v,
          Fields RecordSchema
rs [Shape]
_ <- Shape
s,
          Just Int
j <- Text -> [Text] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
elemIndex Text
f [Text]
fields,
          Just [v]
ls <- (Tree v -> Maybe v) -> [Tree v] -> Maybe [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 Tree v -> Maybe v
forall {a}. Tree a -> Maybe a
leafName [Tree v]
ts ->
            Direction Word16
-> [v]
-> [Mem]
-> ANormal Reference v
-> ANormal Reference v
-> ANormal Reference v
forall v ref.
Var v =>
Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
TLets Direction Word16
d [v
u] [Mem]
ms (Func Reference v -> [v] -> ANormal Reference v
forall v ref. Var v => Func ref v -> [v] -> ANormal ref v
TApp (RecordSchema -> Func Reference v
forall ref v. RecordSchema -> Func ref v
FRec RecordSchema
rs) (Int -> [v] -> [v]
forall a. Int -> [a] -> [a]
take Int
j [v]
ls [v] -> [v] -> [v]
forall a. [a] -> [a] -> [a]
++ v
x v -> [v] -> [v]
forall a. a -> [a] -> [a]
: Int -> [v] -> [v]
forall a. Int -> [a] -> [a]
drop (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [v]
ls)) (ANormal Reference v -> ANormal Reference v)
-> StateT (Sources v) Identity (ANormal Reference v)
-> StateT (Sources v) Identity (ANormal Reference v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ANormal Reference v
-> StateT (Sources v) Identity (ANormal Reference v)
go ANormal Reference v
bo
      -- a let whose binding was a field read of the variable binds a
      -- parameter now: its uses take the parameter, with no let to run
      TLets Direction Word16
d [v
x] [Mem]
ms ANormal Reference v
bn ANormal Reference v
bo ->
        ANormal Reference v
-> StateT (Sources v) Identity (ANormal Reference v)
go ANormal Reference v
bn StateT (Sources v) Identity (ANormal Reference v)
-> (ANormal Reference v
    -> StateT (Sources v) Identity (ANormal Reference v))
-> StateT (Sources v) Identity (ANormal Reference v)
forall a b.
StateT (Sources v) Identity a
-> (a -> StateT (Sources v) Identity b)
-> StateT (Sources v) Identity b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          TVar v
u -> do
            (Sources v -> Sources v) -> StateT (Sources v) Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (v -> Either v (Built v) -> Sources v -> Sources v
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert v
x (v -> Either v (Built v)
forall a b. a -> Either a b
Left v
u))
            ANormal Reference v
-> StateT (Sources v) Identity (ANormal Reference v)
go (Map v v -> ANormal Reference v -> ANormal Reference v
forall v (f :: * -> * -> *).
(Var v, Bifunctor f, Bifoldable f) =>
Map v v -> Term f v -> Term f v
ABTN.renames (v -> v -> Map v v
forall k a. k -> a -> Map k a
M.singleton v
x v
u) ANormal Reference v
bo)
          ANormal Reference v
bn' -> Direction Word16
-> [v]
-> [Mem]
-> ANormal Reference v
-> ANormal Reference v
-> ANormal Reference v
forall v ref.
Var v =>
Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
TLets Direction Word16
d [v
x] [Mem]
ms ANormal Reference v
bn' (ANormal Reference v -> ANormal Reference v)
-> StateT (Sources v) Identity (ANormal Reference v)
-> StateT (Sources v) Identity (ANormal Reference v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ANormal Reference v
-> StateT (Sources v) Identity (ANormal Reference v)
go ANormal Reference v
bo
      ABTN.TAbs v
u ANormal Reference v
b -> v -> ANormal Reference v -> ANormal Reference v
forall v (f :: * -> * -> *). Var v => v -> Term f v -> Term f v
ABTN.TAbs v
u (ANormal Reference v -> ANormal Reference v)
-> StateT (Sources v) Identity (ANormal Reference v)
-> StateT (Sources v) Identity (ANormal Reference v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ANormal Reference v
-> StateT (Sources v) Identity (ANormal Reference v)
go ANormal Reference v
b
      ABTN.TTm ANormalF Reference v (ANormal Reference v)
f -> ANormalF Reference v (ANormal Reference v) -> ANormal Reference v
forall v (f :: * -> * -> *).
(Var v, Bifoldable f) =>
f v (Term f v) -> Term f v
ABTN.TTm (ANormalF Reference v (ANormal Reference v) -> ANormal Reference v)
-> StateT
     (Sources v) Identity (ANormalF Reference v (ANormal Reference v))
-> StateT (Sources v) Identity (ANormal Reference v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ANormal Reference v
 -> StateT (Sources v) Identity (ANormal Reference v))
-> ANormalF Reference v (ANormal Reference v)
-> StateT
     (Sources v) Identity (ANormalF Reference v (ANormal Reference 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) -> ANormalF Reference v a -> f (ANormalF Reference v b)
traverse ANormal Reference v
-> StateT (Sources v) Identity (ANormal Reference v)
go ANormalF Reference v (ANormal Reference v)
f
    -- an arm's variables take the tree's names, so that a product the
    -- body still names after the rewrite has one name wherever it was
    -- matched ('rebuild'); each is recorded as a source of the tree's
    -- name
    arm :: [(v, Tree v)]
-> ANormal Reference v
-> StateT (Sources v) Identity (ANormal Reference v)
arm [(v, Tree v)]
pairs ANormal Reference v
bo = do
      let renaming :: Map v v
renaming = [(v, v)] -> Map v v
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(v
f, Tree v -> v
forall v. Tree v -> v
treeName Tree v
tr) | (v
f, Tree v
tr) <- [(v, Tree v)]
pairs]
      (Sources v -> Sources v) -> StateT (Sources v) Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (Sources v -> Sources v -> Sources v
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union ((v -> Either v (Built v)) -> Map v v -> Sources v
forall a b k. (a -> b) -> Map k a -> Map k b
M.map v -> Either v (Built v)
forall a b. a -> Either a b
Left Map v v
renaming))
      ANormal Reference v
bo' <- (ANormal Reference v
 -> (v, Tree v)
 -> StateT (Sources v) Identity (ANormal Reference v))
-> ANormal Reference v
-> [(v, Tree v)]
-> StateT (Sources v) Identity (ANormal Reference v)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\ANormal Reference v
acc (v
_, Tree v
tr) -> v
-> Tree v
-> ANormal Reference v
-> StateT (Sources v) Identity (ANormal Reference v)
forall v.
Var v =>
v
-> Tree v
-> ANormal Reference v
-> State (Sources v) (ANormal Reference v)
rewrite (Tree v -> v
forall v. Tree v -> v
treeName Tree v
tr) Tree v
tr ANormal Reference v
acc) (Map v v -> ANormal Reference v -> ANormal Reference v
forall v (f :: * -> * -> *).
(Var v, Bifunctor f, Bifoldable f) =>
Map v v -> Term f v -> Term f v
ABTN.renames Map v v
renaming ANormal Reference v
bo) [(v, Tree v)]
pairs
      ANormal Reference v
-> StateT (Sources v) Identity (ANormal Reference v)
go ANormal Reference v
bo'
    leafName :: Tree a -> Maybe a
leafName (Leaf a
l) = a -> Maybe a
forall a. a -> Maybe a
Just a
l
    leafName Tree a
_ = Maybe a
forall a. Maybe a
Nothing

freshVar :: (Var v) => Set.Set v -> (v, Set.Set v)
freshVar :: forall v. Var v => Set v -> (v, Set v)
freshVar Set v
avoid = (v
u, v -> Set v -> Set v
forall a. Ord a => a -> Set a -> Set a
Set.insert v
u Set v
avoid)
  where
    u :: v
u = Set v -> v -> v
forall v. Var v => Set v -> v -> v
Var.freshIn Set v
avoid (Type -> v
forall v. Var v => Type -> v
Var.typed Type
Var.Eta)

-- | A lambda as the emitter sees it ('prepareBody' on its body). Applied to
-- every lambda of a batch, so that the sites the analyses saw are the
-- sites emitted; the analyses pass no setting and get no join points.
prepare :: (Var v) => Batch -> Maybe (JoinSetting v) -> SuperNormal Reference v -> Prepared v
prepare :: forall v.
Var v =>
Batch
-> Maybe (JoinSetting v) -> SuperNormal Reference v -> Prepared v
prepare Batch
b Maybe (JoinSetting v)
js (Lambda [Mem]
ccs (TAbss [v]
vs Term (ANormalF Reference) v
bd)) = SuperNormal Reference v
-> Sources v -> Map [v] [v] -> Set v -> [Join v] -> Prepared v
forall v.
SuperNormal Reference v
-> Sources v -> Map [v] [v] -> Set v -> [Join v] -> Prepared v
Prepared ([Mem] -> Term (ANormalF Reference) v -> SuperNormal Reference v
forall ref v. [Mem] -> ANormal ref v -> SuperNormal ref v
Lambda [Mem]
ccs ([v] -> Term (ANormalF Reference) v -> Term (ANormalF Reference) v
forall v (f :: * -> * -> *). Var v => [v] -> Term f v -> Term f v
TAbss [v]
vs Term (ANormalF Reference) v
bd')) Sources v
srcs Map [v] [v]
lets Set v
dead [Join v]
joins
  where
    (Term (ANormalF Reference) v
bd', Sources v
srcs, Map [v] [v]
lets, Set v
dead, [Join v]
joins) = Batch
-> Maybe Shape
-> Maybe (Joining v)
-> Term (ANormalF Reference) v
-> (Term (ANormalF Reference) v, Sources v, Map [v] [v], Set v,
    [Join v])
forall v.
Var v =>
Batch
-> Maybe Shape
-> Maybe (Joining v)
-> ANormal Reference v
-> (ANormal Reference v, Sources v, Map [v] [v], Set v, [Join v])
prepareBody Batch
b Maybe Shape
forall a. Maybe a
Nothing (Maybe (JoinSetting v)
-> Term (ANormalF Reference) v -> Sources v -> Maybe (Joining v)
forall v.
Var v =>
Maybe (JoinSetting v)
-> ANormal Reference v -> Sources v -> Maybe (Joining v)
joiningOf Maybe (JoinSetting v)
js Term (ANormalF Reference) v
bd Sources v
forall k a. Map k a
M.empty) Term (ANormalF Reference) v
bd

-- | What 'splitJoins' needs of an unoptimized body, given the setting and
-- the sources recorded before the body is prepared again (the second
-- pass of 'specialise').
joiningOf :: (Var v) => Maybe (JoinSetting v) -> ANormal Reference v -> Sources v -> Maybe (Joining v)
joiningOf :: forall v.
Var v =>
Maybe (JoinSetting v)
-> ANormal Reference v -> Sources v -> Maybe (Joining v)
joiningOf Maybe (JoinSetting v)
js ANormal Reference v
bd Sources v
srcs = (\JoinSetting v
setting -> JoinSetting v -> Sites v -> ConLets v -> Sources v -> Joining v
forall v.
JoinSetting v -> Sites v -> ConLets v -> Sources v -> Joining v
Joining JoinSetting v
setting (ConLets v -> Set v -> ANormal Reference v -> Sites v
forall v.
Var v =>
ConLets v -> Set v -> ANormal Reference v -> Sites v
caseSites ConLets v
cons (JoinSetting v -> Set v
forall v. JoinSetting v -> Set v
jsLocals JoinSetting v
setting) ANormal Reference v
base) ConLets v
cons Sources v
srcs) (JoinSetting v -> Joining v)
-> Maybe (JoinSetting v) -> Maybe (Joining v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (JoinSetting v)
js
  where
    base :: ANormal Reference v
base = ANormal Reference v -> ANormal Reference v
forall v. Var v => ANormal Reference v -> ANormal Reference v
uniqueBody ANormal Reference v
bd
    cons :: ConLets v
cons = ANormal Reference v -> ConLets v
forall v. Var v => ANormal Reference v -> ConLets v
conLets ANormal Reference v
base

-- | A body as the analyses and the emitter see it, or as it is while
-- every rewrite is off: every binder with a name of its own
-- ('uniqueBinders'); each match on a product or sum built in the body
-- made its arm ('knownCons'); calls sunk into the branches that build
-- their argument ('sinkCalls'); each let of a call of a function with a
-- product-return worker that takes the result apart made a let of the
-- leaves over a call of the worker ('useReturns'); for the body of a
-- product-return worker, given its shape, each tail that builds the
-- product made a yield of the leaves ('yieldLeaves'); each call that
-- builds one of its callee's patterns made a call of that specialisation
-- with the fields ('useSpecs'); the product lets nothing uses any more
-- left out ('dropDead'); and last, given a setting, each match on the
-- value of a match moved into its tails, with its arms as join points
-- ('splitJoins'). The tails go before the calls, so that a tail call of a
-- function with both kinds of worker takes the one that yields the
-- leaves, which is what the worker's own count needs; the join points go
-- last so that every other rewrite sees one body. With the body come the
-- sources of its variables, the lets the rewrite changed ('Prepared')
-- and the join points.
prepareBody :: (Var v) => Batch -> Maybe Shape -> Maybe (Joining v) -> ANormal Reference v -> (ANormal Reference v, Sources v, M.Map [v] [v], Set.Set v, [Join v])
prepareBody :: forall v.
Var v =>
Batch
-> Maybe Shape
-> Maybe (Joining v)
-> ANormal Reference v
-> (ANormal Reference v, Sources v, Map [v] [v], Set v, [Join v])
prepareBody Batch
b Maybe Shape
ms Maybe (Joining v)
mj ANormal Reference v
bd
  | Bool -> Bool
not Bool
rewritesOn = (ANormal Reference v -> ANormal Reference v
forall v. Var v => ANormal Reference v -> ANormal Reference v
uniqueBody ANormal Reference v
bd, Sources v
forall k a. Map k a
M.empty, Map [v] [v]
forall k a. Map k a
M.empty, Set v
forall a. Set a
Set.empty, [])
  | Bool -> Bool
not ([v] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [v]
loose) =
      [Word]
-> String
-> (ANormal Reference v, Sources v, Map [v] [v], Set v, [Join v])
forall a. HasCallStack => [Word] -> String -> a
internalBug [] (String
 -> (ANormal Reference v, Sources v, Map [v] [v], Set v, [Join v]))
-> String
-> (ANormal Reference v, Sources v, Map [v] [v], Set v, [Join v])
forall a b. (a -> b) -> a -> b
$
        String
"prepareBody: the rewrites left "
          String -> ShowS
forall a. [a] -> [a] -> [a]
++ [v] -> String
forall a. Show a => a -> String
show [v]
loose
          String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" unbound; unoptimized body:\n"
          String -> ShowS
forall a. [a] -> [a] -> [a]
++ ANormal Reference v -> String
forall a. Show a => a -> String
show (ANormal Reference v -> ANormal Reference v
forall v. Var v => ANormal Reference v -> ANormal Reference v
uniqueBody ANormal Reference v
bd)
          String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"\nrewritten body:\n"
          String -> ShowS
forall a. [a] -> [a] -> [a]
++ ANormal Reference v -> String
forall a. Show a => a -> String
show ANormal Reference v
main
          String -> ShowS
forall a. [a] -> [a] -> [a]
++ [String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [String
"\njoin point:\n" String -> ShowS
forall a. [a] -> [a] -> [a]
++ SuperNormal Reference v -> String
forall a. Show a => a -> String
show (Join v -> SuperNormal Reference v
forall v. Join v -> SuperNormal Reference v
joinLambda Join v
jn) | Join v
jn <- [Join v]
joins]
  | Bool
otherwise = (ANormal Reference v -> ANormal Reference v
forall v. Var v => ANormal Reference v -> ANormal Reference v
freshLets ANormal Reference v
main, [Sources v] -> Sources v
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
M.unions [Sources v
srcsK, Sources v
srcs, Sources v
srcsJ], Map [v] [v]
lets, Set v
dead, (Join v -> Join v) -> [Join v] -> [Join v]
forall a b. (a -> b) -> [a] -> [b]
map Join v -> Join v
forall v. Var v => Join v -> Join v
freshJoin [Join v]
joins)
  where
    -- a rewritten body, or a join point under its parameters, is closed
    -- where the unoptimized body is
    scope :: Set v
scope = ANormal Reference v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
ABTN.freeVars ANormal Reference v
bd
    loose :: [v]
loose =
      Set v -> [v]
forall a. Set a -> [a]
Set.toList (Set v -> [v]) -> ([Set v] -> Set v) -> [Set v] -> [v]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Set v] -> Set v
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions ([Set v] -> [v]) -> [Set v] -> [v]
forall a b. (a -> b) -> a -> b
$
        (ANormal Reference v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
ABTN.freeVars ANormal Reference v
main Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
`Set.difference` Set v
scope)
          Set v -> [Set v] -> [Set v]
forall a. a -> [a] -> [a]
: [ ANormal Reference v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
ABTN.freeVars ANormal Reference v
jb Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
`Set.difference` (Set v
scope Set v -> Set v -> Set v
forall a. Semigroup a => a -> a -> a
<> [v] -> Set v
forall a. Ord a => [a] -> Set a
Set.fromList [v]
js)
            | Join v
jn <- [Join v]
joins,
              Lambda [Mem]
_ (TAbss [v]
js ANormal Reference v
jb) <- [Join v -> SuperNormal Reference v
forall v. Join v -> SuperNormal Reference v
joinLambda Join v
jn]
            ]
    (ANormal Reference v
known, Sources v
srcsK)
      | Bool
knownConstructorOn = State (Sources v) (ANormal Reference v)
-> Sources v -> (ANormal Reference v, Sources v)
forall s a. State s a -> s -> (a, s)
runState (ANormal Reference v -> State (Sources v) (ANormal Reference v)
forall v.
Var v =>
ANormal Reference v -> State (Sources v) (ANormal Reference v)
knownCons (ANormal Reference v -> ANormal Reference v
forall v. Var v => ANormal Reference v -> ANormal Reference v
uniqueBody ANormal Reference v
bd)) Sources v
forall k a. Map k a
M.empty
      | Bool
otherwise = (ANormal Reference v -> ANormal Reference v
forall v. Var v => ANormal Reference v -> ANormal Reference v
uniqueBody ANormal Reference v
bd, Sources v
forall k a. Map k a
M.empty)
    (ANormal Reference v
rets, Sources v
srcs, Map [v] [v]
lets) = Batch
-> ANormal Reference v
-> (ANormal Reference v, Sources v, Map [v] [v])
forall v.
Var v =>
Batch
-> ANormal Reference v
-> (ANormal Reference v, Sources v, Map [v] [v])
useReturns Batch
b (ANormal Reference v -> ANormal Reference v
forall v. Var v => ANormal Reference v -> ANormal Reference v
sinkCalls ANormal Reference v
known)
    body :: ANormal Reference v
body = ANormal Reference v -> ANormal Reference v
forall v. Var v => ANormal Reference v -> ANormal Reference v
dropDead (Batch -> ANormal Reference v -> ANormal Reference v
forall v.
Var v =>
Batch -> ANormal Reference v -> ANormal Reference v
useSpecs Batch
b ((ANormal Reference v -> ANormal Reference v)
-> (Shape -> ANormal Reference v -> ANormal Reference v)
-> Maybe Shape
-> ANormal Reference v
-> ANormal Reference v
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ANormal Reference v -> ANormal Reference v
forall a. a -> a
id (Batch -> Shape -> ANormal Reference v -> ANormal Reference v
forall v.
Var v =>
Batch -> Shape -> ANormal Reference v -> ANormal Reference v
yieldLeaves Batch
b) Maybe Shape
ms ANormal Reference v
rets))
    (ANormal Reference v
main, [Join v]
joins, Sources v
srcsJ, Set v
dead) = case Maybe (Joining v)
mj of
      Just Joining v
j | Bool
knownConstructorOn, Maybe Shape
Nothing <- Maybe Shape
ms -> Joining v
-> Sources v
-> ANormal Reference v
-> (ANormal Reference v, [Join v], Sources v, Set v)
forall v.
Var v =>
Joining v
-> Sources v
-> ANormal Reference v
-> (ANormal Reference v, [Join v], Sources v, Set v)
splitJoins Joining v
j (Sources v
srcsK Sources v -> Sources v -> Sources v
forall a. Semigroup a => a -> a -> a
<> Sources v
srcs) ANormal Reference v
body
      Maybe (Joining v)
_ -> (ANormal Reference v
body, [], Sources v
forall k a. Map k a
M.empty, Set v
forall a. Set a
Set.empty)

-- | The body with a number of its own for every indirect let: a let
-- 'splitJoins' copied into several tails carries its number into each
-- copy, and two sections of a combinator cannot share one, so every copy
-- after the first gets a number above the body's. The first keeps the
-- unoptimized number, which is what pairs sections when nothing else
-- changed ('deoptOf').
freshLets :: (Var v) => ANormal Reference v -> ANormal Reference v
freshLets :: forall v. Var v => ANormal Reference v -> ANormal Reference v
freshLets ANormal Reference v
bd = State (Set Word16, Word16) (ANormal Reference v)
-> (Set Word16, Word16) -> ANormal Reference v
forall s a. State s a -> s -> a
evalState (ANormal Reference v
-> State (Set Word16, Word16) (ANormal Reference v)
forall {f :: * -> *} {v} {ref}.
(Var v, MonadState (Set Word16, Word16) f) =>
Term (ANormalF ref) v -> f (Term (ANormalF ref) v)
go ANormal Reference v
bd) (Set Word16
forall a. Set a
Set.empty, Word16
top Word16 -> Word16 -> Word16
forall a. Num a => a -> a -> a
+ Word16
1)
  where
    top :: Word16
top = [Word16] -> Word16
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (Word16
0 Word16 -> [Word16] -> [Word16]
forall a. a -> [a] -> [a]
: [Word16
w | TLets (Indirect Word16
w) [v]
_ [Mem]
_ ANormal Reference v
_ ANormal Reference v
_ <- ANormal Reference v -> [ANormal Reference v]
forall v. ANormal Reference v -> [ANormal Reference v]
universe ANormal Reference v
bd])
    go :: Term (ANormalF ref) v -> f (Term (ANormalF ref) v)
go (ABTN.Term Set v
_ (ABTN.Abs v
u Term (ANormalF ref) v
b)) = v -> Term (ANormalF ref) v -> Term (ANormalF ref) v
forall v (f :: * -> * -> *). Var v => v -> Term f v -> Term f v
ABTN.TAbs v
u (Term (ANormalF ref) v -> Term (ANormalF ref) v)
-> f (Term (ANormalF ref) v) -> f (Term (ANormalF ref) v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Term (ANormalF ref) v -> f (Term (ANormalF ref) v)
go Term (ANormalF ref) v
b
    go (ABTN.Term Set v
_ (ABTN.Tm ANormalF ref v (Term (ANormalF ref) v)
f)) = case ANormalF ref v (Term (ANormalF ref) v)
f of
      ALet (Indirect Word16
w) [Mem]
ccs Term (ANormalF ref) v
bn Term (ANormalF ref) v
bo -> do
        (Set Word16
seen, Word16
next) <- f (Set Word16, Word16)
forall s (m :: * -> *). MonadState s m => m s
get
        Word16
w' <-
          if Word16 -> Set Word16 -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Word16
w Set Word16
seen
            then Word16
next Word16 -> f () -> f Word16
forall a b. a -> f b -> f a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ (Set Word16, Word16) -> f ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Set Word16
seen, Word16
next Word16 -> Word16 -> Word16
forall a. Num a => a -> a -> a
+ Word16
1)
            else Word16
w Word16 -> f () -> f Word16
forall a b. a -> f b -> f a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ (Set Word16, Word16) -> f ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Word16 -> Set Word16 -> Set Word16
forall a. Ord a => a -> Set a -> Set a
Set.insert Word16
w Set Word16
seen, Word16
next)
        ANormalF ref v (Term (ANormalF ref) v) -> Term (ANormalF ref) v
forall v (f :: * -> * -> *).
(Var v, Bifoldable f) =>
f v (Term f v) -> Term f v
ABTN.TTm (ANormalF ref v (Term (ANormalF ref) v) -> Term (ANormalF ref) v)
-> f (ANormalF ref v (Term (ANormalF ref) v))
-> f (Term (ANormalF ref) v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Direction Word16
-> [Mem]
-> Term (ANormalF ref) v
-> Term (ANormalF ref) v
-> ANormalF ref v (Term (ANormalF ref) v)
forall ref v e.
Direction Word16 -> [Mem] -> e -> e -> ANormalF ref v e
ALet (Word16 -> Direction Word16
forall a. a -> Direction a
Indirect Word16
w') [Mem]
ccs (Term (ANormalF ref) v
 -> Term (ANormalF ref) v -> ANormalF ref v (Term (ANormalF ref) v))
-> f (Term (ANormalF ref) v)
-> f (Term (ANormalF ref) v
      -> ANormalF ref v (Term (ANormalF ref) v))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Term (ANormalF ref) v -> f (Term (ANormalF ref) v)
go Term (ANormalF ref) v
bn f (Term (ANormalF ref) v -> ANormalF ref v (Term (ANormalF ref) v))
-> f (Term (ANormalF ref) v)
-> f (ANormalF ref v (Term (ANormalF ref) v))
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Term (ANormalF ref) v -> f (Term (ANormalF ref) v)
go Term (ANormalF ref) v
bo)
      ANormalF ref v (Term (ANormalF ref) v)
_ -> ANormalF ref v (Term (ANormalF ref) v) -> Term (ANormalF ref) v
forall v (f :: * -> * -> *).
(Var v, Bifoldable f) =>
f v (Term f v) -> Term f v
ABTN.TTm (ANormalF ref v (Term (ANormalF ref) v) -> Term (ANormalF ref) v)
-> f (ANormalF ref v (Term (ANormalF ref) v))
-> f (Term (ANormalF ref) v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Term (ANormalF ref) v -> f (Term (ANormalF ref) v))
-> ANormalF ref v (Term (ANormalF ref) v)
-> f (ANormalF ref v (Term (ANormalF 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) -> ANormalF ref v a -> f (ANormalF ref v b)
traverse Term (ANormalF ref) v -> f (Term (ANormalF ref) v)
go ANormalF ref v (Term (ANormalF ref) v)
f

freshJoin :: (Var v) => Join v -> Join v
freshJoin :: forall v. Var v => Join v -> Join v
freshJoin (Join (Lambda [Mem]
ccs (TAbss [v]
vs Term (ANormalF Reference) v
bd)) Maybe [v]
cut Sources v
srcs) = SuperNormal Reference v -> Maybe [v] -> Sources v -> Join v
forall v.
SuperNormal Reference v -> Maybe [v] -> Sources v -> Join v
Join ([Mem] -> Term (ANormalF Reference) v -> SuperNormal Reference v
forall ref v. [Mem] -> ANormal ref v -> SuperNormal ref v
Lambda [Mem]
ccs ([v] -> Term (ANormalF Reference) v -> Term (ANormalF Reference) v
forall v (f :: * -> * -> *). Var v => [v] -> Term f v -> Term f v
TAbss [v]
vs (Term (ANormalF Reference) v -> Term (ANormalF Reference) v
forall v. Var v => ANormal Reference v -> ANormal Reference v
freshLets Term (ANormalF Reference) v
bd))) Maybe [v]
cut Sources v
srcs

-- | The body with every match on a variable bound to a constructor
-- application or a record literal replaced by its arm, the arm's
-- variables renamed to the fields and recorded as their sources (the
-- argument side's 'rewrite', over a tree whose leaves are the fields);
-- a record update of such a record becomes a literal. The let stays
-- until 'dropDead' finds nothing else uses it. A match whose arms do not
-- include the constructor (a default arm would take it) is left as it is.
-- The body of a let is done before the let, since the arm of a later
-- product's match is what names an earlier product's variable (the
-- pair inside a tuple).
knownCons :: (Var v) => ANormal Reference v -> State (Sources v) (ANormal Reference v)
knownCons :: forall v.
Var v =>
ANormal Reference v -> State (Sources v) (ANormal Reference v)
knownCons = ANormal Reference v
-> StateT (Sources v) Identity (ANormal Reference v)
forall v.
Var v =>
ANormal Reference v -> State (Sources v) (ANormal Reference v)
go
  where
    go :: ANormal Reference b
-> StateT (Sources b) Identity (ANormal Reference b)
go ANormal Reference b
tm = case ANormal Reference b
tm of
      TLets Direction Word16
d [b
u] [Mem]
ms ANormal Reference b
bn ANormal Reference b
bo
        | TApp (FCon Reference
r CTag
t) [b]
args <- ANormal Reference b
bn -> Direction Word16
-> [b]
-> [Mem]
-> ANormal Reference b
-> ANormal Reference b
-> ANormal Reference b
forall v ref.
Var v =>
Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
TLets Direction Word16
d [b
u] [Mem]
ms ANormal Reference b
bn (ANormal Reference b -> ANormal Reference b)
-> StateT (Sources b) Identity (ANormal Reference b)
-> StateT (Sources b) Identity (ANormal Reference b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (b
-> Tree b
-> ANormal Reference b
-> StateT (Sources b) Identity (ANormal Reference b)
forall v.
Var v =>
v
-> Tree v
-> ANormal Reference v
-> State (Sources v) (ANormal Reference v)
rewrite b
u (b -> Shape -> [Tree b] -> Tree b
forall v. v -> Shape -> [Tree v] -> Tree v
Node b
u (Reference -> CTag -> [Shape] -> Shape
Con Reference
r CTag
t (Shape
Plain Shape -> [b] -> [Shape]
forall a b. a -> [b] -> [a]
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ [b]
args)) ((b -> Tree b) -> [b] -> [Tree b]
forall a b. (a -> b) -> [a] -> [b]
map b -> Tree b
forall v. v -> Tree v
Leaf [b]
args)) (ANormal Reference b
 -> StateT (Sources b) Identity (ANormal Reference b))
-> StateT (Sources b) Identity (ANormal Reference b)
-> StateT (Sources b) Identity (ANormal Reference b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ANormal Reference b
-> StateT (Sources b) Identity (ANormal Reference b)
go ANormal Reference b
bo)
        | TApp (FRec RecordSchema
rs) [b]
args <- ANormal Reference b
bn -> Direction Word16
-> [b]
-> [Mem]
-> ANormal Reference b
-> ANormal Reference b
-> ANormal Reference b
forall v ref.
Var v =>
Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
TLets Direction Word16
d [b
u] [Mem]
ms ANormal Reference b
bn (ANormal Reference b -> ANormal Reference b)
-> StateT (Sources b) Identity (ANormal Reference b)
-> StateT (Sources b) Identity (ANormal Reference b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (b
-> Tree b
-> ANormal Reference b
-> StateT (Sources b) Identity (ANormal Reference b)
forall v.
Var v =>
v
-> Tree v
-> ANormal Reference v
-> State (Sources v) (ANormal Reference v)
rewrite b
u (b -> Shape -> [Tree b] -> Tree b
forall v. v -> Shape -> [Tree v] -> Tree v
Node b
u (RecordSchema -> [Shape] -> Shape
Fields RecordSchema
rs (Shape
Plain Shape -> [b] -> [Shape]
forall a b. a -> [b] -> [a]
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ [b]
args)) ((b -> Tree b) -> [b] -> [Tree b]
forall a b. (a -> b) -> [a] -> [b]
map b -> Tree b
forall v. v -> Tree v
Leaf [b]
args)) (ANormal Reference b
 -> StateT (Sources b) Identity (ANormal Reference b))
-> StateT (Sources b) Identity (ANormal Reference b)
-> StateT (Sources b) Identity (ANormal Reference b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ANormal Reference b
-> StateT (Sources b) Identity (ANormal Reference b)
go ANormal Reference b
bo)
      ABTN.Term Set b
_ (ABTN.Abs b
v ANormal Reference b
e) -> b -> ANormal Reference b -> ANormal Reference b
forall v (f :: * -> * -> *). Var v => v -> Term f v -> Term f v
ABTN.TAbs b
v (ANormal Reference b -> ANormal Reference b)
-> StateT (Sources b) Identity (ANormal Reference b)
-> StateT (Sources b) Identity (ANormal Reference b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ANormal Reference b
-> StateT (Sources b) Identity (ANormal Reference b)
go ANormal Reference b
e
      ABTN.Term Set b
_ (ABTN.Tm ANormalF Reference b (ANormal Reference b)
f) -> ANormalF Reference b (ANormal Reference b) -> ANormal Reference b
forall v (f :: * -> * -> *).
(Var v, Bifoldable f) =>
f v (Term f v) -> Term f v
ABTN.TTm (ANormalF Reference b (ANormal Reference b) -> ANormal Reference b)
-> StateT
     (Sources b) Identity (ANormalF Reference b (ANormal Reference b))
-> StateT (Sources b) Identity (ANormal Reference b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ANormal Reference b
 -> StateT (Sources b) Identity (ANormal Reference b))
-> ANormalF Reference b (ANormal Reference b)
-> StateT
     (Sources b) Identity (ANormalF Reference b (ANormal Reference 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) -> ANormalF Reference b a -> f (ANormalF Reference b b)
traverse ANormal Reference b
-> StateT (Sources b) Identity (ANormal Reference b)
go ANormalF Reference b (ANormal Reference b)
f

-- | The sites of a body where a match on the value of a match can be
-- moved into the tails that build the value ('site'), by the let's
-- binder: for each arm of the outer match (by constructor number, Nothing
-- for the default arm), the variables its body uses that are bound
-- outside it, and the same for what follows the match when its value is
-- let-bound, the group's local functions aside. Recorded on the
-- unoptimized body, so that a join point's parameters cover every
-- variable the unoptimized code reads, which is what lets a continuation
-- captured inside the join point be written out as the unoptimized
-- frames: a variable the join point's frame cannot rebuild is then one
-- the unoptimized code never reads there ('splitJoins', 'deoptOf').
type Sites v = M.Map v (M.Map (Maybe Int) (Set.Set v), Set.Set v)

caseSites :: (Var v) => ConLets v -> Set.Set v -> ANormal Reference v -> Sites v
caseSites :: forall v.
Var v =>
ConLets v -> Set v -> ANormal Reference v -> Sites v
caseSites ConLets v
cons Set v
locals ANormal Reference v
bd = [(v, (Map (Maybe Int) (Set v), Set v))]
-> Map v (Map (Maybe Int) (Set v), Set v)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(v
p, (EnumMap CTag ([Mem], ANormal Reference v)
-> Maybe (ANormal Reference v) -> Map (Maybe Int) (Set v)
arms EnumMap CTag ([Mem], ANormal Reference v)
cs Maybe (ANormal Reference v)
df, Maybe (Direction Word16, v, [Mem], ANormal Reference v) -> Set v
restFree Maybe (Direction Word16, v, [Mem], ANormal Reference v)
kont)) | Just (Site v
p ANormal Reference v
_ Reference
_ EnumMap CTag ([Mem], ANormal Reference v)
cs Maybe (ANormal Reference v)
df Maybe (Direction Word16, v, [Mem], ANormal Reference v)
kont) <- (ANormal Reference v -> Maybe (Site v))
-> [ANormal Reference v] -> [Maybe (Site v)]
forall a b. (a -> b) -> [a] -> [b]
map (ConLets v -> ANormal Reference v -> Maybe (Site v)
forall v.
Var v =>
ConLets v -> ANormal Reference v -> Maybe (Site v)
site ConLets v
cons) (ANormal Reference v -> [ANormal Reference v]
forall v. ANormal Reference v -> [ANormal Reference v]
universe ANormal Reference v
bd)]
  where
    arms :: EnumMap CTag ([Mem], ANormal Reference v)
-> Maybe (ANormal Reference v) -> Map (Maybe Int) (Set v)
arms EnumMap CTag ([Mem], ANormal Reference v)
cs Maybe (ANormal Reference v)
df =
      [(Maybe Int, Set v)] -> Map (Maybe Int) (Set v)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Maybe Int, Set v)] -> Map (Maybe Int) (Set v))
-> [(Maybe Int, Set v)] -> Map (Maybe Int) (Set v)
forall a b. (a -> b) -> a -> b
$
        [(Int -> Maybe Int
forall a. a -> Maybe a
Just (CTag -> Int
forall k. EnumKey k => k -> Int
EC.keyToInt CTag
t), [v] -> ANormal Reference v -> Set v
free [v]
fs ANormal Reference v
bo) | (CTag
t, ([Mem]
_, TAbss [v]
fs ANormal Reference v
bo)) <- EnumMap CTag ([Mem], ANormal Reference v)
-> [(CTag, ([Mem], ANormal Reference v))]
forall k a. EnumKey k => EnumMap k a -> [(k, a)]
EC.mapToList EnumMap CTag ([Mem], ANormal Reference v)
cs]
          [(Maybe Int, Set v)]
-> [(Maybe Int, Set v)] -> [(Maybe Int, Set v)]
forall a. [a] -> [a] -> [a]
++ [(Maybe Int
forall a. Maybe a
Nothing, [v] -> ANormal Reference v -> Set v
free [] ANormal Reference v
d) | Just ANormal Reference v
d <- [Maybe (ANormal Reference v)
df]]
    restFree :: Maybe (Direction Word16, v, [Mem], ANormal Reference v) -> Set v
restFree = Set v
-> ((Direction Word16, v, [Mem], ANormal Reference v) -> Set v)
-> Maybe (Direction Word16, v, [Mem], ANormal Reference v)
-> Set v
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Set v
forall a. Set a
Set.empty (\(Direction Word16
_, v
x, [Mem]
_, ANormal Reference v
rest) -> [v] -> ANormal Reference v -> Set v
free [v
x] ANormal Reference v
rest)
    free :: [v] -> ANormal Reference v -> Set v
free [v]
fs ANormal Reference v
bo = (ANormal Reference v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
ABTN.freeVars ANormal Reference v
bo Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
`Set.difference` [v] -> Set v
forall a. Ord a => [a] -> Set a
Set.fromList [v]
fs) Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
`Set.difference` Set v
locals

-- | A let whose binding's every tail builds a constructor of one type
-- ('tailCon') and whose body is a data match on the let's variable, used
-- nowhere else: the variable, the binding, the type, the arms and the
-- default arm, and what takes the match's value: nothing when the match
-- is the let's body, so its value is the enclosing section's, or the let
-- binding it (its direction, variable and convention) and the rest of
-- the body, which is how normalization leaves a block whose last
-- expression is the match.
data Site v = Site v (ANormal Reference v) Reference (EC.EnumMap CTag ([Mem], ANormal Reference v)) (Maybe (ANormal Reference v)) (Maybe (Direction Word16, v, [Mem], ANormal Reference v))

site :: (Var v) => ConLets v -> ANormal Reference v -> Maybe (Site v)
site :: forall v.
Var v =>
ConLets v -> ANormal Reference v -> Maybe (Site v)
site ConLets v
cons ANormal Reference v
tm = case ANormal Reference v
tm of
  TLets Direction Word16
_ [v
p] [Mem]
_ ANormal Reference v
bn ANormal Reference v
bo
    | Just (Reference
r, EnumMap CTag ([Mem], ANormal Reference v)
cs, Maybe (ANormal Reference v)
df, Maybe (Direction Word16, v, [Mem], ANormal Reference v)
kont) <- v
-> ANormal Reference v
-> Maybe
     (Reference, EnumMap CTag ([Mem], ANormal Reference v),
      Maybe (ANormal Reference v),
      Maybe (Direction Word16, v, [Mem], ANormal Reference v))
forall {b} {ref}.
Var b =>
b
-> ANormal ref b
-> Maybe
     (ref, EnumMap CTag ([Mem], ANormal ref b), Maybe (ANormal ref b),
      Maybe (Direction Word16, b, [Mem], ANormal ref b))
outer v
p ANormal Reference v
bo,
      v -> ANormal Reference v -> Int
forall v. Var v => v -> ANormal Reference v -> Int
occurrences v
p ANormal Reference v
bo Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1,
      (ANormal Reference v -> Bool) -> [ANormal Reference v] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Maybe (CTag, [v]) -> Bool
forall a. Maybe a -> Bool
isJust (Maybe (CTag, [v]) -> Bool)
-> (ANormal Reference v -> Maybe (CTag, [v]))
-> ANormal Reference v
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ConLets v -> Reference -> ANormal Reference v -> Maybe (CTag, [v])
forall v.
Var v =>
ConLets v -> Reference -> ANormal Reference v -> Maybe (CTag, [v])
tailCon ConLets v
cons Reference
r) (ANormal Reference v -> [ANormal Reference v]
forall v. Var v => ANormal Reference v -> [ANormal Reference v]
tailsOf ANormal Reference v
bn) ->
        Site v -> Maybe (Site v)
forall a. a -> Maybe a
Just (v
-> ANormal Reference v
-> Reference
-> EnumMap CTag ([Mem], ANormal Reference v)
-> Maybe (ANormal Reference v)
-> Maybe (Direction Word16, v, [Mem], ANormal Reference v)
-> Site v
forall v.
v
-> ANormal Reference v
-> Reference
-> EnumMap CTag ([Mem], ANormal Reference v)
-> Maybe (ANormal Reference v)
-> Maybe (Direction Word16, v, [Mem], ANormal Reference v)
-> Site v
Site v
p ANormal Reference v
bn Reference
r EnumMap CTag ([Mem], ANormal Reference v)
cs Maybe (ANormal Reference v)
df Maybe (Direction Word16, v, [Mem], ANormal Reference v)
kont)
  ANormal Reference v
_ -> Maybe (Site v)
forall a. Maybe a
Nothing
  where
    outer :: b
-> ANormal ref b
-> Maybe
     (ref, EnumMap CTag ([Mem], ANormal ref b), Maybe (ANormal ref b),
      Maybe (Direction Word16, b, [Mem], ANormal ref b))
outer b
p = \case
      TMatch b
p' (MatchData ref
r EnumMap CTag ([Mem], ANormal ref b)
cs Maybe (ANormal ref b)
df) | b
p' b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== b
p -> (ref, EnumMap CTag ([Mem], ANormal ref b), Maybe (ANormal ref b),
 Maybe (Direction Word16, b, [Mem], ANormal ref b))
-> Maybe
     (ref, EnumMap CTag ([Mem], ANormal ref b), Maybe (ANormal ref b),
      Maybe (Direction Word16, b, [Mem], ANormal ref b))
forall a. a -> Maybe a
Just (ref
r, EnumMap CTag ([Mem], ANormal ref b)
cs, Maybe (ANormal ref b)
df, Maybe (Direction Word16, b, [Mem], ANormal ref b)
forall a. Maybe a
Nothing)
      TLets Direction Word16
d [b
x] [Mem]
ms (TMatch b
p' (MatchData ref
r EnumMap CTag ([Mem], ANormal ref b)
cs Maybe (ANormal ref b)
df)) ANormal ref b
rest | b
p' b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== b
p -> (ref, EnumMap CTag ([Mem], ANormal ref b), Maybe (ANormal ref b),
 Maybe (Direction Word16, b, [Mem], ANormal ref b))
-> Maybe
     (ref, EnumMap CTag ([Mem], ANormal ref b), Maybe (ANormal ref b),
      Maybe (Direction Word16, b, [Mem], ANormal ref b))
forall a. a -> Maybe a
Just (ref
r, EnumMap CTag ([Mem], ANormal ref b)
cs, Maybe (ANormal ref b)
df, (Direction Word16, b, [Mem], ANormal ref b)
-> Maybe (Direction Word16, b, [Mem], ANormal ref b)
forall a. a -> Maybe a
Just (Direction Word16
d, b
x, [Mem]
ms, ANormal ref b
rest))
      ANormal ref b
_ -> Maybe
  (ref, EnumMap CTag ([Mem], ANormal ref b), Maybe (ANormal ref b),
   Maybe (Direction Word16, b, [Mem], ANormal ref b))
forall a. Maybe a
Nothing

-- | The tails of a binding, through lets, names and the arms of matches
-- (what 'sinkCalls' sees through too).
tailsOf :: (Var v) => ANormal Reference v -> [ANormal Reference v]
tailsOf :: forall v. Var v => ANormal Reference v -> [ANormal Reference v]
tailsOf = \case
  TMatch v
_ Branched Reference (ANormal Reference v)
bs -> (ANormal Reference v -> [ANormal Reference v])
-> [ANormal Reference v] -> [ANormal Reference v]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ANormal Reference v -> [ANormal Reference v]
forall v. Var v => ANormal Reference v -> [ANormal Reference v]
tailsOf (Branched Reference (ANormal Reference v) -> [ANormal Reference v]
forall a. Branched Reference a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Branched Reference (ANormal Reference v)
bs)
  TLets Direction Word16
_ [v]
_ [Mem]
_ ANormal Reference v
_ ANormal Reference v
bo -> ANormal Reference v -> [ANormal Reference v]
forall v. Var v => ANormal Reference v -> [ANormal Reference v]
tailsOf ANormal Reference v
bo
  TName v
_ Either Reference v
_ [v]
_ ANormal Reference v
bo -> ANormal Reference v -> [ANormal Reference v]
forall v. Var v => ANormal Reference v -> [ANormal Reference v]
tailsOf ANormal Reference v
bo
  ABTN.Term Set v
_ (ABTN.Abs v
_ ANormal Reference v
b) -> ANormal Reference v -> [ANormal Reference v]
forall v. Var v => ANormal Reference v -> [ANormal Reference v]
tailsOf ANormal Reference v
b
  ANormal Reference v
tm -> [ANormal Reference v
tm]

-- | The body with a function applied at each tail.
sinkTails :: (Var v) => (ANormal Reference v -> ANormal Reference v) -> ANormal Reference v -> ANormal Reference v
sinkTails :: forall v.
Var v =>
(ANormal Reference v -> ANormal Reference v)
-> ANormal Reference v -> ANormal Reference v
sinkTails ANormal Reference v -> ANormal Reference v
k = \case
  TMatch v
v Branched Reference (ANormal Reference v)
bs -> v
-> Branched Reference (ANormal Reference v) -> ANormal Reference v
forall v ref.
Var v =>
v -> Branched ref (ANormal ref v) -> ANormal ref v
TMatch v
v ((ANormal Reference v -> ANormal Reference v)
-> Branched Reference (ANormal Reference v)
-> Branched Reference (ANormal Reference v)
forall a b.
(a -> b) -> Branched Reference a -> Branched Reference b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((ANormal Reference v -> ANormal Reference v)
-> ANormal Reference v -> ANormal Reference v
forall v.
Var v =>
(ANormal Reference v -> ANormal Reference v)
-> ANormal Reference v -> ANormal Reference v
sinkTails ANormal Reference v -> ANormal Reference v
k) Branched Reference (ANormal Reference v)
bs)
  TLets Direction Word16
d [v]
us [Mem]
ms ANormal Reference v
bn ANormal Reference v
bo -> Direction Word16
-> [v]
-> [Mem]
-> ANormal Reference v
-> ANormal Reference v
-> ANormal Reference v
forall v ref.
Var v =>
Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
TLets Direction Word16
d [v]
us [Mem]
ms ANormal Reference v
bn ((ANormal Reference v -> ANormal Reference v)
-> ANormal Reference v -> ANormal Reference v
forall v.
Var v =>
(ANormal Reference v -> ANormal Reference v)
-> ANormal Reference v -> ANormal Reference v
sinkTails ANormal Reference v -> ANormal Reference v
k ANormal Reference v
bo)
  TName v
v Either Reference v
f [v]
as ANormal Reference v
bo -> v
-> Either Reference v
-> [v]
-> ANormal Reference v
-> ANormal Reference v
forall v ref.
Var v =>
v -> Either ref v -> [v] -> ANormal ref v -> ANormal ref v
TName v
v Either Reference v
f [v]
as ((ANormal Reference v -> ANormal Reference v)
-> ANormal Reference v -> ANormal Reference v
forall v.
Var v =>
(ANormal Reference v -> ANormal Reference v)
-> ANormal Reference v -> ANormal Reference v
sinkTails ANormal Reference v -> ANormal Reference v
k ANormal Reference v
bo)
  ABTN.Term Set v
_ (ABTN.Abs v
u ANormal Reference v
b) -> v -> ANormal Reference v -> ANormal Reference v
forall v (f :: * -> * -> *). Var v => v -> Term f v -> Term f v
ABTN.TAbs v
u ((ANormal Reference v -> ANormal Reference v)
-> ANormal Reference v -> ANormal Reference v
forall v.
Var v =>
(ANormal Reference v -> ANormal Reference v)
-> ANormal Reference v -> ANormal Reference v
sinkTails ANormal Reference v -> ANormal Reference v
k ANormal Reference v
b)
  ANormal Reference v
tm -> ANormal Reference v -> ANormal Reference v
k ANormal Reference v
tm

-- | A tail that builds a constructor of the type, or names a variable
-- bound to one: its tag and its fields.
tailCon :: (Var v) => ConLets v -> Reference -> ANormal Reference v -> Maybe (CTag, [v])
tailCon :: forall v.
Var v =>
ConLets v -> Reference -> ANormal Reference v -> Maybe (CTag, [v])
tailCon ConLets v
cons Reference
r = \case
  TApp (FCon Reference
r' CTag
t) [v]
args | Reference
r' Reference -> Reference -> Bool
forall a. Eq a => a -> a -> Bool
== Reference
r -> (CTag, [v]) -> Maybe (CTag, [v])
forall a. a -> Maybe a
Just (CTag
t, [v]
args)
  TVar v
u | Just (BuiltCon Reference
r' CTag
t [v]
fs) <- v -> ConLets v -> Maybe (Built v)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup v
u ConLets v
cons, Reference
r' Reference -> Reference -> Bool
forall a. Eq a => a -> a -> Bool
== Reference
r -> (CTag, [v]) -> Maybe (CTag, [v])
forall a. a -> Maybe a
Just (CTag
t, [v]
fs)
  ANormal Reference v
_ -> Maybe (CTag, [v])
forall a. Maybe a
Nothing

-- | The body with each site ('site') that the unoptimized body had
-- ('caseSites') rewritten: the let is gone, and each tail of its binding
-- becomes the outer arm for the constructor it builds, with the arm's
-- variables bound to the fields, followed, when the match's value was
-- let-bound, by a let of that variable over the rest of the body. An arm
-- that is a single expression, or that only one tail builds the
-- constructor of, is copied into the tail (the argument side's 'rewrite'
-- does the renaming and records the arm's variables' sources, as for an
-- exploded parameter); any other arm, and the rest when it is more than
-- a single expression, becomes a join point, made once, that the tail
-- calls with the fields (the match's value, for the rest) and the other
-- variables it uses: the variables the unoptimized code used there, each
-- as what stands for it now (renamed, or the fields of a product the body
-- no longer names, through the sources and the unoptimized products), and
-- the variables it uses now. The let's variable is dead everywhere its
-- frame could hold it, and is recorded so. A site whose tails would need
-- more join points than the lambda has room for is left as it is. The
-- join points come out in the order they are made, the k-th at
-- 'joinIndex' k; a site inside an arm or the rest makes its join points
-- before the arm does. With the body come the sources the copied arms
-- recorded and the dead variables.
splitJoins :: (Var v) => Joining v -> Sources v -> ANormal Reference v -> (ANormal Reference v, [Join v], Sources v, Set.Set v)
splitJoins :: forall v.
Var v =>
Joining v
-> Sources v
-> ANormal Reference v
-> (ANormal Reference v, [Join v], Sources v, Set v)
splitJoins (Joining (JoinSetting Word64
owner Reference
grpr Set v
locals) Sites v
sites ConLets v
bcons Sources v
srcs0) Sources v
srcs1 ANormal Reference v
bd = (ANormal Reference v
bd', [Join v] -> [Join v]
forall a. [a] -> [a]
reverse [Join v]
joins, Sources v
srcsOut, Set v
dead)
  where
    cons :: ConLets v
cons = ANormal Reference v -> ConLets v
forall v. Var v => ANormal Reference v -> ConLets v
conLets ANormal Reference v
bd
    srcs :: Sources v
srcs = Sources v
srcs0 Sources v -> Sources v -> Sources v
forall a. Semigroup a => a -> a -> a
<> Sources v
srcs1
    inScope :: Set v
inScope = [Set v] -> Set v
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions (ANormal Reference v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
ABTN.freeVars ANormal Reference v
bd Set v -> [Set v] -> [Set v]
forall a. a -> [a] -> [a]
: (ANormal Reference v -> Set v) -> [ANormal Reference v] -> [Set v]
forall a b. (a -> b) -> [a] -> [b]
map ANormal Reference v -> Set v
forall v. Ord v => ANormal Reference v -> Set v
boundVars (ANormal Reference v -> [ANormal Reference v]
forall v. ANormal Reference v -> [ANormal Reference v]
universe ANormal Reference v
bd))
    (ANormal Reference v
bd', ([Join v]
joins, Sources v
srcsOut, Set v
dead)) = State ([Join v], Sources v, Set v) (ANormal Reference v)
-> ([Join v], Sources v, Set v)
-> (ANormal Reference v, ([Join v], Sources v, Set v))
forall s a. State s a -> s -> (a, s)
runState (Maybe [v]
-> ANormal Reference v
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
go Maybe [v]
forall a. Maybe a
Nothing ANormal Reference v
bd) ([], Sources v
forall k a. Map k a
M.empty, Set v
forall a. Set a
Set.empty)
    go :: Maybe [v]
-> ANormal Reference v
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
go Maybe [v]
cut ANormal Reference v
tm = case ANormal Reference v
tm of
      TLets Direction Word16
d [v
p] [Mem]
ms ANormal Reference v
bn ANormal Reference v
bo
        | Just (Site v
_ ANormal Reference v
_ Reference
r EnumMap CTag ([Mem], ANormal Reference v)
cs Maybe (ANormal Reference v)
df Maybe (Direction Word16, v, [Mem], ANormal Reference v)
kont) <- ConLets v -> ANormal Reference v -> Maybe (Site v)
forall v.
Var v =>
ConLets v -> ANormal Reference v -> Maybe (Site v)
site ConLets v
cons ANormal Reference v
tm,
          Just (Map (Maybe Int) (Set v)
arms, Set v
restFree) <- v -> Sites v -> Maybe (Map (Maybe Int) (Set v), Set v)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup v
p Sites v
sites,
          let built :: [CTag]
built = [CTag
t | ANormal Reference v
tl <- ANormal Reference v -> [ANormal Reference v]
forall v. Var v => ANormal Reference v -> [ANormal Reference v]
tailsOf ANormal Reference v
bn, Just (CTag
t, [v]
_) <- [ConLets v -> Reference -> ANormal Reference v -> Maybe (CTag, [v])
forall v.
Var v =>
ConLets v -> Reference -> ANormal Reference v -> Maybe (CTag, [v])
tailCon ConLets v
cons Reference
r ANormal Reference v
tl]]
              tags :: [CTag]
tags = [CTag] -> [CTag]
forall a. Ord a => [a] -> [a]
nubOrd [CTag]
built
              armOf :: CTag -> Maybe (Maybe Int, [v], ANormal Reference v)
armOf CTag
t = case CTag
-> EnumMap CTag ([Mem], ANormal Reference v)
-> Maybe ([Mem], ANormal Reference v)
forall k a. EnumKey k => k -> EnumMap k a -> Maybe a
EC.lookup CTag
t EnumMap CTag ([Mem], ANormal Reference v)
cs of
                Just ([Mem]
_, TAbss [v]
fs ANormal Reference v
ab) -> (Maybe Int, [v], ANormal Reference v)
-> Maybe (Maybe Int, [v], ANormal Reference v)
forall a. a -> Maybe a
Just (Int -> Maybe Int
forall a. a -> Maybe a
Just (CTag -> Int
forall k. EnumKey k => k -> Int
EC.keyToInt CTag
t), [v]
fs, ANormal Reference v
ab)
                Maybe ([Mem], ANormal Reference v)
Nothing -> (\ANormal Reference v
ab -> (Maybe Int
forall a. Maybe a
Nothing, [], ANormal Reference v
ab)) (ANormal Reference v -> (Maybe Int, [v], ANormal Reference v))
-> Maybe (ANormal Reference v)
-> Maybe (Maybe Int, [v], ANormal Reference v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (ANormal Reference v)
df
              -- an arm is copied when it is a leaf, or when one tail
              -- reaches it (the default arm is reached by every
              -- constructor without an arm of its own)
              tailsTo :: Maybe Int -> Int
tailsTo Maybe Int
k = [()] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [() | CTag
t <- [CTag]
built, Just (Maybe Int
k', [v]
_, ANormal Reference v
_) <- [CTag -> Maybe (Maybe Int, [v], ANormal Reference v)
armOf CTag
t], Maybe Int
k' Maybe Int -> Maybe Int -> Bool
forall a. Eq a => a -> a -> Bool
== Maybe Int
k]
              copied :: Maybe Int -> ANormal Reference v -> Bool
copied Maybe Int
k ANormal Reference v
ab = ANormal Reference v -> Bool
forall {v}. ANormal Reference v -> Bool
leaf ANormal Reference v
ab Bool -> Bool -> Bool
|| Maybe Int -> Int
tailsTo Maybe Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1
              -- the arms of the outer match are inside the let that takes
              -- its value, when there is one
              armCut :: Maybe [v]
armCut = Maybe [v]
-> ((Direction Word16, v, [Mem], ANormal Reference v) -> Maybe [v])
-> Maybe (Direction Word16, v, [Mem], ANormal Reference v)
-> Maybe [v]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Maybe [v]
cut (\(Direction Word16
_, v
x, [Mem]
_, ANormal Reference v
_) -> [v] -> Maybe [v]
forall a. a -> Maybe a
Just [v
x]) Maybe (Direction Word16, v, [Mem], ANormal Reference v)
kont
              restJoin :: Int
restJoin = case Maybe (Direction Word16, v, [Mem], ANormal Reference v)
kont of
                Just (Direction Word16
_, v
_, [Mem]
_, ANormal Reference v
rest) | Bool -> Bool
not (ANormal Reference v -> Bool
forall {v}. ANormal Reference v -> Bool
leaf ANormal Reference v
rest) -> Int
1
                Maybe (Direction Word16, v, [Mem], ANormal Reference v)
_ -> Int
0,
          Just [(Maybe Int, [v], ANormal Reference v)]
targets <- (CTag -> Maybe (Maybe Int, [v], ANormal Reference v))
-> [CTag] -> Maybe [(Maybe Int, [v], ANormal Reference 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 CTag -> Maybe (Maybe Int, [v], ANormal Reference v)
armOf [CTag]
tags -> do
            ([Join v]
made, Sources v
_, Set v
_) <- StateT
  ([Join v], Sources v, Set v) Identity ([Join v], Sources v, Set v)
forall s (m :: * -> *). MonadState s m => m s
get
            let needed :: Int
needed = [Maybe Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Maybe Int] -> [Maybe Int]
forall a. Ord a => [a] -> [a]
nubOrd [Maybe Int
k | (Maybe Int
k, [v]
_, ANormal Reference v
ab) <- [(Maybe Int, [v], ANormal Reference v)]
targets, Bool -> Bool
not (Maybe Int -> ANormal Reference v -> Bool
copied Maybe Int
k ANormal Reference v
ab)]) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
restJoin
            if [Join v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Join v]
made Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
needed Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxJoins
              then State ([Join v], Sources v, Set v) (ANormal Reference v)
fallback
              else do
                (([Join v], Sources v, Set v) -> ([Join v], Sources v, Set v))
-> StateT ([Join v], Sources v, Set v) Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (\([Join v]
js, Sources v
ss, Set v
ds) -> ([Join v]
js, Sources v
ss, v -> Set v -> Set v
forall a. Ord a => a -> Set a -> Set a
Set.insert v
p Set v
ds))
                -- the rest of the body as a join point taking the match's
                -- value, unless it is a single expression, copied
                Maybe (ANormal Reference v)
restCode <- case Maybe (Direction Word16, v, [Mem], ANormal Reference v)
kont of
                  Just (Direction Word16
_, v
x, [Mem]
_, ANormal Reference v
rest)
                    | Bool -> Bool
not (ANormal Reference v -> Bool
forall {v}. ANormal Reference v -> Bool
leaf ANormal Reference v
rest) -> do
                        ANormal Reference v
rest' <- Maybe [v]
-> ANormal Reference v
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
go Maybe [v]
cut ANormal Reference v
rest
                        (Int, [v])
j <- Maybe [v]
-> Sources v
-> [v]
-> Set v
-> ANormal Reference v
-> StateT ([Join v], Sources v, Set v) Identity (Int, [v])
newJoin Maybe [v]
cut Sources v
forall k a. Map k a
M.empty [v
x] Set v
restFree ANormal Reference v
rest'
                        pure (ANormal Reference v -> Maybe (ANormal Reference v)
forall a. a -> Maybe a
Just (Func Reference v -> [v] -> ANormal Reference v
forall v ref. Var v => Func ref v -> [v] -> ANormal ref v
TApp (Reference -> Word64 -> Func Reference v
forall ref v. ref -> Word64 -> Func ref v
FWorker Reference
grpr (Word64 -> Int -> Word64
joinIndex Word64
owner ((Int, [v]) -> Int
forall a b. (a, b) -> a
fst (Int, [v])
j))) (v
x v -> [v] -> [v]
forall a. a -> [a] -> [a]
: (Int, [v]) -> [v]
forall a b. (a, b) -> b
snd (Int, [v])
j)))
                    | Bool
otherwise -> Maybe (ANormal Reference v)
-> StateT
     ([Join v], Sources v, Set v) Identity (Maybe (ANormal Reference v))
forall a. a -> StateT ([Join v], Sources v, Set v) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ANormal Reference v -> Maybe (ANormal Reference v)
forall a. a -> Maybe a
Just ANormal Reference v
rest)
                  Maybe (Direction Word16, v, [Mem], ANormal Reference v)
Nothing -> Maybe (ANormal Reference v)
-> StateT
     ([Join v], Sources v, Set v) Identity (Maybe (ANormal Reference v))
forall a. a -> StateT ([Join v], Sources v, Set v) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (ANormal Reference v)
forall a. Maybe a
Nothing
                -- one join point per arm that is neither a leaf nor built
                -- by one tail only
                Map (Maybe Int) (Int, [v])
numbered <- (Map (Maybe Int) (Int, [v])
 -> (Maybe Int, [v], ANormal Reference v)
 -> StateT
      ([Join v], Sources v, Set v) Identity (Map (Maybe Int) (Int, [v])))
-> Map (Maybe Int) (Int, [v])
-> [(Maybe Int, [v], ANormal Reference v)]
-> StateT
     ([Join v], Sources v, Set v) Identity (Map (Maybe Int) (Int, [v]))
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (v
-> Reference
-> Map (Maybe Int) (Set v)
-> Maybe [v]
-> Map (Maybe Int) (Int, [v])
-> (Maybe Int, [v], ANormal Reference v)
-> StateT
     ([Join v], Sources v, Set v) Identity (Map (Maybe Int) (Int, [v]))
joinFor v
p Reference
r Map (Maybe Int) (Set v)
arms Maybe [v]
armCut) Map (Maybe Int) (Int, [v])
forall k a. Map k a
M.empty [(Maybe Int, [v], ANormal Reference v)
tg | tg :: (Maybe Int, [v], ANormal Reference v)
tg@(Maybe Int
k, [v]
_, ANormal Reference v
ab) <- [(Maybe Int, [v], ANormal Reference v)]
targets, Bool -> Bool
not (Maybe Int -> ANormal Reference v -> Bool
copied Maybe Int
k ANormal Reference v
ab)]
                -- an arm copied into its one tail, renamed by 'rewrite'
                -- (the default arm binds nothing and is copied as it is)
                let copyArm :: CTag
-> [v]
-> Maybe Int
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
copyArm CTag
t [v]
args = \case
                      Just Int
_ -> do
                        let (ANormal Reference v
code, Sources v
armSrcs) = State (Sources v) (ANormal Reference v)
-> Sources v -> (ANormal Reference v, Sources v)
forall s a. State s a -> s -> (a, s)
runState (v
-> Tree v
-> ANormal Reference v
-> State (Sources v) (ANormal Reference v)
forall v.
Var v =>
v
-> Tree v
-> ANormal Reference v
-> State (Sources v) (ANormal Reference v)
rewrite v
p (v -> Shape -> [Tree v] -> Tree v
forall v. v -> Shape -> [Tree v] -> Tree v
Node v
p (Reference -> CTag -> [Shape] -> Shape
Con Reference
r CTag
t (Shape
Plain Shape -> [v] -> [Shape]
forall a b. a -> [b] -> [a]
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ [v]
args)) ((v -> Tree v) -> [v] -> [Tree v]
forall a b. (a -> b) -> [a] -> [b]
map v -> Tree v
forall v. v -> Tree v
Leaf [v]
args)) (v
-> Branched Reference (ANormal Reference v) -> ANormal Reference v
forall v ref.
Var v =>
v -> Branched ref (ANormal ref v) -> ANormal ref v
TMatch v
p (Reference
-> EnumMap CTag ([Mem], ANormal Reference v)
-> Maybe (ANormal Reference v)
-> Branched Reference (ANormal Reference v)
forall ref e.
ref -> EnumMap CTag ([Mem], e) -> Maybe e -> Branched ref e
MatchData Reference
r EnumMap CTag ([Mem], ANormal Reference v)
cs Maybe (ANormal Reference v)
df))) Sources v
forall k a. Map k a
M.empty
                        (([Join v], Sources v, Set v) -> ([Join v], Sources v, Set v))
-> StateT ([Join v], Sources v, Set v) Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (\([Join v]
js, Sources v
ss, Set v
ds) -> ([Join v]
js, Sources v
ss Sources v -> Sources v -> Sources v
forall a. Semigroup a => a -> a -> a
<> Sources v
armSrcs, Set v
ds))
                        pure ANormal Reference v
code
                      Maybe Int
Nothing -> ANormal Reference v
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
forall a. a -> StateT ([Join v], Sources v, Set v) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ANormal Reference v
-> Maybe (ANormal Reference v) -> ANormal Reference v
forall a. a -> Maybe a -> a
fromMaybe ([Word] -> String -> ANormal Reference v
forall a. HasCallStack => [Word] -> String -> a
internalBug [] String
"splitJoins: no default arm") Maybe (ANormal Reference v)
df)
                Map CTag (ANormal Reference v)
onceCodes <-
                  [(CTag, ANormal Reference v)] -> Map CTag (ANormal Reference v)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
                    ([(CTag, ANormal Reference v)] -> Map CTag (ANormal Reference v))
-> StateT
     ([Join v], Sources v, Set v) Identity [(CTag, ANormal Reference v)]
-> StateT
     ([Join v], Sources v, Set v)
     Identity
     (Map CTag (ANormal Reference v))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [StateT
   ([Join v], Sources v, Set v) Identity (CTag, ANormal Reference v)]
-> StateT
     ([Join v], Sources v, Set v) Identity [(CTag, ANormal Reference v)]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence
                      [ (,) CTag
t (ANormal Reference v -> (CTag, ANormal Reference v))
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
-> StateT
     ([Join v], Sources v, Set v) Identity (CTag, ANormal Reference v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> CTag
-> [v]
-> Maybe Int
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
copyArm CTag
t [v]
args Maybe Int
k
                      | (CTag
t, (Maybe Int
k, [v]
_, ANormal Reference v
ab)) <- [CTag]
-> [(Maybe Int, [v], ANormal Reference v)]
-> [(CTag, (Maybe Int, [v], ANormal Reference v))]
forall a b. [a] -> [b] -> [(a, b)]
zip [CTag]
tags [(Maybe Int, [v], ANormal Reference v)]
targets,
                        Maybe Int -> ANormal Reference v -> Bool
copied Maybe Int
k ANormal Reference v
ab,
                        Bool -> Bool
not (ANormal Reference v -> Bool
forall {v}. ANormal Reference v -> Bool
leaf ANormal Reference v
ab),
                        ANormal Reference v
tl <- Int -> [ANormal Reference v] -> [ANormal Reference v]
forall a. Int -> [a] -> [a]
take Int
1 [ANormal Reference v
tl | ANormal Reference v
tl <- ANormal Reference v -> [ANormal Reference v]
forall v. Var v => ANormal Reference v -> [ANormal Reference v]
tailsOf ANormal Reference v
bn, Just (CTag
t', [v]
_) <- [ConLets v -> Reference -> ANormal Reference v -> Maybe (CTag, [v])
forall v.
Var v =>
ConLets v -> Reference -> ANormal Reference v -> Maybe (CTag, [v])
tailCon ConLets v
cons Reference
r ANormal Reference v
tl], CTag
t' CTag -> CTag -> Bool
forall a. Eq a => a -> a -> Bool
== CTag
t],
                        Just (CTag
_, [v]
args) <- [ConLets v -> Reference -> ANormal Reference v -> Maybe (CTag, [v])
forall v.
Var v =>
ConLets v -> Reference -> ANormal Reference v -> Maybe (CTag, [v])
tailCon ConLets v
cons Reference
r ANormal Reference v
tl]
                      ]
                let at :: ANormal Reference v -> ANormal Reference v
at ANormal Reference v
tl
                      | Just (CTag
t, [v]
args) <- ConLets v -> Reference -> ANormal Reference v -> Maybe (CTag, [v])
forall v.
Var v =>
ConLets v -> Reference -> ANormal Reference v -> Maybe (CTag, [v])
tailCon ConLets v
cons Reference
r ANormal Reference v
tl,
                        Just (Maybe Int
k, [v]
fs, ANormal Reference v
ab) <- CTag -> Maybe (Maybe Int, [v], ANormal Reference v)
armOf CTag
t =
                          let arm :: ANormal Reference v
arm
                                | ANormal Reference v -> Bool
forall {v}. ANormal Reference v -> Bool
leaf ANormal Reference v
ab = Map v v -> ANormal Reference v -> ANormal Reference v
forall v (f :: * -> * -> *).
(Var v, Bifunctor f, Bifoldable f) =>
Map v v -> Term f v -> Term f v
ABTN.renames ([(v, v)] -> Map v v
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([v] -> [v] -> [(v, v)]
forall a b. [a] -> [b] -> [(a, b)]
zip [v]
fs [v]
args)) ANormal Reference v
ab
                                | Maybe Int -> ANormal Reference v -> Bool
copied Maybe Int
k ANormal Reference v
ab = ANormal Reference v
-> Maybe (ANormal Reference v) -> ANormal Reference v
forall a. a -> Maybe a -> a
fromMaybe ([Word] -> String -> ANormal Reference v
forall a. HasCallStack => [Word] -> String -> a
internalBug [] String
"splitJoins: an arm without its copy") (CTag
-> Map CTag (ANormal Reference v) -> Maybe (ANormal Reference v)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup CTag
t Map CTag (ANormal Reference v)
onceCodes)
                                | Bool
otherwise = case Maybe Int -> Map (Maybe Int) (Int, [v]) -> Maybe (Int, [v])
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Maybe Int
k Map (Maybe Int) (Int, [v])
numbered of
                                    Just (Int
j, [v]
extra) -> Func Reference v -> [v] -> ANormal Reference v
forall v ref. Var v => Func ref v -> [v] -> ANormal ref v
TApp (Reference -> Word64 -> Func Reference v
forall ref v. ref -> Word64 -> Func ref v
FWorker Reference
grpr (Word64 -> Int -> Word64
joinIndex Word64
owner Int
j)) ([v]
args [v] -> [v] -> [v]
forall a. [a] -> [a] -> [a]
++ [v]
extra)
                                    Maybe (Int, [v])
Nothing -> [Word] -> String -> ANormal Reference v
forall a. HasCallStack => [Word] -> String -> a
internalBug [] String
"splitJoins: an arm without a join point"
                           in case (Maybe (Direction Word16, v, [Mem], ANormal Reference v)
kont, Maybe (ANormal Reference v)
restCode) of
                                (Just (Direction Word16
d', v
x, [Mem]
ms', ANormal Reference v
_), Just ANormal Reference v
rest) -> Direction Word16
-> [v]
-> [Mem]
-> ANormal Reference v
-> ANormal Reference v
-> ANormal Reference v
forall v ref.
Var v =>
Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
TLets Direction Word16
d' [v
x] [Mem]
ms' ANormal Reference v
arm ANormal Reference v
rest
                                (Maybe (Direction Word16, v, [Mem], ANormal Reference v),
 Maybe (ANormal Reference v))
_ -> ANormal Reference v
arm
                      | Bool
otherwise = ANormal Reference v
tl
                -- the binding takes the let's place, so its lets are inside
                -- whatever the let was inside, and the arms it now ends in
                -- are walked for their own sites
                Maybe [v]
-> ANormal Reference v
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
go Maybe [v]
cut ((ANormal Reference v -> ANormal Reference v)
-> ANormal Reference v -> ANormal Reference v
forall v.
Var v =>
(ANormal Reference v -> ANormal Reference v)
-> ANormal Reference v -> ANormal Reference v
sinkTails ANormal Reference v -> ANormal Reference v
at ANormal Reference v
bn)
        where
          fallback :: State ([Join v], Sources v, Set v) (ANormal Reference v)
fallback = Direction Word16
-> [v]
-> [Mem]
-> ANormal Reference v
-> ANormal Reference v
-> ANormal Reference v
forall v ref.
Var v =>
Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
TLets Direction Word16
d [v
p] [Mem]
ms (ANormal Reference v -> ANormal Reference v -> ANormal Reference v)
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
-> StateT
     ([Join v], Sources v, Set v)
     Identity
     (ANormal Reference v -> ANormal Reference v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe [v]
-> ANormal Reference v
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
go (Direction Word16 -> [v] -> Maybe [v] -> Maybe [v]
forall {a} {a}. Direction a -> a -> Maybe a -> Maybe a
cutFor Direction Word16
d [v
p] Maybe [v]
cut) ANormal Reference v
bn StateT
  ([Join v], Sources v, Set v)
  Identity
  (ANormal Reference v -> ANormal Reference v)
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
forall a b.
StateT ([Join v], Sources v, Set v) Identity (a -> b)
-> StateT ([Join v], Sources v, Set v) Identity a
-> StateT ([Join v], Sources v, Set v) Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe [v]
-> ANormal Reference v
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
go Maybe [v]
cut ANormal Reference v
bo
      TLets Direction Word16
d [v]
us [Mem]
ms ANormal Reference v
bn ANormal Reference v
bo -> Direction Word16
-> [v]
-> [Mem]
-> ANormal Reference v
-> ANormal Reference v
-> ANormal Reference v
forall v ref.
Var v =>
Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
TLets Direction Word16
d [v]
us [Mem]
ms (ANormal Reference v -> ANormal Reference v -> ANormal Reference v)
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
-> StateT
     ([Join v], Sources v, Set v)
     Identity
     (ANormal Reference v -> ANormal Reference v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe [v]
-> ANormal Reference v
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
go (Direction Word16 -> [v] -> Maybe [v] -> Maybe [v]
forall {a} {a}. Direction a -> a -> Maybe a -> Maybe a
cutFor Direction Word16
d [v]
us Maybe [v]
cut) ANormal Reference v
bn StateT
  ([Join v], Sources v, Set v)
  Identity
  (ANormal Reference v -> ANormal Reference v)
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
forall a b.
StateT ([Join v], Sources v, Set v) Identity (a -> b)
-> StateT ([Join v], Sources v, Set v) Identity a
-> StateT ([Join v], Sources v, Set v) Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe [v]
-> ANormal Reference v
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
go Maybe [v]
cut ANormal Reference v
bo
      TShift {} -> ANormal Reference v
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
forall a. a -> StateT ([Join v], Sources v, Set v) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ANormal Reference v
tm
      ABTN.Term Set v
_ (ABTN.Abs v
u ANormal Reference v
e) -> v -> ANormal Reference v -> ANormal Reference v
forall v (f :: * -> * -> *). Var v => v -> Term f v -> Term f v
ABTN.TAbs v
u (ANormal Reference v -> ANormal Reference v)
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe [v]
-> ANormal Reference v
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
go Maybe [v]
cut ANormal Reference v
e
      ABTN.Term Set v
_ (ABTN.Tm ANormalF Reference v (ANormal Reference v)
f) -> ANormalF Reference v (ANormal Reference v) -> ANormal Reference v
forall v (f :: * -> * -> *).
(Var v, Bifoldable f) =>
f v (Term f v) -> Term f v
ABTN.TTm (ANormalF Reference v (ANormal Reference v) -> ANormal Reference v)
-> StateT
     ([Join v], Sources v, Set v)
     Identity
     (ANormalF Reference v (ANormal Reference v))
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ANormal Reference v
 -> State ([Join v], Sources v, Set v) (ANormal Reference v))
-> ANormalF Reference v (ANormal Reference v)
-> StateT
     ([Join v], Sources v, Set v)
     Identity
     (ANormalF Reference v (ANormal Reference 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) -> ANormalF Reference v a -> f (ANormalF Reference v b)
traverse (Maybe [v]
-> ANormal Reference v
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
go Maybe [v]
cut) ANormalF Reference v (ANormal Reference v)
f
    -- the binding of an indirect let has the let's frame around it
    cutFor :: Direction a -> a -> Maybe a -> Maybe a
cutFor Direction a
d a
us Maybe a
cut = case Direction a
d of
      Indirect a
_ -> a -> Maybe a
forall a. a -> Maybe a
Just a
us
      Direction a
Direct -> Maybe a
cut
    leaf :: ANormal Reference v -> Bool
leaf ANormal Reference v
ab = [ANormal Reference v] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (ANormal Reference v -> [ANormal Reference v]
forall v. ANormal Reference v -> [ANormal Reference v]
children ANormal Reference v
ab)
    -- the join point for an arm, unless already made: numbered when its
    -- body, with the join points of the sites inside it, is done
    joinFor :: v
-> Reference
-> Map (Maybe Int) (Set v)
-> Maybe [v]
-> Map (Maybe Int) (Int, [v])
-> (Maybe Int, [v], ANormal Reference v)
-> StateT
     ([Join v], Sources v, Set v) Identity (Map (Maybe Int) (Int, [v]))
joinFor v
p Reference
r Map (Maybe Int) (Set v)
arms Maybe [v]
cut Map (Maybe Int) (Int, [v])
acc (Maybe Int
k, [v]
fs, ANormal Reference v
ab)
      | Maybe Int -> Map (Maybe Int) (Int, [v]) -> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.member Maybe Int
k Map (Maybe Int) (Int, [v])
acc = Map (Maybe Int) (Int, [v])
-> StateT
     ([Join v], Sources v, Set v) Identity (Map (Maybe Int) (Int, [v]))
forall a. a -> StateT ([Join v], Sources v, Set v) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Map (Maybe Int) (Int, [v])
acc
      | Bool
otherwise = do
          ANormal Reference v
ab' <- Maybe [v]
-> ANormal Reference v
-> State ([Join v], Sources v, Set v) (ANormal Reference v)
go Maybe [v]
cut ANormal Reference v
ab
          let psrcs :: Sources v
psrcs = case Maybe Int
k of
                Just Int
t -> v -> Either v (Built v) -> Sources v
forall k a. k -> a -> Map k a
M.singleton v
p (Built v -> Either v (Built v)
forall a b. b -> Either a b
Right (Reference -> CTag -> [v] -> Built v
forall v. Reference -> CTag -> [v] -> Built v
BuiltCon Reference
r (Int -> CTag
forall k. EnumKey k => Int -> k
EC.intToKey Int
t) [v]
fs))
                Maybe Int
Nothing -> Sources v
forall k a. Map k a
M.empty
          (Int, [v])
j <- Maybe [v]
-> Sources v
-> [v]
-> Set v
-> ANormal Reference v
-> StateT ([Join v], Sources v, Set v) Identity (Int, [v])
newJoin Maybe [v]
cut Sources v
psrcs [v]
fs (Set v -> Maybe Int -> Map (Maybe Int) (Set v) -> Set v
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Set v
forall a. Set a
Set.empty Maybe Int
k Map (Maybe Int) (Set v)
arms) ANormal Reference v
ab'
          pure (Maybe Int
-> (Int, [v])
-> Map (Maybe Int) (Int, [v])
-> Map (Maybe Int) (Int, [v])
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Maybe Int
k (Int, [v])
j Map (Maybe Int) (Int, [v])
acc)
    -- a join point whose parameters are the given variables and then what
    -- stands for the unoptimized variables and what the body uses now
    newJoin :: Maybe [v]
-> Sources v
-> [v]
-> Set v
-> ANormal Reference v
-> StateT ([Join v], Sources v, Set v) Identity (Int, [v])
newJoin Maybe [v]
cut Sources v
psrcs [v]
fs Set v
base ANormal Reference v
body = do
      ([Join v]
js, Sources v
ss, Set v
ds) <- StateT
  ([Join v], Sources v, Set v) Identity ([Join v], Sources v, Set v)
forall s (m :: * -> *). MonadState s m => m s
get
      let extra :: [v]
extra =
            Set v -> [v]
forall a. Set a -> [a]
Set.toAscList (Set v -> [v]) -> Set v -> [v]
forall a b. (a -> b) -> a -> b
$
              ([Set v] -> Set v
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions ((v -> Set v) -> [v] -> [Set v]
forall a b. (a -> b) -> [a] -> [b]
map (Sources v -> v -> Set v
expand (Sources v
srcs Sources v -> Sources v -> Sources v
forall a. Semigroup a => a -> a -> a
<> Sources v
ss)) (Set v -> [v]
forall a. Set a -> [a]
Set.toList Set v
base)) Set v -> Set v -> Set v
forall a. Semigroup a => a -> a -> a
<> ANormal Reference v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
ABTN.freeVars ANormal Reference v
body)
                Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
`Set.difference` [v] -> Set v
forall a. Ord a => [a] -> Set a
Set.fromList [v]
fs
                Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
`Set.difference` Set v
locals
          params :: [v]
params = [v]
fs [v] -> [v] -> [v]
forall a. [a] -> [a] -> [a]
++ [v]
extra
      ([Join v], Sources v, Set v)
-> StateT ([Join v], Sources v, Set v) Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (SuperNormal Reference v -> Maybe [v] -> Sources v -> Join v
forall v.
SuperNormal Reference v -> Maybe [v] -> Sources v -> Join v
Join ([Mem] -> ANormal Reference v -> SuperNormal Reference v
forall ref v. [Mem] -> ANormal ref v -> SuperNormal ref v
Lambda ((v -> Mem) -> [v] -> [Mem]
forall a b. (a -> b) -> [a] -> [b]
map (Mem -> v -> Mem
forall a b. a -> b -> a
const Mem
BX) [v]
params) ([v] -> ANormal Reference v -> ANormal Reference v
forall v (f :: * -> * -> *). Var v => [v] -> Term f v -> Term f v
TAbss [v]
params ANormal Reference v
body)) Maybe [v]
cut Sources v
psrcs Join v -> [Join v] -> [Join v]
forall a. a -> [a] -> [a]
: [Join v]
js, Sources v
ss, Set v
ds)
      pure ([Join v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Join v]
js Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, [v]
extra)
    -- what stands for a variable of the unoptimized code: itself, what it
    -- was renamed to, or the fields of the product it was; through the
    -- sources given and those recorded so far, since an arm copied into
    -- its one tail at an enclosing site renames its variables, and a
    -- site inside it is reached after
    expand :: Sources v -> v -> Set v
expand Sources v
ss v
w = [v] -> Set v
forall a. Ord a => [a] -> Set a
Set.fromList (Sources v -> Int -> v -> [v]
ex Sources v
ss (Int
0 :: Int) v
w)
    ex :: Sources v -> Int -> v -> [v]
ex Sources v
ss Int
depth v
w
      | Int
depth Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
64 = [Word] -> String -> [v]
forall a. HasCallStack => [Word] -> String -> a
internalBug [] (String -> [v]) -> String -> [v]
forall a b. (a -> b) -> a -> b
$ String
"splitJoins: the source of " String -> ShowS
forall a. [a] -> [a] -> [a]
++ v -> String
forall a. Show a => a -> String
show v
w String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" refers back to itself"
      | v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member v
w Set v
dead0 = []
      | Just (Left v
u) <- v -> Sources v -> Maybe (Either v (Built v))
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup v
w Sources v
ss = Sources v -> Int -> v -> [v]
ex Sources v
ss (Int
depth Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) v
u
      | Just (Right Built v
b) <- v -> Sources v -> Maybe (Either v (Built v))
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup v
w Sources v
ss = (v -> [v]) -> [v] -> [v]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Sources v -> Int -> v -> [v]
ex Sources v
ss (Int
depth Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) (Built v -> [v]
forall v. Built v -> [v]
builtFields Built v
b)
      | Just Built v
b <- v -> ConLets v -> Maybe (Built v)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup v
w ConLets v
bcons = [v
w | v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member v
w Set v
inScope] [v] -> [v] -> [v]
forall a. [a] -> [a] -> [a]
++ (v -> [v]) -> [v] -> [v]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Sources v -> Int -> v -> [v]
ex Sources v
ss (Int
depth Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) (Built v -> [v]
forall v. Built v -> [v]
builtFields Built v
b)
      | v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member v
w Set v
inScope = [v
w]
      | Bool
otherwise = [Word] -> String -> [v]
forall a. HasCallStack => [Word] -> String -> a
internalBug [] (String -> [v]) -> String -> [v]
forall a b. (a -> b) -> a -> b
$ String
"splitJoins: nothing stands for " String -> ShowS
forall a. [a] -> [a] -> [a]
++ v -> String
forall a. Show a => a -> String
show v
w String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" of the unoptimized code"
    -- the variables of sites, dead wherever a frame could hold them
    dead0 :: Set v
dead0 = Sites v -> Set v
forall k a. Map k a -> Set k
M.keysSet Sites v
sites

-- | The body with a name of its own for every binder and an index of its
-- own for every indirect let. Normalization reuses a name under
-- shadowing (the blank variables), and 'ABTN.freshen' gives a name back
-- once the binder's scope is over, while everything here maps variables
-- by name. Only a binder already bound elsewhere in the body is renamed,
-- so that a body whose binders are already unique comes back as it is:
-- the unoptimized body is renamed the same way, and the names are what
-- pairs its lets with the optimized body's ('Prepared'). The let indices
-- number the sections of the emitted combinator, and the ones
-- normalization assigned repeat once the ANF optimizer has inlined a
-- body whose lets carry their own ('uniqueLets'), which is also why the
-- renaming is applied with every rewrite off.
uniqueBody :: (Var v) => ANormal Reference v -> ANormal Reference v
uniqueBody :: forall v. Var v => ANormal Reference v -> ANormal Reference v
uniqueBody ANormal Reference v
bd = ANormal Reference v -> ANormal Reference v
forall v. Var v => ANormal Reference v -> ANormal Reference v
uniqueLets (State (Set v, Set v) (ANormal Reference v)
-> (Set v, Set v) -> ANormal Reference v
forall s a. State s a -> s -> a
evalState (Map v v
-> ANormal Reference v
-> State (Set v, Set v) (ANormal Reference v)
forall v.
Var v =>
Map v v
-> ANormal Reference v
-> State (Set v, Set v) (ANormal Reference v)
uniqueBinders Map v v
forall k a. Map k a
M.empty ANormal Reference v
bd) ([Set v] -> Set v
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions (ANormal Reference v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
ABTN.freeVars ANormal Reference v
bd Set v -> [Set v] -> [Set v]
forall a. a -> [a] -> [a]
: (ANormal Reference v -> Set v) -> [ANormal Reference v] -> [Set v]
forall a b. (a -> b) -> [a] -> [b]
map ANormal Reference v -> Set v
forall v. Ord v => ANormal Reference v -> Set v
boundVars (ANormal Reference v -> [ANormal Reference v]
forall v. ANormal Reference v -> [ANormal Reference v]
universe ANormal Reference v
bd)), Set v
forall a. Set a
Set.empty))

-- | The indirect lets numbered uniquely. A continuation names the
-- sections it resumes by these numbers, in this process and in any other,
-- so a let keeps the number the normalizer gave it, which is what every
-- version writes for a body the optimizer left alone; only a let whose
-- number an earlier let of the body already has (the optimizer inlines
-- bodies that carry their own numbers) gets a new one, the next above
-- every number in the body, in order of appearance, binding before body,
-- so that the same body always gets the same numbers.
uniqueLets :: (Var v) => ANormal Reference v -> ANormal Reference v
uniqueLets :: forall v. Var v => ANormal Reference v -> ANormal Reference v
uniqueLets ANormal Reference v
bd = State (Set Word16, Word16) (ANormal Reference v)
-> (Set Word16, Word16) -> ANormal Reference v
forall s a. State s a -> s -> a
evalState (ANormal Reference v
-> State (Set Word16, Word16) (ANormal Reference v)
forall {f :: * -> *} {v} {ref}.
(Var v, MonadState (Set Word16, Word16) f) =>
Term (ANormalF ref) v -> f (Term (ANormalF ref) v)
go ANormal Reference v
bd) (Set Word16
forall a. Set a
Set.empty, Word16
top)
  where
    top :: Word16
top = [Word16] -> Word16
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (Word16
0 Word16 -> [Word16] -> [Word16]
forall a. a -> [a] -> [a]
: [Word16
w | TLets (Indirect Word16
w) [v]
_ [Mem]
_ ANormal Reference v
_ ANormal Reference v
_ <- ANormal Reference v -> [ANormal Reference v]
forall v. ANormal Reference v -> [ANormal Reference v]
universe ANormal Reference v
bd]) :: Word16
    go :: Term (ANormalF ref) v -> f (Term (ANormalF ref) v)
go (ABTN.Term Set v
_ (ABTN.Abs v
u Term (ANormalF ref) v
b)) = v -> Term (ANormalF ref) v -> Term (ANormalF ref) v
forall v (f :: * -> * -> *). Var v => v -> Term f v -> Term f v
ABTN.TAbs v
u (Term (ANormalF ref) v -> Term (ANormalF ref) v)
-> f (Term (ANormalF ref) v) -> f (Term (ANormalF ref) v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Term (ANormalF ref) v -> f (Term (ANormalF ref) v)
go Term (ANormalF ref) v
b
    go (ABTN.Term Set v
_ (ABTN.Tm ANormalF ref v (Term (ANormalF ref) v)
f)) = case ANormalF ref v (Term (ANormalF ref) v)
f of
      ALet (Indirect Word16
w) [Mem]
ccs Term (ANormalF ref) v
bn Term (ANormalF ref) v
bo -> do
        Word16
w' <- ((Set Word16, Word16) -> (Word16, (Set Word16, Word16)))
-> f Word16
forall a.
((Set Word16, Word16) -> (a, (Set Word16, Word16))) -> f a
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state \(Set Word16
seen, Word16
next) ->
          if Word16 -> Set Word16 -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Word16
w Set Word16
seen
            then (Word16
next Word16 -> Word16 -> Word16
forall a. Num a => a -> a -> a
+ Word16
1, (Word16 -> Set Word16 -> Set Word16
forall a. Ord a => a -> Set a -> Set a
Set.insert (Word16
next Word16 -> Word16 -> Word16
forall a. Num a => a -> a -> a
+ Word16
1) Set Word16
seen, Word16
next Word16 -> Word16 -> Word16
forall a. Num a => a -> a -> a
+ Word16
1))
            else (Word16
w, (Word16 -> Set Word16 -> Set Word16
forall a. Ord a => a -> Set a -> Set a
Set.insert Word16
w Set Word16
seen, Word16
next))
        ANormalF ref v (Term (ANormalF ref) v) -> Term (ANormalF ref) v
forall v (f :: * -> * -> *).
(Var v, Bifoldable f) =>
f v (Term f v) -> Term f v
ABTN.TTm (ANormalF ref v (Term (ANormalF ref) v) -> Term (ANormalF ref) v)
-> f (ANormalF ref v (Term (ANormalF ref) v))
-> f (Term (ANormalF ref) v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Direction Word16
-> [Mem]
-> Term (ANormalF ref) v
-> Term (ANormalF ref) v
-> ANormalF ref v (Term (ANormalF ref) v)
forall ref v e.
Direction Word16 -> [Mem] -> e -> e -> ANormalF ref v e
ALet (Word16 -> Direction Word16
forall a. a -> Direction a
Indirect Word16
w') [Mem]
ccs (Term (ANormalF ref) v
 -> Term (ANormalF ref) v -> ANormalF ref v (Term (ANormalF ref) v))
-> f (Term (ANormalF ref) v)
-> f (Term (ANormalF ref) v
      -> ANormalF ref v (Term (ANormalF ref) v))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Term (ANormalF ref) v -> f (Term (ANormalF ref) v)
go Term (ANormalF ref) v
bn f (Term (ANormalF ref) v -> ANormalF ref v (Term (ANormalF ref) v))
-> f (Term (ANormalF ref) v)
-> f (ANormalF ref v (Term (ANormalF ref) v))
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Term (ANormalF ref) v -> f (Term (ANormalF ref) v)
go Term (ANormalF ref) v
bo)
      ANormalF ref v (Term (ANormalF ref) v)
_ -> ANormalF ref v (Term (ANormalF ref) v) -> Term (ANormalF ref) v
forall v (f :: * -> * -> *).
(Var v, Bifoldable f) =>
f v (Term f v) -> Term f v
ABTN.TTm (ANormalF ref v (Term (ANormalF ref) v) -> Term (ANormalF ref) v)
-> f (ANormalF ref v (Term (ANormalF ref) v))
-> f (Term (ANormalF ref) v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Term (ANormalF ref) v -> f (Term (ANormalF ref) v))
-> ANormalF ref v (Term (ANormalF ref) v)
-> f (ANormalF ref v (Term (ANormalF 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) -> ANormalF ref v a -> f (ANormalF ref v b)
traverse Term (ANormalF ref) v -> f (Term (ANormalF ref) v)
go ANormalF ref v (Term (ANormalF ref) v)
f

-- | The state is the names in use, and the binders seen so far.
uniqueBinders :: (Var v) => M.Map v v -> ANormal Reference v -> State (Set.Set v, Set.Set v) (ANormal Reference v)
uniqueBinders :: forall v.
Var v =>
Map v v
-> ANormal Reference v
-> State (Set v, Set v) (ANormal Reference v)
uniqueBinders Map v v
rn (ABTN.Term Set v
_ (ABTN.Abs v
u Term (ANormalF Reference) v
body)) = do
  v
u' <- v -> StateT (Set v, Set v) Identity v
forall v. Var v => v -> State (Set v, Set v) v
fresh v
u
  v -> Term (ANormalF Reference) v -> Term (ANormalF Reference) v
forall v (f :: * -> * -> *). Var v => v -> Term f v -> Term f v
ABTN.TAbs v
u' (Term (ANormalF Reference) v -> Term (ANormalF Reference) v)
-> State (Set v, Set v) (Term (ANormalF Reference) v)
-> State (Set v, Set v) (Term (ANormalF Reference) v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map v v
-> Term (ANormalF Reference) v
-> State (Set v, Set v) (Term (ANormalF Reference) v)
forall v.
Var v =>
Map v v
-> ANormal Reference v
-> State (Set v, Set v) (ANormal Reference v)
uniqueBinders (v -> v -> Map v v -> Map v v
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert v
u v
u' Map v v
rn) Term (ANormalF Reference) v
body
uniqueBinders Map v v
rn (ABTN.Term Set v
_ (ABTN.Tm ANormalF Reference v (Term (ANormalF Reference) v)
f)) =
  ANormalF Reference v (Term (ANormalF Reference) v)
-> Term (ANormalF Reference) v
forall v (f :: * -> * -> *).
(Var v, Bifoldable f) =>
f v (Term f v) -> Term f v
ABTN.TTm (ANormalF Reference v (Term (ANormalF Reference) v)
 -> Term (ANormalF Reference) v)
-> StateT
     (Set v, Set v)
     Identity
     (ANormalF Reference v (Term (ANormalF Reference) v))
-> State (Set v, Set v) (Term (ANormalF Reference) v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Term (ANormalF Reference) v
 -> State (Set v, Set v) (Term (ANormalF Reference) v))
-> ANormalF Reference v (Term (ANormalF Reference) v)
-> StateT
     (Set v, Set v)
     Identity
     (ANormalF Reference v (Term (ANormalF Reference) 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) -> ANormalF Reference v a -> f (ANormalF Reference v b)
traverse (Map v v
-> Term (ANormalF Reference) v
-> State (Set v, Set v) (Term (ANormalF Reference) v)
forall v.
Var v =>
Map v v
-> ANormal Reference v
-> State (Set v, Set v) (ANormal Reference v)
uniqueBinders Map v v
rn) ((v -> v)
-> ANormalF Reference v (Term (ANormalF Reference) v)
-> ANormalF Reference v (Term (ANormalF Reference) v)
forall a b c.
(a -> b) -> ANormalF Reference a c -> ANormalF Reference b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (\v
v -> v -> v -> Map v v -> v
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault v
v v
v Map v v
rn) ANormalF Reference v (Term (ANormalF Reference) v)
f)

fresh :: (Var v) => v -> State (Set.Set v, Set.Set v) v
fresh :: forall v. Var v => v -> State (Set v, Set v) v
fresh v
u = do
  (Set v
used, Set v
seen) <- StateT (Set v, Set v) Identity (Set v, Set v)
forall s (m :: * -> *). MonadState s m => m s
get
  if v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member v
u Set v
seen
    then do
      let u' :: v
u' = Set v -> v -> v
forall v. Var v => Set v -> v -> v
Var.freshIn Set v
used v
u
      (Set v, Set v) -> StateT (Set v, Set v) Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (v -> Set v -> Set v
forall a. Ord a => a -> Set a -> Set a
Set.insert v
u' Set v
used, v -> Set v -> Set v
forall a. Ord a => a -> Set a -> Set a
Set.insert v
u' Set v
seen)
      pure v
u'
    else do
      (Set v, Set v) -> StateT (Set v, Set v) Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Set v
used, v -> Set v -> Set v
forall a. Ord a => a -> Set a -> Set a
Set.insert v
u Set v
seen)
      pure v
u

-- | A call whose argument is the value of a match, @let s = match c with
-- arms in f .. s ..@ with @s@ used nowhere else, moved into the arms, so
-- that the branch that builds a constructor calls with it: the let's body
-- is the call, so each arm's tail becomes the call with its tail variable
-- for @s@. Only when every tail is a variable (a built constructor's let
-- ends that way): a tail that is an expression would need a let of its
-- own, with an index, to name it.
sinkCalls :: (Var v) => ANormal Reference v -> ANormal Reference v
sinkCalls :: forall v. Var v => ANormal Reference v -> ANormal Reference v
sinkCalls ANormal Reference v
tm = case ANormal Reference v
tm of
  TLets Direction Word16
_ [v
s] [Mem]
_ bn :: ANormal Reference v
bn@(TMatch v
_ Branched Reference (ANormal Reference v)
_) call :: ANormal Reference v
call@(TApp Func Reference v
f [v]
args)
    | v
s v -> [v] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [v]
args,
      v -> ANormal Reference v -> Int
forall v. Var v => v -> ANormal Reference v -> Int
occurrences v
s ANormal Reference v
call Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1,
      Just ANormal Reference v
sunk <- ANormal Reference v
-> (v -> ANormal Reference v) -> Maybe (ANormal Reference v)
forall {t} {ref}.
Var t =>
ANormal ref t -> (t -> ANormal ref t) -> Maybe (ANormal ref t)
sink ANormal Reference v
bn (\v
u -> Func Reference v -> [v] -> ANormal Reference v
forall v ref. Var v => Func ref v -> [v] -> ANormal ref v
TApp Func Reference v
f ((v -> v) -> [v] -> [v]
forall a b. (a -> b) -> [a] -> [b]
map (\v
a -> if v
a v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
s then v
u else v
a) [v]
args)) ->
        ANormal Reference v -> ANormal Reference v
forall v. Var v => ANormal Reference v -> ANormal Reference v
sinkCalls ANormal Reference v
sunk
  ABTN.Term Set v
_ (ABTN.Abs v
u ANormal Reference v
b) -> v -> ANormal Reference v -> ANormal Reference v
forall v (f :: * -> * -> *). Var v => v -> Term f v -> Term f v
ABTN.TAbs v
u (ANormal Reference v -> ANormal Reference v
forall v. Var v => ANormal Reference v -> ANormal Reference v
sinkCalls ANormal Reference v
b)
  ABTN.Term Set v
_ (ABTN.Tm ANormalF Reference v (ANormal Reference v)
f) -> ANormalF Reference v (ANormal Reference v) -> ANormal Reference v
forall v (f :: * -> * -> *).
(Var v, Bifoldable f) =>
f v (Term f v) -> Term f v
ABTN.TTm ((ANormal Reference v -> ANormal Reference v)
-> ANormalF Reference v (ANormal Reference v)
-> ANormalF Reference v (ANormal Reference v)
forall a b.
(a -> b) -> ANormalF Reference v a -> ANormalF Reference v b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ANormal Reference v -> ANormal Reference v
forall v. Var v => ANormal Reference v -> ANormal Reference v
sinkCalls ANormalF Reference v (ANormal Reference v)
f)
  where
    sink :: ANormal ref t -> (t -> ANormal ref t) -> Maybe (ANormal ref t)
sink ANormal ref t
e t -> ANormal ref t
k = case ANormal ref t
e of
      TMatch t
v Branched ref (ANormal ref t)
bs -> t -> Branched ref (ANormal ref t) -> ANormal ref t
forall v ref.
Var v =>
v -> Branched ref (ANormal ref v) -> ANormal ref v
TMatch t
v (Branched ref (ANormal ref t) -> ANormal ref t)
-> Maybe (Branched ref (ANormal ref t)) -> Maybe (ANormal ref t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ANormal ref t -> Maybe (ANormal ref t))
-> Branched ref (ANormal ref t)
-> Maybe (Branched ref (ANormal ref t))
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 t -> (t -> ANormal ref t) -> Maybe (ANormal ref t)
`sink` t -> ANormal ref t
k) Branched ref (ANormal ref t)
bs
      TLets Direction Word16
d [t]
us [Mem]
ms ANormal ref t
bn ANormal ref t
bo -> Direction Word16
-> [t] -> [Mem] -> ANormal ref t -> ANormal ref t -> ANormal ref t
forall v ref.
Var v =>
Direction Word16
-> [v] -> [Mem] -> ANormal ref v -> ANormal ref v -> ANormal ref v
TLets Direction Word16
d [t]
us [Mem]
ms ANormal ref t
bn (ANormal ref t -> ANormal ref t)
-> Maybe (ANormal ref t) -> Maybe (ANormal ref t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ANormal ref t -> (t -> ANormal ref t) -> Maybe (ANormal ref t)
sink ANormal ref t
bo t -> ANormal ref t
k
      TName t
v Either ref t
fn [t]
as ANormal ref t
bo -> t -> Either ref t -> [t] -> ANormal ref t -> ANormal ref t
forall v ref.
Var v =>
v -> Either ref v -> [v] -> ANormal ref v -> ANormal ref v
TName t
v Either ref t
fn [t]
as (ANormal ref t -> ANormal ref t)
-> Maybe (ANormal ref t) -> Maybe (ANormal ref t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ANormal ref t -> (t -> ANormal ref t) -> Maybe (ANormal ref t)
sink ANormal ref t
bo t -> ANormal ref t
k
      TVar t
u -> ANormal ref t -> Maybe (ANormal ref t)
forall a. a -> Maybe a
Just (t -> ANormal ref t
k t
u)
      ABTN.Term Set t
_ (ABTN.Abs t
u ANormal ref t
b) -> t -> ANormal ref t -> ANormal ref t
forall v (f :: * -> * -> *). Var v => v -> Term f v -> Term f v
ABTN.TAbs t
u (ANormal ref t -> ANormal ref t)
-> Maybe (ANormal ref t) -> Maybe (ANormal ref t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ANormal ref t -> (t -> ANormal ref t) -> Maybe (ANormal ref t)
sink ANormal ref t
b t -> ANormal ref t
k
      ANormal ref t
_ -> Maybe (ANormal ref t)
forall a. Maybe a
Nothing

-- | The specialisation a call site goes to, with its number: the first
-- pattern the site's arguments offer everything of.
siteSpec :: (Var v) => ConLets v -> [Pattern] -> [v] -> Maybe (Int, Pattern)
siteSpec :: forall v.
Var v =>
ConLets v -> [[Shape]] -> [v] -> Maybe (Int, [Shape])
siteSpec ConLets v
cons [[Shape]]
pats [v]
args = [(Int, [Shape])] -> Maybe (Int, [Shape])
forall a. [a] -> Maybe a
listToMaybe [(Int
k, [Shape]
pat) | (Int
k, [Shape]
pat) <- [Int] -> [[Shape]] -> [(Int, [Shape])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
1 ..] [[Shape]]
pats, [Shape] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Shape]
pat Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
args, [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Shape -> Shape -> Bool) -> [Shape] -> [Shape] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Shape -> Shape -> Bool
refines [Shape]
site [Shape]
pat)]
  where
    site :: [Shape]
site = (v -> Shape) -> [v] -> [Shape]
forall a b. (a -> b) -> [a] -> [b]
map (ConLets v -> Int -> v -> Shape
forall v. Var v => ConLets v -> Int -> v -> Shape
siteShape ConLets v
cons Int
0) [v]
args

-- | The variables a call passes to a worker, in the worker's parameter
-- order: a plain argument as itself, an exploded one as the field
-- variables of the product let it names, down to the leaves. Nothing when
-- the site does not offer what the pattern asks ('siteSpec' chose it, so
-- it does).
explodedArgs :: (Var v) => ConLets v -> Pattern -> [v] -> Maybe [v]
explodedArgs :: forall v. Var v => ConLets v -> [Shape] -> [v] -> Maybe [v]
explodedArgs ConLets v
cons [Shape]
pat [v]
args
  | [Shape] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Shape]
pat Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
args = [[v]] -> [v]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[v]] -> [v]) -> Maybe [[v]] -> Maybe [v]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Shape, v) -> Maybe [v]) -> [(Shape, v)] -> Maybe [[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 ((Shape -> v -> Maybe [v]) -> (Shape, v) -> Maybe [v]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Shape -> v -> Maybe [v]
go) ([Shape] -> [v] -> [(Shape, v)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Shape]
pat [v]
args)
  | Bool
otherwise = Maybe [v]
forall a. Maybe a
Nothing
  where
    go :: Shape -> v -> Maybe [v]
go Shape
Plain v
a = [v] -> Maybe [v]
forall a. a -> Maybe a
Just [v
a]
    go Shape
s v
a
      | Just Built v
b <- v -> ConLets v -> Maybe (Built v)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup v
a ConLets v
cons, Shape -> Built v -> Bool
forall v. Shape -> Built v -> Bool
matchesBuilt Shape
s Built v
b = [[v]] -> [v]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[v]] -> [v]) -> Maybe [[v]] -> Maybe [v]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Shape, v) -> Maybe [v]) -> [(Shape, v)] -> Maybe [[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 ((Shape -> v -> Maybe [v]) -> (Shape, v) -> Maybe [v]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Shape -> v -> Maybe [v]
go) ([Shape] -> [v] -> [(Shape, v)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Shape -> [Shape]
subShapes Shape
s) (Built v -> [v]
forall v. Built v -> [v]
builtFields Built v
b))
      | Bool
otherwise = Maybe [v]
forall a. Maybe a
Nothing

-- The boring helpers.

lambdas :: SuperGroup Reference v -> [SuperNormal Reference v]
lambdas :: forall v. SuperGroup Reference v -> [SuperNormal Reference v]
lambdas (Rec [(v, SuperNormal Reference v)]
bs SuperNormal Reference v
e) = SuperNormal Reference v
e SuperNormal Reference v
-> [SuperNormal Reference v] -> [SuperNormal Reference v]
forall a. a -> [a] -> [a]
: ((v, SuperNormal Reference v) -> SuperNormal Reference v)
-> [(v, SuperNormal Reference v)] -> [SuperNormal Reference v]
forall a b. (a -> b) -> [a] -> [b]
map (v, SuperNormal Reference v) -> SuperNormal Reference v
forall a b. (a, b) -> b
snd [(v, SuperNormal Reference v)]
bs

lambdaBody :: (Var v) => SuperNormal Reference v -> ANormal Reference v
lambdaBody :: forall v. Var v => SuperNormal Reference v -> ANormal Reference v
lambdaBody (Lambda [Mem]
_ (TAbss [v]
_ Term (ANormalF Reference) v
bd)) = Term (ANormalF Reference) v
bd

-- | A term and every term inside it.
-- | Every subterm, the term first; built with an accumulator, since the
-- analyses walk it once per let of a body and the nested appends of a
-- 'concatMap' cost the body's size times its depth per walk (a doc
-- literal of a few thousand lets took minutes, 2026-11-03).
universe :: ANormal Reference v -> [ANormal Reference v]
universe :: forall v. ANormal Reference v -> [ANormal Reference v]
universe ANormal Reference v
t0 = ANormal Reference v
-> [ANormal Reference v] -> [ANormal Reference v]
forall {v}.
ANormal Reference v
-> [ANormal Reference v] -> [ANormal Reference v]
go ANormal Reference v
t0 []
  where
    go :: ANormal Reference v
-> [ANormal Reference v] -> [ANormal Reference v]
go ANormal Reference v
t [ANormal Reference v]
acc = ANormal Reference v
t ANormal Reference v
-> [ANormal Reference v] -> [ANormal Reference v]
forall a. a -> [a] -> [a]
: (ANormal Reference v
 -> [ANormal Reference v] -> [ANormal Reference v])
-> [ANormal Reference v]
-> [ANormal Reference v]
-> [ANormal Reference v]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ANormal Reference v
-> [ANormal Reference v] -> [ANormal Reference v]
go [ANormal Reference v]
acc (ANormal Reference v -> [ANormal Reference v]
forall v. ANormal Reference v -> [ANormal Reference v]
children ANormal Reference v
t)

children :: ANormal Reference v -> [ANormal Reference v]
children :: forall v. ANormal Reference v -> [ANormal Reference v]
children (ABTN.Term Set v
_ (ABTN.Abs v
_ Term (ANormalF Reference) v
b)) = [Term (ANormalF Reference) v
b]
children (ABTN.Term Set v
_ (ABTN.Tm ANormalF Reference v (Term (ANormalF Reference) v)
f)) = ANormalF Reference v (Term (ANormalF Reference) v)
-> [Term (ANormalF Reference) v]
forall a. ANormalF Reference v a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList ANormalF Reference v (Term (ANormalF Reference) v)
f

boundVars :: (Ord v) => ANormal Reference v -> Set.Set v
boundVars :: forall v. Ord v => ANormal Reference v -> Set v
boundVars (ABTN.Term Set v
_ (ABTN.Abs v
u Term (ANormalF Reference) v
_)) = v -> Set v
forall a. a -> Set a
Set.singleton v
u
boundVars Term (ANormalF Reference) v
_ = Set v
forall a. Set a
Set.empty

-- | How many times a variable is used: its free occurrences, in every
-- position a variable can take.
occurrences :: (Var v) => v -> ANormal Reference v -> Int
occurrences :: forall v. Var v => v -> ANormal Reference v -> Int
occurrences v
v = Sum Int -> Int
forall a. Sum a -> a
getSum (Sum Int -> Int)
-> (ANormal Reference v -> Sum Int) -> ANormal Reference v -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v
-> (ANormalF Reference v (ANormal Reference v) -> Sum Int)
-> ANormal Reference v
-> Sum Int
forall v m.
(Var v, Monoid m) =>
v
-> (ANormalF Reference v (ANormal Reference v) -> m)
-> ANormal Reference v
-> m
foldMentions v
v (Int -> Sum Int
forall a. a -> Sum a
Sum (Int -> Sum Int)
-> (ANormalF Reference v (ANormal Reference v) -> Int)
-> ANormalF Reference v (ANormal Reference v)
-> Sum Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([v] -> Int)
-> (ANormalF Reference v (ANormal Reference v) -> [v])
-> ANormalF Reference v (ANormal Reference v)
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (v -> Bool) -> [v] -> [v]
forall a. (a -> Bool) -> [a] -> [a]
filter (v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
v) ([v] -> [v])
-> (ANormalF Reference v (ANormal Reference v) -> [v])
-> ANormalF Reference v (ANormal Reference v)
-> [v]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ANormalF Reference v (ANormal Reference v) -> [v]
forall {ref} {a} {e}. ANormalF ref a e -> [a]
nodeVars)
  where
    nodeVars :: ANormalF ref a e -> [a]
nodeVars ANormalF ref a e
f = case ANormalF ref a e
f of
      AVar a
u -> [a
u]
      AApp Func ref a
fn [a]
args -> Func ref a -> [a]
forall {ref} {a}. Func ref a -> [a]
funcVars Func ref a
fn [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
args
      AName (Right a
u) [a]
as e
_ -> a
u a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
as
      AName (Left ref
_) [a]
as e
_ -> [a]
as
      AMatch a
u Branched ref e
_ -> [a
u]
      AHnd [ref]
_ a
nh Maybe a
ah e
_ -> a
nh a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Maybe a -> [a]
forall a. Maybe a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Maybe a
ah
      AFrc a
u -> [a
u]
      ADiscard a
u -> [a
u]
      ALocal a
u e
_ -> [a
u]
      AUpdate Bool
_ a
u a
x -> [a
u, a
x]
      ABox a
u -> [a
u]
      ALet {} -> []
      AShift {} -> []
      ALit Lit ref
_ -> []
      ABLit Lit ref
_ -> []
    funcVars :: Func ref a -> [a]
funcVars = \case
      FVar a
u -> [a
u]
      FCont a
u -> [a
u]
      Func ref a
_ -> []

-- | A fold over the nodes of a body that can mention a variable: only
-- the subterms whose free variables hold it are entered, which every
-- term records ('ABTN.freeVars'), so the walk costs the paths to the
-- variable's uses, not the body. The analyses here ask this once per
-- let, which over a body of thousands of lets is what made them
-- quadratic (Paul, 2026-11-03).
foldMentions :: (Var v, Monoid m) => v -> (ANormalF Reference v (ANormal Reference v) -> m) -> ANormal Reference v -> m
foldMentions :: forall v m.
(Var v, Monoid m) =>
v
-> (ANormalF Reference v (ANormal Reference v) -> m)
-> ANormal Reference v
-> m
foldMentions v
v ANormalF Reference v (ANormal Reference v) -> m
at = ANormal Reference v -> m
go
  where
    go :: ANormal Reference v -> m
go ANormal Reference v
tm
      | Bool -> Bool
not (v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member v
v (ANormal Reference v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
ABTN.freeVars ANormal Reference v
tm)) = m
forall a. Monoid a => a
mempty
      | Bool
otherwise = case ANormal Reference v
tm of
          ABTN.Term Set v
_ (ABTN.Abs v
_ ANormal Reference v
b) -> ANormal Reference v -> m
go ANormal Reference v
b
          ABTN.Term Set v
_ (ABTN.Tm ANormalF Reference v (ANormal Reference v)
f) -> ANormalF Reference v (ANormal Reference v) -> m
at ANormalF Reference v (ANormal Reference v)
f m -> m -> m
forall a. Semigroup a => a -> a -> a
<> (ANormal Reference v -> m)
-> ANormalF Reference v (ANormal Reference v) -> m
forall m a. Monoid m => (a -> m) -> ANormalF Reference v a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ANormal Reference v -> m
go ANormalF Reference v (ANormal Reference v)
f

-- | Whether a body uses a variable at all.
mentions :: (Var v) => v -> ANormal Reference v -> Bool
mentions :: forall v. Var v => v -> ANormal Reference v -> Bool
mentions v
v ANormal Reference v
tm = v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member v
v (ANormal Reference v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
ABTN.freeVars ANormal Reference v
tm)