Add VoteBallotBoxesWrapper

Author: sethdusek

core/src/box_kind/ballot_box.rs

@@ -1,5 +1,3 @@
-use std::convert::TryInto;
-
 use crate::{
     contracts::ballot::{BallotContract, BallotContractError},
     oracle_config::{BallotBoxWrapperParameters, CastBallotBoxVoteParameters},
@@ -108,11 +106,7 @@ impl BallotBoxWrapper {
                 .get_register(NonMandatoryRegisterId::R6.into())
                 .ok_or(BallotBoxError::NoPoolBoxAddressInR6)?
                 .try_extract_into::<Digest32>()?;
-            let pb: Digest32 = base16::decode(pool_box_address_hash)
-                .unwrap()
-                .try_into()
-                .unwrap();
-            if pb != register_pool_box_address_hash {
+            if *pool_box_address_hash != register_pool_box_address_hash {
                 warn!("Pool box address in R6 register differs to config. Could be due to vote.");
             }
 
@@ -153,6 +147,76 @@ pub struct BallotBoxWrapperInputs<'a> {
     pub update_nft_token_id: &'a TokenId,
 }
 
+/// A Ballot Box with vote parameters guaranteed to be set
+#[derive(Clone, Debug)]
+pub struct VoteBallotBoxWrapper {
+    ergo_box: ErgoBox,
+    vote_parameters: CastBallotBoxVoteParameters,
+    contract: BallotContract,
+}
+
+impl VoteBallotBoxWrapper {
+    pub fn new(ergo_box: ErgoBox, inputs: BallotBoxWrapperInputs) -> Result<Self, BallotBoxError> {
+        let ballot_token_id = &ergo_box
+            .tokens
+            .as_ref()
+            .ok_or(BallotBoxError::NoBallotToken)?
+            .get(0)
+            .ok_or(BallotBoxError::NoBallotToken)?
+            .token_id;
+        if *ballot_token_id != *inputs.ballot_token_id {
+            return Err(BallotBoxError::UnknownBallotTokenId);
+        }
+
+        if ergo_box
+            .get_register(NonMandatoryRegisterId::R4.into())
+            .ok_or(BallotBoxError::NoGroupElementInR4)?
+            .try_extract_into::<EcPoint>()
+            .is_err()
+        {
+            return Err(BallotBoxError::NoGroupElementInR4);
+        }
+        if ergo_box
+            .get_register(NonMandatoryRegisterId::R5.into())
+            .ok_or(BallotBoxError::NoUpdateBoxCreationHeightInR5)?
+            .try_extract_into::<i32>()
+            .is_err()
+        {
+            return Err(BallotBoxError::NoUpdateBoxCreationHeightInR5);
+        }
+
+        let pool_box_address_hash = ergo_box
+            .get_register(NonMandatoryRegisterId::R6.into())
+            .ok_or(BallotBoxError::NoPoolBoxAddressInR6)?
+            .try_extract_into::<Digest32>()?;
+
+        let reward_token_id = ergo_box
+            .get_register(NonMandatoryRegisterId::R7.into())
+            .ok_or(BallotBoxError::NoRewardTokenIdInR7)?
+            .try_extract_into::<TokenId>()?;
+        let reward_token_quantity = ergo_box
+            .get_register(NonMandatoryRegisterId::R8.into())
+            .ok_or(BallotBoxError::NoRewardTokenQuantityInR8)?
+            .try_extract_into::<i64>()? as u32;
+
+        let contract = BallotContract::from_ergo_tree(ergo_box.ergo_tree.clone(), inputs.into())?;
+        let vote_parameters = CastBallotBoxVoteParameters {
+            pool_box_address_hash,
+            reward_token_id,
+            reward_token_quantity,
+        };
+        Ok(Self {
+            ergo_box,
+            contract,
+            vote_parameters,
+        })
+    }
+
+    pub fn vote_parameters(&self) -> &CastBallotBoxVoteParameters {
+        &self.vote_parameters
+    }
+}
+
 impl BallotBox for BallotBoxWrapper {
     fn contract(&self) -> &BallotContract {
         &self.contract
@@ -186,6 +250,39 @@ impl BallotBox for BallotBoxWrapper {
     }
 }
 
+impl BallotBox for VoteBallotBoxWrapper {
+    fn contract(&self) -> &BallotContract {
+        &self.contract
+    }
+
+    fn ballot_token(&self) -> Token {
+        self.ergo_box
+            .tokens
+            .as_ref()
+            .unwrap()
+            .get(0)
+            .unwrap()
+            .clone()
+    }
+
+    fn min_storage_rent(&self) -> u64 {
+        self.contract.min_storage_rent()
+    }
+
+    fn ballot_token_owner(&self) -> ProveDlog {
+        self.ergo_box
+            .get_register(NonMandatoryRegisterId::R4.into())
+            .unwrap()
+            .try_extract_into::<EcPoint>()
+            .unwrap()
+            .into()
+    }
+
+    fn get_box(&self) -> &ErgoBox {
+        &self.ergo_box
+    }
+}
+
 #[allow(clippy::too_many_arguments)]
 pub fn make_local_ballot_box_candidate(
     contract: &BallotContract,

core/src/cli_commands/update_pool.rs

@@ -11,7 +11,6 @@ use ergo_lib::{
         ergo_box::{box_value::BoxValue, ErgoBox, NonMandatoryRegisterId},
         token::Token,
     },
-    ergotree_ir::mir::constant::Constant,
     ergotree_ir::serialization::SigmaSerializable,
     wallet::{
         box_selector::{BoxSelection, BoxSelector, BoxSelectorError, SimpleBoxSelector},
@@ -24,13 +23,13 @@ use log::{error, info};
 use std::convert::TryInto;
 
 use crate::{
-    box_kind::{make_pool_box_candidate, BallotBox, BallotBoxWrapper, PoolBox, PoolBoxWrapper},
+    box_kind::{make_pool_box_candidate, BallotBox, PoolBox, PoolBoxWrapper, VoteBallotBoxWrapper},
     cli_commands::ergo_explorer_transaction_link,
     contracts::pool::PoolContract,
     contracts::pool::PoolContractInputs,
     node_interface::{current_block_height, get_wallet_status, sign_and_submit_transaction},
-    oracle_config::{OracleConfig, ORACLE_CONFIG},
-    oracle_state::{BallotBoxesSource, OraclePool, PoolBoxSource, StageError, UpdateBoxSource},
+    oracle_config::{CastBallotBoxVoteParameters, OracleConfig, ORACLE_CONFIG},
+    oracle_state::{OraclePool, PoolBoxSource, StageError, UpdateBoxSource, VoteBallotBoxesSource},
     wallet::WalletDataSource,
 };
 use derive_more::From;
@@ -190,7 +189,7 @@ fn display_update_diff(
 #[allow(clippy::too_many_arguments)]
 fn build_update_pool_box_tx(
     pool_box_source: &dyn PoolBoxSource,
-    ballot_boxes: &dyn BallotBoxesSource,
+    ballot_boxes: &dyn VoteBallotBoxesSource,
     wallet: &dyn WalletDataSource,
     update_box: &dyn UpdateBoxSource,
     new_pool_contract: PoolContract,
@@ -201,13 +200,18 @@ fn build_update_pool_box_tx(
     let update_box = update_box.get_update_box()?;
     let min_votes = update_box.min_votes();
     let old_pool_box = pool_box_source.get_pool_box()?;
-    let pool_box_hash = Constant::from(blake2b256_hash(
+    let pool_box_hash = blake2b256_hash(
         &new_pool_contract
             .ergo_tree()
             .sigma_serialize_bytes()
             .unwrap(),
-    ));
+    );
     let reward_tokens = new_reward_tokens.unwrap_or_else(|| old_pool_box.reward_token());
+    let vote_parameters = CastBallotBoxVoteParameters {
+        pool_box_address_hash: pool_box_hash,
+        reward_token_id: reward_tokens.token_id.clone(),
+        reward_token_quantity: *reward_tokens.amount.as_u64() as u32, // TODO: Change vote parameters to i64
+    };
     // 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
@@ -224,28 +228,10 @@ fn build_update_pool_box_tx(
     sorted_ballot_boxes.reverse();
 
     let mut votes_cast = 0;
-    let vote_ballot_boxes: Vec<BallotBoxWrapper> = ballot_boxes
+    let vote_ballot_boxes: Vec<VoteBallotBoxWrapper> = ballot_boxes
         .get_ballot_boxes()?
         .into_iter()
-        .filter(|ballot_box| {
-            let ballot_box = ballot_box.get_box();
-            ballot_box
-                .additional_registers
-                .get(NonMandatoryRegisterId::R5)
-                == Some(&update_box.get_box().creation_info().0.into())
-                && ballot_box
-                    .additional_registers
-                    .get(NonMandatoryRegisterId::R6)
-                    == Some(&pool_box_hash)
-                && ballot_box
-                    .additional_registers
-                    .get(NonMandatoryRegisterId::R7)
-                    == Some(&reward_tokens.token_id.clone().into())
-                && ballot_box
-                    .additional_registers
-                    .get(NonMandatoryRegisterId::R8)
-                    == Some(&(*reward_tokens.amount.as_u64() as i64).into())
-        })
+        .filter(|ballot_box| *ballot_box.vote_parameters() == vote_parameters)
         .scan(&mut votes_cast, |votes_cast, ballot_box| {
             **votes_cast += *ballot_box.ballot_token().amount.as_u64();
             Some(ballot_box)
@@ -378,8 +364,8 @@ mod tests {
 
     use crate::{
         box_kind::{
-            make_local_ballot_box_candidate, make_pool_box_candidate, BallotBoxWrapper,
-            PoolBoxWrapper, PoolBoxWrapperInputs, UpdateBoxWrapper, UpdateBoxWrapperInputs,
+            make_local_ballot_box_candidate, make_pool_box_candidate, PoolBoxWrapper,
+            PoolBoxWrapperInputs, UpdateBoxWrapper, UpdateBoxWrapperInputs, VoteBallotBoxWrapper,
         },
         contracts::{
             ballot::{BallotContract, BallotContractInputs},
@@ -496,11 +482,7 @@ mod tests {
             let secret = DlogProverInput::random();
             let ballot_box_parameters = BallotBoxWrapperParameters {
                 contract_parameters: ballot_contract_parameters.clone(),
-                vote_parameters: Some(crate::oracle_config::CastBallotBoxVoteParameters {
-                    pool_box_address_hash: base16::encode_lower(&pool_box_hash),
-                    reward_token_id: new_reward_tokens.token_id.clone(),
-                    reward_token_quantity: *new_reward_tokens.amount.as_u64() as u32,
-                }),
+                vote_parameters: None,
                 ballot_token_owner_address: AddressEncoder::new(
                     ballot_contract_parameters.p2s.network(),
                 )
@@ -526,7 +508,7 @@ mod tests {
                 ErgoBox::from_box_candidate(&ballot_box_candidate, force_any_val::<TxId>(), 0)
                     .unwrap();
             ballot_boxes.push(
-                BallotBoxWrapper::new(
+                VoteBallotBoxWrapper::new(
                     ballot_box,
                     crate::box_kind::BallotBoxWrapperInputs {
                         parameters: &ballot_box_parameters,
@@ -594,7 +576,6 @@ mod tests {
             change_address,
         )
         .unwrap();
-        println!("{}", serde_json::to_string(&update_tx.spending_tx).unwrap());
 
         wallet.sign_transaction(update_tx, &ctx, None).unwrap();
     }

core/src/oracle_config.rs

@@ -8,7 +8,10 @@ use crate::{
     datapoint_source::{DataPointSource, ExternalScript, PredefinedDataPointSource},
 };
 use anyhow::anyhow;
-use ergo_lib::ergotree_ir::chain::{address::NetworkPrefix, token::TokenId};
+use ergo_lib::{
+    ergo_chain_types::Digest32,
+    ergotree_ir::chain::{address::NetworkPrefix, token::TokenId},
+};
 use log::LevelFilter;
 use serde::{Deserialize, Serialize};
 
@@ -48,9 +51,9 @@ pub struct BallotBoxWrapperParameters {
     pub ballot_token_owner_address: String,
 }
 
-#[derive(Debug, Clone, Deserialize, Serialize)]
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
 pub struct CastBallotBoxVoteParameters {
-    pub pool_box_address_hash: String,
+    pub pool_box_address_hash: Digest32,
     pub reward_token_id: TokenId,
     pub reward_token_quantity: u32,
 }