-- editorconfig-checker-disable-file
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings     #-}

module PlutusIR.Analysis.RetainedSize
    ( RetainedSize (..)
    , AstSize (..)
    , termRetentionMap
    , annotateWithRetainedSize
    ) where

import PlutusPrelude

import PlutusIR.Analysis.Dependencies
import PlutusIR.AstSize
import PlutusIR.Core

import PlutusCore qualified as PLC
import PlutusCore.Builtin as PLC
import PlutusCore.Name.Unique

import Algebra.Graph qualified as C
import Algebra.Graph.ToGraph
import Control.Lens
import Data.Graph.Dom (domTree)
import Data.IntMap.Strict (IntMap)
import Data.IntMap.Strict qualified as IntMap
import Data.Tree
import PlutusCore.Name.UniqueMap qualified as UMap
import PlutusIR.Analysis.Builtins
import PlutusIR.Analysis.VarInfo

{- Note [Retained size analysis]
WARNING: everything in this module assumes global uniqueness of variables.

For calculating the size that each binding retains we use a classic graph algorithm for dominance.
The @algebraic-graphs@ library does not provide one out of the box, so we convert an algebraic graph
to a @fgl@-based representation and use the @dom-lt@ library for calculating the dominator tree
(note that @fgl@ itself has an algorithm for dominance [1], but that was only found out after
everything below was implemented, plus that algorithm does not return a 'Tree', so it's not as
convenient to use).

The procedure for computing the retained size of each binding is two-staged:

1. First we compute how much size each let-binding in a term retains directly.
   For example, a 'TermBind' directly retains only the type signature and the body of the binding.
2. Then we compute the dominator tree for the dependency graph of the term and annotate it with the
   retained size of each binding by adding up the directly retained size of the binding computed at
   stage one and the size retained by all the dependencies of the binding

We don't make any assumptions as to whether, say, a constructor of a data type directly retains the
entire data type (its other constructors, the kind signature of the data type etc). Instead, we
specify that each constructor directly retains only its type signature, each parameter directly
retains only its kind signature etc and we let dependency and dominator tree analysis figure out
whether the constructor in fact retains the entire data type by virtue of the constructor being
connected (or not connected) with other parts of the data type. This has a funny effect: if only
one constructor of the data type is actually used, then it's closer to the root of the dependency
graph and so every other part of the data type is considered a dependency (due to them all being
interconnected) and so that constructor gets annotated with the size of the whole data type, while
every other constructor is only annotated with the size that it directly retains (i.e. the size of
its type signature).

Our main motivation for implementing retained size analysis is to calculate what bindings retain
most size, so that we can prioritize certain optimization passes, however we mostly care about the
size of the final untyped program and given that we do take types and kinds into account during
retained size analysis, results do not translate directly to the untyped setting. So beware of this
pitfall.

Even the most efficient dominator tree algorithms are still not linear and our programs are quite
huge, so before jumping into computing the dominator tree we filter the dependency graph and remove
from it everything that does not directly retain any size, i.e. everything that is not a let-binding.
For example, the dependency graph contains lambda-bound variables and we filter them out as they're
irrelevant for computing dominance. So we're relying on the fact that nodes that don't retain any
size are also not interesting for dominator analysis. Which might be wrong. An example is needed.

[1] https://hackage.haskell.org/package/fgl-5.7.0.3/docs/Data-Graph-Inductive-Query-Dominators.html#v:dom
-}

{- Note [Handling the root]
The dominator tree algorithm works on the assumption that each node is an 'Int'. It's easy to
convert the 'Unique' of a variable to an 'Int', but the dominator tree algorithm needs a root to
walk from and that root also has to be an 'Int' and our uniques start from zero, so we made an
awkward decision to assign the root @-1@ as its index.

Given that the root appears in the dominator tree, we also need to specify the size that it directly
retains. But there doesn't seem to be a sensible way of doing that. Should it be

    rootSize :: Term tyname name uni fun ann -> Size
    rootSize (Let _ _ _ term) = rootSize term
    rootSize term             = termAstSize term

? But what about bindings inside the final non-let term? However we don't need that directly
retained size of the root for anything that is not "don't throw an error on encountering the root
node that does not correspond to any binding" and since we only annotate bindings, the size retained
by the root is not supposed to appear in the final annotated term anyway, so we just make the root
directly retain some ridiculously huge _negative_ number. Now if it ends up being in the final term,
we know there's a bug somewhere and if it doesn't, we don't care about it.
-}

data RetainedSize
    = Retains AstSize
    | NotARetainer
    deriving stock (Int -> RetainedSize -> ShowS
[RetainedSize] -> ShowS
RetainedSize -> String
(Int -> RetainedSize -> ShowS)
-> (RetainedSize -> String)
-> ([RetainedSize] -> ShowS)
-> Show RetainedSize
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> RetainedSize -> ShowS
showsPrec :: Int -> RetainedSize -> ShowS
$cshow :: RetainedSize -> String
show :: RetainedSize -> String
$cshowList :: [RetainedSize] -> ShowS
showList :: [RetainedSize] -> ShowS
Show)

instance Pretty RetainedSize where
    pretty :: forall ann. RetainedSize -> Doc ann
