{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings     #-}

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

import PlutusPrelude

import PlutusIR.Analysis.Dependencies
import PlutusIR.Analysis.Size
import PlutusIR.Core

import PlutusCore qualified as PLC
import PlutusCore.Builtin (ToBuiltinMeaning)
import PlutusCore.Name

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

{- 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             = termSize 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 Size
    | 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
showList :: [RetainedSize] -> ShowS
$cshowList :: [RetainedSize] -> ShowS
show :: RetainedSize -> String
$cshow :: RetainedSize -> String
showsPrec :: Int -> RetainedSize -> ShowS
$cshowsPrec :: Int -> RetainedSize -> ShowS
Show)

instance Pretty RetainedSize where
    pretty :: RetainedSize -> Doc ann
pretty (Retains Size
size) = Doc ann
"$" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Size -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty Size
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 Size)

lookupSize :: Int -> DirectionRetentionMap -> Size
lookupSize :: Int -> DirectionRetentionMap -> Size
lookupSize Int
i (DirectionRetentionMap IntMap Size
ss) = IntMap Size
ss IntMap Size -> Int -> Size
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, Size)
annotateWithSizes :: DirectionRetentionMap -> Tree Int -> Tree (Int, Size)
annotateWithSizes DirectionRetentionMap
sizeInfo = Tree Int -> Tree (Int, Size)
go where
    go :: Tree Int -> Tree (Int, Size)
go (Node Int
i [Tree Int]
ts) = (Int, Size) -> [Tree (Int, Size)] -> Tree (Int, Size)
forall a. a -> [Tree a] -> Tree a
Node (Int
i, Size
sizeI) [Tree (Int, Size)]
rs where
        rs :: [Tree (Int, Size)]
rs = (Tree Int -> Tree (Int, Size)) -> [Tree Int] -> [Tree (Int, Size)]
forall a b. (a -> b) -> [a] -> [b]
map Tree Int -> Tree (Int, Size)
go [Tree Int]
ts
        sizeI :: Size
sizeI = Int -> DirectionRetentionMap -> Size
lookupSize Int
i DirectionRetentionMap
sizeInfo Size -> Size -> Size
forall a. Semigroup a => a -> a -> a
<> (Tree (Int, Size) -> Size) -> [Tree (Int, Size)] -> Size
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((Int, Size) -> Size
forall a b. (a, b) -> b
snd ((Int, Size) -> Size)
-> (Tree (Int, Size) -> (Int, Size)) -> Tree (Int, Size) -> Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree (Int, Size) -> (Int, Size)
forall a. Tree a -> a
rootLabel) [Tree (Int, Size)]
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 (IntMap IntSet -> Rooted)
-> (Graph Node -> IntMap IntSet) -> Graph Node -> Rooted
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Graph Int -> IntMap IntSet
forall t. (ToGraph t, ToVertex t ~ Int) => t -> IntMap IntSet
adjacencyIntMap (Graph Int -> IntMap IntSet)
-> (Graph Node -> Graph Int) -> Graph Node -> IntMap IntSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Node -> Int) -> Graph Node -> Graph Int
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 Size
depsRetentionMap :: DirectionRetentionMap -> Graph Node -> IntMap Size
depsRetentionMap DirectionRetentionMap
sizeInfo = [(Int, Size)] -> IntMap Size
forall a. [(Int, a)] -> IntMap a
IntMap.fromList ([(Int, Size)] -> IntMap Size)
-> (Graph Node -> [(Int, Size)]) -> Graph Node -> IntMap Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree (Int, Size) -> [(Int, Size)]
forall a. Tree a -> [a]
flatten (Tree (Int, Size) -> [(Int, Size)])
-> (Graph Node -> Tree (Int, Size)) -> Graph Node -> [(Int, Size)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DirectionRetentionMap -> Tree Int -> Tree (Int, Size)
annotateWithSizes DirectionRetentionMap
sizeInfo (Tree Int -> Tree (Int, Size))
-> (Graph Node -> Tree Int) -> Graph Node -> Tree (Int, Size)
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 -> UniqueMap Unique Size
bindingSize :: Binding tyname name uni fun ann -> UniqueMap Unique Size
bindingSize (TermBind ann
_ Strictness
_ VarDecl tyname name uni fun ann
var Term tyname name uni fun ann
term) =
    VarDecl tyname name uni fun ann
-> Size -> UniqueMap Unique Size -> UniqueMap Unique Size
forall name unique1 unique2 a.
(HasUnique name unique1, Coercible unique2 Unique) =>
name -> a -> UniqueMap unique2 a -> UniqueMap unique2 a
insertByNameIndex VarDecl tyname name uni fun ann
var (VarDecl tyname name uni fun ann -> Size
forall tyname name (uni :: * -> *) fun ann.
VarDecl tyname name uni fun ann -> Size
varDeclSize VarDecl tyname name uni fun ann
var Size -> Size -> Size
forall a. Semigroup a => a -> a -> a
<> Term tyname name uni fun ann -> Size
forall tyname name (uni :: * -> *) fun ann.
Term tyname name uni fun ann -> Size
termSize Term tyname name uni fun ann
term) UniqueMap Unique Size
forall a. Monoid a => a
mempty
bindingSize (TypeBind ann
_ TyVarDecl tyname ann
tyVar Type tyname uni ann
ty) =
    TyVarDecl tyname ann
-> Size -> UniqueMap Unique Size -> UniqueMap Unique Size
forall name unique1 unique2 a.
(HasUnique name unique1, Coercible unique2 Unique) =>
name -> a -> UniqueMap unique2 a -> UniqueMap unique2 a
insertByNameIndex TyVarDecl tyname ann
tyVar (TyVarDecl tyname ann -> Size
forall tyname ann. TyVarDecl tyname ann -> Size
tyVarDeclSize TyVarDecl tyname ann
tyVar Size -> Size -> Size
forall a. Semigroup a => a -> a -> a
<> Type tyname uni ann -> Size
forall tyname (uni :: * -> *) ann. Type tyname uni ann -> Size
typeSize Type tyname uni ann
ty) UniqueMap Unique Size
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 fun ann]
constrs))
    = TyVarDecl tyname ann
