improve efficiency of treasury withdrawal effect script

agora/Agora/Effect/TreasuryWithdrawal.hs

@@ -15,24 +15,27 @@ module Agora.Effect.TreasuryWithdrawal (
 
 import Agora.Effect (makeEffect)
 import Agora.Plutarch.Orphans ()
+import Agora.Utils (pdelete)
 import Plutarch.Api.V1 (
   PCredential,
   PValue,
   ptuple,
  )
-import Plutarch.Api.V1.Value (pnormalize)
 import Plutarch.Api.V2 (
   AmountGuarantees (Positive),
   KeyGuarantees (Sorted),
   PTuple,
+  PTxInInfo,
+  PTxOut,
   PValidator,
  )
 import Plutarch.DataRepr (
   DerivePConstantViaData (DerivePConstantViaData),
   PDataFields,
  )
-import Plutarch.Extra.ScriptContext (pfindTxInByTxOutRef, pisPubKey)
-import "liqwid-plutarch-extra" Plutarch.Extra.TermCont (pguardC, pletC, pletFieldsC, pmatchC)
+import Plutarch.Extra.Field (pletAllC)
+import Plutarch.Extra.ScriptContext (pisPubKey)
+import "liqwid-plutarch-extra" Plutarch.Extra.TermCont (pguardC, pletC, pletFieldsC)
 import Plutarch.Lift (PConstantDecl, PUnsafeLiftDecl (PLifted))
 import PlutusLedgerApi.V1.Credential (Credential)
 import PlutusLedgerApi.V1.Value (CurrencySymbol, Value)
@@ -129,69 +132,57 @@ instance PTryFrom PData PTreasuryWithdrawalDatum
 -}
 treasuryWithdrawalValidator :: forall {s :: S}. CurrencySymbol -> Term s PValidator
 treasuryWithdrawalValidator currSymbol = makeEffect currSymbol $
