make it possible for delegatee to vote with delegated and own stakes

agora-specs/Sample/Proposal/Vote.hs

@@ -18,6 +18,7 @@ module Sample.Proposal.Vote (
   mkTestTree,
   mkValidOwnerVoteBundle,
   mkValidDelegateeVoteBundle,
+  delegateeVoteWithOwnAndDelegatedStakeBundle,
   transparentAssets,
   transactionNotAuthorized,
   voteForNonexistentOutcome,
@@ -99,6 +100,7 @@ newtype VoteParameters = VoteParameters {voteFor :: ResultTag}
 
 data StakeParameters = StakeParameters
   { numStakes :: Integer
+  , mixInDelegateeAsOwner :: Bool
   , stakeInputParameters :: StakeInputParameters
   , stakeOutputParameters :: StakeOutputParameters
   }
@@ -257,6 +259,16 @@ vote params =
       stakeRedeemer =
         mkStakeRedeemer params.stakeParameters.stakeOutputParameters
 
+      mixOwner i datum =
+        if params.stakeParameters.mixInDelegateeAsOwner
+          && i == 2
+          then
+            datum
+              { owner = PubKeyCredential delegatee
+              , delegatedTo = Nothing
+              }
+          else datum
+
       stakeBuilder :: b
       stakeBuilder =
         foldMap
@@ -266,7 +278,7 @@ vote params =
                     mconcat
                       [ script stakeValidatorHash
                       , withValue stakeInputValue
-                      , withInlineDatum stakeInputDatum
+                      , withInlineDatum $ mixOwner i stakeInputDatum
                       , withRedeemer stakeRedeemer
                       , withRef $ mkStakeRef numProposals' i
                       ]
@@ -277,7 +289,7 @@ vote params =
                         mconcat
                           [ script stakeValidatorHash
                           , withValue stakeOutputValue
-                          , withInlineDatum stakeOutputDatum
+                          , withInlineDatum $ mixOwner i stakeOutputDatum
                           ]
                 ]
           )
@@ -420,6 +432,7 @@ mkValidOwnerVoteBundle stakes =
     , stakeParameters =
         StakeParameters
           { numStakes = stakes
+          , mixInDelegateeAsOwner = False
           , stakeInputParameters =
               StakeInputParameters
                 { perStakeGTs = (def :: ProposalThresholds).vote
@@ -453,6 +466,16 @@ mkValidDelegateeVoteBundle stakes =
               }
         }
 
+delegateeVoteWithOwnAndDelegatedStakeBundle :: ParameterBundle
+delegateeVoteWithOwnAndDelegatedStakeBundle =
+  let template = mkValidDelegateeVoteBundle 5
+   in template
+        { stakeParameters =
+            template.stakeParameters
+              { mixInDelegateeAsOwner = True
+              }
+        }
+
 ownerVoteWithSignleStake :: ParameterBundle
 ownerVoteWithSignleStake = mkValidOwnerVoteBundle 1
 

agora-specs/Spec/Proposal.hs

@@ -160,6 +160,10 @@ specs =
                   "transparent non-GT tokens"
                   Vote.transparentAssets
                   (Vote.Validity True True)
+              , Vote.mkTestTree
+                  "Delegatee vote with own and delegated stakes in one tx"
+                  Vote.delegateeVoteWithOwnAndDelegatedStakeBundle
+                  (Vote.Validity True True)
               ]
           , group
               "illegal"

agora/Agora/Stake/Scripts.hs

@@ -79,6 +79,7 @@ import Plutarch.Extra.Functor (PFunctor (pfmap))
 import "liqwid-plutarch-extra" Plutarch.Extra.List (pfindJust, pmapMaybe)
 import Plutarch.Extra.Maybe (
   passertPJust,
+  pdjust,
   pfromJust,
   pfromMaybe,
   pjust,
@@ -341,7 +342,7 @@ mkStakeValidator impl sstSymbol pstClass gtClass =
 
     authorizedBy <- pletC $ pauthorizedBy # authorizationContext txInfoF
 
-    PPair allHaveSameOwner allHaveSameDelegatee <-
+    PPair allHaveSameOwner allHaveSameOrOwnedByDelegatee <-
       pmatchC $
         pfoldr
           # plam
@@ -354,11 +355,15 @@ mkStakeValidator impl sstSymbol pstClass gtClass =
                           allHaveSameOwner
                             #&& dF.owner
                             #== firstStakeInputDatumF.owner
-                        allHaveSameDelegatee' =
-                          allHaveSameDelegatee
-                            #&& dF.delegatedTo
-                            #== firstStakeInputDatumF.delegatedTo
-                     in pcon $ PPair allHaveSameOwner' allHaveSameDelegatee'
+                        allHaveSameOrOwnedByDelegatee' =
+                          let delegated =
+                                dF.delegatedTo #== firstStakeInputDatumF.delegatedTo
+                              ownedByDelagtee =
+                                pdata (pdjust # dF.owner)
+                                  #== firstStakeInputDatumF.delegatedTo
+                           in allHaveSameDelegatee
+                                #&& (delegated #|| ownedByDelagtee)
+                     in pcon $ PPair allHaveSameOwner' allHaveSameOrOwnedByDelegatee'
             )
           # pcon (PPair (pconstant True) (pconstant True))
           # restOfStakeInputDatums
@@ -369,7 +374,7 @@ mkStakeValidator impl sstSymbol pstClass gtClass =
             # firstStakeInputDatumF.owner
 
         delegateSignsTransaction =
-          allHaveSameDelegatee
+          allHaveSameOrOwnedByDelegatee
             #&& pmaybeData
             # pconstant False
             # plam ((authorizedBy #) . pfromData)