feat: add mustNotSpendUtxosWhere balancer constraint

src/Contract/BalanceTxConstraints.purs

@@ -4,8 +4,10 @@ module Contract.BalanceTxConstraints (module BalanceTxConstraints) where
 
 import Ctl.Internal.BalanceTx.Constraints
   ( BalanceTxConstraintsBuilder
+  , UtxoPredicate
   , mustGenChangeOutsWithMaxTokenQuantity
   , mustNotSpendUtxoWithOutRef
+  , mustNotSpendUtxosWhere
   , mustNotSpendUtxosWithOutRefs
   , mustSendChangeToAddress
   , mustSendChangeWithDatum

src/Internal/BalanceTx/BalanceTx.purs

@@ -26,6 +26,7 @@ import Ctl.Internal.BalanceTx.Collateral
 import Ctl.Internal.BalanceTx.Constraints
   ( BalanceTxConstraintsBuilder
   , _nonSpendableInputs
+  , _nonSpendableInputsPredicates
   )
 import Ctl.Internal.BalanceTx.Constraints
   ( _changeAddress
@@ -108,6 +109,9 @@ import Ctl.Internal.Cardano.Types.Transaction
   , _withdrawals
   , _witnessSet
   )
+import Ctl.Internal.Cardano.Types.TransactionUnspentOutput
+  ( TransactionUnspentOutput(TransactionUnspentOutput)
+  )
 import Ctl.Internal.Cardano.Types.Value
   ( AssetClass
   , Coin(Coin)
@@ -157,7 +161,7 @@ import Data.Array.NonEmpty as NEA
 import Data.Bifunctor (lmap)
 import Data.BigInt (BigInt)
 import Data.Either (Either, hush, note)
-import Data.Foldable (fold, foldMap, foldr, length, null, sum)
+import Data.Foldable (any, fold, foldMap, foldr, length, null, or, sum)
 import Data.Function (on)
 import Data.Lens.Getter ((^.))
 import Data.Lens.Setter ((%~), (.~), (?~))
@@ -275,13 +279,16 @@ balanceTxWithConstraints transaction extraUtxos constraintsBuilder = do
 setTransactionCollateral :: Address -> Transaction -> BalanceTxM Transaction
 setTransactionCollateral changeAddr transaction = do
   nonSpendableSet <- asksConstraints _nonSpendableInputs
+  nonSpendableInputsPredicates <- asksConstraints _nonSpendableInputsPredicates
   collateral <- do
     rawCollateral <- liftEitherContract $ note CouldNotGetCollateral <$>
       Wallet.getWalletCollateral
     -- filter out UTxOs that are set as non-spendable in the balancer constraints
     let
-      isSpendable = not <<< flip Set.member nonSpendableSet <<< _.input <<<
-        unwrap
+      isSpendable =
+        \(TransactionUnspentOutput { input, output }) ->
+          not (Set.member input nonSpendableSet) &&
+            not (any (\p -> p input output) nonSpendableInputsPredicates)
     pure $ Array.filter isSpendable rawCollateral
   addTxCollateralReturn collateral (addTxCollateral collateral transaction)
     changeAddr
@@ -356,44 +363,49 @@ runBalancer p = do
         liftContract $ Wallet.getWalletCollateral <#>
           fold >>> map (unwrap >>> _.input) >>> Set.fromFoldable
       else mempty
-    asksConstraints Constraints._nonSpendableInputs <#>
-      append nonSpendableCollateralInputs >>>
-        \nonSpendableInputs ->
-          foldr
-            ( \(oref /\ output) acc ->
-                let
-                  hasInlineDatum :: Boolean
-                  hasInlineDatum = case (unwrap output).datum of
-                    OutputDatum _ -> true
-                    _ -> false
-
-                  hasScriptRef :: Boolean
-                  hasScriptRef = isJust (unwrap output).scriptRef
-
-                  spendable :: Boolean
-                  spendable = not $ Set.member oref nonSpendableInputs ||
-                    Set.member oref
-                      ( p.transaction ^. _transaction <<< _body <<<
-                          _referenceInputs
-                      )
-
-                  validInContext :: Boolean
-                  validInContext = not $ txHasPlutusV1 &&
-                    (hasInlineDatum || hasScriptRef)
-                in
-                  case spendable, validInContext of
-                    true, true -> acc
-                      { spendable = Map.insert oref output acc.spendable }
-                    true, false -> acc
-                      { invalidInContext = Map.insert oref output
-                          acc.invalidInContext
-                      }
-                    _, _ -> acc
-            )
-            { spendable: Map.empty
-            , invalidInContext: Map.empty
-            }
-            (Map.toUnfoldable p.utxos :: Array _)
+    constraints <- unwrap <$> asks _.constraints
+    let
+      nonSpendableInputs =
+        constraints.nonSpendableInputs <> nonSpendableCollateralInputs
+    pure $ foldr
+      ( \(oref /\ output) acc ->
+          let
+            hasInlineDatum :: Boolean
+            hasInlineDatum = case (unwrap output).datum of
+              OutputDatum _ -> true
+              _ -> false
+
+            hasScriptRef :: Boolean
+            hasScriptRef = isJust (unwrap output).scriptRef
+
+            spendable :: Boolean
+            spendable = not $ or
+              [ Set.member oref nonSpendableInputs
+              , any (\p -> p oref output)
+                  constraints.nonSpendableInputsPredicates
+              , Set.member oref
+                  ( p.transaction ^. _transaction <<< _body <<<
+                      _referenceInputs
+                  )
+              ]
+
+            validInContext :: Boolean
+            validInContext = not $ txHasPlutusV1 &&
+              (hasInlineDatum || hasScriptRef)
+          in
+            case spendable, validInContext of
+              true, true -> acc
+                { spendable = Map.insert oref output acc.spendable }
+              true, false -> acc
+                { invalidInContext = Map.insert oref output
+                    acc.invalidInContext
+                }
+              _, _ -> acc
+      )
+      { spendable: Map.empty
+      , invalidInContext: Map.empty
+      }
+      (Map.toUnfoldable p.utxos :: Array _)
 
   mainLoop :: BalancerState UnindexedTx -> BalanceTxM FinalizedTransaction
   mainLoop = worker <<< PrebalanceTx

src/Internal/BalanceTx/Constraints.purs

@@ -1,8 +1,10 @@
 module Ctl.Internal.BalanceTx.Constraints
   ( BalanceTxConstraints(BalanceTxConstraints)
   , BalanceTxConstraintsBuilder(BalanceTxConstraintsBuilder)
+  , UtxoPredicate
   , buildBalanceTxConstraints
   , mustGenChangeOutsWithMaxTokenQuantity
+  , mustNotSpendUtxosWhere
   , mustNotSpendUtxosWithOutRefs
   , mustNotSpendUtxoWithOutRef
   , mustSendChangeToAddress
@@ -16,6 +18,7 @@ module Ctl.Internal.BalanceTx.Constraints
   , _changeDatum
   , _maxChangeOutputTokenQuantity
   , _nonSpendableInputs
+  , _nonSpendableInputsPredicates
   , _selectionStrategy
   , _srcAddresses
   ) where
@@ -25,15 +28,20 @@ import Prelude
 import Ctl.Internal.BalanceTx.CoinSelection
   ( SelectionStrategy(SelectionStrategyOptimal)
   )
+import Ctl.Internal.Cardano.Types.Transaction (TransactionOutput)
 import Ctl.Internal.Plutus.Conversion
   ( fromPlutusAddress
   , fromPlutusAddressWithNetworkTag
+  , toPlutusTxOutputWithRefScript
   )
 import Ctl.Internal.Plutus.Types.Address
   ( Address
   , AddressWithNetworkTag(AddressWithNetworkTag)
   ) as Plutus
-import Ctl.Internal.Plutus.Types.Transaction (UtxoMap) as Plutus
+import Ctl.Internal.Plutus.Types.Transaction
+  ( TransactionOutputWithRefScript
+  , UtxoMap
+  ) as Plutus
 import Ctl.Internal.Serialization.Address (Address, NetworkId)
 import Ctl.Internal.Types.OutputDatum (OutputDatum)
 import Ctl.Internal.Types.Transaction (TransactionInput)
@@ -45,7 +53,7 @@ import Data.Lens.Iso.Newtype (_Newtype)
 import Data.Lens.Record (prop)
 import Data.Lens.Setter (appendOver, set, setJust)
 import Data.Map (empty) as Map
-import Data.Maybe (Maybe(Nothing))
+import Data.Maybe (Maybe(Nothing), fromMaybe)
 import Data.Newtype (class Newtype, over2, unwrap, wrap)
 import Data.Set (Set)
 import Data.Set (singleton) as Set
@@ -55,6 +63,7 @@ newtype BalanceTxConstraints = BalanceTxConstraints
   { additionalUtxos :: Plutus.UtxoMap
   , maxChangeOutputTokenQuantity :: Maybe BigInt
   , nonSpendableInputs :: Set TransactionInput
+  , nonSpendableInputsPredicates :: Array (UtxoPredicate TransactionOutput)
   , srcAddresses :: Maybe (Array Address)
   , changeAddress :: Maybe Address
   , changeDatum :: Maybe OutputDatum
@@ -63,6 +72,8 @@ newtype BalanceTxConstraints = BalanceTxConstraints
 
 derive instance Newtype BalanceTxConstraints _
 
+type UtxoPredicate (output :: Type) = TransactionInput -> output -> Boolean
+
 _additionalUtxos :: Lens' BalanceTxConstraints Plutus.UtxoMap
 _additionalUtxos = _Newtype <<< prop (Proxy :: Proxy "additionalUtxos")
 
@@ -73,6 +84,11 @@ _maxChangeOutputTokenQuantity =
 _nonSpendableInputs :: Lens' BalanceTxConstraints (Set TransactionInput)
 _nonSpendableInputs = _Newtype <<< prop (Proxy :: Proxy "nonSpendableInputs")
 
+_nonSpendableInputsPredicates
+  :: Lens' BalanceTxConstraints (Array (UtxoPredicate TransactionOutput))
+_nonSpendableInputsPredicates =
+  _Newtype <<< prop (Proxy :: Proxy "nonSpendableInputsPredicates")
+
 _srcAddresses :: Lens' BalanceTxConstraints (Maybe (Array Address))
 _srcAddresses = _Newtype <<< prop (Proxy :: Proxy "srcAddresses")
 
@@ -104,6 +120,7 @@ buildBalanceTxConstraints = applyFlipped defaultConstraints <<< unwrap
     { additionalUtxos: Map.empty
     , maxChangeOutputTokenQuantity: Nothing
     , nonSpendableInputs: mempty
+    , nonSpendableInputsPredicates: mempty
     , srcAddresses: Nothing
     , changeDatum: Nothing
     , changeAddress: Nothing
@@ -166,6 +183,20 @@ mustNotSpendUtxosWithOutRefs = wrap <<< appendOver _nonSpendableInputs
 mustNotSpendUtxoWithOutRef :: TransactionInput -> BalanceTxConstraintsBuilder
 mustNotSpendUtxoWithOutRef = mustNotSpendUtxosWithOutRefs <<< Set.singleton
 
+-- | Tells the balancer not to spend UTxO's based on the given predicate.
+-- | Note that `mustNotSpendUtxosWhere` constraints are stacked when specified
+-- | multiple times, and utxos are tested against each predicate. The order of
+-- | specifying multiple `mustNotSpendUtxosWhere` constraints does NOT affect
+-- | the resulting set.
+mustNotSpendUtxosWhere
+  :: UtxoPredicate Plutus.TransactionOutputWithRefScript
+  -> BalanceTxConstraintsBuilder
+mustNotSpendUtxosWhere p =
+  wrap $ appendOver _nonSpendableInputsPredicates
+    ( Array.singleton \oref out ->
+        fromMaybe false $ p oref <$> toPlutusTxOutputWithRefScript out
+    )
+
 -- | Tells the balancer to use the provided UTxO set when evaluating script
 -- | execution units (sets `additionalUtxoSet` of Ogmios `EvaluateTx`).
 -- | Note that you need to use `unspentOutputs` lookup to make these UTxO's
@@ -176,4 +207,3 @@ mustUseAdditionalUtxos = wrap <<< set _additionalUtxos
 -- | Tells the balancer to use the given strategy for coin selection.
 mustUseCoinSelectionStrategy :: SelectionStrategy -> BalanceTxConstraintsBuilder
 mustUseCoinSelectionStrategy = wrap <<< set _selectionStrategy
-