Merge pull request #254 from tweag/fc/cleanup-wallet-module

`Wallet` module makeover

cooked-validators/src/Cooked/Wallet.hs

@@ -4,101 +4,117 @@
 {-# LANGUAGE TypeSynonymInstances #-}
 {-# OPTIONS_GHC -Wno-orphans #-}
 
+-- | This module defines convenient wrappers for mock chain wallets (around
+-- Plutus mock wallets) and initial distributions (that is the initial state
+-- associating a list of UTxOs with some initial values to each known wallet).
+-- It also exposes a convenient API to construct wallets and distributions,
+-- manipulate them, and fetch information (such as public/private keys and
+-- staking keys).
 module Cooked.Wallet where
 
-import qualified Cardano.Api as Api
-import qualified Cardano.Crypto.Wallet as CWCrypto
-import Control.Arrow
+import qualified Cardano.Api as Cardano
+import qualified Cardano.Crypto.Wallet as Cardano
 import Data.Default
 import Data.Function (on)
-import qualified Data.Map.Strict as M
-import Data.Maybe
+import Data.Map.Strict (Map)
+import qualified Data.Map.Strict as Map
 import qualified Ledger.Address as Pl
-import qualified Ledger.CardanoWallet as CW
+import qualified Ledger.CardanoWallet as Pl
+import qualified Ledger.Credential as Pl
 import qualified Ledger.Crypto as Pl
+import qualified Ledger.Scripts as Pl
 import qualified Ledger.Tx as Pl
-import qualified Ledger.Tx.CardanoAPI as Pl
+import qualified Ledger.Tx.CardanoAPI.Internal as Pl
 import qualified Plutus.Script.Utils.Ada as Pl
-import qualified Plutus.Script.Utils.Scripts as Pl
-import qualified Plutus.V1.Ledger.Api as Pl hiding (TxOut)
-import qualified Plutus.V1.Ledger.Value as Pl
-import qualified Plutus.V2.Ledger.Api as PV2
-import Unsafe.Coerce (unsafeCoerce)
+import qualified Plutus.Script.Utils.Value as Pl
+import qualified Plutus.V2.Ledger.Tx as Pl2
+import qualified PlutusTx as Pl
+import Unsafe.Coerce
 
 -- * MockChain Wallets
 
 -- $mockchainwallets
 --
--- Because the mock wallets from the plutus-apps changes somewhat often, we will
+-- Because the mock wallets from the plutus-apps change somewhat often, we
 -- provide our own wrapper on top of them to make sure that we can easily deal
 -- changes from Plutus.
 
-type Wallet = CW.MockWallet
+type Wallet = Pl.MockWallet
 
-type PrivateKey = CWCrypto.XPrv
+type PrivateKey = Cardano.XPrv
 
 instance Eq Wallet where
- (==) = (==) `on` CW.mwWalletId
+ (==) = (==) `on` Pl.mwWalletId
 
 instance Ord Wallet where
- compare = compare `on` CW.mwWalletId
+ compare = compare `on` Pl.mwWalletId
 
+-- | All the wallets corresponding to known Plutus mock wallets.
 knownWallets :: [Wallet]
-knownWallets = CW.knownMockWallets
+knownWallets = Pl.knownMockWallets
 
+-- | Wallet corresponding to a given wallet number (or wallet ID)
 wallet :: Int -> Wallet
 wallet j
  | j > 0 && j <= 10 = let i = j - 1 in knownWallets !! i
- | otherwise = CW.fromWalletNumber (CW.WalletNumber $ fromIntegral j)
+ | otherwise = Pl.fromWalletNumber (Pl.WalletNumber $ fromIntegral j)
 
+-- | Retrieves the id of the known wallet that corresponds to a public key
+-- hash, if any.
+--
+-- @walletPKHashToId (walletPKHash (wallet 3)) == Just 3@
 walletPKHashToId :: Pl.PubKeyHash -> Maybe Int
-walletPKHashToId = flip M.lookup walletPKHashToIdMap
+walletPKHashToId = flip Map.lookup walletPKHashToIdMap
  where
- walletPKHashToIdMap = M.fromList . flip zip [1 ..] . map walletPKHash $ knownWallets
+ walletPKHashToIdMap =
+ Map.fromList . flip zip [1 ..] . map walletPKHash $ knownWallets
 
+-- | Retrieves a wallet publik key (PK)
 walletPK :: Wallet -> Pl.PubKey
-walletPK = Pl.unPaymentPubKey . CW.paymentPubKey
+walletPK = Pl.unPaymentPubKey . Pl.paymentPubKey
 
+-- | Retrieves a wallet's public staking key (PK), if any
 walletStakingPK :: Wallet -> Maybe Pl.PubKey
 walletStakingPK = fmap Pl.toPublicKey . walletStakingSK
 
+-- | Retrieves a wallet's public key hash
 walletPKHash :: Wallet -> Pl.PubKeyHash
 walletPKHash = Pl.pubKeyHash . walletPK
 
+-- | Retrieves a wallet's public staking key hash, if any
 walletStakingPKHash :: Wallet -> Maybe Pl.PubKeyHash
 walletStakingPKHash = fmap Pl.pubKeyHash . walletStakingPK
 
+-- | Retrieves a wallet's address
 walletAddress :: Wallet -> Pl.Address
 walletAddress w =
  Pl.Address
  (Pl.PubKeyCredential $ walletPKHash w)
  (Pl.StakingHash . Pl.PubKeyCredential <$> walletStakingPKHash w)
 
-walletSK :: CW.MockWallet -> PrivateKey
-walletSK = Pl.unPaymentPrivateKey . CW.paymentPrivateKey
-
--- Massive hack to be able to open a MockPrivateKey; this is needed because
--- the constructor and accessor to MockPrivateKey are not exported. Hence,
--- we make an isomorphic datatype, unsafeCoerce to this datatype then extract
--- whatever we need from it.
-newtype HACK = HACK {please :: CWCrypto.XPrv}
+-- | Retrieves a wallet private key (secret key SK)
+walletSK :: Pl.MockWallet -> PrivateKey
+walletSK = Pl.unPaymentPrivateKey . Pl.paymentPrivateKey
 
--- | Don't use this; its a hack and will be deprecated once we have time
--- to make a PR into plutus exporting the things we need. If you use this anyway,
--- make sure that you only apply it to @MockPrivateKey@; the function is polymorphic
--- because @MockPrivateKey@ is not exported either; having a dedicated function makes
--- it easy to test that this works: check the @Cooked.WalletSpec@ test module.
-hackUnMockPrivateKey :: a -> CWCrypto.XPrv
-hackUnMockPrivateKey = please . unsafeCoerce
+-- FIXME Massive hack to be able to open a MockPrivateKey; this is needed
+-- because the constructor and accessor to MockPrivateKey are not exported.
+-- Hence, we make an isomorphic datatype, unsafeCoerce to this datatype then
+-- extract whatever we need from it.
+newtype HACK = HACK Cardano.XPrv
 
+-- | Retrieves a wallet's private staking key (secret key SK), if any
 walletStakingSK :: Wallet -> Maybe PrivateKey
-walletStakingSK = fmap hackUnMockPrivateKey . CW.mwStakeKey
-
-toPKHMap :: [Wallet] -> M.Map Pl.PubKeyHash Wallet
-toPKHMap ws = M.fromList [(walletPKHash w, w) | w <- ws]
-
--- * Signs a transaction
-
+walletStakingSK = fmap hackUnMockPrivateKey . Pl.mwStakeKey
+ where
+ -- Don't move this function to outside the where clause; its a hack and
+ -- will be deprecated once we have time to make a PR into plutus exporting
+ -- the things we need. If you use this anyway, make sure that you only
+ -- apply it to @MockPrivateKey@; the function is polymorphic because
+ -- @MockPrivateKey@ is not exported either
+ hackUnMockPrivateKey :: a -> Cardano.XPrv
+ hackUnMockPrivateKey x = let HACK y = unsafeCoerce x in y
+
+-- | Signs a transaction
 txAddSignature :: Wallet -> Pl.Tx -> Pl.Tx
 txAddSignature w = Pl.addSignature' (walletSK w)
 
@@ -110,11 +126,11 @@ txAddSignature w = Pl.addSignature' (walletSK w)
 -- the underlying plutus definitions to make it easer when we have
 -- to plug our own, if we ever have the need
 
--- | Describes the initial distribution of /UTxOs/ per wallet. This is important since
---  transaction validation must specify a /collateral/, hence, wallets must posses more
---  than one UTxO to begin with in order to execute a transaction and have some collateral
---  option. The @txCollateral@ is transfered to the node operator in case the transaction
--- fails to validate.
+-- | Describes the initial distribution of /UTxOs/ per wallet. This is
+-- important since transaction validation must specify a /collateral/, hence,
+-- wallets must posses more than one UTxO to begin with in order to execute a
+-- transaction and have some collateral option. The @txCollateral@ is
+-- transfered to the node operator in case the transaction fails to validate.
 --
 -- An initial distribution defined by:
 --
@@ -126,92 +142,90 @@ txAddSignature w = Pl.addSignature' (walletSK w)
 -- > , (wallet 3 , [Pl.lovelaceValueOf 10000000 <> permanentValue "XYZ" 10])
 -- > ]
 --
--- Specifies a starting state where @wallet 1@ contains two /UTxOs/, one with 42 Ada
--- and one with 2 Ada and one "TOK" token; @wallet 2@ contains a single /UTxO/ with 10 Ada and
--- @wallet 3@ has 10 Ada and a permanent value. Check #quickvalues for more on quick
--- and permanent values. (Remember: 1 Ada = 1000000 Lovelace)
+-- Specifies a starting state where @wallet 1@ contains two /UTxOs/, one with
+-- 42 Ada and one with 2 Ada and one "TOK" token; @wallet 2@ contains a single
+-- /UTxO/ with 10 Ada and @wallet 3@ has 10 Ada and a permanent value. Check
+-- #quickvalues for more on quick and permanent values. (Remember: 1 Ada =
+-- 1000000 Lovelace)
 --
