share stake check code among redeemers

agora-specs/Sample/Proposal.hs

@@ -18,7 +18,7 @@ module Sample.Proposal (
   TransitionParameters (..),
   advanceFinishedPropsoal,
   advanceProposalInsufficientVotes,
-  advancePropsoalWithsStake,
+  advancePropsoalWithInvalidOutputStake,
 ) where
 
 import Agora.Governor (GovernorDatum (..))
@@ -403,9 +403,12 @@ mkTransitionTxInfo ::
   ProposalStartingTime ->
   -- | Valid time range of the transaction.
   POSIXTimeRange ->
+  -- | Add a unchanged stake or not.
+  Bool ->
   TxInfo
-mkTransitionTxInfo from to effects votes startingTime validTime =
+mkTransitionTxInfo from to effects votes startingTime validTime shouldAddUnchangedStake =
   let pst = Value.singleton proposalPolicySymbol "" 1
+      sst = Value.assetClassValue stakeAssetClass 1
 
       proposalInputDatum :: ProposalDatum
       proposalInputDatum =
@@ -426,11 +429,48 @@ mkTransitionTxInfo from to effects votes startingTime validTime =
           { status = to
           }
 
+      stakeOwner = signer
+      stakedAmount = 200
+
+      existingLocks :: [ProposalLock]
+      existingLocks =
+        [ ProposalLock (ResultTag 0) (ProposalId 0)
+        , ProposalLock (ResultTag 2) (ProposalId 1)
+        ]
+
+      stakeInputDatum :: StakeDatum
+      stakeInputDatum =
+        StakeDatum
+          { stakedAmount = Tagged stakedAmount
+          , owner = stakeOwner
+          , lockedBy = existingLocks
+          }
+
+      stakeOutputDatum :: StakeDatum
+      stakeOutputDatum = stakeInputDatum
+
+      stakeBuilder :: BaseBuilder
+      stakeBuilder =
+        if shouldAddUnchangedStake
+          then
+            mconcat
+              [ input $
+                  script stakeValidatorHash
+                    . withValue sst
+                    . withDatum stakeInputDatum
+                    . withTxId (txOutRefId stakeRef)
+              , output $
+                  script stakeValidatorHash
+                    . withValue (sst <> minAda)
+                    . withDatum stakeOutputDatum
+              ]
+          else mempty
+
       builder :: BaseBuilder
       builder =
         mconcat
           [ txId "95ba4015e30aef16a3461ea97a779f814aeea6b8009d99a94add4b8293be737a"
-          , signedWith signer
+          , signedWith stakeOwner
           , timeRange validTime
           , input $
               script proposalValidatorHash
@@ -442,13 +482,19 @@ mkTransitionTxInfo from to effects votes startingTime validTime =
                 . withValue (pst <> minAda)
                 . withDatum proposalOutputDatum
           ]
-   in buildTxInfoUnsafe builder
+   in buildTxInfoUnsafe $ builder <> stakeBuilder
+
+-- | Wrapper around 'advanceProposalSuccess'', with valid stake.
+advanceProposalSuccess :: TransitionParameters -> TxInfo
+advanceProposalSuccess ps = advanceProposalSuccess' ps True
 
 {- | Create a valid 'TxInfo' that advances a proposal, given the parameters.
+     The second parameter determines wherther valid stake should be included.
+
      Note that 'TransitionParameters.initialProposalStatus' should not be 'Finished'.
 -}
-advanceProposalSuccess :: TransitionParameters -> TxInfo
+advanceProposalSuccess' :: TransitionParameters -> Bool -> TxInfo
-advanceProposalSuccess params =
+advanceProposalSuccess' params =
   let -- Status of the output proposal.
       toStatus :: ProposalStatus
       toStatus = case params.initialProposalStatus of
@@ -615,6 +661,7 @@ advanceProposalFailureTimeout params =
         votes
         params.proposalStartingTime
         timeRange
+        True
 
 -- | An invalid 'TxInfo' that tries to advance a 'VotingReady' proposal without sufficient votes.
 advanceProposalInsufficientVotes :: TxInfo
@@ -643,6 +690,7 @@ advanceProposalInsufficientVotes =
         votes
         (ProposalStartingTime proposalStartingTime)
         timeRange
+        True
 
 -- | An invalid 'TxInfo' that tries to advance a 'Finished' proposal.
 advanceFinishedPropsoal :: TxInfo
@@ -675,19 +723,21 @@ advanceFinishedPropsoal =
         outcome0WinningVotes
         (ProposalStartingTime 0)
         timeRange
+        True
 