pretty (Retains AstSize
size) = Doc ann
"$" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> AstSize -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. AstSize -> Doc ann
pretty AstSize
size Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
"$"
    pretty RetainedSize
NotARetainer   = Doc ann
forall a. Monoid a => a
mempty

-- See Note [Handling the root].
-- | The 'Int' index of the root.
rootInt :: Int
rootInt :: Int
rootInt = -Int
1

-- See Note [Handling the root].
nodeToInt :: Node -> Int
nodeToInt :: Node -> Int
nodeToInt (Variable (PLC.Unique Int
i)) = Int
i
nodeToInt Node
Root                      = Int
rootInt

-- | A mapping from the index of a binding to what it directly retains.
newtype DirectionRetentionMap = DirectionRetentionMap (IntMap AstSize)

lookupSize :: Int -> DirectionRetentionMap -> AstSize
lookupSize :: Int -> DirectionRetentionMap -> AstSize
lookupSize Int
i (DirectionRetentionMap IntMap AstSize
ss) = IntMap AstSize
ss IntMap AstSize -> Int -> AstSize
forall a. IntMap a -> Int -> a
IntMap.! Int
i

-- | Annotate the dominator tree with the retained size of each entry. The retained size is computed
-- as the size directly retained by the binding plus the size of all its dependencies.
annotateWithSizes :: DirectionRetentionMap -> Tree Int -> Tree (Int, AstSize)
annotateWithSizes :: DirectionRetentionMap -> Tree Int -> Tree (Int, AstSize)
annotateWithSizes DirectionRetentionMap
sizeInfo = Tree Int -> Tree (Int, AstSize)
go where
    go :: Tree Int -> Tree (Int, AstSize)
go (Node Int
i [Tree Int]
ts) = (Int, AstSize) -> [Tree (Int, AstSize)] -> Tree (Int, AstSize)
forall a. a -> [Tree a] -> Tree a
Node (Int
i, AstSize
sizeI) [Tree (Int, AstSize)]
rs where
        rs :: [Tree (Int, AstSize)]
rs = (Tree Int -> Tree (Int, AstSize))
-> [Tree Int] -> [Tree (Int, AstSize)]
forall a b. (a -> b) -> [a] -> [b]
map Tree Int -> Tree (Int, AstSize)
go [Tree Int]
ts
        sizeI :: AstSize
sizeI = Int -> DirectionRetentionMap -> AstSize
lookupSize Int
i DirectionRetentionMap
sizeInfo AstSize -> AstSize -> AstSize
forall a. Semigroup a => a -> a -> a
<> (Tree (Int, AstSize) -> AstSize)
-> [Tree (Int, AstSize)] -> AstSize
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((Int, AstSize) -> AstSize
forall a b. (a, b) -> b
snd ((Int, AstSize) -> AstSize)
-> (Tree (Int, AstSize) -> (Int, AstSize))
-> Tree (Int, AstSize)
-> AstSize
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree (Int, AstSize) -> (Int, AstSize)
forall a. Tree a -> a
rootLabel) [Tree (Int, AstSize)]
rs

-- | Compute the dominator tree of a graph.
toDomTree :: C.Graph Node -> Tree Int
toDomTree :: Graph Node -> Tree Int
toDomTree = Rooted -> Tree Int
domTree (Rooted -> Tree Int)
-> (Graph Node -> Rooted) -> Graph Node -> Tree Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (,) Int
rootInt (Graph -> Rooted) -> (Graph Node -> Graph) -> Graph Node -> Rooted
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Graph Int -> Graph
forall t. (ToGraph t, ToVertex t ~ Int) => t -> Graph
adjacencyIntMap (Graph Int -> Graph)
-> (Graph Node -> Graph Int) -> Graph Node -> Graph
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Node -> Int) -> Graph Node -> Graph Int
forall a b. (a -> b) -> Graph a -> Graph b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Node -> Int
nodeToInt

