-- | This module provides some 'Tweak's that add or remove inputs and outputs
-- from transactions. Some also operate on the minted value.
module Cooked.Tweak.AddInputsAndOutputs
  ( ensureInputTweak,
    addInputTweak,
    removeInputTweak,
    ensureOutputTweak,
    addOutputTweak,
    removeOutputTweak,
    addMintTweak,
    removeMintTweak,
  )
where

import Control.Monad
import Cooked.Skeleton
import Cooked.Tweak.Common
import Data.List
import Data.Map qualified as Map
import Optics.Core
import Plutus.Script.Utils.Scripts qualified as Script
import PlutusLedgerApi.V3 qualified as Api

-- * Adding and removing transaction inputs

-- | Ensure that a given 'Api.TxOutRef' is being spent with a given
-- 'TxSkelRedeemer'. The return value will be @Just@ the added data, if anything
-- changed.
ensureInputTweak :: (MonadTweak m) => Api.TxOutRef -> TxSkelRedeemer -> m (Maybe (Api.TxOutRef, TxSkelRedeemer))
ensureInputTweak :: forall (m :: * -> *).
MonadTweak m =>
TxOutRef -> TxSkelRedeemer -> m (Maybe (TxOutRef, TxSkelRedeemer))
ensureInputTweak TxOutRef
oref TxSkelRedeemer
howConsumed = do
  Map TxOutRef TxSkelRedeemer
presentInputs <- Optic' A_Lens NoIx TxSkel (Map TxOutRef TxSkelRedeemer)
-> m (Map TxOutRef TxSkelRedeemer)
forall (m :: * -> *) k (is :: IxList) a.
(MonadTweak m, Is k A_Getter) =>
Optic' k is TxSkel a -> m a
viewTweak Optic' A_Lens NoIx TxSkel (Map TxOutRef TxSkelRedeemer)
txSkelInsL
  if Map TxOutRef TxSkelRedeemer
presentInputs Map TxOutRef TxSkelRedeemer -> TxOutRef -> Maybe TxSkelRedeemer
forall k a. Ord k => Map k a -> k -> Maybe a
Map.!? TxOutRef
oref Maybe TxSkelRedeemer -> Maybe TxSkelRedeemer -> Bool
forall a. Eq a => a -> a -> Bool
== TxSkelRedeemer -> Maybe TxSkelRedeemer
forall a. a -> Maybe a
Just TxSkelRedeemer
howConsumed
    then Maybe (TxOutRef, TxSkelRedeemer)
-> m (Maybe (TxOutRef, TxSkelRedeemer))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (TxOutRef, TxSkelRedeemer)
forall a. Maybe a
Nothing
    else do
      Optic' A_Lens NoIx TxSkel (Map TxOutRef TxSkelRedeemer)
-> (Map TxOutRef TxSkelRedeemer -> Map TxOutRef TxSkelRedeemer)
-> m ()
forall (m :: * -> *) k (is :: IxList) a.
(MonadTweak m, Is k A_Setter) =>
Optic' k is TxSkel a -> (a -> a) -> m ()
overTweak Optic' A_Lens NoIx TxSkel (Map TxOutRef TxSkelRedeemer)
txSkelInsL (TxOutRef
-> TxSkelRedeemer
-> Map TxOutRef TxSkelRedeemer
-> Map TxOutRef TxSkelRedeemer
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert TxOutRef
oref TxSkelRedeemer
howConsumed)
      Maybe (TxOutRef, TxSkelRedeemer)
-> m (Maybe (TxOutRef, TxSkelRedeemer))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (TxOutRef, TxSkelRedeemer)
 -> m (Maybe (TxOutRef, TxSkelRedeemer)))
-> Maybe (TxOutRef, TxSkelRedeemer)
-> m (Maybe (TxOutRef, TxSkelRedeemer))
forall a b. (a -> b) -> a -> b
$ (TxOutRef, TxSkelRedeemer) -> Maybe (TxOutRef, TxSkelRedeemer)
forall a. a -> Maybe a
Just (TxOutRef
oref, TxSkelRedeemer
howConsumed)

-- | Add an input to a transaction. If the given 'Api.TxOutRef' is already being
-- consumed by the transaction, fail.
addInputTweak :: (MonadTweak m) => Api.TxOutRef -> TxSkelRedeemer -> m ()
addInputTweak :: forall (m :: * -> *).
MonadTweak m =>
TxOutRef -> TxSkelRedeemer -> m ()
addInputTweak TxOutRef
oref TxSkelRedeemer
howConsumed = do
  Map TxOutRef TxSkelRedeemer
