{-# LANGUAGE PatternSynonyms #-}

-- | The wide calling convention's escape analysis, as a pass over the ANF
-- (Unison.Runtime.Wide; design.md "The wide calling convention").
--
-- A variable bound by a let whose value arrives by a yield (a call, a
-- match, a handle: everything 'letYields' says) is /wide/: it occupies
-- 'wideWidth' slots, and when the yield delivered a small constructor (at
-- most 'wideK' fields) its value slot holds a witness and the slots below
-- it the fields, with nothing allocated. So is a let of one variable
-- binding a small constructor application, which the emitter packs into
-- the slots ('wideCon'), when the body matches on the variable or returns
-- it somewhere ('usedWide'): those are the uses the slots serve, and a
-- constructor only ever passed on is packed as a heap value at once, since
-- packing it wide and boxing it before the call would be the same
-- allocation plus the box. Such a variable can be matched on and returned
-- as it is. Where it /escapes/ (an argument of a call, a constructor, a
-- record, a primitive or a request; a captured variable of a closure; the
-- value of an affine update) it has to be an ordinary heap value, which
-- a box makes it, in place: after that the variable holds a boxed value
-- for the rest of the path.
--
-- 'wideBody' marks each such binder 'WD' and places a 'TBox' before the
-- first escape of the variable on each path: once per path however many
-- times the variable escapes after, and not at all on a path where it
-- never does, so a value that is only matched on costs no allocation. The
-- result says everything the emitter (Unison.Runtime.MCode) needs: it lays
-- a 'WD' binder out wide, emits a 'Box' for each 'TBox', checks the two
-- invariants ('WD' only on a let of one variable that yields or builds a
-- small constructor; a 'TBox' only of a wide variable in scope) and
-- decides nothing itself. The pass runs last over every body the emitter
-- sees, the unoptimized copy of a rewritten lambda included, so a box is
-- at the same frame in both.
--
-- With the convention off the yielding binders are still marked, since
-- 'wideWidth' is one then and a wide variable is laid out like any other,
-- and no box is placed.
module Unison.Runtime.MCode.Wide
  ( wideBody,
    letYields,
    wideCon,
  )
where

import Data.Set (Set)
import Data.Set qualified as Set
import Unison.ABT.Normalized qualified as ABTN
import Unison.Reference (Reference)
import Unison.Runtime.ANF
  ( ANormal,
    Direction (..),
    Func (..),
    Mem (..),
    pattern TApp,
    pattern TBLit,
    pattern TBox,
    pattern TDiscard,
    pattern TLets,
    pattern TLit,
    pattern TMatch,
    pattern TName,
    pattern TUpdate,
    pattern TVar,
  )
import Unison.Runtime.Wide (wideK, wideOn)
import Unison.Var (Var)

-- | Whether a let's binding is emitted as a section of its own, whose
-- value arrives by a yield, rather than as an instruction (a literal, a
-- constructor, a record, a record update, a primitive, a discard, an
-- affine update, a box): the complement of the instruction cases of the
-- emitter's emitLet.
letYields :: (Var v) => ANormal Reference v -> Bool
letYields :: forall v. Var v => ANormal Reference v -> Bool
letYields = \case
  TLit Lit Reference
_ -> Bool
False
  TBLit Lit Reference
_ -> Bool
False
  TApp (FCon Reference
_ CTag
_) [v]
_ -> Bool
False
  TApp (FRec RecordSchema
_) [v]
_ -> Bool
False
  TApp (FRecSet FieldName
_) [v]
_ -> Bool
False
  TApp (FPrim Either POp ForeignFunc
_) [v]
_ -> Bool
False
  TDiscard v
_ -> Bool
False
  TUpdate {} -> Bool
False
  TBox v
_ -> Bool
False
  ANormal Reference v
_ -> Bool
True

-- | Whether a let's binding is a small constructor application, which
-- stays in the variable's slots under the convention.
wideCon :: (Var v) => ANormal Reference v -> Bool
wideCon :: forall v. Var v => ANormal Reference v -> Bool
wideCon (TApp (FCon Reference
_ CTag
_) [v]
args) = Bool
wideOn Bool -> Bool -> Bool
&& [v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [v]
args Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
wideK
wideCon ANormal Reference v
_ = Bool
False

-- | Whether a body matches on a variable or returns it: the uses a wide
-- variable serves unboxed. Walks only the subterms that mention the
-- variable, which every term records.
usedWide :: (Var v) => v -> ANormal Reference v -> Bool
usedWide :: forall v. Var v => v -> ANormal Reference v -> Bool
usedWide v
u = ANormal Reference v -> Bool
go
  where
    go :: ANormal Reference v -> Bool
go ANormal Reference v
tm
      | Bool -> Bool
not (v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member v
u (ANormal Reference v -> Set v
forall (f :: * -> * -> *) v. Term f v -> Set v
ABTN.freeVars ANormal Reference v
tm)) = Bool
False
      | Bool
otherwise = case ANormal Reference v
tm of
          TMatch v
v Branched Reference (ANormal Reference v)
bs -> v
v v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
u Bool -> Bool -> Bool
|| (ANormal Reference v -> Bool)
-> Branched Reference (ANormal Reference v) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ANormal Reference v -> Bool
go Branched Reference (ANormal Reference v)
bs
          TVar v
v -> v
v v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
u
          ABTN.Term Set v
_ (ABTN.Abs v
_ ANormal Reference v
b) -> ANormal Reference v -> Bool
go ANormal Reference v
b
          ABTN.Term Set v
_ (ABTN.Tm ANormalF Reference v (ANormal Reference v)
f) -> (ANormal Reference v -> Bool)
-> ANormalF Reference v (ANormal Reference v) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ANormal Reference v -> Bool
go ANormalF Reference v (ANormal Reference v)
f

-- | The body with its wide binders marked and its boxes placed. Linear in
-- the body: one walk, with the set of the wide variables that may still
-- hold an unboxed constructor on the current path, forked at every branch.
wideBody :: (Var v) => ANormal Reference v -> ANormal Reference v
wideBody :: forall v. Var v => ANormal Reference v -> ANormal Reference v
wideBody = Set v -> ANormal Reference v -> ANormal Reference v
forall v.
Var v =>
Set v -> ANormal Reference v -> ANormal Reference v
go Set v
forall a. Set a
Set.empty
  where
    go :: (Var v) => Set v -> ANormal Reference v -> ANormal Reference v
    go :: forall v.
Var v =>
Set v -> ANormal Reference v -> ANormal Reference v
go Set v
un = \case
      -- a parameter or an arm's field is narrow, whatever it shadows
      ABTN.Term Set v
_ (ABTN.Abs v
v 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
v (Set v -> ANormal Reference v -> ANormal Reference v
forall v.
Var v =>
Set v -> ANormal Reference v -> ANormal Reference v
go (v -> Set v -> Set v
forall a. Ord a => a -> Set a -> Set a
Set.delete v
v Set v
un) ANormal Reference v
b)
      TLets Direction Word16
d [v]
us [Mem]
ms ANormal Reference v
bn ANormal Reference v
bo
        -- the binding runs in a section of its own, and boxes what
        -- escapes inside it; the body boxes again, since the binding's
        -- boxes are not its own. One variable is wide; several (the
        -- leaves of a product-return worker's result) are narrow.
        | ANormal Reference v -> Bool
forall v. Var v => ANormal Reference v -> Bool
letYields ANormal Reference v
bn ->
            let ([Mem]
ms', Set v
un') = case [v]
us of
                  [v
u] -> ([Mem
WD], v -> Set v -> Set v
forall a. Ord a => a -> Set a -> Set a
wide v
u ([v] -> Set v -> Set v
forall {t :: * -> *} {a}.
(Foldable t, Ord a) =>
t a -> Set a -> Set a
narrow [v]
us Set v
un))
                  [v]
_ -> ([Mem]
ms, [v] -> Set v -> Set v
forall {t :: * -> *} {a}.
(Foldable t, Ord a) =>
t a -> Set a -> Set a
narrow [v]
us Set v
un)
             in 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' (Set v -> ANormal Reference v -> ANormal Reference v
forall v.
Var v =>
Set v -> ANormal Reference v -> ANormal Reference v
go Set v
un ANormal Reference v
bn) (Set v -> ANormal Reference v -> ANormal Reference v
forall v.
Var v =>
Set v -> ANormal Reference v -> ANormal Reference v
go Set v
un' ANormal Reference v
bo)
        | [v
u] <- [v]
us,
          ANormal Reference v -> Bool
forall v. Var v => ANormal Reference v -> Bool
wideCon ANormal Reference v
bn,
          v -> ANormal Reference v -> Bool
forall v. Var v => v -> ANormal Reference v -> Bool
usedWide v
u ANormal Reference v
bo,
          TApp Func Reference v
_ [v]
args <- ANormal Reference v
bn ->
            Set v
-> [v] -> (Set v -> ANormal Reference v) -> ANormal Reference v
forall {a} {ref}.
Var a =>
Set a -> [a] -> (Set a -> ANormal ref a) -> ANormal ref a
boxes Set v
un [v]
args \Set v
un' -> 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
WD] ANormal Reference v
bn (Set v -> ANormal Reference v -> ANormal Reference v
forall v.
Var v =>
Set v -> ANormal Reference v -> ANormal Reference v
go (v -> Set v -> Set v
forall a. Ord a => a -> Set a -> Set a
wide v
u ([v] -> Set v -> Set v
forall {t :: * -> *} {a}.
(Foldable t, Ord a) =>
t a -> Set a -> Set a
narrow [v]
us Set v
un')) ANormal Reference v
bo)
        -- an instruction: its arguments escape before it
        | TApp Func Reference v
_ [v]
args <- ANormal Reference v
bn ->
            Set v
-> [v] -> (Set v -> ANormal Reference v) -> ANormal Reference v
forall {a} {ref}.
Var a =>
Set a -> [a] -> (Set a -> ANormal ref a) -> ANormal ref a
boxes Set v
un [v]
args \Set v
un' -> 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 (Set v -> ANormal Reference v -> ANormal Reference v
forall v.
Var v =>
Set v -> ANormal Reference v -> ANormal Reference v
go ([v] -> Set v -> Set v
forall {t :: * -> *} {a}.
(Foldable t, Ord a) =>
t a -> Set a -> Set a
narrow [v]
us Set v
un') ANormal Reference v
bo)
        | TUpdate Bool
_ v
_ v
x <- ANormal Reference v
bn ->
            Set v
-> [v] -> (Set v -> ANormal Reference v) -> ANormal Reference v
forall {a} {ref}.
Var a =>
Set a -> [a] -> (Set a -> ANormal ref a) -> ANormal ref a
boxes Set v
un [v
x] \Set v
un' -> 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 (Set v -> ANormal Reference v -> ANormal Reference v
forall v.
Var v =>
Set v -> ANormal Reference v -> ANormal Reference v
go ([v] -> Set v -> Set v
forall {t :: * -> *} {a}.
(Foldable t, Ord a) =>
t a -> Set a -> Set a
narrow [v]
us Set v
un') 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 (Set v -> ANormal Reference v -> ANormal Reference v
forall v.
Var v =>
Set v -> ANormal Reference v -> ANormal Reference v
go Set v
un ANormal Reference v
bn) (Set v -> ANormal Reference v -> ANormal Reference v
forall v.
Var v =>
Set v -> ANormal Reference v -> ANormal Reference v
go ([v] -> Set v -> Set v
forall {t :: * -> *} {a}.
(Foldable t, Ord a) =>
t a -> Set a -> Set a
narrow [v]
us Set v
un) ANormal Reference v
bo)
      TName v
u Either Reference v
f [v]
args ANormal Reference v
bo ->
        Set v
-> [v] -> (Set v -> ANormal Reference v) -> ANormal Reference v
forall {a} {ref}.
Var a =>
Set a -> [a] -> (Set a -> ANormal ref a) -> ANormal ref a
boxes Set v
un [v]
args \Set v
un' -> 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
u Either Reference v
f [v]
args (Set v -> ANormal Reference v -> ANormal Reference v
forall v.
Var v =>
Set v -> ANormal Reference v -> ANormal Reference v
go (v -> Set v -> Set v
forall a. Ord a => a -> Set a -> Set a
Set.delete v
u Set v
un') ANormal Reference v
bo)
      TApp Func Reference v
f [v]
args -> Set v
-> [v] -> (Set v -> ANormal Reference v) -> ANormal Reference v
forall {a} {ref}.
Var a =>
Set a -> [a] -> (Set a -> ANormal ref a) -> ANormal ref a
boxes Set v
un [v]
args \Set v
_ -> 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]
args
      -- a match, a handle, a shift, a local, a force, a variable, a
      -- literal, a discard: nothing escapes, and the walk goes on into
      -- each child with the same set
      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 (Set v -> ANormal Reference v -> ANormal Reference v
forall v.
Var v =>
Set v -> ANormal Reference v -> ANormal Reference v
go Set v
un) ANormalF Reference v (ANormal Reference v)
f)
    wide :: a -> Set a -> Set a
wide a
u Set a
un
      | Bool
wideOn = a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
Set.insert a
u Set a
un
      | Bool
otherwise = Set a
un
    narrow :: t a -> Set a -> Set a
narrow t a
us Set a
un = (a -> Set a -> Set a) -> Set a -> t a -> Set a
forall a b. (a -> b -> b) -> b -> t a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
Set.delete Set a
un t a
us
    -- a box before the node for each argument still unboxed, once each
    boxes :: Set a -> [a] -> (Set a -> ANormal ref a) -> ANormal ref a
boxes Set a
un [] Set a -> ANormal ref a
k = Set a -> ANormal ref a
k Set a
un
    boxes Set a
un (a
a : [a]
as) Set a -> ANormal ref a
k
      | a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member a
a Set a
un = Direction Word16
-> [a] -> [Mem] -> ANormal ref a -> ANormal ref a -> ANormal ref a
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 [] [] (a -> ANormal ref a
forall v ref. Var v => v -> ANormal ref v
TBox a
a) (Set a -> [a] -> (Set a -> ANormal ref a) -> ANormal ref a
boxes (a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
Set.delete a
a Set a
un) [a]
as Set a -> ANormal ref a
k)
      | Bool
otherwise = Set a -> [a] -> (Set a -> ANormal ref a) -> ANormal ref a
boxes Set a
un [a]
as Set a -> ANormal ref a
k