allow voting/retracting votes with multiple stakes

agora/Agora/Plutarch/Orphans.hs

@@ -8,6 +8,8 @@ import Data.Bifunctor (Bifunctor (bimap))
 import Data.Map.Strict qualified as StrictMap
 import Data.Traversable (for)
 import Plutarch.Api.V1 (KeyGuarantees (Sorted), PMap)
+import Plutarch.Num (PNum)
+import Plutarch.SafeMoney (PDiscrete)
 import PlutusTx qualified
 import PlutusTx.AssocMap qualified as AssocMap
 
@@ -74,3 +76,6 @@ instance
       isSorted [] = True
       isSorted [_] = True
       isSorted (x : y : xs) = x < y && isSorted (y : xs)
+
+-- | @since 1.0.0
+deriving anyclass instance PNum (PDiscrete tag)

agora/Agora/Proposal.hs

@@ -617,6 +617,8 @@ newtype PProposalVotes (s :: S)
       PlutusType
     , -- | @since 0.1.0
       PIsData
+    , -- | @since 1.0.0
+      PShow
     )
 
 -- | @since 0.2.0

agora/Agora/Proposal/Scripts.hs

@@ -29,17 +29,16 @@ import Agora.Proposal.Time (
  )
 import Agora.Scripts (AgoraScripts, governorSTSymbol, proposalSTSymbol, stakeSTAssetClass)
 import Agora.Stake (
-  PProposalLock (PVoted),
   PStakeDatum,
   pextractVoteOption,
   pgetStakeRole,
-  pisCreator,
   pisIrrelevant,
   pisPureCreator,
   pisVoter,
  )
 import Agora.Utils (
   plistEqualsBy,
+  pmapMaybe,
  )
 import Plutarch.Api.V1 (PCredential)
 import Plutarch.Api.V1.AssocMap (plookup)
@@ -73,6 +72,7 @@ import Plutarch.Extra.ScriptContext (
   pisTokenSpent,
   ptryFromOutputDatum,
  )
