{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE StrictData            #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE UndecidableInstances  #-}

module UntypedPlutusCore.Evaluation.Machine.Cek.ExBudgetMode
    ( ExBudgetMode (..)
    , CountingSt (..)
    , CekExTally (..)
    , TallyingSt (..)
    , RestrictingSt (..)
    , Hashable
    , counting
    , enormousBudget
    , tallying
    , restricting
    , restrictingEnormous
    )
where

import PlutusPrelude

import UntypedPlutusCore.Evaluation.Machine.Cek.Internal

import PlutusCore.Evaluation.Machine.ExBudget
import PlutusCore.Evaluation.Machine.ExMemory (ExCPU (..), ExMemory (..))
import PlutusCore.Evaluation.Machine.Exception

import Control.Lens (ifoldMap)
import Control.Monad.Except
import Data.HashMap.Monoidal as HashMap
import Data.Hashable (Hashable)
import Data.List (intersperse)
import Data.Map.Strict qualified as Map
import Data.Primitive.PrimArray
import Data.STRef
import Data.SatInt
import Data.Semigroup.Generic
import Prettyprinter
import Text.PrettyBy (IgnorePrettyConfig (..))

-- | Construct an 'ExBudgetMode' out of a function returning a value of the budgeting state type.
-- The value then gets added to the current state via @(<>)@.
monoidalBudgeting
    :: Monoid cost => (ExBudgetCategory fun -> ExBudget -> cost) -> ExBudgetMode cost uni fun
monoidalBudgeting :: (ExBudgetCategory fun -> ExBudget -> cost)
-> ExBudgetMode cost uni fun
monoidalBudgeting ExBudgetCategory fun -> ExBudget -> cost
toCost = (forall s. ST s (ExBudgetInfo cost uni fun s))
-> ExBudgetMode cost uni fun
forall cost (uni :: * -> *) fun.
(forall s. ST s (ExBudgetInfo cost uni fun s))
-> ExBudgetMode cost uni fun
ExBudgetMode ((forall s. ST s (ExBudgetInfo cost uni fun s))
 -> ExBudgetMode cost uni fun)
-> (forall s. ST s (ExBudgetInfo cost uni fun s))
-> ExBudgetMode cost uni fun
forall a b. (a -> b) -> a -> b
$ do
    STRef s cost
costRef <- cost -> ST s (STRef s cost)
forall a s. a -> ST s (STRef s a)
newSTRef cost
forall a. Monoid a => a
mempty
    STRef s ExBudget
budgetRef <- ExBudget -> ST s (STRef s ExBudget)
forall a s. a -> ST s (STRef s a)
newSTRef ExBudget
forall a. Monoid a => a
mempty
    let spend :: ExBudgetCategory fun -> ExBudget -> CekM uni fun s ()
spend ExBudgetCategory fun
key ExBudget
budgetToSpend = ST s () -> CekM uni fun s ()
forall (uni :: * -> *) fun s a. ST s a -> CekM uni fun s a
CekM (ST s () -> CekM uni fun s ()) -> ST s () -> CekM uni fun s ()
forall a b. (a -> b) -> a -> b
$ do
            STRef s cost -> (cost -> cost) -> ST s ()
forall s a. STRef s a -> (a -> a) -> ST s ()
modifySTRef' STRef s cost
costRef (cost -> cost -> cost
forall a. Semigroup a => a -> a -> a
<> ExBudgetCategory fun -> ExBudget -> cost
toCost ExBudgetCategory fun
key ExBudget
budgetToSpend)
            STRef s ExBudget -> (ExBudget -> ExBudget) -> ST s ()
forall s a. STRef s a -> (a -> a) -> ST s ()
modifySTRef' STRef s ExBudget
budgetRef (ExBudget -> ExBudget -> ExBudget
forall a. Semigroup a => a -> a -> a
<> ExBudget
budgetToSpend)
        spender :: CekBudgetSpender uni fun s
spender = (ExBudgetCategory fun -> ExBudget -> CekM uni fun s ())
-> CekBudgetSpender uni fun s
forall (uni :: * -> *) fun s.
(ExBudgetCategory fun -> ExBudget -> CekM uni fun s ())
-> CekBudgetSpender uni fun s
CekBudgetSpender ExBudgetCategory fun -> ExBudget -> CekM uni fun s ()
forall (uni :: * -> *) fun.
ExBudgetCategory fun -> ExBudget -> CekM uni fun s ()
spend
        cumulative :: ST s ExBudget
cumulative = STRef s ExBudget -> ST s ExBudget
forall s a. STRef s a -> ST s a
readSTRef STRef s ExBudget
budgetRef
        final :: ST s cost
final = STRef s cost -> ST s cost
forall s a. STRef s a -> ST s a
readSTRef STRef s cost
costRef
    ExBudgetInfo cost uni fun s -> ST s (ExBudgetInfo cost uni fun s)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExBudgetInfo cost uni fun s -> ST s (ExBudgetInfo cost uni fun s))