--- Check the corresponding @Default InitialDistribution@ instance for the default value.
-newtype InitialDistribution = InitialDistribution {distribution :: M.Map Wallet [Pl.Value]}
+-- Check the corresponding @Default InitialDistribution@ instance for the
+-- default value.
+newtype InitialDistribution = InitialDistribution
+ { unInitialDistribution :: Map Wallet [Pl.Value]
+ }
  deriving (Eq, Show)
 
--- | An initial distribution is valid if all utxos being created contain at least
--- a minimum amount of Ada: 'minAda'.
-validInitialDistribution :: InitialDistribution -> Bool
-validInitialDistribution = all (all hasMinAda . snd) . M.toList . distribution
- where
- hasMinAda vl = minAda `Pl.leq` vl
-
- -- the actual minimal value for a Tx output varies, here we'll just make
- -- sure that there are at least two Ada, which is sufficient for most
- -- "simple" cases.
- minAda :: Pl.Value
- minAda = Pl.lovelaceValueOf 2_000_000
-
 instance Semigroup InitialDistribution where
- (InitialDistribution i) <> (InitialDistribution j) = InitialDistribution $ M.unionWith (<>) i j
+ (InitialDistribution i) <> (InitialDistribution j) =
+ InitialDistribution $ Map.unionWith (<>) i j
 
 instance Monoid InitialDistribution where
