{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase       #-}
{-# LANGUAGE NamedFieldPuns   #-}
{-# LANGUAGE RecordWildCards  #-}

module Wallet.Rollup
    ( doAnnotateBlockchain
    , initialRollup
    , annotateBlockchain
    , Rollup
    -- * Chain event fold
    , initialState
    , handleChainEvent
    , getAnnotatedTransactions
    ) where

import Cardano.Api qualified as C
import Cardano.Node.Emulator.Chain (ChainEvent (..))
import Control.Lens (assign, ifoldr, over, set, use, view, (&), (^.))
import Control.Lens.Combinators (itraverse)
import Control.Monad.State (StateT, evalStateT, runState)
import Data.List (groupBy)
import Data.Map (Map)
import Data.Map qualified as Map
import Ledger (Block, Blockchain, OnChainTx (..), TxIn (TxIn), TxOut, ValidationPhase (..), Value, consumableInputs,
               onChainTxIsValid, outputsProduced, txInRef, txOutRefId, txOutRefIdx, txOutValue, unOnChain)
import Ledger.Tx qualified as Tx
import Wallet.Rollup.Types

------------------------------------------------------------
txInputKey :: TxIn -> TxKey
txInputKey :: TxIn -> TxKey
txInputKey TxIn {TxOutRef
txInRef :: TxOutRef
txInRef :: TxIn -> TxOutRef
txInRef} =
    TxKey :: TxId -> Integer -> TxKey
TxKey
        { _txKeyTxId :: TxId
_txKeyTxId = TxOutRef -> TxId
txOutRefId TxOutRef
txInRef
        , _txKeyTxOutRefIdx :: Integer
_txKeyTxOutRefIdx = TxOutRef -> Integer
txOutRefIdx TxOutRef
txInRef
        }

annotateTransaction ::
       Monad m => SequenceId -> OnChainTx -> StateT Rollup m AnnotatedTx
annotateTransaction :: SequenceId -> OnChainTx -> StateT Rollup m AnnotatedTx
annotateTransaction SequenceId
sequenceId OnChainTx
tx = do
    Map TxKey TxOut
cPreviousOutputs <- Getting (Map TxKey TxOut) Rollup (Map TxKey TxOut)
-> StateT Rollup m (Map TxKey TxOut)
forall s (m :: * -> *) a. MonadState s m => Getting a s a -> m a
use Getting (Map TxKey TxOut) Rollup (Map TxKey TxOut)
Lens' Rollup (Map TxKey TxOut)
previousOutputs
    Map BeneficialOwner Value
cRollingBalances <- Getting
  (Map BeneficialOwner Value) Rollup (Map BeneficialOwner Value)
-> StateT Rollup m (Map BeneficialOwner Value)
forall s (m :: * -> *) a. MonadState s m => Getting a s a -> m a
use Getting
  (Map BeneficialOwner Value) Rollup (Map BeneficialOwner Value)
Lens' Rollup (Map BeneficialOwner Value)
rollingBalances
    [DereferencedInput]
dereferencedInputs <-
        (TxIn -> StateT Rollup m DereferencedInput)
-> [TxIn] -> StateT Rollup m [DereferencedInput]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse
            (\TxIn
txIn ->
                 let key :: TxKey
key = TxIn -> TxKey
txInputKey TxIn
txIn
                  in case TxKey -> Map TxKey TxOut -> Maybe TxOut
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup TxKey
key Map TxKey TxOut
cPreviousOutputs of
                         Just TxOut
txOut -> DereferencedInput -> StateT Rollup m DereferencedInput
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DereferencedInput -> StateT Rollup m DereferencedInput)
-> DereferencedInput -> StateT Rollup m DereferencedInput
forall a b. (a -> b) -> a -> b
$ TxIn -> TxOut -> DereferencedInput
DereferencedInput TxIn
txIn TxOut
txOut
                         Maybe TxOut
Nothing    -> DereferencedInput -> StateT Rollup m DereferencedInput
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DereferencedInput -> StateT Rollup m DereferencedInput)
-> DereferencedInput -> StateT Rollup m DereferencedInput
forall a b. (a -> b) -> a -> b
$ TxKey -> DereferencedInput
InputNotFound TxKey
key)
            (OnChainTx -> [TxIn]
consumableInputs OnChainTx
tx)
    let txId :: TxId
txId = CardanoTx -> TxId
Tx.getCardanoTxId (CardanoTx -> TxId) -> CardanoTx -> TxId
forall a b. (a -> b) -> a -> b
$ OnChainTx -> CardanoTx
unOnChain OnChainTx
tx
        txOuts :: [TxOut]
txOuts = Map TxOutRef TxOut -> [TxOut]
forall k a. Map k a -> [a]
Map.elems (Map TxOutRef TxOut -> [TxOut]) -> Map TxOutRef TxOut -> [TxOut]
forall a b. (a -> b) -> a -> b
$ OnChainTx -> Map TxOutRef TxOut
outputsProduced OnChainTx
tx
        newOutputs :: Map TxKey TxOut
newOutputs =
            (Int -> TxOut -> Map TxKey TxOut -> Map TxKey TxOut)
-> Map TxKey TxOut -> [TxOut] -> Map TxKey TxOut
forall i (f :: * -> *) a b.
FoldableWithIndex i f =>
(i -> a -> b -> b) -> b -> f a -> b
ifoldr
                (\Int
outputIndex ->
                     TxKey -> TxOut -> Map TxKey TxOut -> Map TxKey TxOut
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert
                         TxKey :: TxId -> Integer -> TxKey
TxKey
                             { _txKeyTxId :: TxId
_txKeyTxId = TxId
txId
                             , _txKeyTxOutRefIdx :: Integer
_txKeyTxOutRefIdx = Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
outputIndex
                             })
                Map TxKey TxOut
cPreviousOutputs
                [TxOut]
txOuts
        newBalances :: Map BeneficialOwner Value
newBalances =
            (TxOut -> Map BeneficialOwner Value -> Map BeneficialOwner Value)
-> Map BeneficialOwner Value
-> [TxOut]
-> Map BeneficialOwner Value
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr
                TxOut -> Map BeneficialOwner Value -> Map BeneficialOwner Value
sumAccounts
                Map BeneficialOwner Value
cRollingBalances
                ((ASetter TxOut TxOut (TxOutValue BabbageEra) (TxOutValue BabbageEra)
-> (TxOutValue BabbageEra -> TxOutValue BabbageEra)
-> TxOut
-> TxOut
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter TxOut TxOut (TxOutValue BabbageEra) (TxOutValue BabbageEra)
Lens' TxOut (TxOutValue BabbageEra)
Tx.outValue' TxOutValue BabbageEra -> TxOutValue BabbageEra
forall era. TxOutValue era -> TxOutValue era
negateValue (TxOut -> TxOut)
-> (DereferencedInput -> TxOut) -> DereferencedInput -> TxOut
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DereferencedInput -> TxOut
refersTo (DereferencedInput -> TxOut) -> [DereferencedInput] -> [TxOut]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (DereferencedInput -> Bool)
-> [DereferencedInput] -> [DereferencedInput]
forall a. (a -> Bool) -> [a] -> [a]
filter DereferencedInput -> Bool
isFound [DereferencedInput]
dereferencedInputs) [TxOut] -> [TxOut] -> [TxOut]
forall a. Semigroup a => a -> a -> a
<>
                 [TxOut]
txOuts)
          where
            negateValue :: C.TxOutValue era -> C.TxOutValue era
            negateValue :: TxOutValue era -> TxOutValue era
negateValue  (C.TxOutAdaOnly OnlyAdaSupportedInEra era
wit Lovelace
l) = OnlyAdaSupportedInEra era -> Lovelace -> TxOutValue era
forall era. OnlyAdaSupportedInEra era -> Lovelace -> TxOutValue era
C.TxOutAdaOnly OnlyAdaSupportedInEra era
wit (Lovelace -> Lovelace
forall a. Num a => a -> a
negate Lovelace
l)
            negateValue  (C.TxOutValue MultiAssetSupportedInEra era
wit Value
v)   = MultiAssetSupportedInEra era -> Value -> TxOutValue era
forall era. MultiAssetSupportedInEra era -> Value -> TxOutValue era
C.TxOutValue MultiAssetSupportedInEra era
wit (Value -> Value
C.negateValue Value
v)
        sumAccounts ::
               TxOut -> Map BeneficialOwner Value -> Map BeneficialOwner Value
        sumAccounts :: TxOut -> Map BeneficialOwner Value -> Map BeneficialOwner Value
sumAccounts TxOut
txOut =
            (Maybe Value -> Maybe Value)
-> BeneficialOwner
-> Map BeneficialOwner Value
-> Map BeneficialOwner Value
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
Map.alter Maybe Value -> Maybe Value
sumBalances (TxOut -> BeneficialOwner
toBeneficialOwner TxOut
txOut)
          where
            sumBalances :: Maybe Value -> Maybe Value
            sumBalances :: Maybe Value -> Maybe Value
sumBalances Maybe Value
Nothing         = Value -> Maybe Value
forall a. a -> Maybe a
Just (TxOut -> Value
txOutValue TxOut
txOut)
            sumBalances (Just Value
oldValue) = Value -> Maybe Value
forall a. a -> Maybe a
Just (Value
oldValue Value -> Value -> Value
forall a. Semigroup a => a -> a -> a
<> TxOut -> Value
txOutValue TxOut
txOut)
    ASetter Rollup Rollup (Map TxKey TxOut) (Map TxKey TxOut)
-> Map TxKey TxOut -> StateT Rollup m ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
assign ASetter Rollup Rollup (Map TxKey TxOut) (Map TxKey TxOut)
Lens' Rollup (Map TxKey TxOut)
previousOutputs Map TxKey TxOut
newOutputs
    ASetter
  Rollup
  Rollup
  (Map BeneficialOwner Value)
  (Map BeneficialOwner Value)
-> Map BeneficialOwner Value -> StateT Rollup m ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
assign ASetter
  Rollup
  Rollup
  (Map BeneficialOwner Value)
  (Map BeneficialOwner Value)
Lens' Rollup (Map BeneficialOwner Value)
rollingBalances Map BeneficialOwner Value
newBalances
    AnnotatedTx -> StateT Rollup m AnnotatedTx
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AnnotatedTx -> StateT Rollup m AnnotatedTx)
-> AnnotatedTx -> StateT Rollup m AnnotatedTx
forall a b. (a -> b) -> a -> b
$
        AnnotatedTx :: SequenceId