-{- | An illegal 'TxInfo' that tries to use 'AdvanceProposal' with a stake.
+{- | An illegal 'TxInfo' that tries to output a changed stake with 'AdvanceProposal'.
-     From the perspective of stake validator, the transition is valid,
+     From the perspective of stake validator, the transition is totally valid,
        so the proposal validator should reject this.
 -}
-advancePropsoalWithsStake :: TxInfo
+advancePropsoalWithInvalidOutputStake :: TxInfo
-advancePropsoalWithsStake =
+advancePropsoalWithInvalidOutputStake =
   let templateTxInfo =
-        advanceProposalSuccess
+        advanceProposalSuccess'
           TransitionParameters
             { initialProposalStatus = VotingReady
             , proposalStartingTime = ProposalStartingTime 0
             }
+          False
 
       ---
       -- Now we create a new lock on an arbitrary stake

agora-specs/Spec/Proposal.hs

@@ -51,7 +51,7 @@ import Sample.Proposal qualified as Proposal (
   advanceProposalFailureTimeout,
   advanceProposalInsufficientVotes,
   advanceProposalSuccess,
-  advancePropsoalWithsStake,
+  advancePropsoalWithInvalidOutputStake,
   cosignProposal,
   proposalCreation,
   proposalRef,
@@ -352,7 +352,7 @@ specs =
               )
               AdvanceProposal
               ( ScriptContext
-                  Proposal.advancePropsoalWithsStake
+                  Proposal.advancePropsoalWithInvalidOutputStake
                   (Spending Proposal.proposalRef)
               )
           ]

agora/Agora/Proposal/Scripts.hs

@@ -54,6 +54,7 @@ import Plutarch.Extra.Record (mkRecordConstr, (.&), (.=))
 import Plutarch.Extra.TermCont (
   pguardC,
   pletC,
+  pletFieldsC,
   pmatchC,
   ptryFromC,
  )
@@ -184,29 +185,37 @@ proposalValidator proposal =
     valueSpent <- pletC $ pvalueSpent # txInfoF.inputs
     spentST <- pletC $ psymbolValueOf # stCurrencySymbol #$ valueSpent
 
-    let AssetClass (stakeSym, stakeTn) = proposal.stakeSTAssetClass
-    stakeSTAssetClass <-
-      pletC $ passetClass # pconstant stakeSym # pconstant stakeTn
-    spentStakeST <-
-      pletC $ passetClassValueOf # valueSpent # stakeSTAssetClass
-
     signedBy <- pletC $ ptxSignedBy # txInfoF.signatories
 