+import Plutarch.Extra.Sum (PSum (PSum))
 import "liqwid-plutarch-extra" Plutarch.Extra.TermCont (
   pguardC,
   pletC,
@@ -80,6 +80,7 @@ import "liqwid-plutarch-extra" Plutarch.Extra.TermCont (
   pmatchC,
   ptryFromC,
  )
+import Plutarch.Extra.Traversable (pfoldMap)
 import Plutarch.Extra.Value (psymbolValueOf)
 import Plutarch.SafeMoney (PDiscrete (PDiscrete))
 import Plutarch.Unsafe (punsafeCoerce)
@@ -154,13 +155,9 @@ data PWitnessMultipleStakeContext (s :: S) = PWitnessMultipleStakeContext
 instance DerivePlutusType PWitnessMultipleStakeContext where
   type DPTStrat _ = PlutusTypeScott
 
-{- | Validation context for redeemers which need to modify a single stake.
+-- | @since 1.0.0
-
+newtype PStakeInputsContext (s :: S) = PStakeInputsContext
-     @since 1.0.0
+  { inputStakes :: Term s (PList PStakeDatum)
--}
-data PSpendSingleStakeContext (s :: S) = PSpendSingleStakeContext
-  { inputStake :: Term s PStakeDatum
-  , outputStake :: Term s PStakeDatum
   }
   deriving stock
     ( -- | @since 1.0.0
@@ -172,8 +169,8 @@ data PSpendSingleStakeContext (s :: S) = PSpendSingleStakeContext
     )
 
 -- | @since 1.0.0
-instance DerivePlutusType PSpendSingleStakeContext where
+instance DerivePlutusType PStakeInputsContext where
-  type DPTStrat _ = PlutusTypeScott
+  type DPTStrat _ = PlutusTypeNewtype
 
 {- | The validator for Proposals.
 
@@ -256,66 +253,59 @@ proposalValidator as maximumCosigners =
     --  * has an PST
     --  * has the same proposal id as the proposal input
     --
-    -- We match the proposal id here so that we can support multiple
+    -- We can handle only one proposal under current design.
-    --  proposal inputs in one thansaction.
     proposalOutputDatum <-
       pletC $
-        passertPJust
+        passertPJust # "proposal input should present"
-          # "Own output should be present"
+          #$ pfindJust
-            #$ pfindJust
+            # plam
-          # plam
+              ( flip pletAll $ \outputF ->
-            ( flip pletAll $ \outputF ->
+                  let pstSymbol = pconstant $ proposalSTSymbol as
-                let pstSymbol = pconstant $ proposalSTSymbol as
 
-                    isProposalUTxO =
+                      isProposalUTxO =
-                      foldl1
+                        foldl1
-                        (#&&)
+                          (#&&)
-                        [ ptraceIfFalse "Own by proposal validator" $
+                          [ ptraceIfFalse "Own by proposal validator" $
-                            outputF.address #== proposalInputF.address
+                              outputF.address #== proposalInputF.address
-                        , ptraceIfFalse "Has proposal ST" $
+                          , ptraceIfFalse "Has proposal ST" $
-                            psymbolValueOf # pstSymbol # outputF.value #== 1
+                              psymbolValueOf # pstSymbol # outputF.value #== 1
-                        ]
+                          ]
 
-                    handleProposalUTxO = unTermCont $ do
+                      handleProposalUTxO =
-                      -- Using inline datum to avoid O(n^2) lookup.
+                        -- Using inline datum to avoid O(n^2) lookup.
-                      datum <-
+                        pfromData $
-                        pletC $
+                          ptrace "Resolve proposal datum" $
-                          pfromData $
+                            pfromOutputDatum @(PAsData PProposalDatum)
-                            ptrace "Resolve proposal datum" $
+                              # outputF.datum
-                              pfromOutputDatum @(PAsData PProposalDatum)
+                              # txInfoF.datums
-                                # outputF.datum
+                   in pif
-                                # txInfoF.datums
+                        isProposalUTxO
-
+                        (pjust # handleProposalUTxO)
-                      pure $
+                        pnothing
-                        pif
+              )
-                          ( pfield @"proposalId" # pto datum
+            # pfromData txInfoF.outputs
-                              #== proposalInputDatumF.proposalId
-                          )
-                          (pjust # datum)
-                          pnothing
-                 in pif
-                      isProposalUTxO
-                      handleProposalUTxO
-                      pnothing
-            )
-          # pfromData txInfoF.outputs
 
     --------------------------------------------------------------------------
+    let AssetClass (sstSymbol, sstName) = stakeSTAssetClass as
 
-    -- Handle stake input/output.
+    sstAssetClass <-
+      pletC $
+        passetClass
+          # pconstant sstSymbol
+          # pconstant sstName
+
+    -- Handle stake inputs/outputs.
 
     -- Reslove stake datum if the given UTxO is a stake UTxO.
     getStakeDatum :: Term _ (PTxOut :--> PMaybe PStakeDatum) <-
       pletC $
         plam $
           flip (pletFields @'["value", "datum"]) $ \txOutF ->
-            let AssetClass (stakeSym, _) = stakeSTAssetClass as
+            let isStakeUTxO =
-
-                isStakeUTxO =
                   -- A stake UTxO is a UTxO that carries SST.
-                  psymbolValueOf
+                  passetClassValueOf
-                    # pconstant stakeSym
                     # txOutF.value
+                    # sstAssetClass
                     #== 1
 
                 stake =
@@ -328,6 +318,30 @@ proposalValidator as maximumCosigners =
                       # txInfoF.datums
              in pif isStakeUTxO (pjust # stake) pnothing
 
+    spendStakes' :: Term _ ((PStakeInputsContext :--> PUnit) :--> PUnit) <-
+      pletC $
+        plam $ \val -> unTermCont $ do
+          let stakeInputs =
+                pmapMaybe
+                  # (pfield @"resolved" #>>> getStakeDatum)
+                  # pfromData txInfoF.inputs
+
+              ctx = pcon $ PStakeInputsContext stakeInputs
+
+          pguardC "No stake burnt" $
+            passetClassValueOf # txInfoF.mint # sstAssetClass #== 0
+
+          pure $ val # ctx
+
+    let spendStakes ::
+          ( PStakeInputsContext _ ->
+            TermCont _ ()
+          ) ->
+          Term _ POpaque
+        spendStakes c = popaque $
+          spendStakes' #$ plam $ \sctx ->
+            unTermCont $ pmatchC sctx >>= c >> pure (pconstant ())
+
     -- Witness stakes in reference inputs.
     witnessStakes' ::
       Term
@@ -388,34 +402,6 @@ proposalValidator as maximumCosigners =
           witnessStakes' #$ plam $ \sctxF ->
             unTermCont $ pmatchC sctxF >>= c >> pure (pconstant ())
 
-    -- We don't need to explicitly ensure that there's only one stake in the
-    --   inputs here - the stake validator will do it for us.
-    spendSingleStake' ::
-      Term
-        s
-        ((PSpendSingleStakeContext :--> PUnit) :--> PUnit) <-
-      pletC $
-        let stakeInput =
-              passertPJust # "Stake input should present" #$ pfindJust
-                # ((pfield @"resolved" @_ @PTxInInfo) #>>> getStakeDatum)
-                # txInfoF.inputs
-
-            stakeOutput =
-              passertPJust # "Stake output should present"
-                #$ pfindJust # getStakeDatum # txInfoF.outputs
-
-            ctx = pcon $ PSpendSingleStakeContext stakeInput stakeOutput
-         in plam (# ctx)
-
-    let spendSingleStake ::
-          ( PSpendSingleStakeContext _ ->
-            TermCont _ ()
-          ) ->
-          Term _ POpaque
-        spendSingleStake c = popaque $
-          spendSingleStake' #$ plam $ \sctx ->
-            unTermCont $ pmatchC sctx >>= c >> pure (pconstant ())
-
     ----------------------------------------------------------------------------
 
     proposalRedeemer <- fst <$> ptryFromC @PProposalRedeemer redeemer
@@ -474,8 +460,26 @@ proposalValidator as maximumCosigners =
 
           ----------------------------------------------------------------------
 
-          PVote r -> spendSingleStake $ \sctxF -> do
+          PVote r -> spendStakes $ \sctxF -> do
-            stakeInF <- pletAllC $ pto sctxF.inputStake
+            let totalStakeAmount =
+                  pto $
+                    pfoldMap
+                      # plam
+                        ( \stake -> unTermCont $ do
+                            stakeF <- pletFieldsC @'["stakedAmount", "lockedBy"] stake
+
+                            pguardC "Same stake shouldn't vote on the same proposal twice" $
+                              pnot
+                                #$ pisVoter
+                                #$ pgetStakeRole
+                                  # proposalInputDatumF.proposalId
+                                  # stakeF.lockedBy
+
+                            pure $ pcon $ PSum $ pfromData stakeF.stakedAmount
+                        )
+                      # sctxF.inputStakes
+
+            -- TODO(Connor): check minimum stake amount?
 
             pguardC "Input proposal must be in VotingReady state" $
               currentStatus #== pconstant VotingReady
@@ -494,10 +498,6 @@ proposalValidator as maximumCosigners =
             pguardC "Vote option should be valid" $
               pisJust #$ plookup # voteFor # voteMap
 
-            -- Ensure that no lock with the current proposal id has been put on the stake.
-            pguardC "Same stake shouldn't vote on the same proposal twice" $
-              pnot #$ pisVoter #$ pgetStakeRole # proposalInputDatumF.proposalId # stakeInF.lockedBy
-
             let -- The amount of new votes should be the 'stakedAmount'.
                 -- Update the vote counter of the proposal, and leave other stuff as is.
                 expectedNewVotes =
@@ -506,7 +506,7 @@ proposalValidator as maximumCosigners =
                       pupdate
                         # plam
                           ( \votes -> unTermCont $ do
-                              PDiscrete v <- pmatchC stakeInF.stakedAmount
+                              PDiscrete v <- pmatchC totalStakeAmount
                               pure $ pcon $ PJust $ votes + (pextract # v)
                           )
                         # voteFor
@@ -528,76 +528,68 @@ proposalValidator as maximumCosigners =
             pguardC "Output proposal should be valid" $
               proposalOutputDatum #== expectedProposalOut
 
-            -- We validate the output stake datum here as well: We need the vote option
+          -- Note that the output stake locks validation now happens in the
-            -- to create a valid 'ProposalLock', however the vote option is encoded
+          -- stake validator.
-            -- in the proposal redeemer, which is invisible for the stake validator.
-
-            let newProposalLock =
-                  mkRecordConstr
-                    PVoted
-                    ( #votedOn .= proposalInputDatumF.proposalId
-                        .& #votedFor .= pdata voteFor
-                    )
-
-                -- Prepend the new lock to existing locks
-                expectedProposalLocks =
-                  pcons
-                    # pdata newProposalLock
-                    # pfromData stakeInF.lockedBy
-
-            pguardC "Output stake should be locked by the proposal" $
-              pfield @"lockedBy" # sctxF.outputStake #== expectedProposalLocks
 
           ----------------------------------------------------------------------
 
-          PUnlock _ -> spendSingleStake $ \sctxF -> do
+          PUnlock _ -> spendStakes $ \sctxF -> do
-            stakeInF <- pletAllC $ pto sctxF.inputStake
+            let expectedVotes =
+                  pfoldl
+                    # plam
+                      ( \votes stake -> unTermCont $ do
+                          stakeF <-
+                            pletFieldsC
+                              @'["stakedAmount", "lockedBy"]
+                              stake
 
-            stakeRole <- pletC $ pgetStakeRole # proposalInputDatumF.proposalId # stakeInF.lockedBy
+                          stakeRole <-
+                            pletC $
+                              pgetStakeRole
+                                # proposalInputDatumF.proposalId
+                                # stakeF.lockedBy
 
-            pguardC "Stake input should be relevant" $
+                          pguardC "Stake input should be relevant" $
-              pnot #$ pisIrrelevant # stakeRole
+                            pnot #$ pisIrrelevant # stakeRole
 
-            retractCount <-
+                          let canRetractVotes =
-              pletC $
+                                pnot #$ pisPureCreator # stakeRole
-                pmatch stakeInF.stakedAmount $ \(PDiscrete v) -> pextract # v
 
-            -- The votes can only change when the proposal still allows voting.
+                              voteCount =
-            let shouldUpdateVotes =
+                                pextract
+                                  #$ pto
+                                  $ pfromData stakeF.stakedAmount
+
+                              newVotes =
+                                pretractVotes
+                                  # (pextractVoteOption # stakeRole)
+                                  # voteCount
+                                  # votes
+
+                          pure $ pif canRetractVotes newVotes votes
+                      )
+                    # proposalInputDatumF.votes
+                    # sctxF.inputStakes
+
+                currentTime' =
+                  passertPJust
+                    # "Should be able to get current time"
+                    # currentTime
+
+                inVotingPeriod =
+                  isVotingPeriod # proposalInputDatumF.timingConfig
+                    # proposalInputDatumF.startingTime
+                    # currentTime'
+
+                -- The votes can only change when the proposal still allows voting.
+                shouldUpdateVotes =
                   currentStatus #== pconstant VotingReady
-                    #&& pnot # (pisPureCreator # stakeRole)
+                    #&& inVotingPeriod
-
-                allowRemovingCreatorLock =
-                  currentStatus #== pconstant Finished
-
-                isCreator = pisCreator # stakeRole
-
-                -- If the stake has been used for creating the proposal,
-                --  the creator lock can only be removed when the proposal
-                --  is finished.
-                --
-                -- In other cases, all the locks related to this
-                --   proposal should be removed.
-                validateOutputLocks = plam $ \locks ->
-                  plet
-                    ( pgetStakeRole # proposalInputDatumF.proposalId # locks
-                    )
-                    $ \newStakeRole ->
-                      pif
-                        (isCreator #&& pnot # allowRemovingCreatorLock)
-                        (pisPureCreator # newStakeRole)
-                        (pisIrrelevant # newStakeRole)
 
             pguardC "Proposal output correct" $
               pif
                 shouldUpdateVotes
                 ( let -- Remove votes and leave other parts of the proposal as it.
-                      expectedVotes =
-                        pretractVotes
-                          # (pextractVoteOption # stakeRole)
-                          # retractCount
-                          # proposalInputDatumF.votes
-
                       expectedProposalOut =
                         mkRecordConstr
                           PProposalDatum
@@ -618,12 +610,6 @@ proposalValidator as maximumCosigners =
                     proposalOutputDatum #== proposalInputDatum
                 )
 
-            -- At last, we ensure that all locks belong to this proposal will be removed.
-            stakeOutputLocks <- pletC $ pfield @"lockedBy" # pto sctxF.outputStake
-
-            pguardC "All relevant locks removed from the stake" $
-              validateOutputLocks # stakeOutputLocks
-
           ----------------------------------------------------------------------
 
           PAdvanceProposal _ -> unTermCont $ do

agora/Agora/Stake.hs

@@ -533,8 +533,8 @@ instance DerivePlutusType PProposalContext where
      @1.0.0
 -}
 data PStakeRedeemerHandlerContext (s :: S) = PStakeRedeemerHandlerContext
-  { stakeInputDatums :: Term s (PBuiltinList PStakeDatum)
+  { stakeInputDatums :: Term s (PList PStakeDatum)
-  , stakeOutputDatums :: Term s (PBuiltinList PStakeDatum)
+  , stakeOutputDatums :: Term s (PList PStakeDatum)
   , redeemerContext :: Term s PStakeRedeemerContext
   , sigContext :: Term s PSigContext
   , proposalContext :: Term s PProposalContext

agora/Agora/Stake/Scripts.hs

@@ -61,6 +61,7 @@ import Agora.Stake.Redeemers (
   ppermitVote,
   pretractVote,
  )
+import Agora.Utils (pmapMaybe)
 import Data.Tagged (Tagged (Tagged))
 import Plutarch.Api.V1 (
   KeyGuarantees (Sorted),
@@ -73,6 +74,7 @@ import Plutarch.Api.V2 (
   AmountGuarantees,
   PMintingPolicy,
   PScriptPurpose (PMinting, PSpending),
+  PTxInInfo,
   PTxInfo,
   PTxOut,
   PValidator,
@@ -85,7 +87,7 @@ import Plutarch.Extra.AssetClass (
 import Plutarch.Extra.Category (PSemigroupoid ((#>>>)))
 import Plutarch.Extra.Field (pletAll)
 import Plutarch.Extra.Functor (PFunctor (pfmap))
-import "liqwid-plutarch-extra" Plutarch.Extra.List (pfindJust, pmapMaybe)
+import "liqwid-plutarch-extra" Plutarch.Extra.List (pfindJust)
 import Plutarch.Extra.Maybe (
   passertPJust,
   pfromMaybe,
@@ -457,6 +459,7 @@ mkStakeValidator
                           )
                         # txInfoF.redeemers
 
+                getContext :: Term _ (PTxInInfo :--> PMaybe PProposalContext)
                 getContext = plam $
                   flip pletAll $ \inInfoF ->
                     pfmap
@@ -469,7 +472,18 @@ mkStakeValidator
                         )
                       #$ getProposalDatum
                       # pfromData inInfoF.resolved
-             in pfindJust # getContext # pfromData txInfoF.inputs
+
+                contexts =
+                  pmapMaybe @PList # getContext # pfromData txInfoF.inputs
+             in precList
+                  ( \_ h t ->
+                      pif
+                        (pnull # t)
+                        (pjust # h)
+                        (ptraceError "Ambiguous proposal")
+                  )
+                  (const pnothing)
+                  # contexts
 
           noProposalContext = pcon PNoProposal
 

agora/Agora/Utils.hs

@@ -26,10 +26,12 @@ module Agora.Utils (
   pdeleteBy,
   pisSingleton,
   pfromSingleton,
+  pmapMaybe,
 ) where
 
 import Plutarch.Api.V1 (PPOSIXTime, PTokenName, PValidatorHash)
 import Plutarch.Api.V2 (PScriptHash)
+import Plutarch.Extra.Category (PCategory (pidentity))
 import Plutarch.Extra.Time (PCurrentTime (PCurrentTime))
 import Plutarch.Unsafe (punsafeCoerce)
 import PlutusLedgerApi.V2 (
@@ -250,3 +252,27 @@ pfromSingleton =
             (ptraceError "More than one element")
       )
       (const $ ptraceError "Empty list")
+
+-- | @since 1.0.0
+pmapMaybe ::
+  forall
+    (listO :: PType -> PType)
+    (b :: PType)
+    (listI :: PType -> PType)
+    (a :: PType)
+    (s :: S).
+  (PIsListLike listI a, PIsListLike listO b) =>
+  Term s ((a :--> PMaybe b) :--> listI a :--> listO b)
+pmapMaybe = phoistAcyclic $
+  plam $ \f ->
+    precList
+      ( \self h t ->
+          pmatch
+            (f # h)
+            ( \case
+                PJust x -> pcons # x
+                PNothing -> pidentity
+            )
+            # (self # t)
+      )
+      (const pnil)