-- | 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,
    mockChainState0From,
    mockChainState0,
  )
where

import Cardano.Api qualified as Cardano
import Cardano.Api.Shelley qualified as Cardano
import Cardano.Ledger.Shelley.API qualified as Shelley
import Cardano.Node.Emulator.Internal.Node qualified as Emulator
import Control.Lens qualified as Lens
import Control.Monad.Except
import Cooked.InitialDistribution
import Cooked.MockChain.BlockChain
import Cooked.MockChain.GenerateTx
import Cooked.MockChain.GenerateTx.Output
import Cooked.MockChain.MinAda
import Cooked.MockChain.UtxoState
import Cooked.Skeleton
import Data.Default
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Ledger.Index qualified as Ledger
import Ledger.Orphans ()
import Ledger.Tx qualified as Ledger
import Ledger.Tx.CardanoAPI qualified as Ledger
import Optics.Core
import Optics.TH
import Plutus.Script.Utils.Scripts qualified as Script
import PlutusLedgerApi.V3 qualified as Api

-- | The state used to run the simulation in 'Cooked.MockChain.Direct'
data MockChainState = MockChainState
  { MockChainState -> Params
mcstParams :: Emulator.Params,
    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 (Versioned Script)
mcstConstitution :: Maybe (Script.Versioned Script.Script)
  }
  deriving (Int -> MockChainState -> ShowS
[MockChainState] -> ShowS
MockChainState -> String
(Int -> MockChainState -> ShowS)
-> (MockChainState -> String)
-> ([MockChainState] -> ShowS)
-> Show MockChainState
forall a.
(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 (Versioned Script)
-> MockChainState
MockChainState Params
forall a. Default a => a
def (Params -> EmulatedLedgerState
Emulator.initialState Params
forall a. Default a => a
def) Map TxOutRef (TxSkelOut, Bool)
forall k a. Map k a
Map.empty Maybe (Versioned Script)
forall a. 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 a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl UtxoState -> (TxOutRef, (TxSkelOut, Bool)) -> UtxoState
extractPayload UtxoState
forall a. Monoid a => a
mempty ([(TxOutRef, (TxSkelOut, Bool))] -> UtxoState)
-> (MockChainState -> [(TxOutRef, (TxSkelOut, Bool))])
-> MockChainState
-> UtxoState
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map TxOutRef (TxSkelOut, Bool) -> [(TxOutRef, (TxSkelOut, Bool))]
forall k a. 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 c a. (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 = TxSkelOut -> Address
txSkelOutAddress TxSkelOut
txSkelOut
          newPayloadSet :: UtxoPayloadSet
newPayloadSet =
            [UtxoPayload] -> UtxoPayloadSet
UtxoPayloadSet
              [ TxOutRef
-> Value
-> Maybe (DatumContent, Bool)
-> Maybe ScriptHash
-> UtxoPayload
UtxoPayload
                  TxOutRef
txOutRef
                  (TxSkelOut -> Value
txSkelOutValue TxSkelOut
txSkelOut)
                  ( case TxSkelOut
txSkelOut TxSkelOut
-> Optic' A_Lens NoIx TxSkelOut TxSkelOutDatum -> TxSkelOutDatum
forall k s (is :: IxList) a.
Is k A_Getter =>
s -> Optic' k is s a -> a
^. Optic' A_Lens NoIx TxSkelOut TxSkelOutDatum
txSkelOutDatumL of
                      TxSkelOutDatum
TxSkelOutNoDatum -> Maybe (DatumContent, Bool)
forall a. Maybe a
Nothing
                      TxSkelOutSomeDatum DatumContent
content DatumKind
Inline -> (DatumContent, Bool) -> Maybe (DatumContent, Bool)
forall a. a -> Maybe a
Just (DatumContent
content, Bool
False)
                      TxSkelOutSomeDatum DatumContent
content DatumKind
_ -> (DatumContent, Bool) -> Maybe (DatumContent, Bool)
forall a. a -> Maybe a
Just (DatumContent
content, Bool
True)
                  )
                  (TxSkelOut -> Maybe ScriptHash
txSkelOutReferenceScriptHash 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 (is :: IxList) s t a b.
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 a. 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 (is :: IxList) s t a b.
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 a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
Map.update (\(TxSkelOut
output, Bool
_) -> (TxSkelOut, Bool) -> Maybe (TxSkelOut, Bool)
forall a. a -> Maybe a
Just (TxSkelOut
output, Bool
False)) TxOutRef
oRef)

-- | This creates the initial 'MockChainState' from an initial distribution by
-- submitting an initial transaction with the appropriate content. The genesis
-- key hash has been taken from
-- https://github.com/input-output-hk/cardano-node/blob/543b267d75d3d448e1940f9ec04b42bd01bbb16b/cardano-api/test/Test/Cardano/Api/Genesis.hs#L60
mockChainState0From :: (MonadBlockChainBalancing m) => InitialDistribution -> m MockChainState
mockChainState0From :: forall (m :: * -> *).
MonadBlockChainBalancing m =>
InitialDistribution -> m MockChainState
mockChainState0From (InitialDistribution [TxSkelOut]
initDist) = do
  Params
params <- m Params
forall (m :: * -> *). MonadBlockChainBalancing m => m Params
getParams
  let genesisKeyHash :: Hash GenesisUTxOKey
genesisKeyHash = KeyHash 'Payment -> Hash GenesisUTxOKey
Cardano.GenesisUTxOKeyHash (KeyHash 'Payment -> Hash GenesisUTxOKey)
-> KeyHash 'Payment -> Hash GenesisUTxOKey
forall a b. (a -> b) -> a -> b
$ Hash ADDRHASH (VerKeyDSIGN DSIGN) -> KeyHash 'Payment
forall (r :: KeyRole).
Hash ADDRHASH (VerKeyDSIGN DSIGN) -> KeyHash r
Shelley.KeyHash Hash ADDRHASH (VerKeyDSIGN DSIGN)
"23d51e91ae5adc7ae801e9de4cd54175fb7464ec2680b25686bbb194"
      inputs :: [(TxIn, BuildTxWith BuildTx (Witness WitCtxTxIn ConwayEra))]
inputs = [(NetworkId -> Hash GenesisUTxOKey -> TxIn
Cardano.genesisUTxOPseudoTxIn (Params -> NetworkId
Emulator.pNetworkId Params
params) Hash GenesisUTxOKey
genesisKeyHash, Witness WitCtxTxIn ConwayEra
-> BuildTxWith BuildTx (Witness WitCtxTxIn ConwayEra)
forall a. a -> BuildTxWith BuildTx a
Cardano.BuildTxWith (Witness WitCtxTxIn ConwayEra
 -> BuildTxWith BuildTx (Witness WitCtxTxIn ConwayEra))
-> Witness WitCtxTxIn ConwayEra
-> BuildTxWith BuildTx (Witness WitCtxTxIn ConwayEra)
forall a b. (a -> b) -> a -> b
$ KeyWitnessInCtx WitCtxTxIn -> Witness WitCtxTxIn ConwayEra
forall witctx era. KeyWitnessInCtx witctx -> Witness witctx era
Cardano.KeyWitness KeyWitnessInCtx WitCtxTxIn
Cardano.KeyWitnessForSpending)]
  [TxSkelOut]
outputsMinAda <- (TxSkelOut -> m TxSkelOut) -> [TxSkelOut] -> m [TxSkelOut]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM TxSkelOut -> m TxSkelOut
forall (m :: * -> *).
MonadBlockChainBalancing m =>
TxSkelOut -> m TxSkelOut
toTxSkelOutWithMinAda [TxSkelOut]
initDist
  [TxOut CtxTx ConwayEra]
outputs <- (TxSkelOut -> m (TxOut CtxTx ConwayEra))
-> [TxSkelOut] -> m [TxOut CtxTx ConwayEra]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM TxSkelOut -> m (TxOut CtxTx ConwayEra)
forall (m :: * -> *).
MonadBlockChainBalancing m =>
TxSkelOut -> m (TxOut CtxTx ConwayEra)
toCardanoTxOut [TxSkelOut]
outputsMinAda
  CardanoTx
cardanoTx <-
    Tx ConwayEra -> CardanoTx
Ledger.CardanoEmulatorEraTx (Tx ConwayEra -> CardanoTx)
-> (TxBody ConwayEra -> Tx ConwayEra)
-> TxBody ConwayEra
-> CardanoTx
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Wallet] -> TxBody ConwayEra -> Tx ConwayEra
txSignersAndBodyToCardanoTx []
      (TxBody ConwayEra -> CardanoTx)
-> m (TxBody ConwayEra) -> m CardanoTx
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ToCardanoError -> m (TxBody ConwayEra))
-> (TxBody ConwayEra -> m (TxBody ConwayEra))
-> Either ToCardanoError (TxBody ConwayEra)
-> m (TxBody ConwayEra)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either
        (MockChainError -> m (TxBody ConwayEra)
forall a. MockChainError -> m a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (MockChainError -> m (TxBody ConwayEra))
-> (ToCardanoError -> MockChainError)
-> ToCardanoError
-> m (TxBody ConwayEra)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ToCardanoError -> MockChainError
MCEToCardanoError String
"generateTx :")
        TxBody ConwayEra -> m (TxBody ConwayEra)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return
        (Params
-> CardanoBuildTx -> Either ToCardanoError (TxBody ConwayEra)
Emulator.createTransactionBody Params
params (CardanoBuildTx -> Either ToCardanoError (TxBody ConwayEra))
-> CardanoBuildTx -> Either ToCardanoError (TxBody ConwayEra)
forall a b. (a -> b) -> a -> b
$ TxBodyContent BuildTx ConwayEra -> CardanoBuildTx
Ledger.CardanoBuildTx (TxBodyContent BuildTx ConwayEra
Ledger.emptyTxBodyContent {Cardano.txOuts = outputs, Cardano.txIns = inputs}))
  let index :: UTxO ConwayEra
index = UtxoIndex -> UTxO ConwayEra
Ledger.fromPlutusIndex (UtxoIndex -> UTxO ConwayEra) -> UtxoIndex -> UTxO ConwayEra
forall a b. (a -> b) -> a -> b
$ Blockchain -> UtxoIndex
Ledger.initialise [[CardanoTx -> OnChainTx
Emulator.unsafeMakeValid CardanoTx
cardanoTx]]
      outputsMap :: Map TxOutRef (TxSkelOut, Bool)
outputsMap =
        [(TxOutRef, (TxSkelOut, Bool))] -> Map TxOutRef (TxSkelOut, Bool)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(TxOutRef, (TxSkelOut, Bool))] -> Map TxOutRef (TxSkelOut, Bool))
-> [(TxOutRef, (TxSkelOut, Bool))]
-> Map TxOutRef (TxSkelOut, Bool)
forall a b. (a -> b) -> a -> b
$
          (TxOutRef -> TxSkelOut -> (TxOutRef, (TxSkelOut, Bool)))
-> [TxOutRef] -> [TxSkelOut] -> [(TxOutRef, (TxSkelOut, Bool))]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
            (\TxOutRef
x TxSkelOut
y -> (TxOutRef
x, (TxSkelOut
y, Bool
True)))
            (TxIn -> TxOutRef
Ledger.fromCardanoTxIn (TxIn -> TxOutRef)
-> ((TxOut, TxIn) -> TxIn) -> (TxOut, TxIn) -> TxOutRef
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TxOut, TxIn) -> TxIn
forall a b. (a, b) -> b
snd ((TxOut, TxIn) -> TxOutRef) -> [(TxOut, TxIn)] -> [TxOutRef]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> CardanoTx -> [(TxOut, TxIn)]
Ledger.getCardanoTxOutRefs CardanoTx
cardanoTx)
            [TxSkelOut]
