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

@@ -24,13 +24,13 @@ import Ctl.Internal.BalanceTx.Collateral.Select (selectCollateral)
 import Ctl.Internal.BalanceTx.Constraints
   ( BalanceTxConstraintsBuilder
   , _collateralUtxos
-  , _nonSpendableInputs
   )
 import Ctl.Internal.BalanceTx.Constraints
   ( _changeAddress
   , _changeDatum
   , _maxChangeOutputTokenQuantity
   , _nonSpendableInputs
+  , _nonSpendableInputsPredicates
   , _selectionStrategy
   , _srcAddresses
   ) as Constraints
@@ -144,14 +144,21 @@ import Data.Array.NonEmpty
 import Data.Array.NonEmpty as NEA
 import Data.Bifunctor (lmap)
 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 ((%~), (.~), (?~))
 import Data.Log.Tag (TagSet, tag, tagSetTag)
 import Data.Log.Tag (fromArray) as TagSet
 import Data.Map (Map)
-import Data.Map (empty, filterKeys, insert, lookup, toUnfoldable, union) as Map
+import Data.Map
+  ( empty
+  , filterWithKey
+  , insert
+  , lookup
+  , toUnfoldable
+  , union
+  ) as Map
 import Data.Maybe (Maybe(Just, Nothing), fromMaybe, isJust, maybe)
 import Data.Newtype (class Newtype, unwrap, wrap)
 import Data.Set (Set)
@@ -264,16 +271,23 @@ balanceTxWithConstraints transaction extraUtxos constraintsBuilder = do
 
 setTransactionCollateral :: Address -> Transaction -> BalanceTxM Transaction
 setTransactionCollateral changeAddr transaction = do
-  nonSpendableSet <- asksConstraints _nonSpendableInputs
+  nonSpendableSet <- asksConstraints Constraints._nonSpendableInputs
+  nonSpendableInputsPredicates <- asksConstraints
+    Constraints._nonSpendableInputsPredicates
   mbCollateralUtxos <- asksConstraints _collateralUtxos
   -- We must filter out UTxOs that are set as non-spendable in the balancer
   -- constraints
-  let isSpendable = not <<< flip Set.member nonSpendableSet
+  let
+    isSpendable = \input output ->
+      not (Set.member input nonSpendableSet) &&
+        not (any (\f -> f input output) nonSpendableInputsPredicates)
   collateral <- case mbCollateralUtxos of
     -- if no collateral utxos are specified, use the wallet, but filter
     -- the unspendable ones
     Nothing -> do
-      let isSpendableUtxo = isSpendable <<< _.input <<< unwrap
+      let
+        isSpendableUtxo = \utxo -> isSpendable (unwrap utxo).input
+          (unwrap utxo).output
       { yes: spendableUtxos, no: filteredUtxos } <-
         Array.partition isSpendableUtxo <$> do
           liftEitherContract $ note CouldNotGetCollateral <$>
@@ -291,7 +305,7 @@ setTransactionCollateral changeAddr transaction = do
       let
         coinsPerUtxoUnit = params.coinsPerUtxoUnit
         maxCollateralInputs = UInt.toInt $ params.maxCollateralInputs
-        utxoMap' = fromPlutusUtxoMap networkId $ Map.filterKeys isSpendable
+        utxoMap' = Map.filterWithKey isSpendable $ fromPlutusUtxoMap networkId
           utxoMap
       mbCollateral <- liftEffect $ map Array.fromFoldable <$>
         selectCollateral coinsPerUtxoUnit maxCollateralInputs utxoMap'
@@ -359,44 +373,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 (\f -> f 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
@@ -18,6 +20,7 @@ module Ctl.Internal.BalanceTx.Constraints
   , _changeDatum
   , _maxChangeOutputTokenQuantity
   , _nonSpendableInputs
+  , _nonSpendableInputsPredicates
   , _selectionStrategy
   , _srcAddresses
   ) where
@@ -27,14 +30,17 @@ 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 (TransactionOutputWithRefScript) as Plutus
 import Ctl.Internal.Plutus.Types.Transaction (UtxoMap)
 import Ctl.Internal.Serialization.Address (Address, NetworkId)
 import Ctl.Internal.Types.OutputDatum (OutputDatum)
@@ -46,7 +52,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(Just, Nothing))
+import Data.Maybe (Maybe(Just, Nothing), fromMaybe)
 import Data.Newtype (class Newtype, over2, unwrap, wrap)
 import Data.Set (Set)
 import Data.Set (singleton) as Set
@@ -58,6 +64,7 @@ newtype BalanceTxConstraints = BalanceTxConstraints
   , collateralUtxos :: Maybe UtxoMap
   , maxChangeOutputTokenQuantity :: Maybe BigInt
   , nonSpendableInputs :: Set TransactionInput
+  , nonSpendableInputsPredicates :: Array (UtxoPredicate TransactionOutput)
   , srcAddresses :: Maybe (Array Address)
   , changeAddress :: Maybe Address
   , changeDatum :: Maybe OutputDatum
@@ -66,6 +73,8 @@ newtype BalanceTxConstraints = BalanceTxConstraints
 
 derive instance Newtype BalanceTxConstraints _
 
+type UtxoPredicate (output :: Type) = TransactionInput -> output -> Boolean
+
 _additionalUtxos :: Lens' BalanceTxConstraints UtxoMap
 _additionalUtxos = _Newtype <<< prop (Proxy :: Proxy "additionalUtxos")
 
@@ -79,6 +88,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")
 
@@ -111,6 +125,7 @@ buildBalanceTxConstraints = applyFlipped defaultConstraints <<< unwrap
     , collateralUtxos: Nothing
     , maxChangeOutputTokenQuantity: Nothing
     , nonSpendableInputs: mempty
+    , nonSpendableInputsPredicates: mempty
     , srcAddresses: Nothing
     , changeDatum: Nothing
     , changeAddress: Nothing
@@ -173,6 +188,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