mod simple;
use std::collections::HashMap;
use bounded_vec::BoundedVec;
use ergotree_ir::chain::ergo_box::box_value::BoxValue;
use ergotree_ir::chain::ergo_box::BoxId;
use ergotree_ir::chain::ergo_box::BoxTokens;
use ergotree_ir::chain::ergo_box::ErgoBox;
use ergotree_ir::chain::ergo_box::ErgoBoxCandidate;
use ergotree_ir::chain::token::Token;
use ergotree_ir::chain::token::TokenAmount;
use ergotree_ir::chain::token::TokenAmountError;
use ergotree_ir::chain::token::TokenId;
pub use simple::*;
use thiserror::Error;
pub type SelectedBoxes<T> = BoundedVec<T, 1, { u16::MAX as usize }>;
#[derive(PartialEq, Eq, Debug, Clone)]
pub struct BoxSelection<T: ErgoBoxAssets> {
pub boxes: SelectedBoxes<T>,
pub change_boxes: Vec<ErgoBoxAssetsData>,
}
pub trait BoxSelector<T: ErgoBoxAssets> {
fn select(
&self,
inputs: Vec<T>,
target_balance: BoxValue,
target_tokens: &[Token],
) -> Result<BoxSelection<T>, BoxSelectorError>;
}
#[derive(Error, PartialEq, Eq, Debug, Clone)]
pub enum BoxSelectorError {
#[error("Not enough coins({0} nanoERGs are missing)")]
NotEnoughCoins(u64),
#[error("Not enough tokens, missing {0:?}")]
NotEnoughTokens(Vec<Token>),
#[error("TokenAmountError: {0:?}")]
TokenAmountError(#[from] TokenAmountError),
#[error("CheckPreservationError: {0:?}")]
CheckPreservation(#[from] CheckPreservationError),
#[error("Not enough coins for change box: {0:?}")]
NotEnoughCoinsForChangeBox(#[from] NotEnoughCoinsForChangeBox),
#[error("Selected inputs out of bounds: {0}")]
SelectedInputsOutOfBounds(usize),
}
pub trait ErgoBoxAssets {
fn value(&self) -> BoxValue;
fn tokens(&self) -> Option<BoxTokens>;
}
#[derive(PartialEq, Eq, Debug, Clone)]
pub struct ErgoBoxAssetsData {
pub value: BoxValue,
pub tokens: Option<BoxTokens>,
}
impl ErgoBoxAssets for ErgoBoxAssetsData {
fn value(&self) -> BoxValue {
self.value
}
fn tokens(&self) -> Option<BoxTokens> {
self.tokens.clone()
}
}
impl ErgoBoxAssets for ErgoBoxCandidate {
fn value(&self) -> BoxValue {
self.value
}
fn tokens(&self) -> Option<BoxTokens> {
self.tokens.clone()
}
}
impl ErgoBoxAssets for ErgoBox {
fn value(&self) -> BoxValue {
self.value
}
fn tokens(&self) -> Option<BoxTokens> {
self.tokens.clone()
}
}
pub trait ErgoBoxId {
fn box_id(&self) -> BoxId;
}
impl ErgoBoxId for ErgoBox {
fn box_id(&self) -> BoxId {
self.box_id()
}
}
pub fn sum_value<T: ErgoBoxAssets>(bs: &[T]) -> u64 {
bs.iter().map(|b| *b.value().as_u64()).sum()
}
pub fn sum_tokens(ts: Option<&[Token]>) -> Result<HashMap<TokenId, TokenAmount>, TokenAmountError> {
let mut res: HashMap<TokenId, TokenAmount> = HashMap::new();
ts.into_iter().flatten().try_for_each(|t| {
if let Some(amt) = res.get_mut(&t.token_id) {
*amt = amt.checked_add(&t.amount)?;
} else {
res.insert(t.token_id, t.amount);
}
Ok(())
})?;
Ok(res)
}
pub fn sum_tokens_from_boxes<T: ErgoBoxAssets>(
bs: &[T],
) -> Result<HashMap<TokenId, TokenAmount>, TokenAmountError> {
let mut res: HashMap<TokenId, TokenAmount> = HashMap::new();
bs.iter().try_for_each(|b| {
b.tokens().into_iter().flatten().try_for_each(|t| {
if let Some(amt) = res.get_mut(&t.token_id) {
*amt = amt.checked_add(&t.amount)?;
} else {
res.insert(t.token_id, t.amount);
}
Ok(())
})
})?;
Ok(res)
}
pub fn sum_tokens_from_hashmaps(
tokens1: HashMap<TokenId, TokenAmount>,
tokens2: HashMap<TokenId, TokenAmount>,
) -> Result<HashMap<TokenId, TokenAmount>, TokenAmountError> {
let mut res: HashMap<TokenId, TokenAmount> = HashMap::new();
tokens1
.into_iter()
.chain(tokens2)
.try_for_each(|(id, t_amt)| {
if let Some(amt) = res.get_mut(&id) {
*amt = amt.checked_add(&t_amt)?;
} else {
res.insert(id, t_amt);
}
Ok(())
})?;
Ok(res)
}
pub fn subtract_tokens(
tokens1: &HashMap<TokenId, TokenAmount>,
tokens2: &HashMap<TokenId, TokenAmount>,
) -> Result<HashMap<TokenId, TokenAmount>, TokenAmountError> {
let mut res: HashMap<TokenId, TokenAmount> = tokens1.clone();
tokens2.iter().try_for_each(|(id, t_amt)| {
if let Some(amt) = res.get_mut(id) {
if amt == t_amt {
res.remove(id);
} else {
*amt = amt.checked_sub(t_amt)?;
}
} else {
return Err(TokenAmountError::OutOfBounds(-(*t_amt.as_u64() as i64)));
}
Ok(())
})?;
Ok(res)
}
#[allow(clippy::unwrap_used, clippy::panic)]
#[cfg(feature = "arbitrary")]
pub mod arbitrary {
use std::ops::Range;
use ergotree_ir::chain::{
ergo_box::{
box_value::{arbitrary::ArbBoxValueRange, BoxValue},
BoxTokens,
},
token::{arbitrary::ArbTokenIdParam, Token},
};
use proptest::{arbitrary::Arbitrary, collection::vec, prelude::*};
use super::ErgoBoxAssetsData;
pub struct ArbTokensParam {
pub token_id_param: ArbTokenIdParam,
pub token_count_range: Range<usize>,
}
impl Default for ArbTokensParam {
fn default() -> Self {
ArbTokensParam {
token_id_param: ArbTokenIdParam::default(),
token_count_range: 0..3,
}
}
}
#[derive(Default)]
pub struct ArbErgoBoxAssetsDataParam {
pub value_range: ArbBoxValueRange,
pub tokens_param: ArbTokensParam,
}
impl From<Range<u64>> for ArbErgoBoxAssetsDataParam {
fn from(r: Range<u64>) -> Self {
ArbErgoBoxAssetsDataParam {
value_range: r.into(),
tokens_param: ArbTokensParam::default(),
}
}
}
impl Arbitrary for ErgoBoxAssetsData {
type Parameters = ArbErgoBoxAssetsDataParam;
fn arbitrary_with(args: Self::Parameters) -> Self::Strategy {
(
any_with::<BoxValue>(args.value_range),
vec(
any_with::<Token>(args.tokens_param.token_id_param),
args.tokens_param.token_count_range,
),
)
.prop_map(|(value, tokens)| Self {
value,
tokens: BoxTokens::from_vec(tokens).ok(),
})
.boxed()
}
type Strategy = BoxedStrategy<Self>;
fn arbitrary() -> Self::Strategy {
Self::arbitrary_with(Default::default())
}
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::panic)]
mod tests {
use ergotree_ir::chain::ergo_box::box_value::BoxValue;
use ergotree_ir::chain::ergo_box::BoxTokens;
use ergotree_ir::chain::token::Token;
use proptest::prelude::*;
use sigma_test_util::force_any_val;
use crate::wallet::box_selector::sum_tokens;
use crate::wallet::box_selector::sum_tokens_from_boxes;
use super::ErgoBoxAssetsData;
#[test]
fn test_sum_tokens_repeating_token_id() {
let token = force_any_val::<Token>();
let b = ErgoBoxAssetsData {
value: BoxValue::SAFE_USER_MIN,
tokens: BoxTokens::from_vec(vec![token.clone(), token.clone()]).ok(),
};
assert_eq!(
u64::from(
*sum_tokens_from_boxes(vec![b.clone(), b].as_slice())
.unwrap()
.get(&token.token_id)
.unwrap()
),
u64::from(token.amount) * 4
);
}
proptest! {
#[test]
fn sum_tokens_eq(b in any::<ErgoBoxAssetsData>()) {
prop_assert_eq!(sum_tokens(b.tokens.as_ref().map(BoxTokens::as_ref)), sum_tokens_from_boxes(vec![b].as_slice()))
}
}
}