-  \_cs (datum' :: Term _ PTreasuryWithdrawalDatum) txOutRef' txInfo' -> unTermCont $ do
-    datum <- pletFieldsC @'["receivers", "treasuries"] datum'
-    txInfo <- pletFieldsC @'["outputs", "inputs"] txInfo'
-    PJust ((pfield @"resolved" #) -> txOut) <- pmatchC $ pfindTxInByTxOutRef # txOutRef' # pfromData txInfo.inputs
-    effInput <- pletFieldsC @'["address", "value"] $ txOut
-    outputValues <-
-      pletC $
-        pmap
-          # plam
-            ( \txOut' -> unTermCont $ do
-                txOut <- pletFieldsC @'["address", "value"] $ txOut'
-                let cred = pfield @"credential" # pfromData txOut.address
-                pure . pdata $ ptuple # cred # txOut.value
-            )
-          # pfromData txInfo.outputs
-    inputValues <-
-      pletC $
-        pmap
-          # plam
-            ( \((pfield @"resolved" #) -> txOut') -> unTermCont $ do
-                txOut <- pletFieldsC @'["address", "value"] $ txOut'
-                let cred = pfield @"credential" # pfromData txOut.address
-                pure . pdata $ ptuple # cred # txOut.value
-            )
-          # txInfo.inputs
-    let ofTreasury =
-          pfilter
-            # plam (\((pfield @"_0" #) . pfromData -> cred) -> pelem # cred # datum.treasuries)
-        sumValues = phoistAcyclic $
-          plam $ \v ->
-            pnormalize
-              #$ pfoldr
-              # plam (\(pfromData . (pfield @"_1" #) -> x) y -> x <> y)
-              # mempty
-              # v
-        treasuryInputValuesSum = sumValues #$ ofTreasury # inputValues
-        treasuryOutputValuesSum = sumValues #$ ofTreasury # outputValues
-        receiverValuesSum = sumValues # datum.receivers
-        -- Constraints
-        outputContentMatchesRecivers =
-          pall # plam (\out -> pelem # out # outputValues)
-            #$ datum.receivers
-        excessShouldBePaidToInputs =
-          treasuryOutputValuesSum <> receiverValuesSum #== treasuryInputValuesSum
-        shouldNotPayToEffect =
-          pnot #$ pany
-            # plam
-              ( \x ->
-                  effInput.address #== pfield @"address" # x
-              )
-            # pfromData txInfo.outputs
-        inputsAreOnlyTreasuriesOrCollateral =
-          pall
-            # plam
-              ( \((pfield @"_0" #) . pfromData -> cred) ->
-                  cred #== pfield @"credential" # effInput.address
-                    #|| pelem # cred # datum.treasuries
-                    #|| pisPubKey # pfromData cred
-              )
-            # inputValues
+  \_cs (datum :: Term _ PTreasuryWithdrawalDatum) effectInputRef txInfo -> unTermCont $ do
+    datumF <- pletAllC datum
+    txInfoF <- pletFieldsC @'["outputs", "inputs"] txInfo
+
+    let validateInput :: Term _ (PTxInInfo :--> PBool)
+        validateInput = plam $ \input -> unTermCont $ do
+          inputF <- pletAllC input
+
+          cred <-
+            pletC $
+              pfield @"credential"
+                #$ pfield @"address" # inputF.resolved
+
+          pure $
+            foldl1
+              (#||)
+              [ ptraceIfTrue "Effect input" $ inputF.outRef #== effectInputRef
+              , ptraceIfTrue "Treasury input" $ pelem # cred # datumF.treasuries
+              , ptraceIfTrue "Collateral input" $ pisPubKey # pfromData cred
+              ]
+
+        validateOutput ::
+          Term
+            _
+            ( PBuiltinList (PAsData (PTuple PCredential (PValue 'Sorted 'Positive)))
+                :--> PTxOut
+                :--> PBuiltinList (PAsData (PTuple PCredential (PValue 'Sorted 'Positive)))
+            )
+        validateOutput = plam $ \receivers output -> unTermCont $ do
+          outputF <- pletFieldsC @'["address", "value"] output
+          cred <- pletC $ pfield @"credential" # pfromData outputF.address
+
+          let credValue = pdata $ ptuple # cred # outputF.value
+
+              shouldSendToTreasury =
+                pif
+                  (pelem # cred # datumF.treasuries)
+                  receivers
+                  (ptraceError "Invalid receiver")
+
+          pure $
+            pmatch (pdelete # credValue # receivers) $ \case
+              PJust updatedReceivers ->
+                ptrace "Receiver output" updatedReceivers
+              PNothing ->
+                ptrace "Treasury output" shouldSendToTreasury
+
+    pguardC "All input are valid" $
+      pall # validateInput # txInfoF.inputs
+
+    pguardC "All receiver get correct output" $
+      pnull #$ pfoldl # validateOutput # datumF.receivers # txInfoF.outputs
 
-    pguardC "Transaction should not pay to effects" shouldNotPayToEffect
-    pguardC "Transaction output does not match receivers" outputContentMatchesRecivers
-    pguardC "Remainders should be returned to the treasury" excessShouldBePaidToInputs
-    pguardC "Transaction should only have treasuries specified in the datum as input" inputsAreOnlyTreasuriesOrCollateral
     pure . popaque $ pconstant ()

agora/Agora/Stake/Redeemers.hs

@@ -48,7 +48,7 @@ import Agora.Stake (
   ),
   pstakeLocked,
  )
-import Agora.Utils (pdeleteBy, pfromSingleton, pisSingleton)
+import Agora.Utils (pfromSingleton, pisSingleton, pmustDeleteBy)
 import Plutarch.Api.V1.Address (PCredential)
 import Plutarch.Api.V2 (PMaybeData)
 import Plutarch.Extra.Field (pletAll, pletAllC)
@@ -88,7 +88,7 @@ pbatchUpdateInputs ::
 pbatchUpdateInputs = phoistAcyclic $
   plam $ \f -> flip pmatch $ \ctxF ->
     pnull #$ pfoldr
-      # (pdeleteBy # f)
+      # (pmustDeleteBy # f)
       # ctxF.stakeOutputDatums
       # ctxF.stakeInputDatums
 

agora/Agora/Utils.hs

@@ -24,6 +24,7 @@ module Agora.Utils (
   pcurrentTimeDuration,
   pdelete,
   pdeleteBy,
+  pmustDeleteBy,
   pisSingleton,
   pfromSingleton,
   pmapMaybe,
@@ -40,7 +41,7 @@ import Plutarch.Api.V2 (PScriptHash, PScriptPurpose)
 import Plutarch.Extra.Applicative (PApplicative (ppure))
 import Plutarch.Extra.Category (PCategory (pidentity))
 import Plutarch.Extra.Functor (PFunctor (PSubcategory, pfmap))
-import Plutarch.Extra.Maybe (pnothing)
+import Plutarch.Extra.Maybe (pjust, pnothing)
 import Plutarch.Extra.Ord (PComparator, POrdering (PLT), pcompareBy, pequateBy)
 import Plutarch.Extra.Time (PCurrentTime (PCurrentTime))
 import Plutarch.Unsafe (punsafeCoerce)
@@ -214,15 +215,31 @@ pcurrentTimeDuration = phoistAcyclic $
 pdelete ::
   forall (a :: PType) (list :: PType -> PType) (s :: S).
   (PEq a, PIsListLike list a) =>
-  Term s (a :--> list a :--> list a)
+  Term s (a :--> list a :--> PMaybe (list a))
 pdelete = phoistAcyclic $ pdeleteBy # plam (#==)
 
 -- | @since 1.0.0
 pdeleteBy ::
   forall (a :: PType) (list :: PType -> PType) (s :: S).
   (PIsListLike list a) =>
-  Term s ((a :--> a :--> PBool) :--> a :--> list a :--> list a)
+  Term s ((a :--> a :--> PBool) :--> a :--> list a :--> PMaybe (list a))
 pdeleteBy = phoistAcyclic $
+  plam $ \f' x -> plet (f' # x) $ \f ->
+    precList
+      ( \self h t ->
+          pif
+            (f # h)
+            (pjust # t)
+            (pfmap # (pcons # h) # (self # t))
+      )
+      (const pnothing)
+
+-- | @since 1.0.0
+pmustDeleteBy ::
+  forall (a :: PType) (list :: PType -> PType) (s :: S).
+  (PIsListLike list a) =>
+  Term s ((a :--> a :--> PBool) :--> a :--> list a :--> list a)
+pmustDeleteBy = phoistAcyclic $
   plam $ \f' x -> plet (f' # x) $ \f ->
     precList
       ( \self h t ->
@@ -231,7 +248,7 @@ pdeleteBy = phoistAcyclic $
             t
             (pcons # h #$ self # t)
       )
-      (const pnil)
+      (const $ ptraceError "Cannot delete element")
 
 {- | / O(1) /.Return true if the given list has only one element.