outputsMinAda
  MockChainState -> m MockChainState
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (MockChainState -> m MockChainState)
-> MockChainState -> m MockChainState
forall a b. (a -> b) -> a -> b
$ Params
-> EmulatedLedgerState
-> Map TxOutRef (TxSkelOut, Bool)
-> Maybe (Versioned Script)
-> MockChainState
MockChainState Params
forall a. Default a => a
def (ASetter
  EmulatedLedgerState
  EmulatedLedgerState
  (UTxO ConwayEra)
  (UTxO ConwayEra)
-> UTxO ConwayEra -> EmulatedLedgerState -> EmulatedLedgerState
forall s t a b. ASetter s t a b -> b -> s -> t
Lens.set ASetter
  EmulatedLedgerState
  EmulatedLedgerState
  (UTxO ConwayEra)
  (UTxO ConwayEra)
Lens' EmulatedLedgerState (UTxO ConwayEra)
Emulator.elsUtxoL UTxO ConwayEra
index (Params -> EmulatedLedgerState
Emulator.initialState Params
forall a. Default a => a
def)) Map TxOutRef (TxSkelOut, Bool)
outputsMap Maybe (Versioned Script)
forall a. Maybe a
Nothing

-- | Same as 'mockChainState0From' with the default 'InitialDistribution'
mockChainState0 :: (MonadBlockChainBalancing m) => m MockChainState
mockChainState0 :: forall (m :: * -> *).
MonadBlockChainBalancing m =>
m MockChainState
mockChainState0 = InitialDistribution -> m MockChainState
forall (m :: * -> *).
MonadBlockChainBalancing m =>
InitialDistribution -> m MockChainState
mockChainState0From InitialDistribution
forall a. Default a => a
def