Treasury Withdrawal Effect: Only check treasury inputs

Tried to make it so that its only checking treasury inputs when
checking if transaction is correctly returning the remainders to treasury.

agora/Agora/Effect/TreasuryWithdrawal.hs

@@ -13,16 +13,17 @@ import GHC.Generics qualified as GHC
 import Generics.SOP (Generic, I (I))
 
 import Agora.Effect (makeEffect)
-import Agora.Utils ( paddValue, passert ) 
+import Agora.Utils (findTxOutByTxOutRef, paddValue, passert)
 import Plutarch (popaque)
-import Plutarch.Api.V1
-    ( PCredential,
-      PValue,
-      PTxOut(PTxOut),
-      PTxInInfo(PTxInInfo),
-      PTuple,
-      ptuple,
-      PValidator )
+import Plutarch.Api.V1 (
+  PCredential,
+  PTuple,
+  PTxInInfo (PTxInInfo),
+  PTxOut (PTxOut),
+  PValidator,
+  PValue,
+  ptuple,
+ )
 
 import Plutarch.DataRepr (
   DerivePConstantViaData (..),
@@ -79,9 +80,11 @@ It can be more flexiable over...
 -}
 treasuryWithdrawalValidator :: forall {s :: S}. CurrencySymbol -> Term s PValidator
 treasuryWithdrawalValidator currSymbol = makeEffect currSymbol $
-  \_cs (datum' :: Term _ PTreasuryWithdrawalDatum) _txOutRef' txInfo' -> P.do
+  \_cs (datum' :: Term _ PTreasuryWithdrawalDatum) txOutRef' txInfo' -> P.do
     receivers <- plet $ pfromData $ pfield @"receivers" # datum'
     txInfo <- pletFields @'["outputs", "inputs"] txInfo'
+    PJust txOut <- pmatch $ findTxOutByTxOutRef # txOutRef' # pfromData txInfo'
+    effInput <- pletFields @'["address", "value"] $ txOut
     let outputValues =
           pmap
             # plam
@@ -94,37 +97,58 @@ treasuryWithdrawalValidator currSymbol = makeEffect currSymbol $
         outputContentMatchesRecivers =
           pall # plam (\out -> pelem # out # outputValues)
             #$ receivers
-        sumValues = 
+        sumValues =
           pfoldr
             # plam
               ( \((pfield @"_1" #) . pfromData -> x) y -> P.do
-                  paddValue # (pfromData x) # y
+                  paddValue # pfromData x # y
               )
-            # (pconstant (mempty :: Value))
+            # pconstant (mempty :: Value)
         inputCred =
           pmap
-            # plam (\inInfo -> P.do
-                     PTxInInfo inInfo' <- pmatch $ pfromData inInfo
-                     PTxOut out <- pmatch $ pfromData $ pfield @"resolved" # inInfo'
-                     let addr = pfromData $ pfield @"address" # out
-                     pfield @"credential" # addr)
+            # plam
+              ( \inInfo -> P.do
+                  PTxInInfo inInfo' <- pmatch $ pfromData inInfo
+                  PTxOut out <- pmatch $ pfromData $ pfield @"resolved" # inInfo'
+                  let addr = pfromData $ pfield @"address" # out
+                  pfield @"credential" # addr
+              )
             # pfromData txInfo.inputs
-        totalInput =
+        totalTreasuryInputs =
           pfoldr
-            # plam (\x' y -> P.do
-                   PTxInInfo x <- pmatch $ pfromData x'
-                   PTxOut out <- pmatch $ pfromData $ pfield @"resolved" # x
-                   paddValue # (pfromData $ pfield @"value" # out) # y)
-            # (pconstant (mempty :: Value))
+            # plam
+              ( \x' y -> P.do
+                  PTxInInfo x <- pmatch $ pfromData x'
+                  PTxOut out <- pmatch $ pfromData $ pfield @"resolved" # x
+                  -- only take ones from treasury
+                  pif
+                    (effInput.address #== pfield @"address" # out)
+                    (paddValue # pfromData (pfield @"value" # out) # y)
+                    y
+              )
+            # pconstant (mempty :: Value)
             # pfromData txInfo.inputs
-        sumOutputsToInputAddr = sumValues #$
-          pfilter
-            # plam (\((pfield @"_0" #) . pfromData -> addr) -> pelem # (addr) # inputCred) # outputValues
-        sumReceivers = sumValues # receivers
--- TODO: Probably need to check/exclude the effect input...
-        excessShouldBePaidToInputs = pdata (paddValue # sumReceivers # sumOutputsToInputAddr) #== pdata totalInput
+        sumOutputsToTreasury =
+          sumValues
+            #$ pfilter
+            # plam
+              ( \((pfield @"_0" #) . pfromData -> addr) ->
+                  pelem # addr # inputCred
+                    #&& pnot # (addr #== pfield @"credential" # effInput.address)
+              )
+            # outputValues
+        -- TODO: Probably need to check/exclude the effect input...
+        excessShouldBePaidToInputs =
+          pdata (paddValue # (sumValues # receivers) # sumOutputsToTreasury) #== pdata totalTreasuryInputs
+        shouldNotPayToEffect =
+          pnot #$ pany
+            # plam
+              ( \x ->
+                  effInput.address #== pfield @"address" # pfromData x
+              )
+            # pfromData txInfo.outputs
 
     passert "Transaction output does not match receivers" outputContentMatchesRecivers
-    passert "Remainders should be returned to the treasury" excessShouldBePaidToInputs
-
+    passert "Transaction should not pay to effects" shouldNotPayToEffect
+    passert "Remainders should be returned to the treasury" excessShouldBePaidToInputs -- We might not need this.
     popaque $ pconstant ()