re-implement the treasury withdrawal effect

.gitignore

@@ -25,3 +25,5 @@ TAGS
 # Haddock files and Hoogle databases 
 haddock 
 hoo
+
+.pre-commit-config.yaml

agora/Agora/Bootstrap.hs

@@ -67,7 +67,7 @@ agoraScripts' conf =
     , envelope "NoOp Validator" noOpValidator
     , envelope "Treasury Withdrawal Validator" treasuryWithdrawalValidator
     , envelope "Mutate Governor Validator" mutateGovernorValidator
-    , envelope "Always Succeeds Policy" $ ((plam $ \_ _ -> popaque $ pcon PUnit) :: Term s PMintingPolicy)
+    , envelope "Always Succeeds Policy" ((plam $ \_ _ -> popaque $ pcon PUnit) :: Term s PMintingPolicy)
     ]
   where
     envelope ::

agora/Agora/Effect/TreasuryWithdrawal.hs

@@ -8,20 +8,17 @@ Description: An Effect that withdraws treasury deposit
 An Effect that withdraws treasury deposit
 -}
 module Agora.Effect.TreasuryWithdrawal (
-  TreasuryWithdrawalDatum (TreasuryWithdrawalDatum),
+  TreasuryWithdrawalDatum (..),
   PTreasuryWithdrawalDatum (PTreasuryWithdrawalDatum),
   treasuryWithdrawalValidator,
 ) where
 
 import Agora.Effect (makeEffect)
 import Agora.SafeMoney (AuthorityTokenTag)
+import Agora.Utils (psubtractSortedValue, puncurryTuple)
 import Generics.SOP qualified as SOP
