Plutus Playground Smart Contract

Playground.hs

import qualified Data.Text                 as T
import           Plutus.Contract  hiding (when)
-- ScriptLookups semigroup is defined based on the standard prelude <> and doesn't like the plutus one
import           PlutusTx.Prelude hiding  ((<>)) 
import qualified PlutusTx                 as PlutusTx
import           Playground.Contract
import           Control.Monad                     (void, when)
import           Ledger                            (Address (..), Slot (Slot), Validator, pubKeyHash, txOutTxDatum, txOutValue, txOutTxOut, TxOut, ValidatorCtx, Value, scriptAddress, PubKeyHash, Datum(..), TxOutTx, PubKeyHash (..), ValidatorCtx (..), validatorHash, txInInfoValue, outValue, TxInfo (..), txInInfoWitness, DatumHash, TxInInfo(..) ) 
import qualified Ledger.Value                      as Value 
import qualified Ledger.Typed.Scripts              as Scripts
import qualified Data.Map                          as Map
import           Data.List                         (groupBy, maximumBy, partition) 
import qualified Data.Set                          as Set 
import qualified Ledger.Ada                        as Ada
import           Data.Maybe                        (catMaybes)
import           Ledger.Constraints                (TxConstraints, mustBeSignedBy, mustPayToTheScript, mustValidateIn, mustPayToPubKey, mustPayToOtherScript, scriptInstanceLookups,SomeLookupsAndConstraints (..), mkSomeTx, unspentOutputs, ScriptLookups(..) )
import qualified Ledger.Contexts                   as Validation
import           Wallet.Emulator.Types             (Wallet, walletPubKey)
import qualified Prelude
import qualified Data.Foldable  as Foldable 
import           Control.Lens

import           Data.Semigroup                   


quorum = 2
oneVote = Ada.lovelaceValueOf 1

-- UTILITY FUNCTIONS

-- onchain

