agora/agora-specs/Sample/Proposal/UnlockStake.hs

module Sample.Proposal.UnlockStake (
  unlockStake,
  StakeRole (..),
  UnlockStakeParameters (..),
  votesTemplate,
  emptyEffectFor,
  mkProposalInputDatum,
  mkStakeInputDatum,
  mkProposalValidatorTestCase,
) where

--------------------------------------------------------------------------------

import Agora.Proposal (
  ProposalDatum (..),
  ProposalId (..),
  ProposalRedeemer (Unlock),
  ProposalStatus (..),
  ProposalVotes (..),
  ResultTag (..),
 )
import Agora.Proposal.Scripts (proposalValidator)
import Agora.Proposal.Time (ProposalStartingTime (ProposalStartingTime))
import Agora.Stake (ProposalLock (ProposalLock), Stake (..), StakeDatum (..))
import Control.Monad (join)
import Data.Coerce (coerce)
import Data.Default.Class (Default (def))
import Data.Tagged (Tagged (..), untag)
import Plutarch.Context (
  BaseBuilder,
  buildTxInfoUnsafe,
  input,
  output,
  script,
  txId,
  withDatum,
  withRefIndex,
  withTxId,
  withValue,
 )
import PlutusLedgerApi.V1 (
  DatumHash,
  ScriptContext (..),
  ScriptPurpose (Spending),
  TxInfo (..),
  TxOutRef (..),
  ValidatorHash,
 )
import PlutusLedgerApi.V1.Value qualified as Value
import PlutusTx.AssocMap qualified as AssocMap
import Sample.Proposal.Shared (proposalTxRef, stakeTxRef, testFunc)
import Sample.Shared (
  minAda,
  proposalPolicySymbol,
  proposalValidatorHash,
  signer,
  stake,
  stakeAssetClass,
  stakeValidatorHash,
 )
import Sample.Shared qualified as Shared
import Test.Specification (SpecificationTree)
import Test.Util (sortValue, updateMap)

--------------------------------------------------------------------------------

-- | The template "shape" that votes of proposals generated by 'mkProposalDatumPair' have.
votesTemplate :: ProposalVotes
votesTemplate =
  ProposalVotes $
    AssocMap.fromList
      [ (ResultTag 0, 0)
      , (ResultTag 1, 0)
      ]

-- | Create empty effects for every result tag given the votes.
emptyEffectFor ::
  ProposalVotes ->
  AssocMap.Map ResultTag (AssocMap.Map ValidatorHash DatumHash)
emptyEffectFor (ProposalVotes vs) =
  AssocMap.fromList $
    map (,AssocMap.empty) (AssocMap.keys vs)

-- | The default vote option that will be used by functions in this module.
defaultVoteFor :: ResultTag
defaultVoteFor = ResultTag 0

-- | The default number of GTs the stake will have.
defaultStakedGTs :: Tagged _ Integer
defaultStakedGTs = Tagged 100000

-- | How a stake has been used on a particular proposal.
data StakeRole
  = -- | The stake was spent to vote for a paraticular option.
    Voter
  | -- | The stake was used to created the proposal.
    Creator
  | -- | The stake has nothing to do with the proposal.
    Irrelevant

-- | Parameters for creating a 'TxOut' that unlocks a stake.
data UnlockStakeParameters = UnlockStakeParameters
  { proposalCount :: Integer
  -- ^ The number of proposals in the 'TxOut'.
  , stakeUsage :: StakeRole
  -- ^ The role of the stake we're unlocking.
  , retractVotes :: Bool
  -- ^ Whether to retract votes or not.
  , proposalStatus :: ProposalStatus
  -- ^ The state of all the proposals.
  }

instance Show UnlockStakeParameters where
  show p =
    let role = case p.stakeUsage of
          Voter -> "voter"
          Creator -> "creator"
          _ -> "irrelevant stake"

        action =
          if p.retractVotes
            then "unlock stake + retract votes"
            else "unlock stake"

        while = show p.proposalStatus

        proposalInfo = mconcat [show p.proposalCount, " proposals"]
     in mconcat [proposalInfo, ", ", role, ", ", action, ", ", while]

-- | Generate some input proposals and their corresponding output proposals.
mkProposals :: UnlockStakeParameters -> [(ProposalDatum, ProposalDatum)]
mkProposals p = forEachProposalId p.proposalCount $ mkProposalDatumPair p

-- | Iterate over the proposal id of every proposal, given the number of proposals.
forEachProposalId :: Integer -> (ProposalId -> a) -> [a]
forEachProposalId 0 _ = error "zero proposal"
forEachProposalId n f = f . ProposalId <$> [0 .. n - 1]

