Skip to content

Commit

Permalink
use strum to generate the enum tags
Browse files Browse the repository at this point in the history
  • Loading branch information
OBorce committed Dec 12, 2024
1 parent e16ad88 commit e9d69d2
Show file tree
Hide file tree
Showing 2 changed files with 80 additions and 76 deletions.
37 changes: 7 additions & 30 deletions trezor-common/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@ pub enum OutputValue {
}

#[derive(Debug, Clone, PartialEq, Eq, Encode, Decode, EnumDiscriminants)]
#[strum_discriminants(name(OutputTimeLockIndex), derive(EnumIter))]
#[strum_discriminants(name(OutputTimeLockIndex), derive(EnumIter, FromPrimitive))]
pub enum OutputTimeLock {
#[codec(index = 0)]
UntilHeight(#[codec(compact)] u64),
Expand Down Expand Up @@ -153,7 +153,8 @@ pub enum IsTokenFreezable {
Yes,
}

#[derive(Debug, Clone, PartialEq, Eq, Encode, Decode)]
#[derive(Debug, Clone, PartialEq, Eq, Encode, Decode, EnumDiscriminants)]
#[strum_discriminants(name(TokenTotalSupplyIndex), derive(EnumIter, FromPrimitive))]
pub enum TokenTotalSupply {
#[codec(index = 0)]
Fixed(Amount), // fixed to a certain amount
Expand All @@ -163,13 +164,6 @@ pub enum TokenTotalSupply {
Unlimited, // limited only by the Amount data type
}

#[derive(FromPrimitive)]
pub enum TokenTotalSupplyIndex {
Fixed = 0,
Lockable = 1,
Unlimited = 2,
}

#[derive(Encode)]
pub struct TokenIssuanceV1 {
pub token_ticker: parity_scale_codec::alloc::vec::Vec<u8>,
Expand Down Expand Up @@ -272,20 +266,15 @@ pub struct H256(pub [u8; 32]);
#[derive(PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Debug, Encode, Decode)]
pub struct HtlcSecretHash(pub [u8; 20]);

#[derive(Debug, Clone, PartialEq, Eq, Encode, Decode, Ord, PartialOrd)]
#[derive(Debug, Clone, PartialEq, Eq, Encode, Decode, Ord, PartialOrd, EnumDiscriminants)]
#[strum_discriminants(name(OutPointSourceIdIndex), derive(EnumIter, FromPrimitive))]
pub enum OutPointSourceId {
#[codec(index = 0)]
Transaction(H256),
#[codec(index = 1)]
BlockReward(H256),
}

#[derive(FromPrimitive)]
pub enum OutPointSourceIdIndex {
Transaction = 0,
BlockReward = 1,
}

#[derive(Debug, Clone, PartialEq, Eq, Encode, Decode, Ord, PartialOrd)]
pub struct UtxoOutPoint {
id: OutPointSourceId,
Expand Down Expand Up @@ -333,7 +322,8 @@ pub enum IsTokenUnfreezable {
type OrderId = H256;
type TokenId = H256;

#[derive(Encode)]
#[derive(Encode, EnumDiscriminants)]
#[strum_discriminants(name(AccountCommandIndex), derive(EnumIter, FromPrimitive))]
pub enum AccountCommand {
// Create certain amount of tokens and add them to circulating supply
#[codec(index = 0)]
Expand Down Expand Up @@ -365,19 +355,6 @@ pub enum AccountCommand {
ChangeTokenMetadataUri(TokenId, parity_scale_codec::alloc::vec::Vec<u8>),
}

#[derive(FromPrimitive)]
pub enum AccountCommandIndex {
MintTokens = 0,
UnmintTokens = 1,
LockTokenSupply = 2,
FreezeToken = 3,
UnfreezeToken = 4,
ChangeTokenAuthority = 5,
ConcludeOrder = 6,
FillOrder = 7,
ChangeTokenMetadataUri = 8,
}

#[derive(Encode)]
pub enum TxInput {
#[codec(index = 0)]
Expand Down
119 changes: 73 additions & 46 deletions trezor-common/src/tests/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,8 @@ use common::{chain, primitives};
use crypto::key::KeyKind;
use parity_scale_codec::{DecodeAll, Encode};
use rstest::rstest;
use test_utils::random::{make_seedable_rng, CryptoRng, Rng, Seed};
use strum::IntoEnumIterator;
use test_utils::random::{make_seedable_rng, CryptoRng, IteratorRandom, Rng, Seed};

impl From<primitives::H256> for crate::H256 {
fn from(value: primitives::H256) -> Self {
Expand Down Expand Up @@ -312,16 +313,18 @@ impl From<chain::TxOutput> for crate::TxOutput {
}

fn make_random_destination(rng: &mut (impl Rng + CryptoRng)) -> chain::Destination {
match rng.gen_range(0..=4) {
0 => chain::Destination::AnyoneCanSpend,
1 => chain::Destination::PublicKey(
match chain::DestinationTag::iter().choose(rng).unwrap() {
chain::DestinationTag::AnyoneCanSpend => chain::Destination::AnyoneCanSpend,
chain::DestinationTag::PublicKey => chain::Destination::PublicKey(
crypto::key::PrivateKey::new_from_rng(rng, KeyKind::Secp256k1Schnorr).1,
),
2 => chain::Destination::ScriptHash(primitives::H256(rng.gen()).into()),
3 => chain::Destination::ClassicMultisig(
chain::DestinationTag::ScriptHash => {
chain::Destination::ScriptHash(primitives::H256(rng.gen()).into())
}
chain::DestinationTag::ClassicMultisig => chain::Destination::ClassicMultisig(
(&crypto::key::PrivateKey::new_from_rng(rng, KeyKind::Secp256k1Schnorr).1).into(),
),
_ => chain::Destination::PublicKeyHash(
chain::DestinationTag::PublicKeyHash => chain::Destination::PublicKeyHash(
(&crypto::key::PrivateKey::new_from_rng(rng, KeyKind::Secp256k1Schnorr).1).into(),
),
}
Expand All @@ -332,36 +335,48 @@ fn make_random_bytes(rng: &mut (impl Rng + CryptoRng)) -> Vec<u8> {
}

fn make_random_account_command(rng: &mut (impl Rng + CryptoRng)) -> chain::AccountCommand {
match rng.gen_range(0..=8) {
0 => chain::AccountCommand::MintTokens(
match crate::AccountCommandIndex::iter().choose(rng).unwrap() {
crate::AccountCommandIndex::MintTokens => chain::AccountCommand::MintTokens(
primitives::H256(rng.gen()).into(),
primitives::Amount::from_atoms(rng.gen()),
),
1 => chain::AccountCommand::UnmintTokens(primitives::H256(rng.gen()).into()),
2 => chain::AccountCommand::LockTokenSupply(primitives::H256(rng.gen()).into()),
3 => chain::AccountCommand::FreezeToken(
crate::AccountCommandIndex::UnmintTokens => {
chain::AccountCommand::UnmintTokens(primitives::H256(rng.gen()).into())
}
crate::AccountCommandIndex::LockTokenSupply => {
chain::AccountCommand::LockTokenSupply(primitives::H256(rng.gen()).into())
}
crate::AccountCommandIndex::FreezeToken => chain::AccountCommand::FreezeToken(
primitives::H256(rng.gen()).into(),
if rng.gen::<bool>() {
chain::tokens::IsTokenUnfreezable::Yes
} else {
chain::tokens::IsTokenUnfreezable::No
},
),
4 => chain::AccountCommand::UnfreezeToken(primitives::H256(rng.gen()).into()),
5 => chain::AccountCommand::ChangeTokenAuthority(
primitives::H256(rng.gen()).into(),
make_random_destination(rng),
),
6 => chain::AccountCommand::ConcludeOrder(primitives::H256(rng.gen()).into()),
7 => chain::AccountCommand::FillOrder(
crate::AccountCommandIndex::UnfreezeToken => {
chain::AccountCommand::UnfreezeToken(primitives::H256(rng.gen()).into())
}
crate::AccountCommandIndex::ChangeTokenAuthority => {
chain::AccountCommand::ChangeTokenAuthority(
primitives::H256(rng.gen()).into(),
make_random_destination(rng),
)
}
crate::AccountCommandIndex::ConcludeOrder => {
chain::AccountCommand::ConcludeOrder(primitives::H256(rng.gen()).into())
}
crate::AccountCommandIndex::FillOrder => chain::AccountCommand::FillOrder(
primitives::H256(rng.gen()).into(),
primitives::Amount::from_atoms(rng.gen()),
make_random_destination(rng),
),
_ => chain::AccountCommand::ChangeTokenMetadataUri(
primitives::H256(rng.gen()).into(),
make_random_bytes(rng),
),
crate::AccountCommandIndex::ChangeTokenMetadataUri => {
chain::AccountCommand::ChangeTokenMetadataUri(
primitives::H256(rng.gen()).into(),
make_random_bytes(rng),
)
}
}
}

Expand Down Expand Up @@ -397,23 +412,31 @@ fn make_random_value(rng: &mut (impl Rng + CryptoRng)) -> chain::output_value::O
}

fn make_random_lock(rng: &mut (impl Rng + CryptoRng)) -> chain::timelock::OutputTimeLock {
match rng.gen_range(0..=3) {
0 => chain::timelock::OutputTimeLock::UntilHeight(primitives::BlockHeight::new(rng.gen())),
1 => chain::timelock::OutputTimeLock::UntilTime(
match crate::OutputTimeLockIndex::iter().choose(rng).unwrap() {
crate::OutputTimeLockIndex::UntilHeight => {
chain::timelock::OutputTimeLock::UntilHeight(primitives::BlockHeight::new(rng.gen()))
}
crate::OutputTimeLockIndex::UntilTime => chain::timelock::OutputTimeLock::UntilTime(
chain::block::timestamp::BlockTimestamp::from_int_seconds(rng.gen()),
),
2 => chain::timelock::OutputTimeLock::ForSeconds(rng.gen()),
_ => chain::timelock::OutputTimeLock::ForBlockCount(rng.gen()),
crate::OutputTimeLockIndex::ForSeconds => {
chain::timelock::OutputTimeLock::ForSeconds(rng.gen())
}
crate::OutputTimeLockIndex::ForBlockCount => {
chain::timelock::OutputTimeLock::ForBlockCount(rng.gen())
}
}
}

fn make_random_token_total_supply(
rng: &mut (impl Rng + CryptoRng),
) -> chain::tokens::TokenTotalSupply {
match rng.gen_range(0..=2) {
0 => chain::tokens::TokenTotalSupply::Unlimited,
1 => chain::tokens::TokenTotalSupply::Lockable,
_ => chain::tokens::TokenTotalSupply::Fixed(primitives::Amount::from_atoms(rng.gen())),
match crate::TokenTotalSupplyIndex::iter().choose(rng).unwrap() {
crate::TokenTotalSupplyIndex::Unlimited => chain::tokens::TokenTotalSupply::Unlimited,
crate::TokenTotalSupplyIndex::Lockable => chain::tokens::TokenTotalSupply::Lockable,
crate::TokenTotalSupplyIndex::Fixed => {
chain::tokens::TokenTotalSupply::Fixed(primitives::Amount::from_atoms(rng.gen()))
}
}
}

Expand Down Expand Up @@ -526,32 +549,36 @@ fn make_random_output(rng: &mut (impl Rng + CryptoRng)) -> chain::TxOutput {
#[trace]
#[case(Seed::from_entropy())]
fn check_input_encodings(#[case] seed: Seed) {
let mut rng = make_seedable_rng(seed);
for _ in 0..1000 {
let mut rng = make_seedable_rng(seed);

let inp = make_random_input(&mut rng);
let inp = make_random_input(&mut rng);

let simple_inp: crate::TxInput = inp.clone().into();
let simple_inp: crate::TxInput = inp.clone().into();

assert_eq!(inp.encode(), simple_inp.encode());
assert_eq!(inp.encode(), simple_inp.encode());

let decoded_simple_inp =
chain::TxInput::decode_all(&mut simple_inp.encode().as_slice()).unwrap();
assert_eq!(decoded_simple_inp, inp);
let decoded_simple_inp =
chain::TxInput::decode_all(&mut simple_inp.encode().as_slice()).unwrap();
assert_eq!(decoded_simple_inp, inp);
}
}

#[rstest]
#[trace]
#[case(Seed::from_entropy())]
fn check_output_encodings(#[case] seed: Seed) {
let mut rng = make_seedable_rng(seed);
for _ in 0..1000 {
let mut rng = make_seedable_rng(seed);

let out = make_random_output(&mut rng);
let out = make_random_output(&mut rng);

let simple_out: crate::TxOutput = out.clone().into();
let simple_out: crate::TxOutput = out.clone().into();

assert_eq!(out.encode(), simple_out.encode());
assert_eq!(out.encode(), simple_out.encode());

let decoded_simple_out =
chain::TxOutput::decode_all(&mut simple_out.encode().as_slice()).unwrap();
assert_eq!(decoded_simple_out, out);
let decoded_simple_out =
chain::TxOutput::decode_all(&mut simple_out.encode().as_slice()).unwrap();
assert_eq!(decoded_simple_out, out);
}
}

0 comments on commit e9d69d2

Please sign in to comment.