-- | Compute the retention map of a graph.
depsRetentionMap :: DirectionRetentionMap -> C.Graph Node -> IntMap AstSize
depsRetentionMap :: DirectionRetentionMap -> Graph Node -> IntMap AstSize
depsRetentionMap DirectionRetentionMap
sizeInfo = [(Int, AstSize)] -> IntMap AstSize
forall a. [(Int, a)] -> IntMap a
IntMap.fromList ([(Int, AstSize)] -> IntMap AstSize)
-> (Graph Node -> [(Int, AstSize)]) -> Graph Node -> IntMap AstSize
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree (Int, AstSize) -> [(Int, AstSize)]
forall a. Tree a -> [a]
flatten (Tree (Int, AstSize) -> [(Int, AstSize)])
-> (Graph Node -> Tree (Int, AstSize))
-> Graph Node
-> [(Int, AstSize)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DirectionRetentionMap -> Tree Int -> Tree (Int, AstSize)
annotateWithSizes DirectionRetentionMap
sizeInfo (Tree Int -> Tree (Int, AstSize))
-> (Graph Node -> Tree Int) -> Graph Node -> Tree (Int, AstSize)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Graph Node -> Tree Int
toDomTree

-- | Construct a 'UniqueMap' having size information for each individual part of a 'Binding'.
bindingSize
    :: (HasUnique tyname TypeUnique, HasUnique name TermUnique)
    => Binding tyname name uni fun ann -> PLC.UniqueMap Unique AstSize
bindingSize :: forall tyname name (uni :: * -> *) fun ann.
(HasUnique tyname TypeUnique, HasUnique name TermUnique) =>
Binding tyname name uni fun ann -> UniqueMap Unique AstSize
bindingSize (TermBind ann
_ Strictness
_ VarDecl tyname name uni ann
var Term tyname name uni fun ann
term) =
    VarDecl tyname name uni ann
-> AstSize -> UniqueMap Unique AstSize -> UniqueMap Unique AstSize
forall name unique1 unique2 a.
(HasUnique name unique1, Coercible unique2 Unique) =>
name -> a -> UniqueMap unique2 a -> UniqueMap unique2 a
UMap.insertByNameIndex VarDecl tyname name uni ann
var (VarDecl tyname name uni ann -> AstSize
forall tyname name (uni :: * -> *) ann.
VarDecl tyname name uni ann -> AstSize
varDeclAstSize VarDecl tyname name uni ann
var AstSize -> AstSize -> AstSize
forall a. Semigroup a => a -> a -> a
<> Term tyname name uni fun ann -> AstSize
forall tyname name (uni :: * -> *) fun ann.
Term tyname name uni fun ann -> AstSize
termAstSize Term tyname name uni fun ann
term) UniqueMap Unique AstSize
forall a. Monoid a => a
mempty
bindingSize (TypeBind ann
_ TyVarDecl tyname ann
tyVar Type tyname uni ann
ty) =
    TyVarDecl tyname ann
-> AstSize -> UniqueMap Unique AstSize -> UniqueMap Unique AstSize
forall name unique1 unique2 a.
(HasUnique name unique1, Coercible unique2 Unique) =>
name -> a -> UniqueMap unique2 a -> UniqueMap unique2 a
UMap.insertByNameIndex TyVarDecl tyname ann
tyVar (TyVarDecl tyname ann -> AstSize
forall tyname ann. TyVarDecl tyname ann -> AstSize
tyVarDeclAstSize TyVarDecl tyname ann
tyVar AstSize -> AstSize -> AstSize
forall a. Semigroup a => a -> a -> a
<> Type tyname uni ann -> AstSize
forall tyname (uni :: * -> *) ann. Type tyname uni ann -> AstSize
typeAstSize Type tyname uni ann
ty) UniqueMap Unique AstSize
forall a. Monoid a => a
mempty
bindingSize (DatatypeBind ann
_ (Datatype ann
_ TyVarDecl tyname ann
dataDecl [TyVarDecl tyname ann]
params name
matchName [VarDecl tyname name uni ann]
constrs))
    = TyVarDecl tyname ann
-> AstSize -> UniqueMap Unique AstSize -> UniqueMap Unique AstSize
forall name unique1 unique2 a.
(HasUnique name unique1, Coercible unique2 Unique) =>
name -> a -> UniqueMap unique2 a -> UniqueMap unique2 a
UMap.insertByNameIndex TyVarDecl tyname ann
dataDecl (TyVarDecl tyname ann -> AstSize
forall tyname ann. TyVarDecl tyname ann -> AstSize
tyVarDeclAstSize TyVarDecl tyname ann
dataDecl)
    (UniqueMap Unique AstSize -> UniqueMap Unique AstSize)
-> (UniqueMap Unique AstSize -> UniqueMap Unique AstSize)
-> UniqueMap Unique AstSize
-> UniqueMap Unique AstSize
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (UniqueMap Unique AstSize
 -> [TyVarDecl tyname ann] -> UniqueMap Unique AstSize)
-> [TyVarDecl tyname ann]
-> UniqueMap Unique AstSize
-> UniqueMap Unique AstSize
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((TyVarDecl tyname ann
 -> UniqueMap Unique AstSize -> UniqueMap Unique AstSize)
-> UniqueMap Unique AstSize
-> [TyVarDecl tyname ann]
-> UniqueMap Unique AstSize
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((TyVarDecl tyname ann
  -> UniqueMap Unique AstSize -> UniqueMap Unique AstSize)
 -> UniqueMap Unique AstSize
 -> [TyVarDecl tyname ann]
 -> UniqueMap Unique AstSize)
-> (TyVarDecl tyname ann
    -> UniqueMap Unique AstSize -> UniqueMap Unique AstSize)
-> UniqueMap Unique AstSize
-> [TyVarDecl tyname ann]
-> UniqueMap Unique AstSize
forall a b. (a -> b) -> a -> b
$ \TyVarDecl tyname ann
param -> TyVarDecl tyname ann
-> AstSize -> UniqueMap Unique AstSize -> UniqueMap Unique AstSize
forall name unique1 unique2 a.
(HasUnique name unique1, Coercible unique2 Unique) =>
name -> a -> UniqueMap unique2 a -> UniqueMap unique2 a
UMap.insertByNameIndex TyVarDecl tyname ann
param (AstSize -> UniqueMap Unique AstSize -> UniqueMap Unique AstSize)
-> AstSize -> UniqueMap Unique AstSize -> UniqueMap Unique AstSize
forall a b. (a -> b) -> a -> b
$ TyVarDecl tyname ann -> AstSize
forall tyname ann. TyVarDecl tyname ann -> AstSize
tyVarDeclAstSize TyVarDecl tyname ann
param) [TyVarDecl tyname ann]
params
    (UniqueMap Unique AstSize -> UniqueMap Unique AstSize)
-> (UniqueMap Unique AstSize -> UniqueMap Unique AstSize)
-> UniqueMap Unique AstSize
-> UniqueMap Unique AstSize
forall b c a. (b -> c) -> (a -> b) -> a -> c
. name
-> AstSize -> UniqueMap Unique AstSize -> UniqueMap Unique AstSize
forall name unique1 unique2 a.
(HasUnique name unique1, Coercible unique2 Unique) =>
name -> a -> UniqueMap unique2 a -> UniqueMap unique2 a
UMap.insertByNameIndex name
matchName (Integer -> AstSize
AstSize Integer
1)
    (UniqueMap Unique AstSize -> UniqueMap Unique AstSize)
-> (UniqueMap Unique AstSize -> UniqueMap Unique AstSize)
-> UniqueMap Unique AstSize
-> UniqueMap Unique AstSize
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (UniqueMap Unique AstSize
 -> [VarDecl tyname name uni ann] -> UniqueMap Unique AstSize)
-> [VarDecl tyname name uni ann]
-> UniqueMap Unique AstSize
-> UniqueMap Unique AstSize
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((VarDecl tyname name uni ann
 -> UniqueMap Unique AstSize -> UniqueMap Unique AstSize)
-> UniqueMap Unique AstSize
-> [VarDecl tyname name uni ann]
-> UniqueMap Unique AstSize
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((VarDecl tyname name uni ann
  -> UniqueMap Unique AstSize -> UniqueMap Unique AstSize)
 -> UniqueMap Unique AstSize
 -> [VarDecl tyname name uni ann]
 -> UniqueMap Unique AstSize)
-> (VarDecl tyname name uni ann
    -> UniqueMap Unique AstSize -> UniqueMap Unique AstSize)
-> UniqueMap Unique AstSize
-> [VarDecl tyname name uni ann]
-> UniqueMap Unique AstSize
forall a b. (a -> b) -> a -> b
$ \VarDecl tyname name uni ann
constr -> VarDecl tyname name uni ann
-> AstSize -> UniqueMap Unique AstSize -> UniqueMap Unique AstSize
forall name unique1 unique2 a.
(HasUnique name unique1, Coercible unique2 Unique) =>
name -> a -> UniqueMap unique2 a -> UniqueMap unique2 a
UMap.insertByNameIndex VarDecl tyname name uni ann
constr (AstSize -> UniqueMap Unique AstSize -> UniqueMap Unique AstSize)
-> AstSize -> UniqueMap Unique AstSize -> UniqueMap Unique AstSize
forall a b. (a -> b) -> a -> b
$ VarDecl tyname name uni ann -> AstSize
forall tyname name (uni :: * -> *) ann.
VarDecl tyname name uni ann -> AstSize
varDeclAstSize VarDecl tyname name uni ann
constr) [VarDecl tyname name uni ann]
constrs
    (UniqueMap Unique AstSize -> UniqueMap Unique AstSize)
-> UniqueMap Unique AstSize -> UniqueMap Unique AstSize
forall a b. (a -> b) -> a -> b
$ UniqueMap Unique AstSize
forall a. Monoid a => a
mempty

-- | Construct a 'UniqueMap' having size information for each individual part of every 'Binding'
-- in a term.
bindingSizes
    :: (HasUnique tyname TypeUnique, HasUnique name TermUnique)
    => Term tyname name uni fun ann -> PLC.UniqueMap Unique AstSize
bindingSizes :: forall tyname name (uni :: * -> *) fun ann.
(HasUnique tyname TypeUnique, HasUnique name TermUnique) =>
Term tyname name uni fun ann -> UniqueMap Unique AstSize
bindingSizes (Let ann
_ Recursivity
_ NonEmpty (Binding tyname name uni fun ann)
binds Term tyname name uni fun ann
term) = (Binding tyname name uni fun ann -> UniqueMap Unique AstSize)
-> NonEmpty (Binding tyname name uni fun ann)
-> UniqueMap Unique AstSize
forall m a. Monoid m => (a -> m) -> NonEmpty a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Binding tyname name uni fun ann -> UniqueMap Unique AstSize
forall tyname name (uni :: * -> *) fun ann.
(HasUnique tyname TypeUnique, HasUnique name TermUnique) =>
Binding tyname name uni fun ann -> UniqueMap Unique AstSize
bindingSize NonEmpty (Binding tyname name uni fun ann)
binds UniqueMap Unique AstSize
-> UniqueMap Unique AstSize -> UniqueMap Unique AstSize
forall a. Semigroup a => a -> a -> a
<> Term tyname name uni fun ann -> UniqueMap Unique AstSize
forall tyname name (uni :: * -> *) fun ann.
(HasUnique tyname TypeUnique, HasUnique name TermUnique) =>
Term tyname name uni fun ann -> UniqueMap Unique AstSize
bindingSizes Term tyname name uni fun ann
term
bindingSizes Term tyname name uni fun ann
term                 = Term tyname name uni fun ann
term Term tyname name uni fun ann
-> Getting
     (UniqueMap Unique AstSize)
     (Term tyname name uni fun ann)
     (UniqueMap Unique AstSize)
-> UniqueMap Unique AstSize
forall s a. s -> Getting a s a -> a
^. (Term tyname name uni fun ann
 -> Const (UniqueMap Unique AstSize) (Term tyname name uni fun ann))
-> Term tyname name uni fun ann
-> Const (UniqueMap Unique AstSize) (Term tyname name uni fun ann)
forall tyname name (uni :: * -> *) fun a (f :: * -> *).
Applicative f =>
(Term tyname name uni fun a -> f (Term tyname name uni fun a))
-> Term tyname name uni fun a -> f (Term tyname name uni fun a)
termSubterms ((Term tyname name uni fun ann
  -> Const (UniqueMap Unique AstSize) (Term tyname name uni fun ann))
 -> Term tyname name uni fun ann
 -> Const (UniqueMap Unique AstSize) (Term tyname name uni fun ann))
-> Getting
     (UniqueMap Unique AstSize)
     (Term tyname name uni fun ann)
     (UniqueMap Unique AstSize)
-> Getting
     (UniqueMap Unique AstSize)
     (Term tyname name uni fun ann)
     (UniqueMap Unique AstSize)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Term tyname name uni fun ann -> UniqueMap Unique AstSize)
