-- editorconfig-checker-disable-file
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE MagicHash            #-}
{-# LANGUAGE TypeApplications     #-}
{-# LANGUAGE TypeOperators        #-}
{-# LANGUAGE UndecidableInstances #-}

module PlutusCore.Evaluation.Machine.ExMemoryUsage
    ( CostRose(..)
    , singletonRose
    , ExMemoryUsage(..)
    , flattenCostRose
    , NumBytesCostedAsNumWords(..)
    , IntegerCostedLiterally(..)
    ) where

import PlutusCore.Crypto.BLS12_381.G1 as BLS12_381.G1
import PlutusCore.Crypto.BLS12_381.G2 as BLS12_381.G2
import PlutusCore.Crypto.BLS12_381.Pairing as BLS12_381.Pairing
import PlutusCore.Data
import PlutusCore.Evaluation.Machine.CostStream
import PlutusCore.Evaluation.Machine.ExMemory

import Data.ByteString qualified as BS
import Data.Functor
import Data.Proxy
import Data.SatInt
import Data.Text qualified as T
import Data.Vector.Strict (Vector)
import Data.Vector.Strict qualified as Vector
import Data.Word
import GHC.Exts (Int (I#))
import GHC.Integer
import GHC.Integer.Logarithms
import GHC.Natural
import GHC.Prim
import Universe

{-
 ************************************************************************************
 *  WARNING: exercise caution when altering the ExMemoryUsage instances here.       *
 *                                                                                  *
 *  The instances defined in this file will be used to calculate script validation  *
 *  costs, and if an instance is changed then any scripts which were deployed when  *
 *  a previous instance was in effect MUST STILL VALIDATE using the new instance.   *
 *  It is unsafe to increase the memory usage of a type because that may increase   *
 *  the resource usage of existing scripts beyond the limits set (and paid for)     *
 *  when they were uploaded to the chain, but because our costing functions are all *
 *  monotone it is safe to decrease memory usage, as long it decreases for *all*    *
 *  possible values of the type.                                                    *
 ************************************************************************************
-}

{- Note [ExMemoryUsage instances for non-constants]
In order to calculate the cost of a built-in function we need to feed the 'ExMemory' of each
argument to the costing function associated with the builtin. For a polymorphic builtin this means
that we need to be able to compute the 'ExMemory' of the AST provided as an argument to the builtin.
How do we do that? Our strategy is:

1. if the AST is a wrapped constant, then calculate the 'ExMemory' of the constant
2. if the AST is something else, return 1

This is pretty reasonable: a polymorphic builtin *is* allowed to check if the AST that it got as an
argument is a constant or not, and if it happens to be a constant, the builtin *is* allowed to use
it whatever way it wishes (see Note [Builtins and Plutus type checking] for details). Hence a
builtin may in fact do something ad hoc for constants and we need to account for this possibility in
the costing machinery.

But if the given AST is not a constant, the builtin can't do anything else with it, hence we simply
return 1, meaning "the costing function can't use this 'ExMemory' in any non-vacuous way".

See 'HasMeaningIn' for a full list of constraints determining what a builtin can do with values.

However for all types of values, except the one used by the production evaluator, we implement
'ExMemoryUsage' as a call to 'error'. Not because other evaluators don't compute costs during
evaluation -- the CK machine for example does in fact compute them (because we share the same
builtins machinery between all the evaluators and we want it to be efficient on the production path,
hence it's easier to optimize it for all evaluators than just for the single production evaluator).
And not because the resulting 'ExBudget' is not forced by an evaluator that doesn't care about
costing -- it still gets forced (for the same reason).

The actual reason why we call 'error' is because at the moment no builtin is supposed to have a
costing function that actually computes the 'ExMemory' of the given AST. Currently, if the builtin
takes an 'Opaque', it's not supposed to actually look inside of it (unlike with 'SomeConstant') and
hence the costing function is supposed to ignore that argument. It is possible that we'll eventually
decide to add such a builtin, so the current approach of throwing an 'error' is a precaution
ensuring that we won't add any weirdness by accident.

We don't call 'error' on the production path, because we don't want this risk in there. A failing
test is fine, a failing reasonable transaction is not and we don't want to risk it, even if it seems
very unlikely that such a failure could slip in.

The way we ignore arguments in costing functions is by computing the 'ExMemory' of each of those
arguments lazily. I.e. a call to 'memoryUsage' can only be forced within a costing function and
never outside of one. We have to do this regardless of all the reasoning above: if we compute
the 'ExMemory' of, say, a list strictly, then a builtin prepending an element to a list will
have the complexity of O(length_of_the_list) (because computing the 'ExMemory' of a list requires
traversing the list), while we of course want it to be O(1).
-}

-- | A lazy tree of costs. Convenient for calculating the costs of values of built-in types, because
-- they may have arbitrary branching (in particular a 'Data' object can contain a list of 'Data'
-- objects inside of it).
--
-- 'CostRose' gets collapsed to a lazy linear structure down the pipeline, so that we can
-- stream the costs to the outside where, say, the CEK machine picks them up one by one and handles
-- somehow (in particular, subtracts from the remaining budget).
data CostRose = CostRose {-# UNPACK #-} !CostingInteger ![CostRose]
    deriving stock (Int -> CostRose -> ShowS
[CostRose] -> ShowS
CostRose -> String
(Int -> CostRose -> ShowS)
-> (CostRose -> String) -> ([CostRose] -> ShowS) -> Show CostRose
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> CostRose -> ShowS
showsPrec :: Int -> CostRose -> ShowS
$cshow :: CostRose -> String
show :: CostRose -> String
$cshowList :: [CostRose] -> ShowS
showList :: [CostRose] -> ShowS
Show)

-- | Create a 'CostRose' containing a single cost.
singletonRose :: CostingInteger -> CostRose
singletonRose :: CostingInteger -> CostRose
singletonRose CostingInteger
cost = CostingInteger -> [CostRose] -> CostRose
CostRose CostingInteger
cost []
{-# INLINE singletonRose #-}

-- See Note [Global local functions].
-- This is one way to define the worker. There are many more, see
-- https://github.com/IntersectMBO/plutus/pull/5239#discussion_r1151197471
-- We chose this one, because it's the simplest (no CPS shenanigans) among the safest (retrieving
-- the next element takes O(1) time in the worst case).
--
-- The algorithm is a variation of the defunctionalization technique (see this post in particular:
-- https://www.joachim-breitner.de/blog/778-Don%e2%80%99t_think,_just_defunctionalize), except we
-- don't want a tail-recursive loop and instead emit costs lazily to the outside (as it's the whole
-- point of the lazy costing approach)
flattenCostRoseGo :: CostRose -> [CostRose] -> CostStream
flattenCostRoseGo :: CostRose -> [CostRose] -> CostStream
flattenCostRoseGo (CostRose CostingInteger
cost1 [CostRose]
forest1) [CostRose]
forest2 =
    case [CostRose]
forest1 of
        -- The current subtree doesn't have its own subtrees.
        [] -> case [CostRose]
forest2 of
            -- No more elements in the entire tree, emit the last cost.
            []                -> CostingInteger -> CostStream
CostLast CostingInteger
cost1
            -- There's at least one unhandled subtree encountered earlier, emit the current cost
            -- and collapse the unhandled subtree.
            CostRose
rose2' : [CostRose]
forest2' -> CostingInteger -> CostStream -> CostStream
CostCons CostingInteger
cost1 (CostStream -> CostStream) -> CostStream -> CostStream
forall a b. (a -> b) -> a -> b
$ CostRose -> [CostRose] -> CostStream
flattenCostRoseGo CostRose
rose2' [CostRose]
forest2'
        -- The current subtree has at least one its own subtree.
        CostRose
rose1' : [CostRose]
forest1' ->
            -- Emit the current cost and handle the list of subtrees of the current subtree.
            CostingInteger -> CostStream -> CostStream
CostCons CostingInteger
cost1 (CostStream -> CostStream) -> CostStream -> CostStream
forall a b. (a -> b) -> a -> b
$ case [CostRose]
forest1' of
                -- Same as the case below, except this one avoids creating a chain of
                -- @[] ++ ([] ++ ([] ++ <...>))@, which would make retrieving the next element an
                -- O(depth_of_the_tree) operation in the worst case.
                [] -> CostRose -> [CostRose] -> CostStream
flattenCostRoseGo CostRose
rose1' [CostRose]
forest2
                -- Add the remaining subtrees of the current subtree (@forest1'@) to the stack of
                -- all other subtrees (@forest2@) and handle the new current subtree (@rose1'@).
                [CostRose]
_  -> CostRose -> [CostRose] -> CostStream
flattenCostRoseGo CostRose
rose1' ([CostRose] -> CostStream) -> [CostRose] -> CostStream
forall a b. (a -> b) -> a -> b
$ [CostRose]
forest1' [CostRose] -> [CostRose] -> [CostRose]
forall a. [a] -> [a] -> [a]
++ [CostRose]
forest2

-- | Collapse a 'CostRose' to a lazy linear stream of costs. Retrieving the next element takes O(1)
-- time in the worst case regardless of the recursion pattern of the given 'CostRose'.
flattenCostRose :: CostRose -> CostStream
flattenCostRose :: CostRose -> CostStream
flattenCostRose (CostRose CostingInteger
cost [])              = CostingInteger -> CostStream
CostLast CostingInteger
cost
flattenCostRose (CostRose CostingInteger
cost (CostRose
rose : [CostRose]
forest)) = CostingInteger -> CostStream -> CostStream
CostCons CostingInteger
cost (CostStream -> CostStream) -> CostStream -> CostStream
forall a b. (a -> b) -> a -> b
$ CostRose -> [CostRose] -> CostStream
flattenCostRoseGo CostRose
rose [CostRose]
forest
{-# INLINE flattenCostRose #-}

class ExMemoryUsage a where
    -- Inlining the implementations of this method gave us a 1-2% speedup.
    memoryUsage :: a -> CostRose

{- Note [Alternative memory usage instances].  The `memoryUsage` function provides
 a measure of the size of an object for costing purposes, the idea being that
 the time taken to execute a builtin, and the memory used to contain its result,
 will depend on the size of the inputs.  The name `memoryUsage` is perhaps a
 misnomer: it was originally supposed to measure (in 64-bit words) the heap
 space required to store an object, but this is not always the correct measure
 to use.  For example, the time taken by `AddInteger` or `MultiplyInteger` will
 depend on the logarithms of the inputs (and the logarithm is proportional to
 the memory occupied by the inputs), and the memory occupied by the result will
 be some function of the memory occupied by the inputs, so for these functions
 the actual memory usage is a sensible size measure.  However, calling
 `replicateByte n b` function allocates a number of bytes which is equal to the
 actual value of `n`, which will be exponentially greater than the memory
 occupied by `n`, so this case the memory usage is not a sensible size measure.
 In most cases the default `memoryUsage` function returns the actual memory
 usage, but to deal with cases like `replicateByte` we occasionally use newtype
 wrappers which define a different size measure (see `IntegerCostedLiterally`
 below).  Polymorphic types require some care though: see Note [ExMemoryUsage
 for polymorphic types].
-}

{- Note [ExMemoryUsage for polymorphic types].  For polymorphic types such as
 `pair, `list`, and `array` DO NOT use newtype wrappers to define alternative
 size measures.  The denotations of functions which take polymorphic arguments
 use `SomeConstant` and this will ignore newtype wrappers and will only use the
 default `memoryUsage` function, which could lead to unexpected results.
 Furthermore, actual memory usage is typically not a good size measure for
 polymorphic arguments: the time taken to process a list, for example, will
 typically depend only on the length of the list and not the size of the
 contents.  Currently all such functions are parametrically polymorphic and only
 manipulate pointers without inspecting the contents of their polymorphic
 arguments, so it is reasonable to use size measures which depend only on the
 surface structure of polymorphic objects. -}

{- We expect that all builtins which involve pairs will be constant cost and so
   their memory usage will never be involved in any computations.  The memory
   usage is set to maxBound so that we'll notice if this assumption is ever
   violated -}
instance ExMemoryUsage (a, b) where
    memoryUsage :: (a, b) -> CostRose
memoryUsage (a, b)
_ = CostingInteger -> CostRose
singletonRose CostingInteger
forall a. Bounded a => a
maxBound
    {-# INLINE memoryUsage #-}

{- Note the the `memoryUsage` of an empty list is zero.  This shouldn't cause any
 problems, but be sure to check that no costing function involving lists can
 return zero for an empty list (or any other input).
-}
{- Calculating the memory usage by processing the entire spine of the list eagerly
 is safe because there's no way to cheaply construct a long list: you either
 make one using repeated mkCons, which is expensive, or return one from a
 builtin, which has to be appropriately expensive too. -}
instance ExMemoryUsage [a] where
  memoryUsage :: [a] -> CostRose
memoryUsage [a]
l = CostingInteger -> CostRose
singletonRose (CostingInteger -> CostRose)
-> (Int -> CostingInteger) -> Int -> CostRose
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> CostingInteger
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> CostRose) -> Int -> CostRose
forall a b. (a -> b) -> a -> b
$ [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
l
  {-# INLINE memoryUsage #-}

{- Note the the `memoryUsage` of an empty array is zero.  This shouldn't cause any
 problems, but be sure to check that no costing function involving arrays can
 return zero for an empty array (or any other input).
-}
instance ExMemoryUsage (Vector a) where
    memoryUsage :: Vector a -> CostRose
memoryUsage Vector a
l = CostingInteger -> CostRose
singletonRose (CostingInteger -> CostRose)
-> (Int -> CostingInteger) -> Int -> CostRose
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> CostingInteger
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> CostRose) -> Int -> CostRose
forall a b. (a -> b) -> a -> b
$ Vector a -> Int
forall a. Vector a -> Int
Vector.length Vector a
l
    {-# INLINE memoryUsage #-}

instance (Closed uni, uni `Everywhere` ExMemoryUsage) => ExMemoryUsage (Some (ValueOf uni)) where
    memoryUsage :: Some (ValueOf uni) -> CostRose
memoryUsage (Some (ValueOf uni (Esc a)
uni a
x)) = Proxy ExMemoryUsage
-> uni (Esc a) -> (ExMemoryUsage a => CostRose) -> CostRose
forall (uni :: * -> *) (constr :: * -> Constraint)
       (proxy :: (* -> Constraint) -> *) a r.
(Closed uni, Everywhere uni constr) =>
proxy constr -> uni (Esc a) -> (constr a => r) -> r
forall (constr :: * -> Constraint)
       (proxy :: (* -> Constraint) -> *) a r.
Everywhere uni constr =>
proxy constr -> uni (Esc a) -> (constr a => r) -> r
bring (forall {k} (t :: k). Proxy t
forall (t :: * -> Constraint). Proxy t
Proxy @ExMemoryUsage) uni (Esc a)
uni (a -> CostRose
forall a. ExMemoryUsage a => a -> CostRose
memoryUsage a
x)
    {-# INLINE memoryUsage #-}

instance ExMemoryUsage () where
    memoryUsage :: () -> CostRose
memoryUsage () = CostingInteger -> CostRose
singletonRose CostingInteger
1
    {-# INLINE memoryUsage #-}

{- | When invoking a built-in function, a value of type `NumBytesCostedAsNumWords`
   can be used transparently as a built-in Integer but with a different size
   measure: see Note [Integral types as Integer].  This is required by the
   `integerToByteString` builtin, which takes an argument `w` specifying the
   width (in bytes) of the output bytestring (zero-padded to the desired size).
   The memory consumed by the function is given by `w`, *not* the size of `w`.
   The `NumBytesCostedAsNumWords` type wraps an Int `w` in a newtype whose
   `ExMemoryUsage` is equal to the number of eight-byte words required to
   contain `w` bytes, allowing its costing function to work properly.  We also
   use this for `replicateByte`.  If this is used to wrap an argument in the
   denotation of a builtin then it *MUST* also be used to wrap the same argument
   in the relevant budgeting benchmark.
-}
newtype NumBytesCostedAsNumWords = NumBytesCostedAsNumWords { NumBytesCostedAsNumWords -> Integer
unNumBytesCostedAsNumWords :: Integer }
instance ExMemoryUsage NumBytesCostedAsNumWords where
    memoryUsage :: NumBytesCostedAsNumWords -> CostRose
memoryUsage (NumBytesCostedAsNumWords Integer
n) = CostingInteger -> CostRose
singletonRose (CostingInteger -> CostRose)
-> (Integer -> CostingInteger) -> Integer -> CostRose
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> CostingInteger
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> CostRose) -> Integer -> CostRose
forall a b. (a -> b) -> a -> b
$ ((Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
8) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1
    {-# INLINE memoryUsage #-}
    -- Note that this uses `fromIntegral`, which will narrow large values to
    -- maxBound::SatInt = 2^63-1.  This shouldn't be a problem for costing because no
    -- realistic input should be that large; however if you're going to use this then be
    -- sure to convince yourself that it's safe.

{- | A wrapper for `Integer`s whose "memory usage" for costing purposes is the
   absolute value of the `Integer`.  This is used for costing built-in functions
   such as `shiftByteString` and `rotateByteString`, where the cost may depend
   on the actual value of the shift argument, not its size.  If this is used to
   wrap an argument in the denotation of a builtin then it *MUST* also be used
   to wrap the same argument in the relevant budgeting benchmark.
-}
newtype IntegerCostedLiterally = IntegerCostedLiterally { IntegerCostedLiterally -> Integer
unIntegerCostedLiterally :: Integer }
instance ExMemoryUsage IntegerCostedLiterally where
    memoryUsage :: IntegerCostedLiterally -> CostRose
memoryUsage (IntegerCostedLiterally Integer
n) = CostingInteger -> CostRose
singletonRose (CostingInteger -> CostRose)
-> (Integer -> CostingInteger) -> Integer -> CostRose
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> CostingInteger
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> CostRose) -> Integer -> CostRose
forall a b. (a -> b) -> a -> b
$ Integer -> Integer
forall a. Num a => a -> a
abs Integer
n
    {-# INLINE memoryUsage #-}
    -- Note that this uses `fromIntegral`, which will narrow large values to
    -- maxBound::SatInt = 2^63-1.  This shouldn't be a problem for costing because no
    -- realistic input should be that large; however if you're going to use this then be
    -- sure to convince yourself that it's safe.

-- | Calculate a 'CostingInteger' for the given 'Integer'.
memoryUsageInteger :: Integer -> CostingInteger
-- integerLog2# is unspecified for 0 (but in practice returns -1)
-- ^ This changed with GHC 9.2: it now returns 0.  It's probably safest if we
-- keep this special case for the time being though.
memoryUsageInteger :: Integer -> CostingInteger
memoryUsageInteger Integer
0 = CostingInteger
1
-- Assume 64 Int
memoryUsageInteger Integer
i = Int -> CostingInteger
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> CostingInteger) -> Int -> CostingInteger
forall a b. (a -> b) -> a -> b
$ Int# -> Int
I# (Integer -> Int#
integerLog2# (Integer -> Integer
forall a. Num a => a -> a
abs Integer
i) Int# -> Int# -> Int#
`quotInt#` Integer -> Int#
integerToInt Integer
64) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
-- So that the produced GHC Core doesn't explode in size, we don't win anything by inlining this
-- function anyway.
{-# NOINLINE memoryUsageInteger #-}

instance ExMemoryUsage Integer where
    memoryUsage :: Integer -> CostRose
memoryUsage Integer
i = CostingInteger -> CostRose
singletonRose (CostingInteger -> CostRose) -> CostingInteger -> CostRose
forall a b. (a -> b) -> a -> b
$ Integer -> CostingInteger
memoryUsageInteger Integer
i
    {-# INLINE memoryUsage #-}

instance ExMemoryUsage Natural where
    -- Same as Integer since we are going via Integer
    memoryUsage :: Natural -> CostRose
memoryUsage Natural
n = Integer -> CostRose
forall a. ExMemoryUsage a => a -> CostRose
memoryUsage (Integer -> CostRose) -> Integer -> CostRose
forall a b. (a -> b) -> a -> b
$ Natural -> Integer
forall a. Integral a => a -> Integer
toInteger Natural
n
    {-# INLINE memoryUsage #-}

instance ExMemoryUsage Word8 where
    memoryUsage :: Word8 -> CostRose
memoryUsage Word8
_ = CostingInteger -> CostRose
singletonRose CostingInteger
1
    {-# INLINE memoryUsage #-}

{- Bytestrings: we want the empty bytestring and bytestrings of length 1-8 to have
   size 1, bytestrings of length 9-16 to have size 2, etc.  Note that (-1)
   `quot` 8 == 0, so the code below gives the correct answer for the empty
   bytestring.  -}
instance ExMemoryUsage BS.ByteString where
    -- Don't use `div` here!  That gives 0 instead of 1 for the empty bytestring.
    memoryUsage :: ByteString -> CostRose
memoryUsage ByteString
bs = CostingInteger -> CostRose
singletonRose (CostingInteger -> CostRose)
-> (Int -> CostingInteger) -> Int -> CostRose
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> CostingInteger
unsafeToSatInt (Int -> CostRose) -> Int -> CostRose
forall a b. (a -> b) -> a -> b
$ ((Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
8) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 where
        n :: Int
n = ByteString -> Int
BS.length ByteString
bs
    {-# INLINE memoryUsage #-}

instance ExMemoryUsage T.Text where
    -- This says that @Text@ allocates 1 'CostingInteger' worth of memory (i.e. 8 bytes) per
    -- character, which is a conservative overestimate (i.e. is safe) regardless of whether @Text@
    -- is UTF16-based (like it used to when we implemented this instance) or UTF8-based (like it is
    -- now).
    --
    -- Note that the @ExMemoryUsage Char@ instance does not affect this one, this is for performance
    -- reasons, since @T.length@ is O(1) unlike @sum . map (memoryUsage @Char) . T.unpack@. We used
    -- to have the latter, but changed it to the former for easy performance gains.
    --
    -- We may want to make this a bit less of an overestimate in future just not to overcharge
    -- users.
    memoryUsage :: Text -> CostRose
memoryUsage = CostingInteger -> CostRose
singletonRose (CostingInteger -> CostRose)
-> (Text -> CostingInteger) -> Text -> CostRose
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> CostingInteger
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> CostingInteger) -> (Text -> Int) -> Text -> CostingInteger
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Int
T.length
    {-# INLINE memoryUsage #-}

instance ExMemoryUsage Int where
    memoryUsage :: Int -> CostRose
memoryUsage Int
_ = CostingInteger -> CostRose
singletonRose CostingInteger
1
    {-# INLINE memoryUsage #-}

-- If you ever change this, also change @ExMemoryUsage T.Text@.
instance ExMemoryUsage Char where
    memoryUsage :: Char -> CostRose
memoryUsage Char
_ = CostingInteger -> CostRose
singletonRose CostingInteger
1
    {-# INLINE memoryUsage #-}

instance ExMemoryUsage Bool where
    memoryUsage :: Bool -> CostRose
memoryUsage Bool
_ = CostingInteger -> CostRose
singletonRose CostingInteger
1
    {-# INLINE memoryUsage #-}

-- | Add two 'CostRose's. We don't make this into a 'Semigroup' instance, because there exist
-- different ways to add two 'CostRose's (e.g. we could optimize the case when one of the roses
-- contains only one element or we can make the function lazy in the second argument). Here we chose
-- the version that is most efficient when the first argument is a statically known constant (we
-- didn't do any benchmarking though, so it may not be the most efficient one) as we need this
-- below.
addConstantRose :: CostRose -> CostRose -> CostRose
addConstantRose :: CostRose -> CostRose -> CostRose
addConstantRose (CostRose CostingInteger
cost1 [CostRose]
forest1) (CostRose CostingInteger
cost2 [CostRose]
forest2) =
    CostingInteger -> [CostRose] -> CostRose
CostRose (CostingInteger
cost1 CostingInteger -> CostingInteger -> CostingInteger
forall a. Num a => a -> a -> a
+ CostingInteger
cost2) ([CostRose]
forest1 [CostRose] -> [CostRose] -> [CostRose]
forall a. [a] -> [a] -> [a]
++ [CostRose]
forest2)
{-# INLINE addConstantRose #-}

{- A naive traversal for size.  This accounts for the number of nodes in a Data
   object and also the sizes of the contents of the nodes.  This is not ideal,
   but it seems to be the best we can do.  At present this only comes into play
   for 'equalsData', which is implemented using the derived implementation of
   '==' (fortunately the costing functions are lazy, so this won't be called for
   things like 'unBData' which have constant costing functions because they only
   have to look at the top node).  The problem is that when we call 'equalsData'
   the comparison will take place entirely in Haskell, so the costing functions
   for the contents of 'I' and 'B' nodes won't be called.  Thus if we just
   counted the number of nodes the sizes of 'I 2' and 'B <huge bytestring>'
   would be the same but they'd take different amounts of time to compare.  It's
   not clear how to trade off the costs of processing a node and processing the
   contents of nodes: the implementation below compromises by charging four
   units per node, but we may wish to revise this after experimentation.
-}
instance ExMemoryUsage Data where
    memoryUsage :: Data -> CostRose
memoryUsage = Data -> CostRose
sizeData where
        dataNodeRose :: CostRose
dataNodeRose = CostingInteger -> CostRose
singletonRose CostingInteger
4
        {-# INLINE dataNodeRose #-}

        sizeData :: Data -> CostRose
sizeData Data
d = CostRose -> CostRose -> CostRose
addConstantRose CostRose
dataNodeRose (CostRose -> CostRose) -> CostRose -> CostRose
forall a b. (a -> b) -> a -> b
$ case Data
d of
            -- TODO: include the size of the tag, but not just yet. See PLT-1176.
            Constr Integer
_ [Data]
l -> CostingInteger -> [CostRose] -> CostRose
CostRose CostingInteger
0 ([CostRose] -> CostRose) -> [CostRose] -> CostRose
forall a b. (a -> b) -> a -> b
$ [Data]
l [Data] -> (Data -> CostRose) -> [CostRose]
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> Data -> CostRose
sizeData
            Map [(Data, Data)]
l      -> CostingInteger -> [CostRose] -> CostRose
CostRose CostingInteger
0 ([CostRose] -> CostRose) -> [CostRose] -> CostRose
forall a b. (a -> b) -> a -> b
$ [(Data, Data)]
l [(Data, Data)] -> ((Data, Data) -> [CostRose]) -> [CostRose]
forall a b. [a] -> (a -> [b]) -> [b]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \(Data
d1, Data
d2) -> [Data
d1, Data
d2] [Data] -> (Data -> CostRose) -> [CostRose]
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> Data -> CostRose
sizeData
            List [Data]
l     -> CostingInteger -> [CostRose] -> CostRose
CostRose CostingInteger
0 ([CostRose] -> CostRose) -> [CostRose] -> CostRose
forall a b. (a -> b) -> a -> b
$ [Data]
l [Data] -> (Data -> CostRose) -> [CostRose]
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> Data -> CostRose
sizeData
            I Integer
n        -> Integer -> CostRose
forall a. ExMemoryUsage a => a -> CostRose
memoryUsage Integer
n
            B ByteString
b        -> ByteString -> CostRose
forall a. ExMemoryUsage a => a -> CostRose
memoryUsage ByteString
b

{- Note [Costing constant-size types]
The memory usage of each of the BLS12-381 types is constant, so we may be able
to optimise things a little by ensuring that we don't re-compute the size of
(say) a G1 element every time one is used. GHC will probably do this anyway, but
we make sure by defining a top level function for each of the size measures and
getting the memoryUsage instances to call those.
-}

g1ElementCost :: CostRose
g1ElementCost :: CostRose
g1ElementCost = CostingInteger -> CostRose
singletonRose (CostingInteger -> CostRose)
-> (Int -> CostingInteger) -> Int -> CostRose
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> CostingInteger
unsafeToSatInt (Int -> CostRose) -> Int -> CostRose
forall a b. (a -> b) -> a -> b
$ Int
BLS12_381.G1.memSizeBytes Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
8
{-# NOINLINE g1ElementCost #-}

instance ExMemoryUsage BLS12_381.G1.Element where
    memoryUsage :: Element -> CostRose
memoryUsage Element
_ = CostRose
g1ElementCost
    -- Should be 18

g2ElementCost :: CostRose
g2ElementCost :: CostRose
g2ElementCost = CostingInteger -> CostRose
singletonRose (CostingInteger -> CostRose)
-> (Int -> CostingInteger) -> Int -> CostRose
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> CostingInteger
unsafeToSatInt (Int -> CostRose) -> Int -> CostRose
forall a b. (a -> b) -> a -> b
$ Int
BLS12_381.G2.memSizeBytes Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
8
{-# NOINLINE g2ElementCost #-}

instance ExMemoryUsage BLS12_381.G2.Element where
    memoryUsage :: Element -> CostRose
memoryUsage Element
_ = CostRose
g2ElementCost
    -- Should be 36

mlResultElementCost :: CostRose
mlResultElementCost :: CostRose
mlResultElementCost = CostingInteger -> CostRose
singletonRose (CostingInteger -> CostRose)
-> (Int -> CostingInteger) -> Int -> CostRose
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> CostingInteger
unsafeToSatInt (Int -> CostRose) -> Int -> CostRose
forall a b. (a -> b) -> a -> b
$ Int
BLS12_381.Pairing.mlResultMemSizeBytes Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
8
{-# NOINLINE mlResultElementCost #-}

instance ExMemoryUsage BLS12_381.Pairing.MlResult where
    memoryUsage :: MlResult -> CostRose
memoryUsage MlResult
_ = CostRose
mlResultElementCost
    -- Should be 72