{-# 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 #-}