-> Getting
     (UniqueMap Unique AstSize)
     (Term tyname name uni fun ann)
     (UniqueMap Unique AstSize)
forall (p :: * -> * -> *) (f :: * -> *) s a.
(Profunctor p, Contravariant f) =>
(s -> a) -> Optic' p f s a
to Term tyname name uni fun ann -> UniqueMap Unique AstSize
forall tyname name (uni :: * -> *) fun ann.
(HasUnique tyname TypeUnique, HasUnique name TermUnique) =>
Term tyname name uni fun ann -> UniqueMap Unique AstSize
bindingSizes

-- | Same as 'bindingSizes' but is wrapped in a newtype and has a bogus entry for the root.
toDirectionRetentionMap
    :: (HasUnique tyname TypeUnique, HasUnique name TermUnique)
    => Term tyname name uni fun ann -> DirectionRetentionMap
toDirectionRetentionMap :: forall tyname name (uni :: * -> *) fun ann.
(HasUnique tyname TypeUnique, HasUnique name TermUnique) =>
Term tyname name uni fun ann -> DirectionRetentionMap
toDirectionRetentionMap Term tyname name uni fun ann
term =
    IntMap AstSize -> DirectionRetentionMap
DirectionRetentionMap (IntMap AstSize -> DirectionRetentionMap)
-> (UniqueMap Unique AstSize -> IntMap AstSize)
-> UniqueMap Unique AstSize
-> DirectionRetentionMap
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstSize -> IntMap AstSize -> IntMap AstSize
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
rootInt AstSize
rootSize (IntMap AstSize -> IntMap AstSize)
-> (UniqueMap Unique AstSize -> IntMap AstSize)
-> UniqueMap Unique AstSize
-> IntMap AstSize
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UniqueMap Unique AstSize -> IntMap AstSize
forall unique a. UniqueMap unique a -> IntMap a
PLC.unUniqueMap (UniqueMap Unique AstSize -> DirectionRetentionMap)
-> UniqueMap Unique AstSize -> DirectionRetentionMap
forall a b. (a -> b) -> a -> b
$ Term tyname name uni fun ann -> UniqueMap Unique AstSize
forall tyname name (uni :: * -> *) fun ann.
(HasUnique tyname TypeUnique, HasUnique name TermUnique) =>
Term tyname name uni fun ann -> UniqueMap Unique AstSize
bindingSizes Term tyname name uni fun ann
term where
        -- See Note [Handling the root].
        rootSize :: AstSize