-> TxId
-> CardanoTx
-> [DereferencedInput]
-> Map BeneficialOwner Value
-> Bool
-> AnnotatedTx
AnnotatedTx
            { SequenceId
sequenceId :: SequenceId
sequenceId :: SequenceId
sequenceId
            , TxId
txId :: TxId
txId :: TxId
txId
            , tx :: CardanoTx
tx = OnChainTx -> CardanoTx
unOnChain OnChainTx
tx
            , [DereferencedInput]
dereferencedInputs :: [DereferencedInput]
dereferencedInputs :: [DereferencedInput]
dereferencedInputs
            , balances :: Map BeneficialOwner Value
balances = Map BeneficialOwner Value
newBalances
            , valid :: Bool
valid = OnChainTx -> Bool
onChainTxIsValid OnChainTx
tx
            }

annotateChainSlot :: Monad m => Int -> Block -> StateT Rollup m [AnnotatedTx]
annotateChainSlot :: Int -> Block -> StateT Rollup m [AnnotatedTx]
annotateChainSlot Int
slotIndex =
    (Int -> OnChainTx -> StateT Rollup m AnnotatedTx)
-> Block -> StateT Rollup m [AnnotatedTx]
forall i (t :: * -> *) (f :: * -> *) a b.
(TraversableWithIndex i t, Applicative f) =>
(i -> a -> f b) -> t a -> f (t b)
itraverse (\Int
txIndex -> SequenceId -> OnChainTx -> StateT Rollup m AnnotatedTx
forall (m :: * -> *).
Monad m =>
SequenceId -> OnChainTx -> StateT Rollup m AnnotatedTx
annotateTransaction SequenceId :: Int -> Int -> SequenceId
SequenceId {Int
txIndex :: Int
slotIndex :: Int
txIndex :: Int
slotIndex :: Int
..})

