-- | This module exposes the internal state in which our direct simulation is
-- run, and functions to update and query it.
module Cooked.MockChain.MockChainState
  ( MockChainState (..),
    mcstParamsL,
    mcstLedgerStateL,
    mcstOutputsL,
    mcstConstitutionL,
    mcstToUtxoState,
    addOutput,
    removeOutput,
  )
where

import Cardano.Node.Emulator.Internal.Node qualified as Emulator
import Cooked.MockChain.UtxoState
import Cooked.Skeleton
import Data.Default
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Ledger.Orphans ()
import Optics.Core
import Optics.TH
import PlutusLedgerApi.V3 qualified as Api

-- | The state used to run the simulation in 'Cooked.MockChain.Direct'
data MockChainState where
  MockChainState ::
    { -- | The parametors of the emulated blockchain
      MockChainState -> Params
mcstParams :: Emulator.Params,
      -- | The ledger state of the emulated blockchain
      MockChainState -> EmulatedLedgerState
mcstLedgerState :: Emulator.EmulatedLedgerState,
      -- | Associates to each 'Api.TxOutRef' the 'TxSkelOut' that produced it,
      -- alongside a boolean to state whether this UTxO is still present in the
      -- index ('True') or has already been consumed ('False').
      MockChainState -> Map TxOutRef (TxSkelOut, Bool)
mcstOutputs :: Map Api.TxOutRef (TxSkelOut, Bool),
      -- | The constitution script to be used with proposals
      MockChainState -> Maybe VScript
mcstConstitution :: Maybe VScript
    } ->
    MockChainState
  deriving (Int -> MockChainState -> ShowS
[MockChainState] -> ShowS
MockChainState -> String
(Int -> MockChainState -> ShowS)
-> (MockChainState -> String)
-> ([MockChainState] -> ShowS)
-> Show MockChainState
forall (a :: OpticKind).
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> MockChainState -> ShowS
showsPrec :: Int -> MockChainState -> ShowS
$cshow :: MockChainState -> String
show :: MockChainState -> String
$cshowList :: [MockChainState] -> ShowS
showList :: [MockChainState] -> ShowS
Show)

-- | A lens to set or get the parameters of the 'MockChainState'
makeLensesFor [("mcstParams", "mcstParamsL")] ''MockChainState

-- | A lens to set or get the ledger state of the 'MockChainState'
makeLensesFor [("mcstLedgerState", "mcstLedgerStateL")] ''MockChainState

-- | A lens to set or get the outputs of the 'MockChainState'
makeLensesFor [("mcstOutputs", "mcstOutputsL")] ''MockChainState

-- | A lens to set or get the constitution script of the 'MockChainState'
makeLensesFor [("mcstConstitution", "mcstConstitutionL")] ''MockChainState

instance Default MockChainState where
  def :: MockChainState
def = Params
-> EmulatedLedgerState
-> Map TxOutRef (TxSkelOut, Bool)
-> Maybe VScript
-> MockChainState
MockChainState Params
forall (a :: OpticKind). Default a => a
def (Params -> EmulatedLedgerState
Emulator.initialState Params
forall (a :: OpticKind). Default a => a
def) Map TxOutRef (TxSkelOut, Bool)
forall (k :: OpticKind) (a :: OpticKind). Map k a
Map.empty Maybe VScript
forall (a :: OpticKind). Maybe a
Nothing

-- | Builds a 'UtxoState' from a 'MockChainState'
mcstToUtxoState :: MockChainState -> UtxoState
mcstToUtxoState :: MockChainState -> UtxoState
mcstToUtxoState =
  (UtxoState -> (TxOutRef, (TxSkelOut, Bool)) -> UtxoState)
-> UtxoState -> [(TxOutRef, (TxSkelOut, Bool))] -> UtxoState
forall (b :: OpticKind) (a :: OpticKind).
(b -> a -> b) -> b -> [a] -> b
forall (t :: OpticKind -> OpticKind) (b :: OpticKind)
       (a :: OpticKind).
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl UtxoState -> (TxOutRef, (TxSkelOut, Bool)) -> UtxoState
extractPayload UtxoState
forall (a :: OpticKind). Monoid a => a
mempty ([(TxOutRef, (TxSkelOut, Bool))] -> UtxoState)
-> (MockChainState -> [(TxOutRef, (TxSkelOut, Bool))])
-> MockChainState
-> UtxoState
forall (b :: OpticKind) (c :: OpticKind) (a :: OpticKind).
(b -> c) -> (a -> b) -> a -> c
. Map TxOutRef (TxSkelOut, Bool) -> [(TxOutRef, (TxSkelOut, Bool))]
forall (k :: OpticKind) (a :: OpticKind). Map k a -> [(k, a)]
Map.toList (Map TxOutRef (TxSkelOut, Bool) -> [(TxOutRef, (TxSkelOut, Bool))])
-> (MockChainState -> Map TxOutRef (TxSkelOut, Bool))
-> MockChainState
-> [(TxOutRef, (TxSkelOut, Bool))]
forall (b :: OpticKind) (c :: OpticKind) (a :: OpticKind).
(b -> c) -> (a -> b) -> a -> c
. MockChainState -> Map TxOutRef (TxSkelOut, Bool)
mcstOutputs
  where
    extractPayload :: UtxoState -> (Api.TxOutRef, (TxSkelOut, Bool)) -> UtxoState
    extractPayload :: UtxoState -> (TxOutRef, (TxSkelOut, Bool)) -> UtxoState
extractPayload UtxoState
utxoState (TxOutRef
txOutRef, (TxSkelOut
txSkelOut, Bool
bool)) =
      let newAddress :: Address