rootSize = Integer -> AstSize
AstSize (- Integer
10 Integer -> Int -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
10::Int))

-- | Check if a 'Node' appears in 'DirectionRetentionMap'.
hasSizeIn :: DirectionRetentionMap -> Node -> Bool
hasSizeIn :: DirectionRetentionMap -> Node -> Bool
hasSizeIn DirectionRetentionMap
_             Node
Root                                   = Bool
True
hasSizeIn (DirectionRetentionMap IntMap AstSize
ss) (Variable (PLC.Unique Int
i)) = Int
i Int -> IntMap AstSize -> Bool
forall a. Int -> IntMap a -> Bool
`IntMap.member` IntMap AstSize
ss

-- | Compute the retention map of a term.
termRetentionMap
    :: (HasUnique tyname TypeUnique, HasUnique name TermUnique, ToBuiltinMeaning uni fun)
    => BuiltinsInfo uni fun
    -> VarsInfo tyname name uni ann
    -> Term tyname name uni fun ann
    -> IntMap AstSize
termRetentionMap :: forall tyname name (uni :: * -> *) fun ann.
(HasUnique tyname TypeUnique, HasUnique name TermUnique,
 ToBuiltinMeaning uni fun) =>
BuiltinsInfo uni fun
-> VarsInfo tyname name uni ann
-> Term tyname name uni fun ann
-> IntMap AstSize
termRetentionMap BuiltinsInfo uni fun
binfo VarsInfo tyname name uni ann
vinfo Term tyname name uni fun ann
term = DirectionRetentionMap -> Graph Node -> IntMap AstSize
depsRetentionMap DirectionRetentionMap
sizeInfo Graph Node
deps where
    sizeInfo :: DirectionRetentionMap