annotateBlockchain :: Monad m => Blockchain -> StateT Rollup m [[AnnotatedTx]]
annotateBlockchain :: Blockchain -> StateT Rollup m [[AnnotatedTx]]
annotateBlockchain = ([[AnnotatedTx]] -> [[AnnotatedTx]])
-> StateT Rollup m [[AnnotatedTx]]
-> StateT Rollup m [[AnnotatedTx]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[AnnotatedTx]] -> [[AnnotatedTx]]
forall a. [a] -> [a]
reverse (StateT Rollup m [[AnnotatedTx]]
 -> StateT Rollup m [[AnnotatedTx]])
-> (Blockchain -> StateT Rollup m [[AnnotatedTx]])
-> Blockchain
-> StateT Rollup m [[AnnotatedTx]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Block -> StateT Rollup m [AnnotatedTx])
-> Blockchain -> StateT Rollup m [[AnnotatedTx]]
forall i (t :: * -> *) (f :: * -> *) a b.
(TraversableWithIndex i t, Applicative f) =>
(i -> a -> f b) -> t a -> f (t b)
itraverse Int -> Block -> StateT Rollup m [AnnotatedTx]
forall (m :: * -> *).
Monad m =>
Int -> Block -> StateT Rollup m [AnnotatedTx]
annotateChainSlot (Blockchain -> StateT Rollup m [[AnnotatedTx]])
-> (Blockchain -> Blockchain)
-> Blockchain
-> StateT Rollup m [[AnnotatedTx]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Blockchain -> Blockchain
forall a. [a] -> [a]
reverse

initialRollup :: Rollup
initialRollup :: Rollup
initialRollup =
    Rollup :: Map TxKey TxOut -> Map BeneficialOwner Value -> Rollup
Rollup {_previousOutputs :: Map TxKey TxOut
_previousOutputs = Map TxKey TxOut
forall k a. Map k a
Map.empty, _rollingBalances :: Map BeneficialOwner Value
_rollingBalances = Map BeneficialOwner Value
forall k a. Map k a
Map.empty}

initialState :: RollupState
initialState :: RollupState
initialState =
    RollupState :: SequenceId -> Rollup -> [AnnotatedTx] -> RollupState
RollupState { _rollup :: Rollup
_rollup = Rollup
initialRollup, _annotatedTransactions :: [AnnotatedTx]
_annotatedTransactions = [], _currentSequenceId :: SequenceId
_currentSequenceId = Int -> Int -> SequenceId
SequenceId Int
0 Int
0 }

doAnnotateBlockchain :: Monad m => Blockchain -> m [[AnnotatedTx]]
doAnnotateBlockchain :: Blockchain -> m [[AnnotatedTx]]
doAnnotateBlockchain Blockchain
blockchain =
    StateT Rollup m [[AnnotatedTx]] -> Rollup -> m [[AnnotatedTx]]
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (Blockchain -> StateT Rollup m [[AnnotatedTx]]
forall (m :: * -> *).
Monad m =>
Blockchain -> StateT Rollup m [[AnnotatedTx]]
annotateBlockchain Blockchain
blockchain) Rollup
initialRollup

getAnnotatedTransactions :: RollupState -> [[AnnotatedTx]]
getAnnotatedTransactions :: RollupState -> [[AnnotatedTx]]
getAnnotatedTransactions = (AnnotatedTx -> AnnotatedTx -> Bool)
-> [AnnotatedTx] -> [[AnnotatedTx]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy ((AnnotatedTx -> Int) -> AnnotatedTx -> AnnotatedTx -> Bool
forall a b. Eq a => (b -> a) -> b -> b -> Bool
equating (SequenceId -> Int
slotIndex (SequenceId -> Int)
-> (AnnotatedTx -> SequenceId) -> AnnotatedTx -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AnnotatedTx -> SequenceId
sequenceId)) ([AnnotatedTx] -> [[AnnotatedTx]])
-> (RollupState -> [AnnotatedTx]) -> RollupState -> [[AnnotatedTx]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [AnnotatedTx] -> [AnnotatedTx]
forall a. [a] -> [a]
reverse ([AnnotatedTx] -> [AnnotatedTx])
-> (RollupState -> [AnnotatedTx]) -> RollupState -> [AnnotatedTx]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Getting [AnnotatedTx] RollupState [AnnotatedTx]
-> RollupState -> [AnnotatedTx]
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting [AnnotatedTx] RollupState [AnnotatedTx]
Lens' RollupState [AnnotatedTx]
annotatedTransactions

handleChainEvent :: RollupState -> ChainEvent -> RollupState
handleChainEvent :: RollupState -> ChainEvent -> RollupState
handleChainEvent RollupState
s = \case
    SlotAdd Slot
_                           -> RollupState
s RollupState -> (RollupState -> RollupState) -> RollupState
forall a b. a -> (a -> b) -> b
& ASetter RollupState RollupState SequenceId SequenceId
-> (SequenceId -> SequenceId) -> RollupState -> RollupState
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter RollupState RollupState SequenceId SequenceId
Lens' RollupState SequenceId
currentSequenceId (ASetter SequenceId SequenceId Int Int
-> Int -> SequenceId -> SequenceId
forall s t a b. ASetter s t a b -> b -> s -> t
set ASetter SequenceId SequenceId Int Int
Lens' SequenceId Int
txIndexL Int
0 (SequenceId -> SequenceId)
-> (SequenceId -> SequenceId) -> SequenceId -> SequenceId
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ASetter SequenceId SequenceId Int Int
-> (Int -> Int) -> SequenceId -> SequenceId
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter SequenceId SequenceId Int Int
Lens' SequenceId Int
slotIndexL Int -> Int
forall a. Enum a => a -> a
succ)
    TxnValidate TxId
_ CardanoTx
tx [Text]
_                  -> RollupState -> OnChainTx -> RollupState
addTx RollupState
s (CardanoTx -> OnChainTx
Valid CardanoTx
tx)
    TxnValidationFail ValidationPhase
Phase2 TxId
_ CardanoTx
tx ValidationError
_ Value
_ [Text]
_ -> RollupState -> OnChainTx -> RollupState
addTx RollupState
s (CardanoTx -> OnChainTx
Invalid CardanoTx
tx)
    ChainEvent
_                                   -> RollupState
s

addTx :: RollupState -> OnChainTx -> RollupState
addTx :: RollupState -> OnChainTx -> RollupState
addTx RollupState
s OnChainTx
tx =
    let (AnnotatedTx
tx', Rollup
newState) = State Rollup AnnotatedTx -> Rollup -> (AnnotatedTx, Rollup)
forall s a. State s a -> s -> (a, s)
runState (SequenceId -> OnChainTx -> State Rollup AnnotatedTx
forall (m :: * -> *).
Monad m =>
SequenceId -> OnChainTx -> StateT Rollup m AnnotatedTx
annotateTransaction (RollupState
s RollupState
-> Getting SequenceId RollupState SequenceId -> SequenceId
forall s a. s -> Getting a s a -> a
^. Getting SequenceId RollupState SequenceId
Lens' RollupState SequenceId
currentSequenceId) OnChainTx
tx) (RollupState
s RollupState -> Getting Rollup RollupState Rollup -> Rollup
forall s a. s -> Getting a s a -> a
^. Getting Rollup RollupState Rollup
Lens' RollupState Rollup
rollup)
    in RollupState
s RollupState -> (RollupState -> RollupState) -> RollupState
forall a b. a -> (a -> b) -> b
& ASetter RollupState RollupState SequenceId SequenceId
-> (SequenceId -> SequenceId) -> RollupState -> RollupState
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter RollupState RollupState SequenceId SequenceId
Lens' RollupState SequenceId
currentSequenceId (ASetter SequenceId SequenceId Int Int
-> (Int -> Int) -> SequenceId -> SequenceId
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter SequenceId SequenceId Int Int
Lens' SequenceId Int
txIndexL Int -> Int
forall a. Enum a => a -> a
succ)
         RollupState -> (RollupState -> RollupState) -> RollupState
forall a b. a -> (a -> b) -> b
& ASetter RollupState RollupState [AnnotatedTx] [AnnotatedTx]
-> ([AnnotatedTx] -> [AnnotatedTx]) -> RollupState -> RollupState
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter RollupState RollupState [AnnotatedTx] [AnnotatedTx]
Lens' RollupState [AnnotatedTx]
annotatedTransactions ((:) AnnotatedTx
tx')
         RollupState -> (RollupState -> RollupState) -> RollupState
forall a b. a -> (a -> b) -> b
& ASetter RollupState RollupState Rollup Rollup
-> Rollup -> RollupState -> RollupState
forall s t a b. ASetter s t a b -> b -> s -> t
set ASetter RollupState RollupState Rollup Rollup
Lens' RollupState Rollup
rollup Rollup
newState

-- https://hackage.haskell.org/package/Cabal-3.2.1.0/docs/src/Distribution.Utils.Generic.html#equating
equating :: Eq a => (b -> a) -> b -> b -> Bool
equating :: (b -> a) -> b -> b -> Bool
equating b -> a
p b
x b
y = b -> a
p b
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== b -> a
p b
y