presentInputs <- Optic' A_Lens NoIx TxSkel (Map TxOutRef TxSkelRedeemer)
-> m (Map TxOutRef TxSkelRedeemer)
forall (m :: * -> *) k (is :: IxList) a.
(MonadTweak m, Is k A_Getter) =>
Optic' k is TxSkel a -> m a
viewTweak Optic' A_Lens NoIx TxSkel (Map TxOutRef TxSkelRedeemer)
txSkelInsL
  Bool -> m ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (TxOutRef -> Map TxOutRef TxSkelRedeemer -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.notMember TxOutRef
oref Map TxOutRef TxSkelRedeemer
presentInputs)
  Optic' A_Lens NoIx TxSkel (Map TxOutRef TxSkelRedeemer)
-> (Map TxOutRef TxSkelRedeemer -> Map TxOutRef TxSkelRedeemer)
-> m ()
forall (m :: * -> *) k (is :: IxList) a.
(MonadTweak m, Is k A_Setter) =>
Optic' k is TxSkel a -> (a -> a) -> m ()
overTweak Optic' A_Lens NoIx TxSkel (Map TxOutRef TxSkelRedeemer)
txSkelInsL (TxOutRef
-> TxSkelRedeemer
-> Map TxOutRef TxSkelRedeemer
-> Map TxOutRef TxSkelRedeemer
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert TxOutRef
oref TxSkelRedeemer
howConsumed)

-- | Remove transaction inputs according to a given predicate. The returned list
-- contains all removed inputs.
removeInputTweak :: (MonadTweak m) => (Api.TxOutRef -> TxSkelRedeemer -> Bool) -> m [(Api.TxOutRef, TxSkelRedeemer)]
removeInputTweak :: forall (m :: * -> *).
MonadTweak m =>
(TxOutRef -> TxSkelRedeemer -> Bool)
-> m [(TxOutRef, TxSkelRedeemer)]
removeInputTweak TxOutRef -> TxSkelRedeemer -> Bool
removePred = do
  Map TxOutRef TxSkelRedeemer
presentInputs <- Optic' A_Lens NoIx TxSkel (Map TxOutRef TxSkelRedeemer)
-> m (Map TxOutRef TxSkelRedeemer)
forall (m :: * -> *) k (is :: IxList) a.
(MonadTweak m, Is k A_Getter) =>
Optic' k is TxSkel a -> m a
viewTweak Optic' A_Lens NoIx TxSkel (Map TxOutRef TxSkelRedeemer)
txSkelInsL
  let (Map TxOutRef TxSkelRedeemer
removed, Map TxOutRef TxSkelRedeemer
kept) = (TxOutRef -> TxSkelRedeemer -> Bool)
-> Map TxOutRef TxSkelRedeemer
-> (Map TxOutRef TxSkelRedeemer, Map TxOutRef TxSkelRedeemer)
forall k a. (k -> a -> Bool) -> Map k a -> (Map k a, Map k a)
Map.partitionWithKey TxOutRef -> TxSkelRedeemer -> Bool
removePred Map TxOutRef TxSkelRedeemer
presentInputs
  Optic' A_Lens NoIx TxSkel (Map TxOutRef TxSkelRedeemer)
-> Map TxOutRef TxSkelRedeemer -> m ()
forall (m :: * -> *) k (is :: IxList) a.
(MonadTweak m, Is k A_Setter) =>
Optic' k is TxSkel a -> a -> m ()
setTweak Optic' A_Lens NoIx TxSkel (Map TxOutRef TxSkelRedeemer)
txSkelInsL Map TxOutRef TxSkelRedeemer
kept
  [(TxOutRef, TxSkelRedeemer)] -> m [(TxOutRef, TxSkelRedeemer)]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ([(TxOutRef, TxSkelRedeemer)] -> m [(TxOutRef, TxSkelRedeemer)])
-> [(TxOutRef, TxSkelRedeemer)] -> m [(TxOutRef, TxSkelRedeemer)]
forall a b. (a -> b) -> a -> b
$ Map TxOutRef TxSkelRedeemer -> [(TxOutRef, TxSkelRedeemer)]
forall k a. Map k a -> [(k, a)]
Map.toList Map TxOutRef TxSkelRedeemer
removed

-- * Adding and removing transaction outputs