sizeInfo = Term tyname name uni fun ann -> DirectionRetentionMap
forall tyname name (uni :: * -> *) fun ann.
(HasUnique tyname TypeUnique, HasUnique name TermUnique) =>
Term tyname name uni fun ann -> DirectionRetentionMap
toDirectionRetentionMap Term tyname name uni fun ann
term
    deps :: Graph Node
deps = (Node -> Bool) -> Graph Node -> Graph Node
forall a. (a -> Bool) -> Graph a -> Graph a
C.induce (DirectionRetentionMap -> Node -> Bool
hasSizeIn DirectionRetentionMap
sizeInfo) (Graph Node -> Graph Node) -> Graph Node -> Graph Node
forall a b. (a -> b) -> a -> b
$ BuiltinsInfo uni fun
-> VarsInfo tyname name uni ann
-> Term tyname name uni fun ann
-> Graph Node
forall g tyname name (uni :: * -> *) fun a.
(DepGraph g, HasUnique tyname TypeUnique,
 HasUnique name TermUnique, ToBuiltinMeaning uni fun) =>
BuiltinsInfo uni fun
-> VarsInfo tyname name uni a -> Term tyname name uni fun a -> g
runTermDeps BuiltinsInfo uni fun
binfo VarsInfo tyname name uni ann
vinfo Term tyname name uni fun ann
term

-- | Apply a function to the annotation of each part of every 'Binding' in a term.
reannotateBindings
    :: (HasUnique name TermUnique, HasUnique tyname TypeUnique)
    => (Unique -> ann -> ann)
    -> Term tyname name uni fun ann
    -> Term tyname name uni fun ann
reannotateBindings :: forall name tyname ann (uni :: * -> *) fun.
(HasUnique name TermUnique, HasUnique tyname TypeUnique) =>
(Unique -> ann -> ann)
-> Term tyname name uni fun ann -> Term tyname name uni fun ann
reannotateBindings Unique -> ann -> ann
f = Term tyname name uni fun ann -> Term tyname name uni fun ann
forall {unique} {unique} {tyname} {name} {uni :: * -> *} {fun}.
(Coercible unique Int, Coercible unique Int,
 HasUnique tyname unique, HasUnique name unique) =>
Term tyname name uni fun ann -> Term tyname name uni fun ann
goTerm where
    -- We don't need these helper functions anywhere else, so we make them into local definitions.
    goVarDecl :: VarDecl tyname name uni ann -> VarDecl tyname name uni ann
goVarDecl (VarDecl ann
ann name
name Type tyname uni ann
ty) = ann -> name -> Type tyname uni ann -> VarDecl tyname name uni ann
forall tyname name (uni :: * -> *) ann.
ann -> name -> Type tyname uni ann -> VarDecl tyname name uni ann
VarDecl (Unique -> ann -> ann
f (name
name name -> Getting Unique name Unique -> Unique
forall s a. s -> Getting a s a -> a
^. Getting Unique name Unique
forall name unique. HasUnique name unique => Lens' name Unique
Lens' name Unique
theUnique) ann
ann) name
name Type tyname uni ann
ty
    goTyVarDecl :: TyVarDecl tyname ann -> TyVarDecl tyname ann
goTyVarDecl (TyVarDecl ann
ann tyname
tyname Kind ann
kind) = ann -> tyname -> Kind ann -> TyVarDecl tyname ann
forall tyname ann.
ann -> tyname -> Kind ann -> TyVarDecl tyname ann
TyVarDecl (Unique -> ann -> ann
f (tyname
tyname tyname -> Getting Unique tyname Unique -> Unique
forall s a. s -> Getting a s a -> a
^. Getting Unique tyname Unique
forall name unique. HasUnique name unique => Lens' name Unique
Lens' tyname Unique
theUnique) ann
ann) tyname
tyname Kind ann
kind
    goDatatype :: Datatype tyname name uni ann -> Datatype tyname name uni ann
goDatatype (Datatype ann
ann TyVarDecl tyname ann
dataTyDecl [TyVarDecl tyname ann]
paramTyDecls name
matchName [VarDecl tyname name uni ann]
constrDecls) =
        ann