-> Size -> UniqueMap Unique Size -> UniqueMap Unique Size
forall name unique1 unique2 a.
(HasUnique name unique1, Coercible unique2 Unique) =>
name -> a -> UniqueMap unique2 a -> UniqueMap unique2 a
insertByNameIndex TyVarDecl tyname ann
dataDecl (TyVarDecl tyname ann -> Size
forall tyname ann. TyVarDecl tyname ann -> Size
tyVarDeclSize TyVarDecl tyname ann
dataDecl)
    (UniqueMap Unique Size -> UniqueMap Unique Size)
-> (UniqueMap Unique Size -> UniqueMap Unique Size)
-> UniqueMap Unique Size
-> UniqueMap Unique Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (UniqueMap Unique Size
 -> [TyVarDecl tyname ann] -> UniqueMap Unique Size)
-> [TyVarDecl tyname ann]
-> UniqueMap Unique Size
-> UniqueMap Unique Size
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((TyVarDecl tyname ann
 -> UniqueMap Unique Size -> UniqueMap Unique Size)
-> UniqueMap Unique Size
-> [TyVarDecl tyname ann]
-> UniqueMap Unique Size
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((TyVarDecl tyname ann
  -> UniqueMap Unique Size -> UniqueMap Unique Size)
 -> UniqueMap Unique Size
 -> [TyVarDecl tyname ann]
 -> UniqueMap Unique Size)
-> (TyVarDecl tyname ann
    -> UniqueMap Unique Size -> UniqueMap Unique Size)
-> UniqueMap Unique Size
-> [TyVarDecl tyname ann]
-> UniqueMap Unique Size
forall a b. (a -> b) -> a -> b
$ \TyVarDecl tyname ann
param -> TyVarDecl tyname ann
-> Size -> UniqueMap Unique Size -> UniqueMap Unique Size
forall name unique1 unique2 a.
(HasUnique name unique1, Coercible unique2 Unique) =>
name -> a -> UniqueMap unique2 a -> UniqueMap unique2 a
insertByNameIndex TyVarDecl tyname ann
param (Size -> UniqueMap Unique Size -> UniqueMap Unique Size)
-> Size -> UniqueMap Unique Size -> UniqueMap Unique Size
forall a b. (a -> b) -> a -> b
$ TyVarDecl tyname ann -> Size
forall tyname ann. TyVarDecl tyname ann -> Size
tyVarDeclSize TyVarDecl tyname ann
param) [TyVarDecl tyname ann]
params
    (UniqueMap Unique Size -> UniqueMap Unique Size)