newAddress = Optic' A_Getter NoIx TxSkelOut Address -> TxSkelOut -> Address
forall (k :: OpticKind) (is :: IxList) (s :: OpticKind)
       (a :: OpticKind).
Is k A_Getter =>
Optic' k is s a -> s -> a
view Optic' A_Getter NoIx TxSkelOut Address
txSkelOutAddressG TxSkelOut
txSkelOut
          newPayloadSet :: UtxoPayloadSet
newPayloadSet =
            [UtxoPayload] -> UtxoPayloadSet
UtxoPayloadSet
              [ TxOutRef
-> Value -> UtxoPayloadDatum -> Maybe ScriptHash -> UtxoPayload
UtxoPayload
                  TxOutRef
txOutRef
                  (Optic' A_Lens NoIx TxSkelOut Value -> TxSkelOut -> Value
forall (k :: OpticKind) (is :: IxList) (s :: OpticKind)
       (a :: OpticKind).
Is k A_Getter =>
Optic' k is s a -> s -> a
view Optic' A_Lens NoIx TxSkelOut Value
txSkelOutValueL TxSkelOut
txSkelOut)
                  ( case TxSkelOut
txSkelOut TxSkelOut
-> Optic' A_Lens NoIx TxSkelOut TxSkelOutDatum -> TxSkelOutDatum
forall (k :: OpticKind) (s :: OpticKind) (is :: IxList)
       (a :: OpticKind).
Is k A_Getter =>
s -> Optic' k is s a -> a
^. Optic' A_Lens NoIx TxSkelOut TxSkelOutDatum
txSkelOutDatumL of
                      TxSkelOutDatum
NoTxSkelOutDatum -> UtxoPayloadDatum
NoUtxoPayloadDatum
                      SomeTxSkelOutDatum dat
content DatumKind
kind -> dat -> Bool -> UtxoPayloadDatum
forall (dat :: OpticKind).
DatumConstrs dat =>
dat -> Bool -> UtxoPayloadDatum
SomeUtxoPayloadDatum dat
content (DatumKind
kind DatumKind -> DatumKind -> Bool
forall (a :: OpticKind). Eq a => a -> a -> Bool
/= DatumKind
Inline)
                  )
                  (Optic' An_AffineFold NoIx TxSkelOut ScriptHash
-> TxSkelOut -> Maybe ScriptHash
forall (k :: OpticKind) (is :: IxList) (s :: OpticKind)
       (a :: OpticKind).
Is k An_AffineFold =>
Optic' k is s a -> s -> Maybe a
preview Optic' An_AffineFold NoIx TxSkelOut ScriptHash
txSkelOutReferenceScriptHashAF TxSkelOut
txSkelOut)
              ]
       in if Bool
bool
            then UtxoState
utxoState {availableUtxos = Map.insertWith (<>) newAddress newPayloadSet (availableUtxos utxoState)}
            else UtxoState
utxoState {consumedUtxos = Map.insertWith (<>) newAddress newPayloadSet (consumedUtxos utxoState)}

-- | Stores an output in a 'MockChainState'
addOutput :: Api.TxOutRef -> TxSkelOut -> MockChainState -> MockChainState
addOutput :: TxOutRef -> TxSkelOut -> MockChainState -> MockChainState
addOutput TxOutRef
oRef TxSkelOut
txSkelOut = Lens' MockChainState (Map TxOutRef (TxSkelOut, Bool))
-> (Map TxOutRef (TxSkelOut, Bool)
    -> Map TxOutRef (TxSkelOut, Bool))
-> MockChainState
-> MockChainState
forall (k :: OpticKind) (is :: IxList) (s :: OpticKind)
       (t :: OpticKind) (a :: OpticKind) (b :: OpticKind).
Is k A_Setter =>
Optic k is s t a b -> (a -> b) -> s -> t
over Lens' MockChainState (Map TxOutRef (TxSkelOut, Bool))
mcstOutputsL (TxOutRef
-> (TxSkelOut, Bool)
-> Map TxOutRef (TxSkelOut, Bool)
-> Map TxOutRef (TxSkelOut, Bool)
forall (k :: OpticKind) (a :: OpticKind).
Ord k =>
k -> a -> Map k a -> Map k a
Map.insert TxOutRef
oRef (TxSkelOut
txSkelOut, Bool
True))

-- | Removes an output from the 'MockChainState'
removeOutput :: Api.TxOutRef -> MockChainState -> MockChainState
removeOutput :: TxOutRef -> MockChainState -> MockChainState
removeOutput TxOutRef
oRef = Lens' MockChainState (Map TxOutRef (TxSkelOut, Bool))
-> (Map TxOutRef (TxSkelOut, Bool)
    -> Map TxOutRef (TxSkelOut, Bool))
-> MockChainState
-> MockChainState
forall (k :: OpticKind) (is :: IxList) (s :: OpticKind)
       (t :: OpticKind) (a :: OpticKind) (b :: OpticKind).
Is k A_Setter =>
Optic k is s t a b -> (a -> b) -> s -> t
over Lens' MockChainState (Map TxOutRef (TxSkelOut, Bool))
mcstOutputsL (((TxSkelOut, Bool) -> Maybe (TxSkelOut, Bool))
-> TxOutRef
-> Map TxOutRef (TxSkelOut, Bool)
-> Map TxOutRef (TxSkelOut, Bool)
forall (k :: OpticKind) (a :: OpticKind).
Ord k =>
(a -> Maybe a) -> k -> Map k a -> Map k a
Map.update (\(TxSkelOut
output, Bool
_) -> (TxSkelOut, Bool) -> Maybe (TxSkelOut, Bool)
forall (a :: OpticKind). a -> Maybe a
Just (TxSkelOut
output, Bool
False)) TxOutRef
oRef)