-> TyVarDecl tyname ann
-> [TyVarDecl tyname ann]
-> name
-> [VarDecl tyname name uni ann]
-> Datatype tyname name uni ann
forall tyname name (uni :: * -> *) a.
a
-> TyVarDecl tyname a
-> [TyVarDecl tyname a]
-> name
-> [VarDecl tyname name uni a]
-> Datatype tyname name uni a
Datatype
            -- We don't have any other suitable place to associate the name of the matcher with an
            -- annotation, so we do it here. Fortunately, the matcher is the only thing that
            -- survives erasure, so this even makes some sense.
            (Unique -> ann -> ann
f (name
matchName name -> Getting Unique name Unique -> Unique
forall s a. s -> Getting a s a -> a
^. Getting Unique name Unique
forall name unique. HasUnique name unique => Lens' name Unique
Lens' name Unique
theUnique) ann
ann)
            (TyVarDecl tyname ann -> TyVarDecl tyname ann
forall {unique} {tyname}.
(Coercible unique Int, HasUnique tyname unique) =>
TyVarDecl tyname ann -> TyVarDecl tyname ann
goTyVarDecl TyVarDecl tyname ann
dataTyDecl)
            (TyVarDecl tyname ann -> TyVarDecl tyname ann
forall {unique} {tyname}.
(Coercible unique Int, HasUnique tyname unique) =>
TyVarDecl tyname ann -> TyVarDecl tyname ann
goTyVarDecl (TyVarDecl tyname ann -> TyVarDecl tyname ann)
-> [TyVarDecl tyname ann] -> [TyVarDecl tyname ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [TyVarDecl tyname ann]
paramTyDecls)
            name
matchName
            (VarDecl tyname name uni ann -> VarDecl tyname name uni ann
forall {unique} {name} {tyname} {uni :: * -> *}.
(Coercible unique Int, HasUnique name unique) =>
VarDecl tyname name uni ann -> VarDecl tyname name uni ann
goVarDecl (VarDecl tyname name uni ann -> VarDecl tyname name uni ann)
-> [VarDecl tyname name uni ann] -> [VarDecl tyname name uni ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [VarDecl tyname name uni ann]
constrDecls)

    -- Note that @goBind@ and @goTerm@ are mutually recursive.
    goBind :: Binding tyname name uni fun ann -> Binding tyname name uni fun ann
goBind (TermBind ann
ann Strictness
str VarDecl tyname name uni ann
var Term tyname name uni fun ann
term) = ann
-> Strictness
-> VarDecl tyname name uni ann
-> Term tyname name uni fun ann
-> Binding tyname name uni fun ann
forall tyname name (uni :: * -> *) fun a.
a
-> Strictness
-> VarDecl tyname name uni a
-> Term tyname name uni fun a
-> Binding tyname name uni fun a
TermBind ann
ann Strictness
str (VarDecl tyname name uni ann -> VarDecl tyname name uni ann
forall {unique} {name} {tyname} {uni :: * -> *}.
(Coercible unique Int, HasUnique name unique) =>
VarDecl tyname name uni ann -> VarDecl tyname name uni ann
goVarDecl VarDecl tyname name uni ann
var) (Term tyname name uni fun ann -> Binding tyname name uni fun ann)
-> Term tyname name uni fun ann -> Binding tyname name uni fun ann
forall a b. (a -> b) -> a -> b
$ Term tyname name uni fun ann -> Term tyname name uni fun ann
goTerm Term tyname name uni fun ann
term
    goBind (TypeBind ann
ann TyVarDecl tyname ann
tyVar Type tyname uni ann
ty)     = ann
-> TyVarDecl tyname ann
-> Type tyname uni ann
-> Binding tyname name uni fun ann
forall tyname name (uni :: * -> *) fun a.
a
-> TyVarDecl tyname a
-> Type tyname uni a
-> Binding tyname name uni fun a
TypeBind ann
ann (TyVarDecl tyname ann -> TyVarDecl tyname ann
forall {unique} {tyname}.
(Coercible unique Int, HasUnique tyname unique) =>
TyVarDecl tyname ann -> TyVarDecl tyname ann
goTyVarDecl TyVarDecl tyname ann
tyVar) Type tyname uni ann
ty
    goBind (DatatypeBind ann
ann Datatype tyname name uni ann
datatype) = ann
-> Datatype tyname name uni ann -> Binding tyname name uni fun ann
forall tyname name (uni :: * -> *) fun a.
a -> Datatype tyname name uni a -> Binding tyname name uni fun a
DatatypeBind ann
ann (Datatype tyname name uni ann -> Binding tyname name uni fun ann)
-> Datatype tyname name uni ann -> Binding tyname name uni fun ann
forall a b. (a -> b) -> a -> b
$ Datatype tyname name uni ann -> Datatype tyname name uni ann
forall {unique} {unique} {tyname} {name} {uni :: * -> *}.
(Coercible unique Int, Coercible unique Int,
 HasUnique tyname unique, HasUnique name unique) =>
Datatype tyname name uni ann -> Datatype tyname name uni ann
goDatatype Datatype tyname name uni ann
datatype

    goTerm :: Term tyname name uni fun ann -> Term tyname name uni fun ann
goTerm (Let ann
ann Recursivity
recy NonEmpty (Binding tyname name uni fun ann)
binds Term tyname name uni fun ann
term) = ann
-> Recursivity
-> NonEmpty (Binding tyname name uni fun ann)
-> Term tyname name uni fun ann
-> Term tyname name uni fun ann
forall tyname name (uni :: * -> *) fun a.
a
-> Recursivity
-> NonEmpty (Binding tyname name uni fun a)
-> Term tyname name uni fun a
-> Term tyname name uni fun a
Let ann
ann Recursivity
recy (Binding tyname name uni fun ann -> Binding tyname name uni fun ann
goBind (Binding tyname name uni fun ann
 -> Binding tyname name uni fun ann)
-> NonEmpty (Binding tyname name uni fun ann)
-> NonEmpty (Binding tyname name uni fun ann)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NonEmpty (Binding tyname name uni fun ann)
binds) (Term tyname name uni fun ann -> Term tyname name uni fun ann)
-> Term tyname name uni fun ann -> Term tyname name uni fun ann
forall a b. (a -> b) -> a -> b
$ Term tyname name uni fun ann -> Term tyname name uni fun ann
goTerm Term tyname name uni fun ann
term
    goTerm Term tyname name uni fun ann
term                      = Term tyname name uni fun ann
term Term tyname name uni fun ann
-> (Term tyname name uni fun ann -> Term tyname name uni fun ann)
-> Term tyname name uni fun ann
forall a b. a -> (a -> b) -> b
& (Term tyname name uni fun ann
 -> Identity (Term tyname name uni fun ann))
-> Term tyname name uni fun ann
-> Identity (Term tyname name uni fun ann)
forall tyname name (uni :: * -> *) fun a (f :: * -> *).
Applicative f =>
(Term tyname name uni fun a -> f (Term tyname name uni fun a))
-> Term tyname name uni fun a -> f (Term tyname name uni fun a)
termSubterms ((Term tyname name uni fun ann
  -> Identity (Term tyname name uni fun ann))
 -> Term tyname name uni fun ann
 -> Identity (Term tyname name uni fun ann))
-> (Term tyname name uni fun ann -> Term tyname name uni fun ann)
-> Term tyname name uni fun ann
-> Term tyname name uni fun ann
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ Term tyname name uni fun ann -> Term tyname name uni fun ann
goTerm

-- Ideally we should have a separate step putting uniques into annotations, so that we can reuse it
-- both here and for scoping analysis.
-- See Note [Retained size analysis]
-- | Annotate each part of every 'Binding' in a term with the size that it retains.
annotateWithRetainedSize
    :: (HasUnique name TermUnique, HasUnique tyname TypeUnique, ToBuiltinMeaning uni fun)
    => BuiltinsInfo uni fun
    -> Term tyname name uni fun ann
    -> Term tyname name uni fun RetainedSize
-- @reannotateBindings@ only processes annotations "associated with" a unique, so it can't change
-- the type. Therefore we need to set all the bindings to an appropriate type beforehand.
annotateWithRetainedSize :: forall name tyname (uni :: * -> *) fun ann.
(HasUnique name TermUnique, HasUnique tyname TypeUnique,
 ToBuiltinMeaning uni fun) =>
BuiltinsInfo uni fun
-> Term tyname name uni fun ann
-> Term tyname name uni fun RetainedSize
annotateWithRetainedSize BuiltinsInfo uni fun
binfo Term tyname name uni fun ann
term = (Unique -> RetainedSize -> RetainedSize)
-> Term tyname name uni fun RetainedSize
-> Term tyname name uni fun RetainedSize
forall name tyname ann (uni :: * -> *) fun.
(HasUnique name TermUnique, HasUnique tyname TypeUnique) =>
(Unique -> ann -> ann)
-> Term tyname name uni fun ann -> Term tyname name uni fun ann
reannotateBindings (Int -> RetainedSize -> RetainedSize
forall {p}. Int -> p -> RetainedSize
upd (Int -> RetainedSize -> RetainedSize)
-> (Unique -> Int) -> Unique -> RetainedSize -> RetainedSize
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Unique -> Int
unUnique) (Term tyname name uni fun RetainedSize
 -> Term tyname name uni fun RetainedSize)
-> Term tyname name uni fun RetainedSize
-> Term tyname name uni fun RetainedSize
forall a b. (a -> b) -> a -> b
$ RetainedSize
NotARetainer RetainedSize
-> Term tyname name uni fun ann
-> Term tyname name uni fun RetainedSize
forall a b.
a -> Term tyname name uni fun b -> Term tyname name uni fun a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Term tyname name uni fun ann
term where
    retentionMap :: IntMap AstSize
retentionMap = BuiltinsInfo uni fun
-> VarsInfo tyname name uni ann
-> Term tyname name uni fun ann
-> IntMap AstSize
forall tyname name (uni :: * -> *) fun ann.
(HasUnique tyname TypeUnique, HasUnique name TermUnique,
 ToBuiltinMeaning uni fun) =>
BuiltinsInfo uni fun
-> VarsInfo tyname name uni ann
-> Term tyname name uni fun ann
-> IntMap AstSize
termRetentionMap BuiltinsInfo uni fun
binfo VarsInfo tyname name uni ann
vinfo Term tyname name uni fun ann
term
    vinfo :: VarsInfo tyname name uni ann
vinfo = Term tyname name uni fun ann -> VarsInfo tyname name uni ann
forall name tyname (uni :: * -> *) fun a.
(HasUnique name TermUnique, HasUnique tyname TypeUnique) =>
Term tyname name uni fun a -> VarsInfo tyname name uni a
termVarInfo Term tyname name uni fun ann
term
    -- If a binding is not in the retention map, then it's still a retainer, just retains zero size.
    upd :: Int -> p -> RetainedSize
upd Int
i p
_ = AstSize -> RetainedSize
Retains (AstSize -> RetainedSize) -> AstSize -> RetainedSize
forall a b. (a -> b) -> a -> b
$ AstSize -> Int -> IntMap AstSize -> AstSize
forall a. a -> Int -> IntMap a -> a
IntMap.findWithDefault (Integer -> AstSize
AstSize Integer
0) Int
i IntMap AstSize
retentionMap