-> (UniqueMap Unique Size -> UniqueMap Unique Size)
-> UniqueMap Unique Size
-> UniqueMap Unique Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. name -> Size -> UniqueMap Unique Size -> UniqueMap Unique Size
forall name unique1 unique2 a.
(HasUnique name unique1, Coercible unique2 Unique) =>
name -> a -> UniqueMap unique2 a -> UniqueMap unique2 a
insertByNameIndex name
matchName (Integer -> Size
Size Integer
1)
    (UniqueMap Unique Size -> UniqueMap Unique Size)
-> (UniqueMap Unique Size -> UniqueMap Unique Size)
-> UniqueMap Unique Size
-> UniqueMap Unique Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (UniqueMap Unique Size
 -> [VarDecl tyname name uni fun ann] -> UniqueMap Unique Size)
-> [VarDecl tyname name uni fun ann]
-> UniqueMap Unique Size
-> UniqueMap Unique Size
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((VarDecl tyname name uni fun ann
 -> UniqueMap Unique Size -> UniqueMap Unique Size)
-> UniqueMap Unique Size
-> [VarDecl tyname name uni fun ann]
-> UniqueMap Unique Size
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((VarDecl tyname name uni fun ann
  -> UniqueMap Unique Size -> UniqueMap Unique Size)
 -> UniqueMap Unique Size
 -> [VarDecl tyname name uni fun ann]
 -> UniqueMap Unique Size)
-> (VarDecl tyname name uni fun ann
    -> UniqueMap Unique Size -> UniqueMap Unique Size)
-> UniqueMap Unique Size
-> [VarDecl tyname name uni fun ann]
-> UniqueMap Unique Size
forall a b. (a -> b) -> a -> b
$ \VarDecl tyname name uni fun ann
constr -> VarDecl tyname name uni fun ann
-> Size -> UniqueMap Unique Size -> UniqueMap Unique Size
forall name unique1 unique2 a.
(HasUnique name unique1, Coercible unique2 Unique) =>
name -> a -> UniqueMap unique2 a -> UniqueMap unique2 a
insertByNameIndex VarDecl tyname name uni fun ann
constr (Size -> UniqueMap Unique Size -> UniqueMap Unique Size)
-> Size -> UniqueMap Unique Size -> UniqueMap Unique Size
forall a b. (a -> b) -> a -> b
$ VarDecl tyname name uni fun ann -> Size
forall tyname name (uni :: * -> *) fun ann.
VarDecl tyname name uni fun ann -> Size
varDeclSize VarDecl tyname name uni fun ann
constr) [VarDecl tyname name uni fun ann]
constrs
    (UniqueMap Unique Size -> UniqueMap Unique Size)
-> UniqueMap Unique Size -> UniqueMap Unique Size
forall a b. (a -> b) -> a -> b
$ UniqueMap Unique Size
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 -> UniqueMap Unique Size
bindingSizes :: Term tyname name uni fun ann -> UniqueMap Unique Size
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 Size)
-> NonEmpty (Binding tyname name uni fun ann)
-> UniqueMap Unique Size
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Binding tyname name uni fun ann -> UniqueMap Unique Size
forall tyname name (uni :: * -> *) fun ann.
(HasUnique tyname TypeUnique, HasUnique name TermUnique) =>
Binding tyname name uni fun ann -> UniqueMap Unique Size
bindingSize NonEmpty (Binding tyname name uni fun ann)
binds UniqueMap Unique Size
-> UniqueMap Unique Size -> UniqueMap Unique Size
forall a. Semigroup a => a -> a -> a
<> Term tyname name uni fun ann -> UniqueMap Unique Size
forall tyname name (uni :: * -> *) fun ann.
(HasUnique tyname TypeUnique, HasUnique name TermUnique) =>
Term tyname name uni fun ann -> UniqueMap Unique Size
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 Size)
     (Term tyname name uni fun ann)
     (UniqueMap Unique Size)