-- | Create the input stake and its corresponding output stake.
mkStakeDatumPair :: UnlockStakeParameters -> (StakeDatum, StakeDatum)
mkStakeDatumPair c =
  let output =
        StakeDatum
          { stakedAmount = defaultStakedGTs
          , owner = signer
          , lockedBy = []
          }

      inputLocks = join $ forEachProposalId c.proposalCount (mkStakeLocks c.stakeUsage)

      input = output {lockedBy = inputLocks}
   in (input, output)
  where
    mkStakeLocks :: StakeRole -> ProposalId -> [ProposalLock]
    mkStakeLocks Voter pid = [ProposalLock defaultVoteFor pid]
    mkStakeLocks Creator pid =
      map (`ProposalLock` pid) $
        AssocMap.keys $ getProposalVotes votesTemplate
    mkStakeLocks _ _ = []

-- | Create the input proposal datum.
mkProposalInputDatum :: UnlockStakeParameters -> ProposalId -> ProposalDatum
mkProposalInputDatum p pid = fst $ mkProposalDatumPair p pid

-- | Create the input stake datum.
mkStakeInputDatum :: UnlockStakeParameters -> StakeDatum
mkStakeInputDatum = fst . mkStakeDatumPair

-- | Create a input proposal and its corresponding output proposal.
mkProposalDatumPair ::
  UnlockStakeParameters ->
  ProposalId ->
  (ProposalDatum, ProposalDatum)
mkProposalDatumPair params pid =
  let owner = signer

      inputVotes = mkInputVotes params.stakeUsage $ untag defaultStakedGTs

      input =
        ProposalDatum
          { proposalId = pid
          , effects = emptyEffectFor votesTemplate
          , status = params.proposalStatus
          , cosigners = [owner]
          , thresholds = def
          , votes = inputVotes
          , timingConfig = def
          , startingTime = ProposalStartingTime 0
          }

      output =
        if params.retractVotes
          then input {votes = votesTemplate}
          else input
   in (input, output)
  where
    -- Assemble the votes of the input proposal based on 'votesTemplate'.
    mkInputVotes ::
      StakeRole ->
      -- The staked amount/votes.
      Integer ->
      ProposalVotes
    mkInputVotes Voter vc =
      ProposalVotes $
        updateMap (Just . const vc) defaultVoteFor $
          getProposalVotes votesTemplate
    mkInputVotes Creator _ =
      ProposalVotes $
        updateMap (Just . const 1000) defaultVoteFor $
          getProposalVotes votesTemplate
    mkInputVotes _ _ = votesTemplate

-- | Create a 'TxInfo' that tries to unlock a stake.
unlockStake :: UnlockStakeParameters -> TxInfo
unlockStake p =
  let pst = Value.singleton proposalPolicySymbol "" 1
      sst = Value.assetClassValue stakeAssetClass 1

      pIODatums = mkProposals p
      (sInDatum, sOutDatum) = mkStakeDatumPair p

      proposals =
        foldMap
          ( \(i, o) ->
              mconcat
                @BaseBuilder
                [ input $
                    script proposalValidatorHash
                      . withValue pst
                      . withDatum i
                      . withTxId proposalTxRef
                      . withRefIndex (coerce i.proposalId + 2)
                , output $
                    script proposalValidatorHash
                      . withValue (sortValue $ pst <> minAda)
                      . withDatum o
                ]
          )
          pIODatums

      stakeValue =
        sortValue $
          mconcat
            [ Value.assetClassValue
                (untag stake.gtClassRef)
                (untag defaultStakedGTs)
            , sst
            , minAda
            ]

      stakes =
        mconcat @BaseBuilder
          [ input $
              script stakeValidatorHash
                . withValue stakeValue
                . withDatum sInDatum
                . withTxId stakeTxRef
                . withRefIndex 1
          , output $
              script stakeValidatorHash
                . withValue stakeValue
                . withDatum sOutDatum
          ]

      builder =
        mconcat @BaseBuilder
          [ txId "388bc0b897b3dadcd479da4c88291de4113a50b72ddbed001faf7fc03f11bc52"
          , proposals
          , stakes
          ]
   in buildTxInfoUnsafe builder

-- | Create a test case that tests the proposal validator's @'Unlock' _@ redeemer.
mkProposalValidatorTestCase :: UnlockStakeParameters -> Bool -> SpecificationTree
mkProposalValidatorTestCase p shouldSucceed =
  let datum = mkProposalInputDatum p $ ProposalId 0
      redeemer = Unlock (ResultTag 0)
      name = show p
      scriptContext =
        ScriptContext
          (unlockStake p)
          (Spending (TxOutRef proposalTxRef 2))
   in testFunc
        shouldSucceed
        name
        (proposalValidator Shared.proposal)
        datum
        redeemer
        scriptContext