-- | 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.Inputs
  ( ensureInputTweak,
    addInputTweak,
    removeInputTweak,
    modifySpendRedeemersOfTypeTweak,
  )
where

import Control.Monad
import Cooked.Skeleton
import Cooked.Tweak.Common
import Data.Map qualified as Map
import Data.Maybe (fromMaybe)
import Optics.Core
import PlutusLedgerApi.V3 qualified as Api
import Type.Reflection (Typeable)

-- | 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

-- | Applies an optional modification to all spend redeemers of type a
modifySpendRedeemersOfTypeTweak :: forall a b m. (Typeable a, RedeemerConstrs b, MonadTweak m) => (a -> Maybe b) -> m ()
modifySpendRedeemersOfTypeTweak :: forall a b (m :: * -> *).
(Typeable a, RedeemerConstrs b, MonadTweak m) =>
(a -> Maybe b) -> m ()
modifySpendRedeemersOfTypeTweak a -> Maybe b
f = do
  [(TxOutRef, TxSkelRedeemer)]
presentInputs <- Map TxOutRef TxSkelRedeemer -> [(TxOutRef, TxSkelRedeemer)]
forall k a. Map k a -> [(k, a)]
Map.toList (Map TxOutRef TxSkelRedeemer -> [(TxOutRef, TxSkelRedeemer)])
-> m (Map TxOutRef TxSkelRedeemer)
-> m [(TxOutRef, TxSkelRedeemer)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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
  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 -> m ())
-> Map TxOutRef TxSkelRedeemer -> m ()
forall a b. (a -> b) -> a -> b
$
    [(TxOutRef, TxSkelRedeemer)] -> Map TxOutRef TxSkelRedeemer
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(TxOutRef, TxSkelRedeemer)] -> Map TxOutRef TxSkelRedeemer)
-> [(TxOutRef, TxSkelRedeemer)] -> Map TxOutRef TxSkelRedeemer
forall a b. (a -> b) -> a -> b
$
      [(TxOutRef, TxSkelRedeemer)]
presentInputs [(TxOutRef, TxSkelRedeemer)]
-> ((TxOutRef, TxSkelRedeemer) -> (TxOutRef, TxSkelRedeemer))
-> [(TxOutRef, TxSkelRedeemer)]
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \(TxOutRef
oRef, TxSkelRedeemer
red) -> (TxOutRef
oRef,) (TxSkelRedeemer -> (TxOutRef, TxSkelRedeemer))
-> (Maybe TxSkelRedeemer -> TxSkelRedeemer)
-> Maybe TxSkelRedeemer
-> (TxOutRef, TxSkelRedeemer)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TxSkelRedeemer -> Maybe TxSkelRedeemer -> TxSkelRedeemer
forall a. a -> Maybe a -> a
fromMaybe TxSkelRedeemer
red (Maybe TxSkelRedeemer -> (TxOutRef, TxSkelRedeemer))
-> Maybe TxSkelRedeemer -> (TxOutRef, TxSkelRedeemer)
forall a b. (a -> b) -> a -> b
$ do
        a
typedRedeemer <- TxSkelRedeemer -> Maybe a
forall a. Typeable a => TxSkelRedeemer -> Maybe a
getTypedRedeemer TxSkelRedeemer
red
        (b -> TxSkelRedeemer -> TxSkelRedeemer
forall redeemer.
RedeemerConstrs redeemer =>
redeemer -> TxSkelRedeemer -> TxSkelRedeemer
`setTypedRedeemer` TxSkelRedeemer
red) (b -> TxSkelRedeemer) -> Maybe b -> Maybe TxSkelRedeemer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> Maybe b
f a
typedRedeemer