add validation logic for unlocking stakes

agora/Agora/Proposal.hs

@@ -33,6 +33,7 @@ module Agora.Proposal (
   pemptyVotesFor,
   pwinner,
   pneutralOption,
+  pretractVotes,
 ) where
 
 --------------------------------------------------------------------------------
@@ -63,7 +64,7 @@ import Plutarch.DataRepr (DerivePConstantViaData (..), PDataFields, PIsDataReprI
 import Plutarch.Extra.List (pnotNull)
 import Plutarch.Extra.Map qualified as PM
 import Plutarch.Extra.Map.Unsorted qualified as PUM
-import Plutarch.Extra.TermCont (pletC)
+import Plutarch.Extra.TermCont (pguardC, pletC)
 import Plutarch.Lift (
   DerivePConstantViaNewtype (..),
   PConstantDecl,
@@ -367,6 +368,24 @@ newtype PProposalVotes (s :: S)
   = PProposalVotes (Term s (PMap 'Unsorted PResultTag PInteger))
   deriving (PlutusType, PIsData) via (DerivePNewtype PProposalVotes (PMap 'Unsorted PResultTag PInteger))
 
+-- | Retract votes given the option and the amount of votes.
+pretractVotes :: Term s (PProposalVotes :--> PResultTag :--> PInteger :--> PProposalVotes)
+pretractVotes = phoistAcyclic $
+  plam $ \votes rt count ->
+    let voteMap :: Term _ (PMap 'Unsorted PResultTag PInteger)
+        voteMap = pto votes
+     in pcon $
+          PProposalVotes $
+            PM.pupdate
+              # plam
+                ( \oldCount -> unTermCont $ do
+                    newCount <- pletC $ oldCount - count
+                    pguardC "Resulting vote count greater or equal to 0" $ 0 #<= newCount
+                    pure $ pcon $ PJust newCount
+                )
+              # rt
+              # voteMap
+
 instance PUnsafeLiftDecl PProposalVotes where type PLifted PProposalVotes = ProposalVotes
 deriving via
   (DerivePConstantViaNewtype ProposalVotes PProposalVotes (PMap 'Unsorted PResultTag PInteger))

agora/Agora/Proposal/Scripts.hs

@@ -17,6 +17,7 @@ import Agora.Proposal (
   PProposalVotes (PProposalVotes),
   Proposal (governorSTAssetClass, stakeSTAssetClass),
   ProposalStatus (..),
+  pretractVotes,
  )
 import Agora.Proposal.Time (
   currentProposalTime,
@@ -25,7 +26,13 @@ import Agora.Proposal.Time (
   isLockingPeriod,
   isVotingPeriod,
  )
-import Agora.Stake (PProposalLock (..), PStakeDatum (..), findStakeOwnedBy)
+import Agora.Stake (
+  PProposalLock (..),
+  PStakeDatum (..),
+  PStakeUsage (..),
+  findStakeOwnedBy,
+  pgetStakeUsage,
+ )
 import Agora.Utils (
   findTxOutByTxOutRef,
   getMintingPolicySymbol,
@@ -183,7 +190,6 @@ proposalValidator proposal =
     let stCurrencySymbol =
           pconstant $ getMintingPolicySymbol (proposalPolicy proposal.governorSTAssetClass)
     valueSpent <- pletC $ pvalueSpent # txInfoF.inputs
-    spentST <- pletC $ psymbolValueOf # stCurrencySymbol #$ valueSpent
 
     signedBy <- pletC $ ptxSignedBy # txInfoF.signatories
 
@@ -225,6 +231,7 @@ proposalValidator proposal =
           # (pfield @"datumHash" # ownOutput)
           # txInfoF.datums
 
+    proposalUnchanged <- pletC $ proposalOut #== proposalDatum
     --------------------------------------------------------------------------
     -- Find the stake input and stake output by SST.
 
@@ -234,8 +241,6 @@ proposalValidator proposal =
     spentStakeST <-
       pletC $ passetClassValueOf # valueSpent # stakeSTAssetClass
 
-    pguardC "ST at inputs must be 1" (spentST #== 1)
-
     let stakeInput =
           pfield @"resolved"
             #$ mustBePJust
@@ -397,7 +402,82 @@ proposalValidator proposal =
 
           pure $ popaque (pconstant ())
         --------------------------------------------------------------------------
-        PUnlock _r -> popaque (pconstant ())
+        PUnlock r -> unTermCont $ do
+          -- At draft stage, the votes should be empty.
+          pguardC "Shouldn't retract votes from a draft propsoal" $
+            pnot #$ proposalF.status #== pconstantData Draft
+
+          -- This is the vote option we're retracting from.
+          retractFrom <- pletC $ pfield @"resultTag" # r
+
+          -- Determine if the input stake is actually locked by this proposal.
+          stakeUsage <- pletC $ pgetStakeUsage # stakeInF.lockedBy # proposalF.proposalId
+
+          pguardC "Stake input relevant" $
+            pmatch stakeUsage $ \case
+              PDidNothing ->
+                ptrace "Not relevant" $
+                  pconstant False
+              PCreated ->
+                ptraceIfFalse "Too early" $
+                  proposalF.status #== pconstantData Finished
+              PVotedFor rt ->
+                ptraceIfFalse "Result tag not match" $
+                  rt #== retractFrom
+
+          -- The count of removing votes is equal to the 'stakeAmount' of input stake.
+          retractCount <-
+            pletC $
+              pmatch stakeInF.stakedAmount $ (\(PDiscrete v) -> pextract # v)
+
+          -- The votes can only change when the proposal still allows voting.
+          let shouldUpdateVotes =
+                proposalF.status #== pconstantData VotingReady
+                  #&& pnot # (pcon PCreated #== stakeUsage)
+
+          pguardC "Proposal output correct" $
+            pif
+              shouldUpdateVotes
+              ( let -- Remove votes and leave other parts of the proposal as it.
+                    expectedVotes = pretractVotes # proposalF.votes # retractFrom # retractCount
+
+                    expectedProposalOut =
+                      mkRecordConstr
+                        PProposalDatum
+                        ( #proposalId .= proposalF.proposalId
+                            .& #effects .= proposalF.effects
+                            .& #status .= proposalF.status
+                            .& #cosigners .= proposalF.cosigners
+                            .& #thresholds .= proposalF.thresholds
+                            .& #votes .= pdata expectedVotes
+                            .& #timingConfig .= proposalF.timingConfig
+                            .& #startingTime .= proposalF.startingTime
+                        )
+                 in ptraceIfFalse "Update votes" $
+                      expectedProposalOut #== proposalOut
+              )
+              -- No change to the proposal is allowed.
+              $ ptraceIfFalse "Proposal unchanged" proposalUnchanged
+
+          -- At last, we ensure that all locks belong to this proposal will be removed.
+          stakeOutputLocks <- pletC $ pfield @"lockedBy" # stakeOut
+
+          let templateStakeOut =
+                mkRecordConstr
+                  PStakeDatum
+                  ( #stakedAmount .= stakeInF.stakedAmount
+                      .& #owner .= stakeInF.owner
+                      .& #lockedBy .= stakeOutputLocks
+                  )
+
+          pguardC "Only locks updated in the output stake" $
+            templateStakeOut #== stakeOut
+
+          pguardC "All relevant locks removed from the stake" $
+            pgetStakeUsage # pfromData stakeOutputLocks
+              # proposalF.proposalId #== pcon PDidNothing
+
+          pure $ popaque (pconstant ())
         --------------------------------------------------------------------------
         PAdvanceProposal _r -> unTermCont $ do
           pguardC "Stake should not change" stakeUnchanged

agora/Agora/Stake.hs

@@ -18,10 +18,12 @@ module Agora.Stake (
   PStakeDatum (..),
   PStakeRedeemer (..),
   PProposalLock (..),
+  PStakeUsage (..),
 
   -- * Utility functions
   stakeLocked,
   findStakeOwnedBy,
+  pgetStakeUsage,
 ) where
 
 --------------------------------------------------------------------------------
@@ -29,7 +31,7 @@ module Agora.Stake (
 import Control.Applicative (Const)
 import Data.Tagged (Tagged (..))
 import GHC.Generics qualified as GHC
-import Generics.SOP (Generic, I (I))
+import Generics.SOP (Generic, HasDatatypeInfo, I (I))
 import Prelude hiding (Num (..))
 
 --------------------------------------------------------------------------------
@@ -56,8 +58,8 @@ import Plutarch.DataRepr (
   PDataFields,
   PIsDataReprInstances (PIsDataReprInstances),
  )
-import Plutarch.Extra.List (pnotNull)
-import Plutarch.Extra.TermCont (pletC, pmatchC)
+import Plutarch.Extra.List (pmapMaybe, pnotNull)
+import Plutarch.Extra.TermCont (pletC, pletFieldsC, pmatchC)
 import Plutarch.Internal (punsafeCoerce)
 import Plutarch.Lift (PConstantDecl, PUnsafeLiftDecl (..))
 import Plutarch.SafeMoney (PDiscrete)
@@ -286,7 +288,7 @@ stakeDatumOwnedBy =
       pletFields @'["owner"] (pto stakeDatum) $ \stakeDatumF ->
         stakeDatumF.owner #== pdata pk
 
--- Does the input have a `Stake` owned by a particular PK?
+-- | Does the input have a `Stake` owned by a particular PK?
 isInputStakeOwnedBy ::
   Term
     _
@@ -299,7 +301,7 @@ isInputStakeOwnedBy =
   plam $ \ac ss datums txInInfo' -> unTermCont $ do
     PTxInInfo ((pfield @"resolved" #) -> txOut) <- pmatchC $ pfromData txInInfo'
     PTxOut txOut' <- pmatchC txOut
-    txOutF <- tcont $ pletFields @'["value", "datumHash"] txOut'
+    txOutF <- pletFieldsC @'["value", "datumHash"] txOut'
     outStakeST <- pletC $ passetClassValueOf # txOutF.value # ac
     pure $
       pmatch txOutF.datumHash $ \case
@@ -312,3 +314,53 @@ isInputStakeOwnedBy =
                 PJust v -> stakeDatumOwnedBy # ss # pfromData (punsafeCoerce v)
             )
             (pcon PFalse)
+
+{- | Represent the usage of a stake on a particular proposal.
+     A stake can be used to either create or vote on a proposal.
+-}
+data PStakeUsage (s :: S)
+  = PVotedFor (Term s PResultTag)
+  | PCreated
+  | PDidNothing
+  deriving stock (GHC.Generic)
+  deriving anyclass (Generic, PlutusType, HasDatatypeInfo, PEq)
+
+{- | / O(n) /.Return the usage of a stake on a particular proposal,
+      given the 'lockedBy' field of a stake and the target proposal.
+-}
+pgetStakeUsage ::
+  Term
+    _
+    ( PBuiltinList (PAsData PProposalLock)
+        :--> PProposalId
+        :--> PStakeUsage
+    )
+pgetStakeUsage = phoistAcyclic $
+  plam $ \locks pid ->
+    let -- All locks from the given proposal.
+        filteredLocks =
+          pmapMaybe
+            # plam
+              ( \lock'@(pfromData -> lock) -> unTermCont $ do
+                  lockF <- pletFieldsC @'["proposalTag"] lock
+
+                  pure $
+                    pif
+                      (lockF.proposalTag #== pid)
+                      (pcon $ PJust lock')
+                      (pcon PNothing)
+              )
+            # locks
+
+        lockCount' = plength # filteredLocks
+     in plet lockCount' $ \lockCount ->
+          pif (lockCount #== 0) (pcon PDidNothing) $
+            pif
+              (lockCount #== 1)
+              ( pcon $
+                  PVotedFor $
+                    pfromData $
+                      pfield @"vote" #$ phead # filteredLocks
+              )
+              -- Note: see the implementation of the governor.
+              (pcon PCreated)

agora/Agora/Stake/Scripts.hs

@@ -287,7 +287,7 @@ stakeValidator stake =
 
           pure $
             pmatch stakeRedeemer $ \case
-              PRetractVotes _ -> unTermCont $ do
+              PRetractVotes l -> unTermCont $ do
                 pguardC
                   "Owner signs this transaction"
                   ownerSignsTransaction
@@ -301,15 +301,22 @@ stakeValidator stake =
                   spentProposalST #== 1
 
                 pguardC "A UTXO must exist with the correct output" $
-                  unTermCont $ do
-                    let valueCorrect = ownOutputValueUnchanged
+                  let expectedLocks = pfield @"locks" # l
 
-                    -- TODO: check output datum is expected.
+                      expectedDatum =
+                        mkRecordConstr
+                          PStakeDatum
+                          ( #stakedAmount .= stakeDatum.stakedAmount
+                              .& #owner .= stakeDatum.owner
+                              .& #lockedBy .= expectedLocks
+                          )
 
-                    pure $
-                      foldl1
+                      valueCorrect = ownOutputValueUnchanged
+                      outputDatumCorrect = stakeOut #== expectedDatum
+                   in foldl1
                         (#&&)
                         [ ptraceIfFalse "valueCorrect" valueCorrect
+                        , ptraceIfFalse "datumCorrect" outputDatumCorrect
                         ]
 
                 pure $ popaque (pconstant ())

bench.csv

@@ -2,28 +2,28 @@ name,cpu,mem,size
 Agora/Effects/Treasury Withdrawal Effect/effect/Simple,289461528,703055,3191
 Agora/Effects/Treasury Withdrawal Effect/effect/Simple with multiple treasuries ,448521458,1070167,3518
 Agora/Effects/Treasury Withdrawal Effect/effect/Mixed Assets,408085321,966048,3383
-Agora/Effects/Governor Mutation Effect/validator/valid new governor datum/governor validator should pass,83758582,229228,7664
+Agora/Effects/Governor Mutation Effect/validator/valid new governor datum/governor validator should pass,83758582,229228,7665
 Agora/Effects/Governor Mutation Effect/validator/valid new governor datum/effect validator should pass,97345575,266935,3358
 Agora/Stake/policy/stakeCreation,43114795,124549,2156
-Agora/Stake/validator/stakeDepositWithdraw deposit,171823342,464745,4144
-Agora/Stake/validator/stakeDepositWithdraw withdraw,171823342,464745,4132
-Agora/Proposal/policy/proposalCreation,23140177,69194,1517
-Agora/Proposal/validator/cosignature/proposal,213692648,591151,5767
-Agora/Proposal/validator/cosignature/stake,115369581,282557,4681
-Agora/Proposal/validator/voting/proposal,167847632,446101,5696
-Agora/Proposal/validator/voting/stake,99545453,256941,4655
-Agora/Proposal/validator/advancing/successfully advance to next state/Draft -> VotringReady,167492034,450393,5594
-Agora/Proposal/validator/advancing/successfully advance to next state/VotingReady -> Locked,166648612,448890,5597
-Agora/Proposal/validator/advancing/successfully advance to next state/Locked -> Finished,168345079,452496,5597
-Agora/Proposal/validator/advancing/successfully advance to failed state: timeout/Draft -> Finished,166362233,447663,5596
-Agora/Proposal/validator/advancing/successfully advance to failed state: timeout/VotingReady -> Finished,164953322,444958,5597
-Agora/Proposal/validator/advancing/successfully advance to failed state: timeout/Locked -> Finished,166084300,447362,5597
+Agora/Stake/validator/stakeDepositWithdraw deposit,171823342,464745,4189
+Agora/Stake/validator/stakeDepositWithdraw withdraw,171823342,464745,4177
+Agora/Proposal/policy/proposalCreation,23140177,69194,1518
+Agora/Proposal/validator/cosignature/proposal,204675349,564476,6565
+Agora/Proposal/validator/cosignature/stake,114125937,284821,4726
+Agora/Proposal/validator/voting/proposal,166129664,437310,6494
+Agora/Proposal/validator/voting/stake,107127768,275725,4700
+Agora/Proposal/validator/advancing/successfully advance to next state/Draft -> VotringReady,162018766,433842,6392
+Agora/Proposal/validator/advancing/successfully advance to next state/VotingReady -> Locked,161175344,432339,6395
+Agora/Proposal/validator/advancing/successfully advance to next state/Locked -> Finished,162871811,435945,6395
+Agora/Proposal/validator/advancing/successfully advance to failed state: timeout/Draft -> Finished,160888965,431112,6394
+Agora/Proposal/validator/advancing/successfully advance to failed state: timeout/VotingReady -> Finished,159480054,428407,6395
+Agora/Proposal/validator/advancing/successfully advance to failed state: timeout/Locked -> Finished,160611032,430811,6395
 Agora/AuthorityToken/singleAuthorityTokenBurned/Correct simple,21017788,55883,806
 Agora/AuthorityToken/singleAuthorityTokenBurned/Correct many inputs,33204186,88241,900
 Agora/Treasury/Validator/Positive/Allows for effect changes,29938856,79744,1390
 Agora/AuthorityToken/singleAuthorityTokenBurned/Correct simple,21017788,55883,806
 Agora/AuthorityToken/singleAuthorityTokenBurned/Correct many inputs,33204186,88241,900
 Agora/Governor/policy/GST minting,43087287,120125,1829
-Agora/Governor/validator/proposal creation,262494214,690689,8180
-Agora/Governor/validator/GATs minting,349283864,932132,8301
-Agora/Governor/validator/mutate governor state,84905433,234687,7765
+Agora/Governor/validator/proposal creation,261928725,689487,8181
+Agora/Governor/validator/GATs minting,352305185,937264,8302
+Agora/Governor/validator/mutate governor state,84905433,234687,7766