-> UniqueMap Unique Size
forall s a. s -> Getting a s a -> a
^. (Term tyname name uni fun ann
 -> Const (UniqueMap Unique Size) (Term tyname name uni fun ann))
-> Term tyname name uni fun ann
-> Const (UniqueMap Unique Size) (Term tyname name uni fun ann)
forall tyname name (uni :: * -> *) fun a.
Traversal'
  (Term tyname name uni fun a) (Term tyname name uni fun a)
termSubterms ((Term tyname name uni fun ann
  -> Const (UniqueMap Unique Size) (Term tyname name uni fun ann))
 -> Term tyname name uni fun ann
 -> Const (UniqueMap Unique Size) (Term tyname name uni fun ann))
-> Getting
     (UniqueMap Unique Size)
     (Term tyname name uni fun ann)
     (UniqueMap Unique Size)
-> Getting
     (UniqueMap Unique Size)
     (Term tyname name uni fun ann)
     (UniqueMap Unique Size)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Term tyname name uni fun ann -> UniqueMap Unique Size)
-> Getting
     (UniqueMap Unique Size)
     (Term tyname name uni fun ann)
     (UniqueMap Unique Size)
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 Size
forall tyname name (uni :: * -> *) fun ann.
(HasUnique tyname TypeUnique, HasUnique name TermUnique) =>
Term tyname name uni fun ann -> UniqueMap Unique Size
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 :: Term tyname name uni fun ann -> DirectionRetentionMap
toDirectionRetentionMap Term tyname name uni fun ann
term =
    IntMap Size -> DirectionRetentionMap
DirectionRetentionMap (IntMap Size -> DirectionRetentionMap)
-> (UniqueMap Unique Size -> IntMap Size)
-> UniqueMap Unique Size
-> DirectionRetentionMap
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Size -> IntMap Size -> IntMap Size
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
rootInt Size
rootSize (IntMap Size -> IntMap Size)
-> (UniqueMap Unique Size -> IntMap Size)
-> UniqueMap Unique Size
-> IntMap Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UniqueMap Unique Size -> IntMap Size
forall unique a. UniqueMap unique a -> IntMap a
unUniqueMap (UniqueMap Unique Size -> DirectionRetentionMap)
-> UniqueMap Unique Size -> DirectionRetentionMap
forall a b. (a -> b) -> a -> b
$ Term tyname name uni fun ann -> UniqueMap Unique Size
forall tyname name (uni :: * -> *) fun ann.
(HasUnique tyname TypeUnique, HasUnique name TermUnique) =>
Term tyname name uni fun ann -> UniqueMap Unique Size
bindingSizes Term tyname name uni fun ann
term where
        -- See Note [Handling the root].
        rootSize :: Size
rootSize = Integer -> Size
Size (- 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 Size
ss) (Variable (PLC.Unique Int
i)) = Int
i Int -> IntMap Size -> Bool
forall a. Int -> IntMap a -> Bool
`IntMap.member` IntMap Size
ss

-- | Compute the retention map of a term.
termRetentionMap
    :: (HasUnique tyname TypeUnique, HasUnique name TermUnique, ToBuiltinMeaning uni fun)
    => Term tyname name uni fun ann -> IntMap Size
termRetentionMap :: Term tyname name uni fun ann -> IntMap Size
termRetentionMap Term tyname name uni fun ann
term = DirectionRetentionMap -> Graph Node -> IntMap Size
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
$ (Graph Node, StrictnessMap) -> Graph Node
forall a b. (a, b) -> a
fst ((Graph Node, StrictnessMap) -> Graph Node)
-> (Graph Node, StrictnessMap) -> Graph Node
forall a b. (a -> b) -> a -> b
$ Term tyname name uni fun ann -> (Graph Node, StrictnessMap)
forall g tyname name (uni :: * -> *) fun a.
(DepGraph g, HasUnique tyname TypeUnique,
 HasUnique name TermUnique, ToBuiltinMeaning uni fun) =>
Term tyname name uni fun a -> (g, StrictnessMap)
runTermDeps 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 :: (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 tyname unique name unique (uni :: * -> *) fun.
(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 fun ann -> VarDecl tyname name uni fun ann
goVarDecl (VarDecl ann
ann name
name Type tyname uni ann
ty) = ann
-> name -> Type tyname uni ann -> VarDecl tyname name uni fun ann
forall k tyname name (uni :: * -> *) (fun :: k) ann.
ann
-> name -> Type tyname uni ann -> VarDecl tyname name uni fun 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
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
theUnique) ann
ann) tyname
tyname Kind ann
kind
    goDatatype :: Datatype tyname name uni fun ann