-    pguardC "ST at inputs must be 1" (spentST #== 1)
-
     currentTime <- pletC $ currentProposalTime # txInfoF.validRange
 
-    -- Filter out own output with own address and PST.
+    -- Own output is an output that
-    -- Delay the evaluation cause in some cases there won't be any continuing output.
+    --  * is sent to the address of the proposal validator
+    --  * has an PST
+    --  * has the same proposal id as the proposal input
+    --
+    -- We match the proposal id here so that we can support multiple
+    --  proposal inputs in one thansaction.
     ownOutput <-
       pletC $
         mustBePJust # "Own output should be present" #$ pfind
           # plam
             ( \input -> unTermCont $ do
-                inputF <- tcont $ pletFields @'["address", "value"] input
+                inputF <- tcont $ pletFields @'["address", "value", "datumHash"] input
+
+                -- TODO: this is highly inefficient: O(n) for every output,
+                --       Maybe we can cache the sorted datum map?
+                let datum =
+                      mustFindDatum' @PProposalDatum
+                        # inputF.datumHash
+                        # txInfoF.datums
+
+                    proposalId = pfield @"proposalId" # datum
+
                 pure $
                   inputF.address #== ownAddress
                     #&& psymbolValueOf # stCurrencySymbol # inputF.value #== 1
+                    #&& proposalId #== proposalF.proposalId
             )
           # pfromData txInfoF.outputs
 
@@ -216,6 +225,46 @@ proposalValidator proposal =
           # (pfield @"datumHash" # ownOutput)
           # txInfoF.datums
 
+    --------------------------------------------------------------------------
+    -- Find the stake input and stake output by SST.
+
+    let AssetClass (stakeSym, stakeTn) = proposal.stakeSTAssetClass
+    stakeSTAssetClass <-
+      pletC $ passetClass # pconstant stakeSym # pconstant stakeTn
+    spentStakeST <-
+      pletC $ passetClassValueOf # valueSpent # stakeSTAssetClass
+
+    pguardC "ST at inputs must be 1" (spentST #== 1)
+
+    let stakeInput =
+          pfield @"resolved"
+            #$ mustBePJust
+            # "Stake input should be present"
+              #$ pfind
+            # plam
+              ( \(pfromData . (pfield @"value" #) . (pfield @"resolved" #) -> value) ->
+                  passetClassValueOf # value # stakeSTAssetClass #== 1
+              )
+            # pfromData txInfoF.inputs
+
+    stakeIn <- pletC $ mustFindDatum' @PStakeDatum # (pfield @"datumHash" # stakeInput) # txInfoF.datums
+    stakeInF <- pletFieldsC @'["stakedAmount", "lockedBy", "owner"] stakeIn
+
+    let stakeOutput =
+          mustBePJust # "Stake output should be present"
+            #$ pfind
+            # plam
+              ( \(pfromData . (pfield @"value" #) -> value) ->
+                  passetClassValueOf # value # stakeSTAssetClass #== 1
+              )
+            # pfromData txInfoF.outputs
+
+    stakeOut <- pletC $ mustFindDatum' @PStakeDatum # (pfield @"datumHash" # stakeOutput) # txInfoF.datums
+
+    stakeUnchanged <- pletC $ stakeIn #== stakeOut
+
+    --------------------------------------------------------------------------
+
     pure $
       pmatch proposalRedeemer $ \case
         PVote r -> unTermCont $ do
@@ -232,23 +281,6 @@ proposalValidator proposal =
           pguardC "Vote option should be valid" $
             pisJust #$ plookup # voteFor # voteMap
 
-          -- Find the input stake, the amount of new votes should be the 'stakedAmount'.
-          let stakeInput =
-                pfield @"resolved"
-                  #$ mustBePJust
-                  # "Stake input should be present"
-                    #$ pfind
-                  # plam
-                    ( \(pfromData . (pfield @"value" #) . (pfield @"resolved" #) -> value) ->
-                        passetClassValueOf # value # stakeSTAssetClass #== 1
-                    )
-                  # pfromData txInfoF.inputs
-
-              stakeIn :: Term _ PStakeDatum
-              stakeIn = mustFindDatum' # (pfield @"datumHash" # stakeInput) # txInfoF.datums
-
-          stakeInF <- tcont $ pletFields @'["stakedAmount", "lockedBy", "owner"] stakeIn
-
           -- Ensure that no lock with the current proposal id has been put on the stake.
           pguardC "Same stake shouldn't vote on the same propsoal twice" $
             pnot #$ pany
@@ -258,7 +290,8 @@ proposalValidator proposal =
                 )
               # pfromData stakeInF.lockedBy
 
-          let -- Update the vote counter of the proposal, and leave other stuff as is.
+          let -- The amount of new votes should be the 'stakedAmount'.
+              -- Update the vote counter of the proposal, and leave other stuff as is.
               expectedNewVotes = pmatch (pfromData proposalF.votes) $ \(PProposalVotes m) ->
                 pcon $
                   PProposalVotes $
@@ -289,18 +322,6 @@ proposalValidator proposal =
           -- to create a valid 'ProposalLock', however the vote option is encoded
           -- in the proposal redeemer, which is invisible for the stake validator.
 
-          let stakeOutput =
-                mustBePJust # "Stake output should be present"
-                  #$ pfind
-                  # plam
-                    ( \(pfromData . (pfield @"value" #) -> value) ->
-                        passetClassValueOf # value # stakeSTAssetClass #== 1
-                    )
-                  # pfromData txInfoF.outputs
-
-              stakeOut :: Term _ PStakeDatum
-              stakeOut = mustFindDatum' # (pfield @"datumHash" # stakeOutput) # txInfoF.datums
-
           let newProposalLock =
                 mkRecordConstr
                   PProposalLock
@@ -325,6 +346,8 @@ proposalValidator proposal =
           pure $ popaque (pconstant ())
         --------------------------------------------------------------------------
         PCosign r -> unTermCont $ do
+          pguardC "Stake should not change" stakeUnchanged
+
           newSigs <- pletC $ pfield @"newCosigners" # r
 
           pguardC "Cosigners are unique" $
@@ -374,13 +397,11 @@ proposalValidator proposal =
 
           pure $ popaque (pconstant ())
         --------------------------------------------------------------------------
-        PUnlock _r ->
+        PUnlock _r -> popaque (pconstant ())
-          popaque (pconstant ())
         --------------------------------------------------------------------------
         PAdvanceProposal _r -> unTermCont $ do
-          pguardC "No stake input is allowed" $ spentStakeST #== 0
+          pguardC "Stake should not change" stakeUnchanged
 
-          currentTime <- pletC $ currentProposalTime # txInfoF.validRange
           proposalOutStatus <- pletC $ pfield @"status" # proposalOut
 
           let -- Only the status of proposals should be updated in this case.

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,7665
+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/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,145357978,397941,5721
+Agora/Proposal/validator/cosignature/proposal,213692648,591151,5767
 Agora/Proposal/validator/cosignature/stake,115369581,282557,4681
-Agora/Proposal/validator/voting/proposal,154824944,415642,5650
+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,94701799,249495,5027
+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,93858377,247992,5030
+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,95554844,251598,5030
+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,93571998,246765,5029
+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,92163087,244060,5030
+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,93294065,246464,5030
+Agora/Proposal/validator/advancing/successfully advance to failed state: timeout/Locked -> Finished,166084300,447362,5597
 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,261928725,689487,8181
+Agora/Governor/validator/proposal creation,262494214,690689,8180
-Agora/Governor/validator/GATs minting,349849353,933334,8302
+Agora/Governor/validator/GATs minting,349283864,932132,8301
-Agora/Governor/validator/mutate governor state,84905433,234687,7766
+Agora/Governor/validator/mutate governor state,84905433,234687,7765