-- | Ensure that a certain output is produced by a transaction. The return value
-- will be @Just@ the added output, when applicable.
ensureOutputTweak :: (MonadTweak m) => TxSkelOut -> m (Maybe TxSkelOut)
ensureOutputTweak :: forall (m :: * -> *).
MonadTweak m =>
TxSkelOut -> m (Maybe TxSkelOut)
ensureOutputTweak TxSkelOut
txSkelOut = do
  [TxSkelOut]
presentOutputs <- Optic' A_Lens NoIx TxSkel [TxSkelOut] -> m [TxSkelOut]
forall (m :: * -> *) k (is :: IxList) a.
(MonadTweak m, Is k A_Getter) =>
Optic' k is TxSkel a -> m a
viewTweak Optic' A_Lens NoIx TxSkel [TxSkelOut]
txSkelOutsL
  if TxSkelOut
txSkelOut TxSkelOut -> [TxSkelOut] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [TxSkelOut]
presentOutputs
    then Maybe TxSkelOut -> m (Maybe TxSkelOut)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe TxSkelOut
forall a. Maybe a
Nothing
    else do
      TxSkelOut -> m ()
forall (m :: * -> *). MonadTweak m => TxSkelOut -> m ()
addOutputTweak TxSkelOut
txSkelOut
      Maybe TxSkelOut -> m (Maybe TxSkelOut)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe TxSkelOut -> m (Maybe TxSkelOut))
-> Maybe TxSkelOut -> m (Maybe TxSkelOut)
forall a b. (a -> b) -> a -> b
$ TxSkelOut -> Maybe TxSkelOut
forall a. a -> Maybe a
Just TxSkelOut
txSkelOut

-- | Add a transaction output, at the end of the current list of outputs, thus
-- retaining the initial outputs order.
addOutputTweak :: (MonadTweak m) => TxSkelOut -> m ()
addOutputTweak :: forall (m :: * -> *). MonadTweak m => TxSkelOut -> m ()
addOutputTweak TxSkelOut
txSkelOut = Optic' A_Lens NoIx TxSkel [TxSkelOut]
-> ([TxSkelOut] -> [TxSkelOut]) -> m ()
forall (m :: * -> *) k (is :: IxList) a.
(MonadTweak m, Is k A_Setter) =>
Optic' k is TxSkel a -> (a -> a) -> m ()
overTweak Optic' A_Lens NoIx TxSkel [TxSkelOut]
txSkelOutsL ([TxSkelOut] -> [TxSkelOut] -> [TxSkelOut]
forall a. [a] -> [a] -> [a]
++ [TxSkelOut
txSkelOut])

-- | Remove transaction outputs according to some predicate. The returned list
-- contains all the removed outputs.
removeOutputTweak :: (MonadTweak m) => (TxSkelOut -> Bool) -> m [TxSkelOut]
removeOutputTweak :: forall (m :: * -> *).
MonadTweak m =>
(TxSkelOut -> Bool) -> m [TxSkelOut]
removeOutputTweak TxSkelOut -> Bool
removePred = do
  [TxSkelOut]
presentOutputs <- Optic' A_Lens NoIx TxSkel [TxSkelOut] -> m [TxSkelOut]
forall (m :: * -> *) k (is :: IxList) a.
(MonadTweak m, Is k A_Getter) =>
Optic' k is TxSkel a -> m a
viewTweak Optic' A_Lens NoIx TxSkel [TxSkelOut]
txSkelOutsL
  let ([TxSkelOut]
removed, [TxSkelOut]
kept) = (TxSkelOut -> Bool) -> [TxSkelOut] -> ([TxSkelOut], [TxSkelOut])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition TxSkelOut -> Bool
removePred [TxSkelOut]
presentOutputs
  Optic' A_Lens NoIx TxSkel [TxSkelOut] -> [TxSkelOut] -> m ()
forall (m :: * -> *) k (is :: IxList) a.
(MonadTweak m, Is k A_Setter) =>
Optic' k is TxSkel a -> a -> m ()
setTweak Optic' A_Lens NoIx TxSkel [TxSkelOut]
txSkelOutsL [TxSkelOut]
kept
  [TxSkelOut] -> m [TxSkelOut]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return [TxSkelOut]
removed

-- * Adding and removing minted values