-> ExBudgetInfo cost uni fun s
-> ST s (ExBudgetInfo cost uni fun s)
forall a b. (a -> b) -> a -> b
$ CekBudgetSpender uni fun s
-> ST s cost -> ST s ExBudget -> ExBudgetInfo cost uni fun s
forall cost (uni :: * -> *) fun s.
CekBudgetSpender uni fun s
-> ST s cost -> ST s ExBudget -> ExBudgetInfo cost uni fun s
ExBudgetInfo CekBudgetSpender uni fun s
forall (uni :: * -> *). CekBudgetSpender uni fun s
spender ST s cost
final ST s ExBudget
cumulative

-- | For calculating the cost of execution by counting up using the 'Monoid' instance of 'ExBudget'.
newtype CountingSt = CountingSt ExBudget
    deriving stock (CountingSt -> CountingSt -> Bool
(CountingSt -> CountingSt -> Bool)
-> (CountingSt -> CountingSt -> Bool) -> Eq CountingSt
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: CountingSt -> CountingSt -> Bool
$c/= :: CountingSt -> CountingSt -> Bool
== :: CountingSt -> CountingSt -> Bool
$c== :: CountingSt -> CountingSt -> Bool
Eq, Int -> CountingSt -> ShowS
[CountingSt] -> ShowS
CountingSt -> String
(Int -> CountingSt -> ShowS)
-> (CountingSt -> String)
-> ([CountingSt] -> ShowS)
-> Show CountingSt
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [CountingSt] -> ShowS
$cshowList :: [CountingSt] -> ShowS
show :: CountingSt -> String
$cshow :: CountingSt -> String
showsPrec :: Int -> CountingSt -> ShowS
$cshowsPrec :: Int -> CountingSt -> ShowS
Show)
    deriving newtype (b -> CountingSt -> CountingSt
NonEmpty CountingSt -> CountingSt
CountingSt -> CountingSt -> CountingSt
(CountingSt -> CountingSt -> CountingSt)
-> (NonEmpty CountingSt -> CountingSt)
-> (forall b. Integral b => b -> CountingSt -> CountingSt)
-> Semigroup CountingSt
forall b. Integral b => b -> CountingSt -> CountingSt
forall a.
(a -> a -> a)
-> (NonEmpty a -> a)
-> (forall b. Integral b => b -> a -> a)
-> Semigroup a
stimes :: b -> CountingSt -> CountingSt
$cstimes :: forall b. Integral b => b -> CountingSt -> CountingSt
sconcat :: NonEmpty CountingSt -> CountingSt
$csconcat :: NonEmpty CountingSt -> CountingSt
<> :: CountingSt -> CountingSt -> CountingSt
$c<> :: CountingSt -> CountingSt -> CountingSt
Semigroup, Semigroup CountingSt
CountingSt
Semigroup CountingSt
-> CountingSt
-> (CountingSt -> CountingSt -> CountingSt)
-> ([CountingSt] -> CountingSt)
-> Monoid CountingSt
[CountingSt] -> CountingSt
CountingSt -> CountingSt -> CountingSt
forall a.
Semigroup a -> a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
mconcat :: [CountingSt] -> CountingSt
$cmconcat :: [CountingSt] -> CountingSt
mappend :: CountingSt -> CountingSt -> CountingSt
$cmappend :: CountingSt -> CountingSt -> CountingSt
mempty :: CountingSt
$cmempty :: CountingSt
$cp1Monoid :: Semigroup CountingSt
Monoid, PrettyBy config, CountingSt -> ()
(CountingSt -> ()) -> NFData CountingSt
forall a. (a -> ()) -> NFData a
rnf :: CountingSt -> ()
$crnf :: CountingSt -> ()
NFData)

instance Pretty CountingSt where
    pretty :: CountingSt -> Doc ann
pretty (CountingSt ExBudget
budget) = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$ Doc ann
"required budget:" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> ExBudget -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty ExBudget
budget Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
line

-- | For calculating the cost of execution.
counting :: ExBudgetMode CountingSt uni fun
counting :: ExBudgetMode CountingSt uni fun
counting = (ExBudgetCategory fun -> ExBudget -> CountingSt)
-> ExBudgetMode CountingSt uni fun
forall cost fun (uni :: * -> *).
Monoid cost =>
(ExBudgetCategory fun -> ExBudget -> cost)
-> ExBudgetMode cost uni fun
monoidalBudgeting ((ExBudgetCategory fun -> ExBudget -> CountingSt)
 -> ExBudgetMode CountingSt uni fun)
-> (ExBudgetCategory fun -> ExBudget -> CountingSt)
-> ExBudgetMode CountingSt uni fun
forall a b. (a -> b) -> a -> b
$ (ExBudget -> CountingSt)
-> ExBudgetCategory fun -> ExBudget -> CountingSt
forall a b. a -> b -> a
const ExBudget -> CountingSt
CountingSt