- mempty = InitialDistribution M.empty
+ mempty = InitialDistribution Map.empty
 
 instance Default InitialDistribution where
- def = InitialDistribution $ M.fromList $ zip knownWallets (repeat $ replicate 10 defLovelace)
+ def =
+ InitialDistribution $
+ Map.fromList $
+ zip knownWallets (repeat $ replicate 10 defLovelace)
  where
  defLovelace = Pl.lovelaceValueOf 100_000_000
 
 -- | Ensures the distribution is valid by adding any missing Ada to all utxos.
 distributionFromList :: [(Wallet, [Pl.Value])] -> InitialDistribution
-distributionFromList = InitialDistribution . M.fromList . fmap (second $ fmap ensureHasMinAda)
- where
- -- the actual minimal value for a Tx output varies, here we'll just make
- -- sure that there are at least two Ada, which is sufficient for most
- -- "simple" cases.
- ensureHasMinAda :: Pl.Value -> Pl.Value
- ensureHasMinAda val = val <> Pl.toValue missingAda
- where
- missingAda = max 0 $ Pl.Lovelace 2_000_000 - Pl.fromValue val
+distributionFromList = InitialDistribution . Map.fromList
 
 -- | Extension of the default initial distribution with additional value in
 -- some wallets.
-initialDistribution' :: [(Wallet, [Pl.Value])] -> InitialDistribution
-initialDistribution' = (def <>) . distributionFromList
+initialDistribution :: [(Wallet, [Pl.Value])] -> InitialDistribution
+initialDistribution = (def <>) . distributionFromList
 