-- | Add a new entry to the 'TxSkelMints' of the transaction skeleton under
-- modification. As this is implemented in terms of 'addToTxSkelMints', the same
-- caveats apply as do to that function!
addMintTweak :: (MonadTweak m) => (Script.Versioned Script.MintingPolicy, TxSkelRedeemer, Api.TokenName, Integer) -> m ()
addMintTweak :: forall (m :: * -> *).
MonadTweak m =>
(Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)
-> m ()
addMintTweak (Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)
mint = Optic' A_Lens NoIx TxSkel TxSkelMints
-> (TxSkelMints -> TxSkelMints) -> m ()
forall (m :: * -> *) k (is :: IxList) a.
(MonadTweak m, Is k A_Setter) =>
Optic' k is TxSkel a -> (a -> a) -> m ()
overTweak Optic' A_Lens NoIx TxSkel TxSkelMints
txSkelMintsL ((TxSkelMints -> TxSkelMints) -> m ())
-> (TxSkelMints -> TxSkelMints) -> m ()
forall a b. (a -> b) -> a -> b
$ (Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)
-> TxSkelMints -> TxSkelMints
addToTxSkelMints (Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)
mint

-- | Remove some entries from the 'TxSkelMints' of a transaction, according to
-- some predicate. The returned list holds the removed entries.
removeMintTweak ::
  (MonadTweak m) =>
  ((Script.Versioned Script.MintingPolicy, TxSkelRedeemer, Api.TokenName, Integer) -> Bool) ->
  m [(Script.Versioned Script.MintingPolicy, TxSkelRedeemer, Api.TokenName, Integer)]
removeMintTweak :: forall (m :: * -> *).
MonadTweak m =>
((Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)
 -> Bool)
-> m [(Versioned MintingPolicy, TxSkelRedeemer, TokenName,
       Integer)]
removeMintTweak (Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)
-> Bool
removePred = do
  [(Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)]
presentMints <- Optic'
  A_Getter
  NoIx
  TxSkel
  [(Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)]
-> m [(Versioned MintingPolicy, TxSkelRedeemer, TokenName,
       Integer)]
forall (m :: * -> *) k (is :: IxList) a.
(MonadTweak m, Is k A_Getter) =>
Optic' k is TxSkel a -> m a
viewTweak (Optic'
   A_Getter
   NoIx
   TxSkel
   [(Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)]
 -> m [(Versioned MintingPolicy, TxSkelRedeemer, TokenName,
        Integer)])
-> Optic'
     A_Getter
     NoIx
     TxSkel
     [(Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)]
-> m [(Versioned MintingPolicy, TxSkelRedeemer, TokenName,
       Integer)]
forall a b. (a -> b) -> a -> b
$ Optic' A_Lens NoIx TxSkel TxSkelMints
txSkelMintsL Optic' A_Lens NoIx TxSkel TxSkelMints
-> Optic
     A_Getter
     NoIx
     TxSkelMints
     TxSkelMints
     [(Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)]
     [(Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)]
-> Optic'
     A_Getter
     NoIx
     TxSkel
     [(Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)]
forall k l m (is :: IxList) (js :: IxList) (ks :: IxList) s t u v a
       b.
(JoinKinds k l m, AppendIndices is js ks) =>
Optic k is s t u v -> Optic l js u v a b -> Optic m ks s t a b
% (TxSkelMints
 -> [(Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)])
-> Optic
     A_Getter
     NoIx
     TxSkelMints
     TxSkelMints
     [(Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)]
     [(Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)]
forall s a. (s -> a) -> Getter s a
to TxSkelMints
-> [(Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)]
txSkelMintsToList
  let ([(Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)]
removed, [(Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)]
kept) = ((Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)
 -> Bool)
-> [(Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)]
-> ([(Versioned MintingPolicy, TxSkelRedeemer, TokenName,
      Integer)],
    [(Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition (Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)
-> Bool
removePred [(Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)]
presentMints
  Optic' A_Lens NoIx TxSkel TxSkelMints -> TxSkelMints -> m ()
forall (m :: * -> *) k (is :: IxList) a.
(MonadTweak m, Is k A_Setter) =>
Optic' k is TxSkel a -> a -> m ()
setTweak Optic' A_Lens NoIx TxSkel TxSkelMints
txSkelMintsL (TxSkelMints -> m ()) -> TxSkelMints -> m ()
forall a b. (a -> b) -> a -> b
$ [(Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)]
-> TxSkelMints
txSkelMintsFromList [(Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)]
kept
  [(Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)]
-> m [(Versioned MintingPolicy, TxSkelRedeemer, TokenName,
       Integer)]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return [(Versioned MintingPolicy, TxSkelRedeemer, TokenName, Integer)]
removed