-import Plutarch.Api.V1 (
+import Plutarch.Api.V1 (PCredential, PCurrencySymbol, PValue)
-  PCredential,
+import Plutarch.Api.V1.Value (pforgetPositive)
-  PCurrencySymbol,
-  PValue,
-  ptuple,
- )
 import Plutarch.Api.V2 (
   AmountGuarantees (Positive),
   KeyGuarantees (Sorted),
@@ -42,11 +39,11 @@ import Plutarch.Extra.IsData (
  )
 import Plutarch.Extra.ScriptContext (pisPubKey)
 import Plutarch.Extra.Tagged (PTagged)
+import Plutarch.Extra.Traversable (pfoldMap)
 import Plutarch.Lift (PConstantDecl, PUnsafeLiftDecl (PLifted))
 import PlutusLedgerApi.V1.Credential (Credential)
 import PlutusLedgerApi.V1.Value (Value)
 import PlutusTx qualified
-import "liqwid-plutarch-extra" Plutarch.Extra.List (pdeleteFirst)
 import "liqwid-plutarch-extra" Plutarch.Extra.TermCont (
   pguardC,
   pletC,
@@ -141,9 +138,11 @@ instance PTryFrom PData (PAsData PTreasuryWithdrawalDatum)
 
        2. Left over assets should be redirected back to Treasury
 
-     It can be more flexiable over...
+    The output order should be:
 
-     - The number of outputs themselves
+      1. Reciever outputs. They should be in the same order as the 'receivers' field of the datum.
+
+      2. Other outputs: treasury outputs, colleteral outputs, etc.
 
      @since 1.0.0
 -}
@@ -156,54 +155,117 @@ treasuryWithdrawalValidator = plam $
       datumF <- pletAllC datum
       txInfoF <- pletFieldsC @'["outputs", "inputs"] txInfo
 
-      let validateInput :: Term _ (PTxInInfo :--> PBool)
+      let
-          validateInput = plam $ \input -> unTermCont $ do
+        -- Validate the input and if it's from one of the treasuries,
-            inputF <- pletAllC input
+        --  return the value.
+        --
+        -- Only effect inputs, treasury inputs and public key inputs are
+        --  allowed.
+        extractTreasuryInputValue ::
+          Term _ (PTxInInfo :--> PValue 'Sorted 'Positive)
+        extractTreasuryInputValue = plam $ \input -> unTermCont $ do
+          inputF <- pletAllC input
+          resolvedF <- pletFieldsC @'["address", "value"] inputF.resolved
 
-            cred <-
+          cred <- pletC $ pfield @"credential" # resolvedF.address
-              pletC $
-                pfield @"credential"
-                  #$ pfield @"address"
-                  # inputF.resolved
 
-            pure $
+          let isEffectInput =
-              foldl1
+                ptraceIfTrue "Effect input" $
-                (#||)
+                  inputF.outRef #== effectInputRef
-                [ ptraceIfTrue "Effect input" $ inputF.outRef #== effectInputRef
+              isTreasuryInput =
-                , ptraceIfTrue "Treasury input" $ pelem # cred # datumF.treasuries
+                ptraceIfTrue "Treasury input" $
-                , ptraceIfTrue "Collateral input" $ pisPubKey # pfromData cred
+                  pelem # pdata cred # datumF.treasuries
-                ]
+              isPubkeyInput =
+                ptraceIfTrue "Pubkey input" $
+                  pisPubKey # cred
+          pure
+            $ pif
+              (isEffectInput #|| isPubkeyInput)
+              mempty
+            $ pif isTreasuryInput resolvedF.value
+            $ ptraceError "Unknown input"
 
-          validateOutput ::
+        treasuryInputAmount =
-            Term
+          pfoldMap
-              _
+            # extractTreasuryInputValue
-              ( PBuiltinList (PAsData (PTuple PCredential (PValue 'Sorted 'Positive)))
+            # txInfoF.inputs
-                  :--> 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
+        sentAmout =
+          pfoldMap
+            # plam ((puncurryTuple # plam (const id) #) . pfromData)
+            # pfromData datumF.receivers
 
-                shouldSendToTreasury =
+        treasuryLeftOverAmount =
-                  pif
+          psubtractSortedValue
-                    (pelem # cred # datumF.treasuries)
+            # treasuryInputAmount
-                    receivers
+            # sentAmout
-                    (ptraceError "Invalid receiver")
 
-            pure $
+        remainingOutputs =
-              pmatch (pdeleteFirst # credValue # receivers) $ \case
+          ptrace "Check receiver outputs" $
-                PJust updatedReceivers ->
+            checkReceiverOutputs
-                  ptrace "Receiver output" updatedReceivers
+              # datumF.receivers
-                PNothing ->
+              # txInfoF.outputs
-                  ptrace "Treasury output" shouldSendToTreasury
 
-      pguardC "All input are valid" $
+        extractTeasuryOutputValue ::
-        pall # validateInput # txInfoF.inputs
+          Term _ (PTxOut :--> PValue 'Sorted 'Positive)
+        extractTeasuryOutputValue = plam $
+          flip (pletFields @'["address", "value"]) $ \outputF ->
+            let cred = pfield @"credential" # outputF.address
 
-      pguardC "All receiver get correct output" $
+                isTreasuryOutput =
-        pnull #$ pfoldl # validateOutput # datumF.receivers # txInfoF.outputs
+                  pelem # cred # datumF.treasuries
+             in pif
+                  isTreasuryOutput
+                  outputF.value
+                  mempty
+
+        -- Return the value if it'll be sent to one of the treasuries.
+        treasuryOutputAmount =
+          pfoldMap
+            # extractTeasuryOutputValue
+            # remainingOutputs
+
+      pguardC "Unused treasury should stay at treasury validators" $
+        treasuryLeftOverAmount #== pforgetPositive treasuryOutputAmount
 
       pure . popaque $ pconstant ()
+  where
+    -- Make sure that all the receivers get the correct payment and return the
+    --  remaining outputs.
+    checkReceiverOutputs ::
+      forall (s :: S).
+      Term
+        s
+        ( PBuiltinList
+            (PAsData (PTuple PCredential (PValue 'Sorted 'Positive)))
+            :--> PBuiltinList PTxOut
+            :--> PBuiltinList PTxOut
+        )
+    checkReceiverOutputs = pfix #$ plam $ \self receivers outputs ->
+      pelimList
+        ( \r rs ->
+            pelimList
+              ( \o os -> pletFields @'["value", "address"] o $ \oF ->
+                  let isValidReceiverOutput =
+                        puncurryTuple
+                          # plam
+                            ( \expCred expVal ->
+                                foldl1
+                                  (#&&)
+                                  [ ptraceIfFalse "Valid credential" $
+                                      expCred #== pfield @"credential" # oF.address
+                                  , ptraceIfFalse "Valid value" $
+                                      expVal #== oF.value
+                                  ]
+                            )
+                          # pfromData r
+                   in pif
+                        isValidReceiverOutput
+                        (self # rs # os)
+                        (ptraceError "Invalid receiver output")
+              )
+              (ptraceError "Unable to exhaust receivers")
+              outputs
+        )
+        outputs
+        receivers

agora/Agora/Utils.hs

@@ -18,22 +18,29 @@ module Agora.Utils (
   ptag,
   puntag,
   phashDatum,
+  puncurryTuple,
+  psubtractSortedValue,
 ) where
 
+import Plutarch.Api.V1 (KeyGuarantees (Sorted))
+import Plutarch.Api.V1.AssocMap (punionWith)
 import Plutarch.Api.V1.Scripts (PDatumHash (PDatumHash))
 import Plutarch.Api.V2 (
-  AmountGuarantees,
+  AmountGuarantees (NoGuarantees),
-  KeyGuarantees,
   PCurrencySymbol,
   PMaybeData (PDNothing),
+  PTuple,
   PValue,
  )
 import Plutarch.Builtin (pforgetData, pserialiseData)
 import Plutarch.Crypto (pblake2b_256)
+import Plutarch.DataRepr (punDataSum)
 import Plutarch.Extra.AssetClass (PAssetClass, PAssetClassData, ptoScottEncoding)
+import Plutarch.Extra.Field (pletAll)
 import Plutarch.Extra.Tagged (PTagged)
 import Plutarch.Extra.Value (psymbolValueOf)
-import Plutarch.Unsafe (punsafeDowncast)
+import Plutarch.Num ((#-))
+import Plutarch.Unsafe (punsafeCoerce, punsafeDowncast)
 import PlutusLedgerApi.V2 (
   Address (Address),
   Credential (ScriptCredential),
@@ -139,3 +146,27 @@ phashDatum =
         . (pserialiseData #)
         . pforgetData
         . pdata
+
+puncurryTuple ::
+  forall (c :: PType) (a :: PType) (b :: PType) (s :: S).
+  (PIsData a, PIsData b) =>
+  Term s ((a :--> b :--> c) :--> PTuple a b :--> c)
+puncurryTuple = phoistAcyclic $
+  plam $
+    \f ((punDataSum #) -> r) ->
+      pletAll r $ \rF -> f # rF._0 # rF._1
+
+psubtractSortedValue ::
+  forall (ag :: AmountGuarantees) (s :: S).
+  Term
+    s
+    ( PValue 'Sorted ag
+        :--> PValue 'Sorted ag
+        :--> PValue 'Sorted 'NoGuarantees
+    )
+psubtractSortedValue = phoistAcyclic $ plam $ \a b ->
+  punsafeCoerce $
+    punionWith
+      # (punionWith # plam (#-))
+      # pto a
+      # pto b