-- | For a detailed report on what costs how much + the same overall budget that 'Counting' gives.
-- The (derived) 'Monoid' instance of 'CekExTally' is the main piece of the machinery.
newtype CekExTally fun = CekExTally (MonoidalHashMap (ExBudgetCategory fun) ExBudget)
    deriving stock (CekExTally fun -> CekExTally fun -> Bool
(CekExTally fun -> CekExTally fun -> Bool)
-> (CekExTally fun -> CekExTally fun -> Bool)
-> Eq (CekExTally fun)
forall fun. Eq fun => CekExTally fun -> CekExTally fun -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: CekExTally fun -> CekExTally fun -> Bool
$c/= :: forall fun. Eq fun => CekExTally fun -> CekExTally fun -> Bool
== :: CekExTally fun -> CekExTally fun -> Bool
$c== :: forall fun. Eq fun => CekExTally fun -> CekExTally fun -> Bool
Eq, (forall x. CekExTally fun -> Rep (CekExTally fun) x)
-> (forall x. Rep (CekExTally fun) x -> CekExTally fun)
-> Generic (CekExTally fun)
forall x. Rep (CekExTally fun) x -> CekExTally fun
forall x. CekExTally fun -> Rep (CekExTally fun) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall fun x. Rep (CekExTally fun) x -> CekExTally fun
forall fun x. CekExTally fun -> Rep (CekExTally fun) x
$cto :: forall fun x. Rep (CekExTally fun) x -> CekExTally fun
$cfrom :: forall fun x. CekExTally fun -> Rep (CekExTally fun) x
Generic, Int -> CekExTally fun -> ShowS
[CekExTally fun] -> ShowS
CekExTally fun -> String
(Int -> CekExTally fun -> ShowS)
-> (CekExTally fun -> String)
-> ([CekExTally fun] -> ShowS)
-> Show (CekExTally fun)
forall fun. Show fun => Int -> CekExTally fun -> ShowS
forall fun. Show fun => [CekExTally fun] -> ShowS
forall fun. Show fun => CekExTally fun -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [CekExTally fun] -> ShowS
$cshowList :: forall fun. Show fun => [CekExTally fun] -> ShowS
show :: CekExTally fun -> String
$cshow :: forall fun. Show fun => CekExTally fun -> String
showsPrec :: Int -> CekExTally fun -> ShowS
$cshowsPrec :: forall fun. Show fun => Int -> CekExTally fun -> ShowS
Show)
    deriving (b -> CekExTally fun -> CekExTally fun
NonEmpty (CekExTally fun) -> CekExTally fun
CekExTally fun -> CekExTally fun -> CekExTally fun
(CekExTally fun -> CekExTally fun -> CekExTally fun)
-> (NonEmpty (CekExTally fun) -> CekExTally fun)
-> (forall b. Integral b => b -> CekExTally fun -> CekExTally fun)
-> Semigroup (CekExTally fun)
forall fun.
(Eq fun, Hashable fun) =>
NonEmpty (CekExTally fun) -> CekExTally fun
forall fun.
(Eq fun, Hashable fun) =>
CekExTally fun -> CekExTally fun -> CekExTally fun
forall fun b.
(Eq fun, Hashable fun, Integral b) =>
b -> CekExTally fun -> CekExTally fun
forall b. Integral b => b -> CekExTally fun -> CekExTally fun
forall a.
(a -> a -> a)
-> (NonEmpty a -> a)
-> (forall b. Integral b => b -> a -> a)
-> Semigroup a
stimes :: b -> CekExTally fun -> CekExTally fun
$cstimes :: forall fun b.
(Eq fun, Hashable fun, Integral b) =>
b -> CekExTally fun -> CekExTally fun
sconcat :: NonEmpty (CekExTally fun) -> CekExTally fun
$csconcat :: forall fun.
(Eq fun, Hashable fun) =>
NonEmpty (CekExTally fun) -> CekExTally fun
<> :: CekExTally fun -> CekExTally fun -> CekExTally fun
$c<> :: forall fun.
(Eq fun, Hashable fun) =>
CekExTally fun -> CekExTally fun -> CekExTally fun
Semigroup, Semigroup (CekExTally fun)
CekExTally fun
Semigroup (CekExTally fun)
-> CekExTally fun
-> (CekExTally fun -> CekExTally fun -> CekExTally fun)
-> ([CekExTally fun] -> CekExTally fun)
-> Monoid (CekExTally fun)
[CekExTally fun] -> CekExTally fun
CekExTally fun -> CekExTally fun -> CekExTally fun
forall fun. (Eq fun, Hashable fun) => Semigroup (CekExTally fun)
forall fun. (Eq fun, Hashable fun) => CekExTally fun
forall fun.
(Eq fun, Hashable fun) =>
[CekExTally fun] -> CekExTally fun
forall fun.
(Eq fun, Hashable fun) =>
CekExTally fun -> CekExTally fun -> CekExTally fun
forall a.
Semigroup a -> a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
mconcat :: [CekExTally fun] -> CekExTally fun
$cmconcat :: forall fun.
(Eq fun, Hashable fun) =>
[CekExTally fun] -> CekExTally fun
mappend :: CekExTally fun -> CekExTally fun -> CekExTally fun
$cmappend :: forall fun.
(Eq fun, Hashable fun) =>
CekExTally fun -> CekExTally fun -> CekExTally fun
mempty :: CekExTally fun
$cmempty :: forall fun. (Eq fun, Hashable fun) => CekExTally fun
$cp1Monoid :: forall fun. (Eq fun, Hashable fun) => Semigroup (CekExTally fun)
Monoid) via (GenericSemigroupMonoid (CekExTally fun))
    deriving anyclass (CekExTally fun -> ()
(CekExTally fun -> ()) -> NFData (CekExTally fun)
forall fun. NFData fun => CekExTally fun -> ()
forall a. (a -> ()) -> NFData a
rnf :: CekExTally fun -> ()
$crnf :: forall fun. NFData fun => CekExTally fun -> ()
NFData)
    deriving (PrettyBy config) via (IgnorePrettyConfig (CekExTally fun))