-> Datatype tyname name uni fun ann
goDatatype (Datatype ann
ann TyVarDecl tyname ann
dataTyDecl [TyVarDecl tyname ann]
paramTyDecls name
matchName [VarDecl tyname name uni fun ann]
constrDecls) =
        ann
-> TyVarDecl tyname ann
-> [TyVarDecl tyname ann]
-> name
-> [VarDecl tyname name uni fun ann]
-> Datatype tyname name uni fun ann
forall tyname name (uni :: * -> *) fun a.
a
-> TyVarDecl tyname a
-> [TyVarDecl tyname a]
-> name
-> [VarDecl tyname name uni fun a]
-> Datatype tyname name uni fun 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
theUnique) ann
ann)
            (TyVarDecl tyname ann -> TyVarDecl tyname ann
forall tyname unique.
HasUnique tyname unique =>
TyVarDecl tyname ann -> TyVarDecl tyname ann
goTyVarDecl TyVarDecl tyname ann
dataTyDecl)
            (TyVarDecl tyname ann -> TyVarDecl tyname ann
forall tyname unique.
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 fun ann -> VarDecl tyname name uni fun ann
forall name unique tyname (uni :: * -> *) fun fun.
HasUnique name unique =>
VarDecl tyname name uni fun ann -> VarDecl tyname name uni fun ann
goVarDecl (VarDecl tyname name uni fun ann
 -> VarDecl tyname name uni fun ann)
-> [VarDecl tyname name uni fun ann]
-> [VarDecl tyname name uni fun ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [VarDecl tyname name uni fun 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 fun ann
var Term tyname name uni fun ann
term) = ann
-> Strictness
-> VarDecl tyname name uni fun 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 fun a
-> Term tyname name uni fun a
-> Binding tyname name uni fun a
TermBind ann
ann Strictness
str (VarDecl tyname name uni fun ann -> VarDecl tyname name uni fun ann
forall name unique tyname (uni :: * -> *) fun fun.
HasUnique name unique =>
VarDecl tyname name uni fun ann -> VarDecl tyname name uni fun ann
goVarDecl VarDecl tyname name uni fun 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 tyname unique.
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 fun ann
datatype) = ann
-> Datatype tyname name uni fun ann
-> Binding tyname name uni fun ann
forall tyname name (uni :: * -> *) fun a.
a
-> Datatype tyname name uni fun a -> Binding tyname name uni fun a
DatatypeBind ann
ann (Datatype tyname name uni fun ann
 -> Binding tyname name uni fun ann)
-> Datatype tyname name uni fun ann
-> Binding tyname name uni fun ann
forall a b. (a -> b) -> a -> b
$ Datatype tyname name uni fun ann
-> Datatype tyname name uni fun ann
forall name unique tyname unique (uni :: * -> *) fun fun.
(HasUnique name unique, HasUnique tyname unique) =>
Datatype tyname name uni fun ann
-> Datatype tyname name uni fun ann
goDatatype Datatype tyname name uni fun 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.
Traversal'
  (Term tyname name uni fun a) (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)
    => 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 :: Term tyname name uni fun ann
-> Term tyname name uni fun RetainedSize
annotateWithRetainedSize 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 (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Term tyname name uni fun ann
term where
    retentionMap :: IntMap Size
retentionMap = Term tyname name uni fun ann -> IntMap Size
forall tyname name (uni :: * -> *) fun ann.
(HasUnique tyname TypeUnique, HasUnique name TermUnique,
 ToBuiltinMeaning uni fun) =>
Term tyname name uni fun ann -> IntMap Size
termRetentionMap 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
_ = Size -> RetainedSize
Retains (Size -> RetainedSize) -> Size -> RetainedSize
forall a b. (a -> b) -> a -> b
$ Size -> Int -> IntMap Size -> Size
forall a. a -> Int -> IntMap a -> a
IntMap.findWithDefault (Integer -> Size
Size Integer
0) Int
i IntMap Size
retentionMap