+-- | Bootstraps an initial transaction resulting in a state where wallets
+-- posess UTxOs fitting a given 'InitialDistribution'
 initialTxFor :: InitialDistribution -> Pl.Tx
-initialTxFor initDist
- | not $ validInitialDistribution initDist =
- error "Not all UTxOs have at least minAda; this initial distribution is unusable"
- | otherwise =
- mempty
- { Pl.txMint = fromRight' . Pl.toCardanoValue $ mconcat (map (mconcat . snd) initDist'),
-  Pl.txOutputs = concatMap (\(w, vs) -> map (initUtxosFor w) vs) initDist'
-  }
+initialTxFor initDist =
+ mempty
+ { Pl.txMint =
+  fromRight'
+  . Pl.toCardanoValue
+   $ mconcat (map (mconcat . snd) initDist'),
+ Pl.txOutputs = concatMap (\(w, vs) -> map (initUtxosFor w) vs) initDist'
+ }
  where
  initUtxosFor w v = toPlTxOut @() (walletAddress w) v Nothing
 
- initDist' = M.toList $ distribution initDist
+ initDist' = Map.toList $ unInitialDistribution initDist
 
  toPlTxOut :: Pl.ToData a => Pl.Address -> Pl.Value -> Maybe a -> Pl.TxOut
  toPlTxOut addr value datum = toPlTxOut' addr value datum'
  where
- datum' = maybe PV2.NoOutputDatum (PV2.OutputDatumHash . Pl.datumHash . Pl.Datum . Pl.toBuiltinData) datum
-
- toPlTxOut' :: Pl.Address -> Pl.Value -> PV2.OutputDatum -> Pl.TxOut
+ datum' =
+ maybe
+ Pl2.NoOutputDatum
+ ( Pl2.OutputDatumHash
+ . Pl.datumHash
+ . Pl.Datum
+ . Pl.toBuiltinData
+ )
+ datum
+
+ toPlTxOut' :: Pl.Address -> Pl.Value -> Pl2.OutputDatum -> Pl.TxOut
  toPlTxOut' addr value datum = Pl.TxOut $ toCardanoTxOut' addr value datum
 
- toCardanoTxOut' :: Pl.Address -> Pl.Value -> PV2.OutputDatum -> Api.TxOut Api.CtxTx Api.BabbageEra
+ toCardanoTxOut' ::
+ Pl.Address ->
+ Pl.Value ->
+ Pl2.OutputDatum ->
+ Cardano.TxOut Cardano.CtxTx Cardano.BabbageEra
  toCardanoTxOut' addr value datum =
  fromRight' $
  Pl.toCardanoTxOut
  theNetworkId
- (PV2.TxOut addr value datum Nothing)
+ (Pl2.TxOut addr value datum Nothing)
 
  fromRight' :: Show e => Either e a -> a
  fromRight' x = case x of
  Left err -> error $ show err
  Right res -> res
 
- theNetworkId :: Api.NetworkId
- theNetworkId = Api.Testnet $ Api.NetworkMagic 42 -- TODO PORT what's magic?
-
-valuesForWallet :: InitialDistribution -> Wallet -> [Pl.Value]
-valuesForWallet d w = fromMaybe [] $ w `M.lookup` distribution d
+ theNetworkId :: Cardano.NetworkId
+ theNetworkId = Cardano.Testnet $ Cardano.NetworkMagic 42 -- TODO PORT what's magic?

examples/tests/AuctionSpec.hs

@@ -54,7 +54,7 @@ bananasIn v = Value.assetClassValueOf v bananaAssetClass
 
 -- | initial distribution s.t. everyone owns five bananas
 testInit :: InitialDistribution
-testInit = initialDistribution' [(i, [Ada.lovelaceValueOf 20_000_000 <> banana 5]) | i <- knownWallets]
+testInit = initialDistribution [(i, [Ada.lovelaceValueOf 20_000_000 <> banana 5]) | i <- knownWallets]
 
 -- * Successful single-trace runs