instance (Show fun, Ord fun) => Pretty (CekExTally fun) where
    pretty :: CekExTally fun -> Doc ann
pretty (CekExTally MonoidalHashMap (ExBudgetCategory fun) ExBudget
m) =
        let om :: Map (ExBudgetCategory fun) ExBudget
om = [(ExBudgetCategory fun, ExBudget)]
-> Map (ExBudgetCategory fun) ExBudget
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(ExBudgetCategory fun, ExBudget)]
 -> Map (ExBudgetCategory fun) ExBudget)
-> [(ExBudgetCategory fun, ExBudget)]
-> Map (ExBudgetCategory fun) ExBudget
forall a b. (a -> b) -> a -> b
$ MonoidalHashMap (ExBudgetCategory fun) ExBudget
-> [(ExBudgetCategory fun, ExBudget)]
forall k a. MonoidalHashMap k a -> [(k, a)]
HashMap.toList MonoidalHashMap (ExBudgetCategory fun) ExBudget
m
        in Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$ [Doc ann] -> Doc ann
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold ([Doc ann
"{ "] [Doc ann] -> [Doc ann] -> [Doc ann]
forall a. Semigroup a => a -> a -> a
<> (Doc ann -> [Doc ann] -> [Doc ann]
forall a. a -> [a] -> [a]
intersperse (Doc ann
forall ann. Doc ann
line Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
"| ") ([Doc ann] -> [Doc ann]) -> [Doc ann] -> [Doc ann]
forall a b. (a -> b) -> a -> b
$ (Doc ann -> Doc ann) -> [Doc ann] -> [Doc ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
group ([Doc ann] -> [Doc ann]) -> [Doc ann] -> [Doc ann]
forall a b. (a -> b) -> a -> b
$
          (ExBudgetCategory fun -> ExBudget -> [Doc ann])
-> Map (ExBudgetCategory fun) ExBudget -> [Doc ann]
forall i (f :: * -> *) m a.
(FoldableWithIndex i f, Monoid m) =>
(i -> a -> m) -> f a -> m
ifoldMap (\ExBudgetCategory fun
k ExBudget
v -> [(ExBudgetCategory fun -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty ExBudgetCategory fun
k Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"causes" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> ExBudget -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty ExBudget
v)]) Map (ExBudgetCategory fun) ExBudget
om) [Doc ann] -> [Doc ann] -> [Doc ann]
forall a. Semigroup a => a -> a -> a
<> [Doc ann
"}"])

data TallyingSt fun = TallyingSt (CekExTally fun) ExBudget
    deriving stock (TallyingSt fun -> TallyingSt fun -> Bool
(TallyingSt fun -> TallyingSt fun -> Bool)
-> (TallyingSt fun -> TallyingSt fun -> Bool)
-> Eq (TallyingSt fun)
forall fun. Eq fun => TallyingSt fun -> TallyingSt fun -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: TallyingSt fun -> TallyingSt fun -> Bool
$c/= :: forall fun. Eq fun => TallyingSt fun -> TallyingSt fun -> Bool
== :: TallyingSt fun -> TallyingSt fun -> Bool
$c== :: forall fun. Eq fun => TallyingSt fun -> TallyingSt fun -> Bool
Eq, Int -> TallyingSt fun -> ShowS
[TallyingSt fun] -> ShowS
TallyingSt fun -> String
(Int -> TallyingSt fun -> ShowS)
-> (TallyingSt fun -> String)
-> ([TallyingSt fun] -> ShowS)
-> Show (TallyingSt fun)
forall fun. Show fun => Int -> TallyingSt fun -> ShowS
forall fun. Show fun => [TallyingSt fun] -> ShowS
forall fun. Show fun => TallyingSt fun -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [TallyingSt fun] -> ShowS
$cshowList :: forall fun. Show fun => [TallyingSt fun] -> ShowS
show :: TallyingSt fun -> String
$cshow :: forall fun. Show fun => TallyingSt fun -> String
showsPrec :: Int -> TallyingSt fun -> ShowS
$cshowsPrec :: forall fun. Show fun => Int -> TallyingSt fun -> ShowS
Show, (forall x. TallyingSt fun -> Rep (TallyingSt fun) x)
-> (forall x. Rep (TallyingSt fun) x -> TallyingSt fun)
-> Generic (TallyingSt fun)
forall x. Rep (TallyingSt fun) x -> TallyingSt fun
forall x. TallyingSt fun -> Rep (TallyingSt fun) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall fun x. Rep (TallyingSt fun) x -> TallyingSt fun
forall fun x. TallyingSt fun -> Rep (TallyingSt fun) x
$cto :: forall fun x. Rep (TallyingSt fun) x -> TallyingSt fun
$cfrom :: forall fun x. TallyingSt fun -> Rep (TallyingSt fun) x
Generic)
    deriving (b -> TallyingSt fun -> TallyingSt fun
NonEmpty (TallyingSt fun) -> TallyingSt fun
TallyingSt fun -> TallyingSt fun -> TallyingSt fun
(TallyingSt fun -> TallyingSt fun -> TallyingSt fun)
-> (NonEmpty (TallyingSt fun) -> TallyingSt fun)
-> (forall b. Integral b => b -> TallyingSt fun -> TallyingSt fun)
-> Semigroup (TallyingSt fun)
forall fun.
(Eq fun, Hashable fun) =>
NonEmpty (TallyingSt fun) -> TallyingSt fun
forall fun.
(Eq fun, Hashable fun) =>
TallyingSt fun -> TallyingSt fun -> TallyingSt fun
forall fun b.
(Eq fun, Hashable fun, Integral b) =>
b -> TallyingSt fun -> TallyingSt fun
forall b. Integral b => b -> TallyingSt fun -> TallyingSt fun
forall a.
(a -> a -> a)
-> (NonEmpty a -> a)
-> (forall b. Integral b => b -> a -> a)
-> Semigroup a
stimes :: b -> TallyingSt fun -> TallyingSt fun
$cstimes :: forall fun b.
(Eq fun, Hashable fun, Integral b) =>
b -> TallyingSt fun -> TallyingSt fun
sconcat :: NonEmpty (TallyingSt fun) -> TallyingSt fun
$csconcat :: forall fun.
(Eq fun, Hashable fun) =>
NonEmpty (TallyingSt fun) -> TallyingSt fun
<> :: TallyingSt fun -> TallyingSt fun -> TallyingSt fun
$c<> :: forall fun.
(Eq fun, Hashable fun) =>
TallyingSt fun -> TallyingSt fun -> TallyingSt fun
Semigroup, Semigroup (TallyingSt fun)
TallyingSt fun
Semigroup (TallyingSt fun)
-> TallyingSt fun
-> (TallyingSt fun -> TallyingSt fun -> TallyingSt fun)
-> ([TallyingSt fun] -> TallyingSt fun)
-> Monoid (TallyingSt fun)
[TallyingSt fun] -> TallyingSt fun
TallyingSt fun -> TallyingSt fun -> TallyingSt fun
forall fun. (Eq fun, Hashable fun) => Semigroup (TallyingSt fun)
forall fun. (Eq fun, Hashable fun) => TallyingSt fun
forall fun.
(Eq fun, Hashable fun) =>
[TallyingSt fun] -> TallyingSt fun
forall fun.
(Eq fun, Hashable fun) =>
TallyingSt fun -> TallyingSt fun -> TallyingSt fun
forall a.
Semigroup a -> a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
mconcat :: [TallyingSt fun] -> TallyingSt fun
$cmconcat :: forall fun.
(Eq fun, Hashable fun) =>
[TallyingSt fun] -> TallyingSt fun
mappend :: TallyingSt fun -> TallyingSt fun -> TallyingSt fun
$cmappend :: forall fun.
(Eq fun, Hashable fun) =>
TallyingSt fun -> TallyingSt fun -> TallyingSt fun
mempty :: TallyingSt fun
$cmempty :: forall fun. (Eq fun, Hashable fun) => TallyingSt fun
$cp1Monoid :: forall fun. (Eq fun, Hashable fun) => Semigroup (TallyingSt fun)
Monoid) via (GenericSemigroupMonoid (TallyingSt fun))
    deriving anyclass (TallyingSt fun -> ()
(TallyingSt fun -> ()) -> NFData (TallyingSt fun)
forall fun. NFData fun => TallyingSt fun -> ()
forall a. (a -> ()) -> NFData a
rnf :: TallyingSt fun -> ()
$crnf :: forall fun. NFData fun => TallyingSt fun -> ()
NFData)
    deriving (PrettyBy config) via (IgnorePrettyConfig (TallyingSt fun))

instance (Show fun, Ord fun) => Pretty (TallyingSt fun) where
    pretty :: TallyingSt fun -> Doc ann
pretty (TallyingSt CekExTally fun
tally ExBudget
budget) = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$ [Doc ann] -> Doc ann
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold
        [ Doc ann
"{ tally: ", CekExTally fun -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty CekExTally fun
tally, Doc ann
forall ann. Doc ann
line
        , Doc ann
"| budget: ", ExBudget -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty ExBudget
budget, Doc ann
forall ann. Doc ann
line
        , Doc ann
"}"
        ]

-- | For a detailed report on what costs how much + the same overall budget that 'Counting' gives.
tallying :: (Eq fun, Hashable fun) => ExBudgetMode (TallyingSt fun) uni fun
tallying :: ExBudgetMode (TallyingSt fun) uni fun
tallying =
    (ExBudgetCategory fun -> ExBudget -> TallyingSt fun)
-> ExBudgetMode (TallyingSt fun) uni fun
forall cost fun (uni :: * -> *).
Monoid cost =>
(ExBudgetCategory fun -> ExBudget -> cost)
-> ExBudgetMode cost uni fun
monoidalBudgeting ((ExBudgetCategory fun -> ExBudget -> TallyingSt fun)
 -> ExBudgetMode (TallyingSt fun) uni fun)
-> (ExBudgetCategory fun -> ExBudget -> TallyingSt fun)
-> ExBudgetMode (TallyingSt fun) uni fun
forall a b. (a -> b) -> a -> b
$ \ExBudgetCategory fun
key ExBudget
budgetToSpend ->
        CekExTally fun -> ExBudget -> TallyingSt fun
forall fun. CekExTally fun -> ExBudget -> TallyingSt fun
TallyingSt (MonoidalHashMap (ExBudgetCategory fun) ExBudget -> CekExTally fun
forall fun.
MonoidalHashMap (ExBudgetCategory fun) ExBudget -> CekExTally fun
CekExTally (MonoidalHashMap (ExBudgetCategory fun) ExBudget -> CekExTally fun)
-> MonoidalHashMap (ExBudgetCategory fun) ExBudget
-> CekExTally fun
forall a b. (a -> b) -> a -> b
$ ExBudgetCategory fun
-> ExBudget -> MonoidalHashMap (ExBudgetCategory fun) ExBudget
forall k a. (Eq k, Hashable k) => k -> a -> MonoidalHashMap k a
singleton ExBudgetCategory fun
key ExBudget
budgetToSpend) ExBudget
budgetToSpend

newtype RestrictingSt = RestrictingSt ExRestrictingBudget
    deriving stock (RestrictingSt -> RestrictingSt -> Bool
(RestrictingSt -> RestrictingSt -> Bool)
-> (RestrictingSt -> RestrictingSt -> Bool) -> Eq RestrictingSt
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: RestrictingSt -> RestrictingSt -> Bool
$c/= :: RestrictingSt -> RestrictingSt -> Bool
== :: RestrictingSt -> RestrictingSt -> Bool
$c== :: RestrictingSt -> RestrictingSt -> Bool
Eq, Int -> RestrictingSt -> ShowS
[RestrictingSt] -> ShowS
RestrictingSt -> String
(Int -> RestrictingSt -> ShowS)
-> (RestrictingSt -> String)
-> ([RestrictingSt] -> ShowS)
-> Show RestrictingSt
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [RestrictingSt] -> ShowS
$cshowList :: [RestrictingSt] -> ShowS
show :: RestrictingSt -> String
$cshow :: RestrictingSt -> String
showsPrec :: Int -> RestrictingSt -> ShowS
$cshowsPrec :: Int -> RestrictingSt -> ShowS
Show)
    deriving newtype (b -> RestrictingSt -> RestrictingSt
NonEmpty RestrictingSt -> RestrictingSt
RestrictingSt -> RestrictingSt -> RestrictingSt
(RestrictingSt -> RestrictingSt -> RestrictingSt)
-> (NonEmpty RestrictingSt -> RestrictingSt)
-> (forall b. Integral b => b -> RestrictingSt -> RestrictingSt)
-> Semigroup RestrictingSt
forall b. Integral b => b -> RestrictingSt -> RestrictingSt
forall a.
(a -> a -> a)
-> (NonEmpty a -> a)
-> (forall b. Integral b => b -> a -> a)
-> Semigroup a
stimes :: b -> RestrictingSt -> RestrictingSt
$cstimes :: forall b. Integral b => b -> RestrictingSt -> RestrictingSt
sconcat :: NonEmpty RestrictingSt -> RestrictingSt
$csconcat :: NonEmpty RestrictingSt -> RestrictingSt
<> :: RestrictingSt -> RestrictingSt -> RestrictingSt
$c<> :: RestrictingSt -> RestrictingSt -> RestrictingSt
Semigroup, Semigroup RestrictingSt
RestrictingSt
Semigroup RestrictingSt
-> RestrictingSt
-> (RestrictingSt -> RestrictingSt -> RestrictingSt)
-> ([RestrictingSt] -> RestrictingSt)
-> Monoid RestrictingSt
[RestrictingSt] -> RestrictingSt
RestrictingSt -> RestrictingSt -> RestrictingSt
forall a.
Semigroup a -> a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
mconcat :: [RestrictingSt] -> RestrictingSt
$cmconcat :: [RestrictingSt] -> RestrictingSt
mappend :: RestrictingSt -> RestrictingSt -> RestrictingSt
$cmappend :: RestrictingSt -> RestrictingSt -> RestrictingSt
mempty :: RestrictingSt
$cmempty :: RestrictingSt
$cp1Monoid :: Semigroup RestrictingSt
Monoid, RestrictingSt -> ()
(RestrictingSt -> ()) -> NFData RestrictingSt
forall a. (a -> ()) -> NFData a
rnf :: RestrictingSt -> ()
$crnf :: RestrictingSt -> ()
NFData)
    deriving anyclass (PrettyBy config)

instance Pretty RestrictingSt where
    pretty :: RestrictingSt -> Doc ann
pretty (RestrictingSt ExRestrictingBudget
budget) = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$ Doc ann
"final budget:" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> ExRestrictingBudget -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty ExRestrictingBudget
budget Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
line

-- | For execution, to avoid overruns.
restricting :: forall uni fun . (PrettyUni uni fun) => ExRestrictingBudget -> ExBudgetMode RestrictingSt uni fun
restricting :: ExRestrictingBudget -> ExBudgetMode RestrictingSt uni fun
restricting (ExRestrictingBudget initB :: ExBudget
initB@(ExBudget ExCPU
cpuInit ExMemory
memInit)) = (forall s. ST s (ExBudgetInfo RestrictingSt uni fun s))
-> ExBudgetMode RestrictingSt uni fun
forall cost (uni :: * -> *) fun.
(forall s. ST s (ExBudgetInfo cost uni fun s))
-> ExBudgetMode cost uni fun
ExBudgetMode ((forall s. ST s (ExBudgetInfo RestrictingSt uni fun s))
 -> ExBudgetMode RestrictingSt uni fun)
-> (forall s. ST s (ExBudgetInfo RestrictingSt uni fun s))
-> ExBudgetMode RestrictingSt uni fun
forall a b. (a -> b) -> a -> b
$ do
    -- We keep the counters in a PrimArray. This is better than an STRef since it stores its contents unboxed.
    --
    -- If we don't specify the element type then GHC has difficulty inferring it, but it's
    -- annoying to specify the monad, since it refers to the 's' which is not in scope.
    MutablePrimArray s SatInt
ref <- Int -> ST s (MutablePrimArray (PrimState (ST s)) SatInt)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray @_ @SatInt Int
2
    let
        cpuIx :: Int
cpuIx = Int
0
        memIx :: Int
memIx = Int
1
        readCpu :: ST s ExCPU
readCpu = Coercible SatInt ExCPU => SatInt -> ExCPU
coerce @_ @ExCPU (SatInt -> ExCPU) -> ST s SatInt -> ST s ExCPU
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutablePrimArray (PrimState (ST s)) SatInt -> Int -> ST s SatInt
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
readPrimArray MutablePrimArray s SatInt
MutablePrimArray (PrimState (ST s)) SatInt
ref Int
cpuIx
        writeCpu :: ExCPU -> ST s ()
writeCpu ExCPU
cpu = MutablePrimArray (PrimState (ST s)) SatInt
-> Int -> SatInt -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray s SatInt
MutablePrimArray (PrimState (ST s)) SatInt
ref Int
cpuIx (SatInt -> ST s ()) -> SatInt -> ST s ()
forall a b. (a -> b) -> a -> b
$ ExCPU -> SatInt
coerce ExCPU
cpu
        readMem :: ST s ExMemory
readMem = Coercible SatInt ExMemory => SatInt -> ExMemory
coerce @_ @ExMemory (SatInt -> ExMemory) -> ST s SatInt -> ST s ExMemory
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutablePrimArray (PrimState (ST s)) SatInt -> Int -> ST s SatInt
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
readPrimArray MutablePrimArray s SatInt
MutablePrimArray (PrimState (ST s)) SatInt
ref Int
memIx
        writeMem :: ExMemory -> ST s ()
writeMem ExMemory
mem = MutablePrimArray (PrimState (ST s)) SatInt
-> Int -> SatInt -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray s SatInt
MutablePrimArray (PrimState (ST s)) SatInt
ref Int
memIx (SatInt -> ST s ()) -> SatInt -> ST s ()
forall a b. (a -> b) -> a -> b
$ ExMemory -> SatInt
coerce ExMemory
mem

    ExCPU -> ST s ()
writeCpu ExCPU
cpuInit
    ExMemory -> ST s ()
writeMem ExMemory
memInit
    let
        spend :: p -> ExBudget -> CekM uni fun s ()
spend p
_ (ExBudget ExCPU
cpuToSpend ExMemory
memToSpend) = do
            ExCPU
cpuLeft <- ST s ExCPU -> CekM uni fun s ExCPU
forall (uni :: * -> *) fun s a. ST s a -> CekM uni fun s a
CekM ST s ExCPU
readCpu
            ExMemory
memLeft <- ST s ExMemory -> CekM uni fun s ExMemory
forall (uni :: * -> *) fun s a. ST s a -> CekM uni fun s a
CekM ST s ExMemory
readMem
            let cpuLeft' :: ExCPU
cpuLeft' = ExCPU
cpuLeft ExCPU -> ExCPU -> ExCPU
forall a. Num a => a -> a -> a
- ExCPU
cpuToSpend
            let memLeft' :: ExMemory
memLeft' = ExMemory
memLeft ExMemory -> ExMemory -> ExMemory
forall a. Num a => a -> a -> a
- ExMemory
memToSpend
            -- Note that even if we throw an out-of-budget error, we still need to record
            -- what the final state was.
            ST s () -> CekM uni fun s ()
forall (uni :: * -> *) fun s a. ST s a -> CekM uni fun s a
CekM (ST s () -> CekM uni fun s ()) -> ST s () -> CekM uni fun s ()
forall a b. (a -> b) -> a -> b
$ ExCPU -> ST s ()
writeCpu ExCPU
cpuLeft'
            ST s () -> CekM uni fun s ()
forall (uni :: * -> *) fun s a. ST s a -> CekM uni fun s a
CekM (ST s () -> CekM uni fun s ()) -> ST s () -> CekM uni fun s ()
forall a b. (a -> b) -> a -> b
$ ExMemory -> ST s ()
writeMem ExMemory
memLeft'
            Bool -> CekM uni fun s () -> CekM uni fun s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (ExCPU
cpuLeft' ExCPU -> ExCPU -> Bool
forall a. Ord a => a -> a -> Bool
< ExCPU
0 Bool -> Bool -> Bool
|| ExMemory
memLeft' ExMemory -> ExMemory -> Bool
forall a. Ord a => a -> a -> Bool
< ExMemory
0) (CekM uni fun s () -> CekM uni fun s ())
-> CekM uni fun s () -> CekM uni fun s ()
forall a b. (a -> b) -> a -> b
$ do
                let budgetLeft :: ExBudget
budgetLeft = ExCPU -> ExMemory -> ExBudget
ExBudget ExCPU
cpuLeft' ExMemory
memLeft'
                AReview
  (EvaluationError CekUserError (MachineError fun))
  (EvaluationError CekUserError (MachineError fun))
-> EvaluationError CekUserError (MachineError fun)
-> Maybe (Term NamedDeBruijn uni fun ())
-> CekM uni fun s ()
forall exc e t term (m :: * -> *) x.
(exc ~ ErrorWithCause e term, MonadError exc m) =>
AReview e t -> t -> Maybe term -> m x
throwingWithCause AReview
  (EvaluationError CekUserError (MachineError fun))
  (EvaluationError CekUserError (MachineError fun))
forall r user internal.
AsEvaluationError r user internal =>
Prism' r (EvaluationError user internal)
_EvaluationError
                    (CekUserError -> EvaluationError CekUserError (MachineError fun)
forall user internal. user -> EvaluationError user internal
UserEvaluationError (CekUserError -> EvaluationError CekUserError (MachineError fun))
-> (ExRestrictingBudget -> CekUserError)
-> ExRestrictingBudget
-> EvaluationError CekUserError (MachineError fun)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExRestrictingBudget -> CekUserError
CekOutOfExError (ExRestrictingBudget
 -> EvaluationError CekUserError (MachineError fun))
-> ExRestrictingBudget
-> EvaluationError CekUserError (MachineError fun)
forall a b. (a -> b) -> a -> b
$ ExBudget -> ExRestrictingBudget
ExRestrictingBudget ExBudget
budgetLeft)
                    Maybe (Term NamedDeBruijn uni fun ())
forall a. Maybe a
Nothing
        spender :: CekBudgetSpender uni fun s
spender = (ExBudgetCategory fun -> ExBudget -> CekM uni fun s ())
-> CekBudgetSpender uni fun s
forall (uni :: * -> *) fun s.
(ExBudgetCategory fun -> ExBudget -> CekM uni fun s ())
-> CekBudgetSpender uni fun s
CekBudgetSpender ExBudgetCategory fun -> ExBudget -> CekM uni fun s ()
forall (uni :: * -> *) fun p.
(Everywhere uni PrettyConst, Typeable uni, Typeable fun,
 Pretty fun, Closed uni, GShow uni) =>
p -> ExBudget -> CekM uni fun s ()
spend
        remaining :: ST s ExBudget
remaining = ExCPU -> ExMemory -> ExBudget
ExBudget (ExCPU -> ExMemory -> ExBudget)
-> ST s ExCPU -> ST s (ExMemory -> ExBudget)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ST s ExCPU
readCpu ST s (ExMemory -> ExBudget) -> ST s ExMemory -> ST s ExBudget
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ST s ExMemory
readMem
        cumulative :: ST s ExBudget
cumulative = do
            ExBudget
r <- ST s ExBudget
remaining
            ExBudget -> ST s ExBudget
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExBudget -> ST s ExBudget) -> ExBudget -> ST s ExBudget
forall a b. (a -> b) -> a -> b
$ ExBudget
initB ExBudget -> ExBudget -> ExBudget
`minusExBudget` ExBudget
r
        final :: ST s RestrictingSt
