make vote-update reward token parameter optional;

use SpecToken for reward tokens in vote-update and update-pool commands;

Author: greenhat

core/src/box_kind/ballot_box.rs

@@ -3,8 +3,7 @@ use crate::{
         BallotContract, BallotContractError, BallotContractInputs, BallotContractParameters,
     },
     oracle_types::BlockHeight,
-    pool_config::CastBallotBoxVoteParameters,
-    spec_token::{BallotTokenId, SpecToken, TokenIdKind, UpdateTokenId},
+    spec_token::{BallotTokenId, RewardTokenId, SpecToken, TokenIdKind, UpdateTokenId},
 };
 use ergo_lib::{
     chain::ergo_box::box_builder::{ErgoBoxCandidateBuilder, ErgoBoxCandidateBuilderError},
@@ -13,7 +12,7 @@ use ergo_lib::{
         chain::{
             address::{Address, AddressEncoderError},
             ergo_box::{box_value::BoxValue, ErgoBox, ErgoBoxCandidate, NonMandatoryRegisterId},
-            token::{Token, TokenId},
+            token::TokenId,
         },
         mir::constant::{TryExtractFromError, TryExtractInto},
         serialization::SigmaSerializationError,
@@ -28,10 +27,6 @@ pub enum BallotBoxError {
     NoBallotToken,
     #[error("ballot box: unknown ballot token id in `TOKENS(0)`")]
     UnknownBallotTokenId,
-    #[error("ballot box: no reward token id in R7 register")]
-    NoRewardTokenIdInR7,
-    #[error("ballot box: no reward token quantity in R8 register")]
-    NoRewardTokenQuantityInR8,
     #[error("ballot box: no group element in R4 register")]
     NoGroupElementInR4,
     #[error("ballot box: unexpected group element in R4 register")]
@@ -48,6 +43,8 @@ pub enum BallotBoxError {
     TryExtractFrom(#[from] TryExtractFromError),
     #[error("ballot box: SigmaSerializationError {0:?}")]
     SigmaSerialization(#[from] SigmaSerializationError),
+    #[error("invalid reward token: {0:?}")]
+    InvalidRewardToken(String),
 }
 
 pub trait BallotBox {
@@ -130,6 +127,13 @@ impl BallotBoxWrapperInputs {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct CastBallotBoxVoteParameters {
+    pub pool_box_address_hash: Digest32,
+    pub reward_token_opt: Option<SpecToken<RewardTokenId>>,
+    pub update_box_creation_height: i32,
+}
+
 /// A Ballot Box with vote parameters guaranteed to be set
 #[derive(Clone, Debug)]
 pub struct VoteBallotBoxWrapper {
@@ -169,21 +173,31 @@ impl VoteBallotBoxWrapper {
             .ok_or(BallotBoxError::NoPoolBoxAddressInR6)?
             .try_extract_into::<Digest32>()?;
 
-        let reward_token_id = ergo_box
+        let reward_token_id_opt = ergo_box
             .get_register(NonMandatoryRegisterId::R7.into())
-            .ok_or(BallotBoxError::NoRewardTokenIdInR7)?
-            .try_extract_into::<TokenId>()?;
-        let reward_token_quantity = ergo_box
+            .map(|c| c.try_extract_into::<TokenId>());
+        let reward_token_quantity_opt = ergo_box
             .get_register(NonMandatoryRegisterId::R8.into())
-            .ok_or(BallotBoxError::NoRewardTokenQuantityInR8)?
-            .try_extract_into::<i64>()? as u64;
+            .map(|c| c.try_extract_into::<i64>());
 
+        let reward_token_opt = match (reward_token_id_opt, reward_token_quantity_opt) {
+            (Some(Ok(reward_token_id)), Some(Ok(reward_token_quantity))) => Some(SpecToken {
+                token_id: RewardTokenId::from_token_id_unchecked(reward_token_id),
+                amount: (reward_token_quantity as u64).try_into().unwrap(),
+            }),
+            (None, None) => None,
+            (id, amt) => {
+                return Err(BallotBoxError::InvalidRewardToken(format!(
+                    "Reward token id {:?} amount {:?}",
+                    id, amt
+                )))
+            }
+        };
         let contract =
             BallotContract::from_ergo_tree(ergo_box.ergo_tree.clone(), &inputs.contract_inputs)?;
         let vote_parameters = CastBallotBoxVoteParameters {
             pool_box_address_hash,
-            reward_token_id,
-            reward_token_quantity,
+            reward_token_opt,
             update_box_creation_height,
         };
         Ok(Self {
@@ -269,7 +283,7 @@ pub fn make_local_ballot_box_candidate(
     update_box_creation_height: BlockHeight,
     ballot_token: SpecToken<BallotTokenId>,
     pool_box_address_hash: Digest32,
-    reward_tokens: Token,
+    reward_token_opt: Option<SpecToken<RewardTokenId>>,
     value: BoxValue,
     creation_height: BlockHeight,
 ) -> Result<ErgoBoxCandidate, ErgoBoxCandidateBuilderError> {
@@ -283,11 +297,16 @@ pub fn make_local_ballot_box_candidate(
         (update_box_creation_height.0 as i32).into(),
     );
     builder.set_register_value(NonMandatoryRegisterId::R6, pool_box_address_hash.into());
-    builder.set_register_value(NonMandatoryRegisterId::R7, reward_tokens.token_id.into());
-    builder.set_register_value(
-        NonMandatoryRegisterId::R8,
-        (*reward_tokens.amount.as_u64() as i64).into(),
-    );
+    if let Some(reward_tokens) = reward_token_opt {
+        builder.set_register_value(
+            NonMandatoryRegisterId::R7,
+            reward_tokens.token_id.token_id().into(),
+        );
+        builder.set_register_value(
+            NonMandatoryRegisterId::R8,
+            (*reward_tokens.amount.as_u64() as i64).into(),
+        );
+    }
     builder.add_token(ballot_token.into());
     builder.build()
 }

core/src/box_kind/pool_box.rs

@@ -4,7 +4,6 @@ use ergo_lib::ergotree_ir::chain::ergo_box::box_value::BoxValue;
 use ergo_lib::ergotree_ir::chain::ergo_box::ErgoBox;
 use ergo_lib::ergotree_ir::chain::ergo_box::ErgoBoxCandidate;
 use ergo_lib::ergotree_ir::chain::ergo_box::NonMandatoryRegisterId;
-use ergo_lib::ergotree_ir::chain::token::Token;
 use ergo_lib::ergotree_ir::mir::constant::TryExtractInto;
 use thiserror::Error;
 
@@ -229,14 +228,14 @@ pub fn make_pool_box_candidate_unchecked(
     datapoint: i64,
     epoch_counter: EpochCounter,
     pool_nft_token: SpecToken<PoolTokenId>,
-    reward_token: Token,
+    reward_token: SpecToken<RewardTokenId>,
     value: BoxValue,
     creation_height: BlockHeight,
 ) -> Result<ErgoBoxCandidate, ErgoBoxCandidateBuilderError> {
     let mut builder = ErgoBoxCandidateBuilder::new(value, contract.ergo_tree(), creation_height.0);
     builder.set_register_value(NonMandatoryRegisterId::R4, datapoint.into());
     builder.set_register_value(NonMandatoryRegisterId::R5, (epoch_counter.0 as i32).into());
     builder.add_token(pool_nft_token.into());
-    builder.add_token(reward_token);
+    builder.add_token(reward_token.into());
     builder.build()
 }

core/src/cli_commands/update_pool.rs

@@ -9,7 +9,6 @@ use ergo_lib::{
     ergotree_ir::chain::{
         address::Address,
         ergo_box::{ErgoBox, NonMandatoryRegisterId},
-        token::Token,
     },
     ergotree_ir::serialization::SigmaSerializable,
     wallet::{
@@ -24,16 +23,17 @@ use std::convert::TryInto;
 
 use crate::{
     box_kind::{
-        make_pool_box_candidate_unchecked, BallotBox, PoolBox, PoolBoxWrapper, VoteBallotBoxWrapper,
+        make_pool_box_candidate_unchecked, BallotBox, CastBallotBoxVoteParameters, PoolBox,
+        PoolBoxWrapper, VoteBallotBoxWrapper,
     },
     cli_commands::ergo_explorer_transaction_link,
     contracts::pool::PoolContract,
     node_interface::{SignTransaction, SubmitTransaction},
     oracle_config::BASE_FEE,
     oracle_state::{OraclePool, PoolBoxSource, StageError, UpdateBoxSource, VoteBallotBoxesSource},
     oracle_types::BlockHeight,
-    pool_config::{CastBallotBoxVoteParameters, PoolConfig, POOL_CONFIG},
-    spec_token::TokenIdKind,
+    pool_config::{PoolConfig, POOL_CONFIG},
+    spec_token::{RewardTokenId, SpecToken, TokenIdKind},
     wallet::{WalletDataError, WalletDataSource},
 };
 use derive_more::From;
@@ -77,7 +77,7 @@ pub fn update_pool(
     tx_signer: &dyn SignTransaction,
     tx_submit: &dyn SubmitTransaction,
     new_pool_box_hash_str: Option<String>,
-    new_reward_tokens: Option<Token>,
+    new_reward_tokens: Option<SpecToken<RewardTokenId>>,
     height: BlockHeight,
 ) -> Result<(), UpdatePoolError> {
     info!("Opening pool_config_updated.yaml");
@@ -137,9 +137,9 @@ fn display_update_diff(
     old_pool_config: &PoolConfig,
     new_pool_config: &PoolConfig,
     old_pool_box: PoolBoxWrapper,
-    new_reward_tokens: Option<Token>,
+    new_reward_tokens: Option<SpecToken<RewardTokenId>>,
 ) {
-    let new_tokens = new_reward_tokens.unwrap_or_else(|| old_pool_box.reward_token().into());
+    let new_tokens = new_reward_tokens.unwrap_or_else(|| old_pool_box.reward_token());
     let new_pool_contract =
         PoolContract::checked_load(&new_pool_config.pool_box_wrapper_inputs.contract_inputs)
             .unwrap();
@@ -189,7 +189,7 @@ fn build_update_pool_box_tx(
     wallet: &dyn WalletDataSource,
     update_box: &dyn UpdateBoxSource,
     new_pool_contract: PoolContract,
-    new_reward_tokens: Option<Token>,
+    new_reward_tokens: Option<SpecToken<RewardTokenId>>,
     height: BlockHeight,
     change_address: Address,
 ) -> Result<TransactionContext<UnsignedTransaction>, UpdatePoolError> {
@@ -202,13 +202,12 @@ fn build_update_pool_box_tx(
             .sigma_serialize_bytes()
             .unwrap(),
     );
-    let reward_tokens = new_reward_tokens.unwrap_or_else(|| old_pool_box.reward_token().into());
     let vote_parameters = CastBallotBoxVoteParameters {
         pool_box_address_hash: pool_box_hash,
-        reward_token_id: reward_tokens.token_id,
-        reward_token_quantity: *reward_tokens.amount.as_u64(),
+        reward_token_opt: new_reward_tokens.clone(),
         update_box_creation_height: update_box.get_box().creation_height as i32,
     };
+    let reward_tokens = new_reward_tokens.unwrap_or_else(|| old_pool_box.reward_token());
     // Find ballot boxes that are voting for the new pool hash
     let mut sorted_ballot_boxes = ballot_boxes.get_ballot_boxes()?;
     // Sort in descending order of ballot token amounts. If two boxes have the same amount of ballot tokens, also compare box value, in case some boxes were incorrectly created below minStorageRent
@@ -278,11 +277,12 @@ fn build_update_pool_box_tx(
     }
 
     let target_balance = *BASE_FEE;
-    let target_tokens = if reward_tokens.token_id != old_pool_box.reward_token().token_id() {
-        vec![reward_tokens.clone()]
-    } else {
-        vec![]
-    };
+    let target_tokens =
+        if reward_tokens.token_id.token_id() != old_pool_box.reward_token().token_id() {
+            vec![reward_tokens.clone().into()]
+        } else {
+            vec![]
+        };
     let box_selector = SimpleBoxSelector::new();
     let selection = box_selector.select(unspent_boxes, target_balance, &target_tokens)?;
     let mut input_boxes = vec![old_pool_box.get_box().clone(), update_box.get_box().clone()];
@@ -321,7 +321,7 @@ fn build_update_pool_box_tx(
         change_address,
     );
 
-    if reward_tokens.token_id != old_pool_box.reward_token().token_id() {
+    if reward_tokens.token_id.token_id() != old_pool_box.reward_token().token_id() {
         tx_builder.set_token_burn_permit(vec![old_pool_box.reward_token().into()]);
     }
 
@@ -379,7 +379,7 @@ mod tests {
             generate_token_ids, make_wallet_unspent_box, BallotBoxesMock, PoolBoxMock,
             UpdateBoxMock, WalletDataMock,
         },
-        spec_token::{RefreshTokenId, SpecToken, TokenIdKind},
+        spec_token::{RefreshTokenId, RewardTokenId, SpecToken, TokenIdKind},
     };
 
     use super::build_update_pool_box_tx;
@@ -405,8 +405,8 @@ mod tests {
             token_id: token_ids.reward_token_id.clone(),
             amount: 1500.try_into().unwrap(),
         };
-        let new_reward_tokens = Token {
-            token_id: force_any_tokenid(),
+        let new_reward_tokens = SpecToken {
+            token_id: RewardTokenId::from_token_id_unchecked(force_any_tokenid()),
             amount: force_any_val(),
         };
         dbg!(&new_reward_tokens);
@@ -498,7 +498,7 @@ mod tests {
                     amount: 1.try_into().unwrap(),
                 },
                 pool_box_hash,
-                new_reward_tokens.clone(),
+                Some(new_reward_tokens.clone()),
                 ballot_contract.min_storage_rent(),
                 height,
             )
@@ -524,7 +524,7 @@ mod tests {
             // create a wallet box with new reward tokens
             secret.public_image(),
             BASE_FEE.checked_mul_u32(4_000_000_000).unwrap(),
-            Some(vec![new_reward_tokens.clone()].try_into().unwrap()),
+            Some(vec![new_reward_tokens.clone().into()].try_into().unwrap()),
         );
         let change_address = AddressEncoder::unchecked_parse_network_address_from_str(
             "9iHyKxXs2ZNLMp9N9gbUT9V8gTbsV7HED1C1VhttMfBUMPDyF7r",