-- 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(..)
    , ListCostedByLength(..)
    , ArrayCostedByLength(..)
    ) 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

instance (ExMemoryUsage a, ExMemoryUsage b) => ExMemoryUsage (a, b) where
    memoryUsage :: (a, b) -> CostRose
memoryUsage (a
a, b
b) = CostingInteger -> [CostRose] -> CostRose
CostRose CostingInteger
1 [a -> CostRose
forall a. ExMemoryUsage a => a -> CostRose
memoryUsage a
a, b -> CostRose
forall a. ExMemoryUsage a => a -> CostRose
memoryUsage b
b]
    {-# 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.

{- | A wrappper for lists whose "memory usage" for costing purposes is just the
   length of the list, ignoring the sizes of the elements. 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 ListCostedByLength a = ListCostedByLength { forall a. ListCostedByLength a -> [a]
unListCostedByLength :: [a] }
instance ExMemoryUsage (ListCostedByLength a) where
    memoryUsage :: ListCostedByLength a -> CostRose
memoryUsage (ListCostedByLength [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 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.

newtype ArrayCostedByLength a = ArrayCostedByLength { forall a. ArrayCostedByLength a -> Vector a
unArrayCostedByLength :: Vector a }
instance ExMemoryUsage (ArrayCostedByLength a) where
    memoryUsage :: ArrayCostedByLength a -> CostRose
memoryUsage (ArrayCostedByLength 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 #-}
    -- 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 #-}

instance ExMemoryUsage a => ExMemoryUsage [a] where
    -- sizeof([a]) = (1 + 3N) words + N * sizeof(v)
    memoryUsage :: [a] -> CostRose
memoryUsage = CostingInteger -> [CostRose] -> CostRose
CostRose CostingInteger
nilCost ([CostRose] -> CostRose) -> ([a] -> [CostRose]) -> [a] -> CostRose
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> CostRose) -> [a] -> [CostRose]
forall a b. (a -> b) -> [a] -> [b]
map (CostRose -> CostRose -> CostRose
addConstantRose CostRose
consRose (CostRose -> CostRose) -> (a -> CostRose) -> a -> CostRose
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> CostRose
forall a. ExMemoryUsage a => a -> CostRose
memoryUsage) where
        -- As per https://wiki.haskell.org/GHC/Memory_Footprint
        nilCost :: CostingInteger
nilCost = CostingInteger
1
        {-# INLINE nilCost #-}
        consRose :: CostRose
consRose = CostingInteger -> CostRose
singletonRose CostingInteger
3
        {-# INLINE consRose #-}
    {-# INLINE memoryUsage #-}

instance ExMemoryUsage a => ExMemoryUsage (Vector a) where
    -- sizeof(Vector v) = (7 + N) words + N * sizeof(v)
    memoryUsage :: Vector a -> CostRose
memoryUsage Vector a
v = CostingInteger -> [CostRose] -> CostRose
CostRose CostingInteger
arrayCost [ a -> CostRose
forall a. ExMemoryUsage a => a -> CostRose
memoryUsage a
a | a
a <- Vector a -> [a]
forall a. Vector a -> [a]
Vector.toList Vector a
v ]
      where
        arrayCost :: SatInt
        arrayCost :: CostingInteger
arrayCost = CostingInteger
7 CostingInteger -> CostingInteger -> CostingInteger
forall a. Num a => a -> a -> a
+ Int -> CostingInteger
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vector a -> Int
forall a. Vector a -> Int
Vector.length Vector a
v)
        {-# INLINE arrayCost #-}
    {-# INLINE memoryUsage #-}

{- Another 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