{-# LANGUAGE BangPatterns #-}

module PlutusCore.Evaluation.Machine.ExBudgetStream
    ( ExBudgetStream(..)
    , sumExBudgetStream
    , zipCostStream
    ) where

import PlutusCore.Evaluation.Machine.CostStream
import PlutusCore.Evaluation.Machine.ExBudget
import PlutusCore.Evaluation.Machine.ExMemory

import Data.Coerce

-- | A lazy stream of 'ExBudget's. Basically @NonEmpty ExBudget@, except the elements are
-- stored strictly.
--
-- The semantics of a stream are those of the 'fold' of its elements. I.e. a stream that is a
-- reordered version of another stream is considered equal to that stream.
--
-- An 'ExBudgetStream' is what one gets by zipping two 'CostStream's (one for CPU, one for memory),
-- which is why the two data types are so similar. The only reason why we don't express both the
-- concepts in terms of a single data type is efficiency, in particular unboxing is crucial for
-- 'CostStream' and we don't care about it in 'ExBudgetStream', because we can't get the spender
-- in the CEK machine to get inlined and so unboxing 'ExBudget' here would only result in boxing it
-- back once it's about to be spent.
data ExBudgetStream
    = ExBudgetLast !ExBudget
    | ExBudgetCons !ExBudget ExBudgetStream
    deriving stock (Int -> ExBudgetStream -> ShowS
[ExBudgetStream] -> ShowS
ExBudgetStream -> String
(Int -> ExBudgetStream -> ShowS)
-> (ExBudgetStream -> String)
-> ([ExBudgetStream] -> ShowS)
-> Show ExBudgetStream
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ExBudgetStream -> ShowS
showsPrec :: Int -> ExBudgetStream -> ShowS
$cshow :: ExBudgetStream -> String
show :: ExBudgetStream -> String
$cshowList :: [ExBudgetStream] -> ShowS
showList :: [ExBudgetStream] -> ShowS
Show)

-- See Note [Global local functions].
sumExBudgetStreamGo :: ExBudget -> ExBudgetStream -> ExBudget
sumExBudgetStreamGo :: ExBudget -> ExBudgetStream -> ExBudget
sumExBudgetStreamGo !ExBudget
acc (ExBudgetLast ExBudget
budget)         = ExBudget
acc ExBudget -> ExBudget -> ExBudget
forall a. Semigroup a => a -> a -> a
<> ExBudget
budget
sumExBudgetStreamGo !ExBudget
acc (ExBudgetCons ExBudget
budget ExBudgetStream
budgets) = ExBudget -> ExBudgetStream -> ExBudget
sumExBudgetStreamGo (ExBudget
acc ExBudget -> ExBudget -> ExBudget
forall a. Semigroup a => a -> a -> a
<> ExBudget
budget) ExBudgetStream
budgets

-- | Add up all the budgets in a 'ExBudgetStream'.
sumExBudgetStream :: ExBudgetStream -> ExBudget
sumExBudgetStream :: ExBudgetStream -> ExBudget
sumExBudgetStream (ExBudgetLast ExBudget
budget0)          = ExBudget
budget0
sumExBudgetStream (ExBudgetCons ExBudget
budget0 ExBudgetStream
budgets0) = ExBudget -> ExBudgetStream -> ExBudget
sumExBudgetStreamGo ExBudget
budget0 ExBudgetStream
budgets0
{-# INLINE sumExBudgetStream #-}

-- | Convert a 'CostStream' to an 'ExBudgetStream' by applying a function to each element.
costToExBudgetStream :: (CostingInteger -> ExBudget) -> CostStream -> ExBudgetStream
costToExBudgetStream :: (CostingInteger -> ExBudget) -> CostStream -> ExBudgetStream
costToExBudgetStream CostingInteger -> ExBudget
f = CostStream -> ExBudgetStream
go where
    go :: CostStream -> ExBudgetStream
go (CostLast CostingInteger
cost)       = ExBudget -> ExBudgetStream
ExBudgetLast (CostingInteger -> ExBudget
f CostingInteger
cost)
    go (CostCons CostingInteger
cost CostStream
costs) = ExBudget -> ExBudgetStream -> ExBudgetStream
ExBudgetCons (CostingInteger -> ExBudget
f CostingInteger
cost) (ExBudgetStream -> ExBudgetStream)
-> ExBudgetStream -> ExBudgetStream
forall a b. (a -> b) -> a -> b
$ CostStream -> ExBudgetStream
go CostStream
costs
{-# INLINE costToExBudgetStream #-}

-- | Convert a 'CostingInteger' representing a CPU cost and a 'CostingInteger' representing a memory
-- cost to an 'ExBudget'.
toExBudget :: CostingInteger -> CostingInteger -> ExBudget
toExBudget :: CostingInteger -> CostingInteger -> ExBudget
toExBudget = (ExCPU -> ExMemory -> ExBudget)
-> CostingInteger -> CostingInteger -> ExBudget
forall a b. Coercible a b => a -> b
coerce ExCPU -> ExMemory -> ExBudget
ExBudget
{-# INLINE toExBudget #-}

-- See Note [Global local functions].
zipCostStreamGo :: CostStream -> CostStream -> ExBudgetStream
zipCostStreamGo :: CostStream -> CostStream -> ExBudgetStream
zipCostStreamGo (CostLast CostingInteger
cpu) (CostLast CostingInteger
mem) =
    ExBudget -> ExBudgetStream
ExBudgetLast (ExBudget -> ExBudgetStream) -> ExBudget -> ExBudgetStream
forall a b. (a -> b) -> a -> b
$ CostingInteger -> CostingInteger -> ExBudget
toExBudget CostingInteger
cpu CostingInteger
mem
zipCostStreamGo (CostLast CostingInteger
cpu) (CostCons CostingInteger
mem CostStream
mems) =
    ExBudget -> ExBudgetStream -> ExBudgetStream
ExBudgetCons (CostingInteger -> CostingInteger -> ExBudget
toExBudget CostingInteger
cpu CostingInteger
mem) (ExBudgetStream -> ExBudgetStream)
-> ExBudgetStream -> ExBudgetStream
forall a b. (a -> b) -> a -> b
$ (CostingInteger -> ExBudget) -> CostStream -> ExBudgetStream
costToExBudgetStream (\CostingInteger
mem' -> CostingInteger -> CostingInteger -> ExBudget
toExBudget CostingInteger
0 CostingInteger
mem') CostStream
mems
zipCostStreamGo (CostCons CostingInteger
cpu CostStream
cpus) (CostLast CostingInteger
mem) =
    ExBudget -> ExBudgetStream -> ExBudgetStream
ExBudgetCons (CostingInteger -> CostingInteger -> ExBudget
toExBudget CostingInteger
cpu CostingInteger
mem) (ExBudgetStream -> ExBudgetStream)
-> ExBudgetStream -> ExBudgetStream
forall a b. (a -> b) -> a -> b
$ (CostingInteger -> ExBudget) -> CostStream -> ExBudgetStream
costToExBudgetStream (\CostingInteger
cpu' -> CostingInteger -> CostingInteger -> ExBudget
toExBudget CostingInteger
cpu' CostingInteger
0) CostStream
cpus
zipCostStreamGo (CostCons CostingInteger
cpu CostStream
cpus) (CostCons CostingInteger
mem CostStream
mems) =
    ExBudget -> ExBudgetStream -> ExBudgetStream
ExBudgetCons (CostingInteger -> CostingInteger -> ExBudget
toExBudget CostingInteger
cpu CostingInteger
mem) (ExBudgetStream -> ExBudgetStream)
-> ExBudgetStream -> ExBudgetStream
forall a b. (a -> b) -> a -> b
$ CostStream -> CostStream -> ExBudgetStream
zipCostStreamGo CostStream
cpus CostStream
mems

-- | Zip two 'CostStream' together (one with CPU costs and the other one with memory costs,
-- respectively) to get an 'ExBudgetStream'. If one is longer than the other, then it's assumed to
-- contain the required amount of zeros for two streams to have the same length (all those zeros
-- \"appear\" in the tail of the stream).
zipCostStream :: CostStream -> CostStream -> ExBudgetStream
zipCostStream :: CostStream -> CostStream -> ExBudgetStream
zipCostStream CostStream
cpus0 CostStream
mems0 = case (CostStream
cpus0, CostStream
mems0) of
    -- See Note [Single-element streams].
    (CostLast CostingInteger
cpu, CostLast CostingInteger
mem) -> ExBudget -> ExBudgetStream
ExBudgetLast (ExBudget -> ExBudgetStream) -> ExBudget -> ExBudgetStream
forall a b. (a -> b) -> a -> b
$ CostingInteger -> CostingInteger -> ExBudget
toExBudget CostingInteger
cpu CostingInteger
mem
    (CostStream, CostStream)
_                            -> CostStream -> CostStream -> ExBudgetStream
zipCostStreamGo CostStream
cpus0 CostStream
mems0
{-# INLINE zipCostStream #-}