final = ExRestrictingBudget -> RestrictingSt
RestrictingSt (ExRestrictingBudget -> RestrictingSt)
-> (ExBudget -> ExRestrictingBudget) -> ExBudget -> RestrictingSt
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExBudget -> ExRestrictingBudget
ExRestrictingBudget (ExBudget -> RestrictingSt) -> ST s ExBudget -> ST s RestrictingSt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ST s ExBudget
remaining
    ExBudgetInfo RestrictingSt uni fun s
-> ST s (ExBudgetInfo RestrictingSt uni fun s)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExBudgetInfo RestrictingSt uni fun s
 -> ST s (ExBudgetInfo RestrictingSt uni fun s))
-> ExBudgetInfo RestrictingSt uni fun s
-> ST s (ExBudgetInfo RestrictingSt uni fun s)
forall a b. (a -> b) -> a -> b
$ CekBudgetSpender uni fun s
-> ST s RestrictingSt
-> ST s ExBudget
-> ExBudgetInfo RestrictingSt uni fun s
forall cost (uni :: * -> *) fun s.
CekBudgetSpender uni fun s
-> ST s cost -> ST s ExBudget -> ExBudgetInfo cost uni fun s
ExBudgetInfo CekBudgetSpender uni fun s
spender ST s RestrictingSt
final ST s ExBudget
cumulative

-- | 'restricting' instantiated at 'enormousBudget'.
restrictingEnormous :: (PrettyUni uni fun) => ExBudgetMode RestrictingSt uni fun
restrictingEnormous :: ExBudgetMode RestrictingSt uni fun
restrictingEnormous = ExRestrictingBudget -> ExBudgetMode RestrictingSt uni fun
forall (uni :: * -> *) fun.
PrettyUni uni fun =>
ExRestrictingBudget -> ExBudgetMode RestrictingSt uni fun
restricting ExRestrictingBudget
enormousBudget