-- the standard Maybe.fromJust and List.sortOn don't work onchain
{-# INLINABLE fromJust #-}
fromJust :: Maybe a -> a
fromJust (Just a) = a

-- Classic Haskell fake quicksort (not a real quicksort, but is good enough)
{-# INLINABLE sort #-}
sort :: (Ord a) => [a] -> [a]
sort [] = []
sort (x:xs) =
  let smallerSorted = sort [a | a <- xs, a <= x]
      biggerSorted = sort [a | a <- xs, a > x]
  in  smallerSorted ++ [x] ++ biggerSorted

{-# INLINABLE datumToData #-}
datumToData :: (PlutusTx.IsData a) => Datum -> Maybe a
datumToData datum = PlutusTx.fromData (getDatum datum)

{-# INLINABLE findExtractData #-}
findExtractData :: DatumHash -> TxInfo -> VoteDatum
findExtractData dh txInfo = fromJust(datumToData (fromJust (Validation.findDatum dh txInfo)))

{-# INLINABLE extractData #-}
extractData :: (PlutusTx.IsData a) => TxOutTx -> Maybe a
extractData txOut = do
                        datum <- txOutTxDatum txOut
                        datumToData datum   

{-# INLINABLE extractVote #-}
extractVote :: TxOutTx -> VoteDatum
extractVote txo = fromJust (extractData txo)

{-# INLINABLE validatorHashOf #-}
validatorHashOf :: TxInInfo -> ValidatorHash
validatorHashOf TxInInfo{txInInfoWitness=Just (vHash, _, _) } = vHash 

{-# INLINABLE datumHashOf #-}
datumHashOf :: TxInInfo -> DatumHash
datumHashOf TxInInfo{txInInfoWitness=Just ( _, _, dHash) } = dHash 


-- offchain
extractWallet :: TxOutTx -> PubKeyHash 
extractWallet tx = votedWallet (extractVote tx)

extractPayout :: TxOutTx -> Integer 
extractPayout tx = payout (extractVote tx)

extractOwner :: TxOutTx -> PubKeyHash 
extractOwner tx = owner (extractVote tx)

findMostVotedGroup :: (a -> a -> Bool) -> [a] -> ([a], Integer)
findMostVotedGroup grouper elements = 
                let 
                    tally = map (\vs -> (vs, length vs)) $ groupBy grouper elements
                    compareTally (_, count1) (_, count2) = count1 `compare` count2
                in
                   maximumBy compareTally tally

pubKeyHashOf :: Wallet -> PubKeyHash
pubKeyHashOf = pubKeyHash . walletPubKey


-- ONCHAIN VALIDATORS

-- Vote script
data VoteDatum = VoteDatum {   
      votedWallet     :: PubKeyHash,
      payout          :: Integer, 
      owner           :: PubKeyHash
    } deriving (Generic, Show)

PlutusTx.makeLift ''VoteDatum
PlutusTx.makeIsDataIndexed ''VoteDatum [('VoteDatum,0)]


data Vote
instance Scripts.ScriptType Vote where
    type instance RedeemerType Vote = ()
    type instance DatumType Vote = VoteDatum


{-# INLINABLE voteScript #-}
voteScript  ::  ValidatorHash -> VoteDatum -> () -> ValidatorCtx -> Bool 
voteScript treasury VoteDatum{owner=voteOwner} _ ctx@ValidatorCtx{valCtxTxInfo=txInfo@TxInfo{txInfoInputs}} =                     
                        let 
                            -- Spending path 1: The vote must be spent with the treasury in the same transaction
                            -- (check that the treasury hash is present in the inputs)
                            collectVotesAction = traceIfFalse " **** Must be spent together with the treasury" $ any (\txInInfo -> validatorHashOf txInInfo == treasury) txInfoInputs 
                            -- or
                            -- Spending path 2: The vote must be returned to the owner
                            returnVoteAction =  traceIfFalse " **** Vote must be paid back to owner" $ Validation.valuePaidTo txInfo voteOwner == oneVote
                        in
                            collectVotesAction || returnVoteAction 

                        

voteScriptInstance :: ValidatorHash -> Scripts.ScriptInstance Vote
voteScriptInstance treasuryHash = Scripts.validator @Vote
    ($$(PlutusTx.compile [|| voteScript ||]) `PlutusTx.applyCode` PlutusTx.liftCode treasuryHash)
    $$(PlutusTx.compile [|| wrap ||]) where
        wrap = Scripts.wrapValidator @VoteDatum @()

voteScriptAddress :: ValidatorHash -> Address
voteScriptAddress treasuryHash = Ledger.scriptAddress (Scripts.validatorScript ( voteScriptInstance treasuryHash))

-- Treasury script
data Treasury
instance Scripts.ScriptType Treasury where
    type instance RedeemerType Treasury = ()
    type instance DatumType Treasury = ()


{-# INLINABLE treasuryScript #-}
treasuryScript  ::  () -> () -> ValidatorCtx -> Bool 
treasuryScript _ _ ctx@ValidatorCtx{valCtxTxInfo=txInfo@TxInfo{txInfoInputs}} =                     
                    let                         

                        -- split the inputs by partition looking for the input that has the hash of the current script (treasury)
                        -- the result should be a list of votes (all the for the same address and amount) and one input from the treasury                
                        ([treasury], allVotes@(aVote:_) ) = partition (\txInInfo -> validatorHashOf txInInfo == Validation.ownHash ctx) txInfoInputs 

                        -- extract the datum from the first vote 
                        aVoteDatum = fromJust (Validation.findDatum (datumHashOf aVote) txInfo)

                        -- validate all the votes have the same voted wallet and same payout
                        compareDatums d1 d2 = 
                                let 
                                    data1 = findExtractData d1 txInfo
                                    data2 = findExtractData d2 txInfo
                                in
                                    votedWallet data1 == votedWallet data2 && payout data1 == payout data2

                        allVotesAreTheSame =  traceIfFalse " **** Not all votes are the same" $ all (True == ) $ 
                            fmap (\voteTxInfo -> compareDatums (datumHashOf (aVote)) (datumHashOf voteTxInfo)  ) allVotes         

                        -- validate the voted wallet is paid with the voted amount
                        votedAddress = votedWallet (fromJust(datumToData aVoteDatum))         
                        votedAmount  = payout (fromJust(datumToData aVoteDatum))   
                        
                        paidToVoted  =  Validation.valuePaidTo txInfo votedAddress 
                        ensureVotedIsPaid = traceIfFalse " **** Voted wallet is not paid the amount" $ paidToVoted == Ada.lovelaceValueOf votedAmount

                        -- validate the value is spent from the treasury
                        remainingInTreasury = Validation.valueLockedBy txInfo (Validation.ownHash ctx)
                        valueInTreasury = txInInfoValue treasury
                        validTreasurySpend = traceIfFalse " **** Treasury spend" $ valueInTreasury == (paidToVoted + remainingInTreasury)

                        -- validate the votes are preserved (input votes == output votes)    
                        outVotes = Validation.scriptOutputsAt ( validatorHashOf aVote) txInfo 
                        -- compare values
                        sameVotesValues = (fmap txInInfoValue allVotes) == (fmap snd outVotes)
                        -- compare datums
                        votesPreserved = traceIfFalse " **** votes preserved" $ sort ( fmap datumHashOf allVotes)  == sort  (fmap fst outVotes)

                        -- validate enough votes    
                        quorumCheck = traceIfFalse " **** Not enough votes" (length allVotes >= quorum )

                    in quorumCheck && allVotesAreTheSame && ensureVotedIsPaid && validTreasurySpend && votesPreserved
                        
                        

treasuryScriptInstance :: Scripts.ScriptInstance Treasury
treasuryScriptInstance  = Scripts.validator @Treasury
    $$(PlutusTx.compile [|| treasuryScript ||])
    $$(PlutusTx.compile [|| wrap ||]) where
        wrap = Scripts.wrapValidator @() @()


treasuryScriptHash :: ValidatorHash
treasuryScriptHash = validatorHash $ Scripts.validatorScript treasuryScriptInstance

treasuryScriptAddress :: Address
treasuryScriptAddress = Ledger.scriptAddress (Scripts.validatorScript treasuryScriptInstance)


-- OFFCHAIN ENDPOINTS

type VotingSchema =
    BlockchainActions
        .\/ Endpoint "1-setup treasury" Integer
        .\/ Endpoint "2-vote" VoteParams
        .\/ Endpoint "3-collect" ()
        .\/ Endpoint "4-return vote" Wallet

-- Initialize the treasury with some value
setupTreasury :: Contract () VotingSchema T.Text ()
setupTreasury = do
                    trasuryAmount <- endpoint @"1-setup treasury" @Integer
                    let            
                        tx = mustPayToTheScript () ( Ada.lovelaceValueOf trasuryAmount) 
                    void (submitTxConstraints treasuryScriptInstance tx)


-- | Parameters for the "vote" endpoint
data VoteParams = VoteParams
    { votedFor :: Wallet
    , amount :: Integer
    }
    deriving stock (Prelude.Eq, Prelude.Show, Generic)
    deriving anyclass (FromJSON, ToJSON, ToSchema, ToArgument)

-- Vote 
vote :: Contract () VotingSchema T.Text ()
vote= do
        VoteParams votedFor amount <- endpoint @"2-vote" @VoteParams
        voter <- pubKeyHash <$> ownPubKey
        let
          votedforPKH = pubKeyHashOf votedFor
          txAddVote = mustPayToTheScript VoteDatum{votedWallet=votedforPKH, payout=amount, owner=voter} oneVote
        void (submitTxConstraints (voteScriptInstance treasuryScriptHash) txAddVote)


-- Tally votes endpoint
-- collect enough votes and spend the voted amount from the treasury to the winning choice
tally :: Contract () VotingSchema T.Text ()
tally = do        
            endpoint @"3-collect" @()
            votesUtxo <- utxoAt (voteScriptAddress treasuryScriptHash)
            treasuryUtxo <- utxoAt treasuryScriptAddress              
            collector <- pubKeyHash <$> ownPubKey
            let
                utxoList = Map.toList votesUtxo
                comparator (_,x) (_,y) = (extractWallet x == extractWallet y) && (extractPayout x == extractPayout y)
                (winningVotes, count) = findMostVotedGroup comparator utxoList
                winningUtxos = Map.fromList winningVotes 
            if count >= quorum then 
                let 
                    ScriptAddress voteScriptHash = voteScriptAddress treasuryScriptHash

                    -- collect from the winning utxos and the treasury
                    txVotesUtxos = collectFromScript winningUtxos ()
                    txInputTreasury = collectFromScript treasuryUtxo () 

                    

                    votedPayout = extractPayout $ snd.head $ winningVotes
                    winningWallet = extractWallet (snd $ head winningVotes) 
                    datum = Datum $ PlutusTx.toData $ VoteDatum{votedWallet=winningWallet, payout=votedPayout,owner=collector}
                                
                    -- pay the voted amount from the treasury (and keep the remainder in the treasury)           
                    totalTreasury = sum $ map (Ada.getLovelace. Ada.fromValue . txOutValue . txOutTxOut . snd) $  Map.toList treasuryUtxo                                       
                    txPayWinner = mustPayToPubKey winningWallet ( Ada.lovelaceValueOf votedPayout)
                    txRepayTreasury = mustPayToOtherScript treasuryScriptHash datum ( Ada.lovelaceValueOf ( totalTreasury - votedPayout ))

                    -- rebuild spent votes
                    rebuildVote utxo =  mustPayToOtherScript  voteScriptHash (fromJust (txOutTxDatum utxo)) oneVote 

                    rebuildVoteTxs = map (rebuildVote.snd) winningVotes
                    txRebuildVotes = Prelude.foldl1 (<>) rebuildVoteTxs
                              
                    -- treasury script constraints          
                    treasuryUtxosConstraint = txInputTreasury <> txPayWinner  <> txRepayTreasury
                    treasuryLookups = (scriptInstanceLookups treasuryScriptInstance) <> (unspentOutputs treasuryUtxo)

                    -- vote script constraints
                    votesUtxosConstraint = txVotesUtxos <> txRebuildVotes
                    votesLookups = (scriptInstanceLookups (voteScriptInstance treasuryScriptHash) ) <> (unspentOutputs votesUtxo)

                    -- Since the treasury and the vote constraints and spends have to be on the same transaction, 
                    -- we can't use the standard submitTxConstraintsXXXX and we need to manually create the transaction  
                    treasurySpend = SomeLookupsAndConstraints treasuryLookups treasuryUtxosConstraint
                    voteSpend = SomeLookupsAndConstraints votesLookups votesUtxosConstraint
                    tx = mkSomeTx [treasurySpend, voteSpend]
                in                     
                  do
                    logInfo @String "Submit TX"
                    void $ do
                        tx <- either (throwError . review _ConstraintResolutionError) pure tx
                        submitUnbalancedTx tx
            else
                do 
                    logInfo @String "NOT ENOUGH VOTES"
                    throwError $ T.pack "Not enough votes"

returnVote :: Contract () VotingSchema T.Text ()
returnVote = do
                    -- Should be 
                    -- voter <- pubKeyHash <$> ownPubKey
                    -- but the simulator doesn't allow multiple actions on wallets                    
                    wallet <- endpoint @"4-return vote" @Wallet
                    voteUtxos <- utxoAt (voteScriptAddress treasuryScriptHash)                     
                    let            
                        voteScript = voteScriptInstance treasuryScriptHash
                        voter = pubKeyHashOf wallet                        
                        votesToReturnUtxos = filter (\txOut -> voter == (extractOwner (snd txOut))) $  Map.toList voteUtxos                        
                        txPayToVoter = Foldable.fold $  map  ( \utxo -> mustPayToPubKey voter (  txOutValue $ txOutTxOut $ snd utxo ) ) votesToReturnUtxos
                        txVotesUtxos = collectFromScript ( Map.fromList votesToReturnUtxos) ()
                        tx = txPayToVoter <> txVotesUtxos
                    logInfo $ votesToReturnUtxos
                    logInfo $ show voter
                    void (submitTxConstraintsSpending voteScript (Map.fromList votesToReturnUtxos) tx)



endpoints :: Contract () VotingSchema T.Text ()
endpoints = setupTreasury `select` vote `select` tally `select` returnVote

mkSchemaDefinitions ''VotingSchema

$(mkKnownCurrencies [])

Simulation.json

[0,[]]