Add flexible asset withdrawal with tests

.codecov.yml

@@ -11,3 +11,6 @@ coverage:
     patch:
       default:
         enabled: no
+
+ignore:
+  - "contracts/infinity-builder/.*"

contracts/infinity-pair/src/execute.rs

@@ -6,12 +6,12 @@ use crate::state::{BondingCurve, PairType, TokenId, INFINITY_GLOBAL, NFT_DEPOSIT
 
 use cosmwasm_std::{
     coin, ensure, ensure_eq, has_coins, Addr, Coin, DepsMut, Env, MessageInfo, Order, StdResult,
-    Uint128,
 };
+use cw721::{Cw721QueryMsg, TokensResponse};
 use cw_utils::{maybe_addr, must_pay, nonpayable};
 use infinity_shared::InfinityError;
 use sg_marketplace_common::address::address_or;
-use sg_marketplace_common::coin::transfer_coin;
+use sg_marketplace_common::coin::transfer_coins;
 use sg_marketplace_common::nft::{only_with_owner_approval, transfer_nft};
 use sg_std::Response;
 
@@ -58,34 +58,56 @@ pub fn handle_execute_msg(
             execute_receive_nft(deps, info, pair, cw721_receive_msg.token_id)
         },
         ExecuteMsg::WithdrawNfts {
+            collection,
             token_ids,
+            asset_recipient,
         } => {
             nonpayable(&info)?;
             only_pair_owner(&info.sender, &pair)?;
-            execute_withdraw_nfts(deps, info, pair, token_ids)
+            execute_withdraw_nfts(
+                deps,
+                info,
+                pair,
+                api.addr_validate(&collection)?,
+                token_ids,
+                maybe_addr(api, asset_recipient)?,
+            )
         },
         ExecuteMsg::WithdrawAnyNfts {
+            collection,
             limit,
+            asset_recipient,
         } => {
             nonpayable(&info)?;
             only_pair_owner(&info.sender, &pair)?;
-            execute_withdraw_any_nfts(deps, info, pair, limit)
+            execute_withdraw_any_nfts(
+                deps,
+                env,
+                info,
+                pair,
+                api.addr_validate(&collection)?,
+                limit,
+                maybe_addr(api, asset_recipient)?,
+            )
         },
         ExecuteMsg::DepositTokens {} => {
             only_pair_owner(&info.sender, &pair)?;
             execute_deposit_tokens(deps, info, env, pair)
         },
         ExecuteMsg::WithdrawTokens {
-            amount,
+            funds,
+            asset_recipient,
         } => {
             nonpayable(&info)?;
             only_pair_owner(&info.sender, &pair)?;
-            execute_withdraw_tokens(deps, info, env, pair, amount)
+            execute_withdraw_tokens(deps, info, env, pair, funds, maybe_addr(api, asset_recipient)?)
         },
-        ExecuteMsg::WithdrawAllTokens {} => {
+        ExecuteMsg::WithdrawAllTokens {
+            asset_recipient,
+        } => {
             nonpayable(&info)?;
             only_pair_owner(&info.sender, &pair)?;
-            execute_withdraw_all_tokens(deps, info, env, pair)
+            execute_withdraw_all_tokens(deps, info, env, pair, maybe_addr(api, asset_recipient)?)
         },
         ExecuteMsg::UpdatePairConfig {
             is_active,
@@ -182,7 +204,9 @@ pub fn execute_withdraw_nfts(
     deps: DepsMut,
     _info: MessageInfo,
     mut pair: Pair,
+    collection: Addr,
     token_ids: Vec<String>,
+    asset_recipient: Option<Addr>,
 ) -> Result<(Pair, Response), ContractError> {
     ensure!(
         !token_ids.is_empty(),
@@ -191,35 +215,44 @@ pub fn execute_withdraw_nfts(
 
     let mut response = Response::new();
 
-    for token_id in &token_ids {
-        ensure!(
-            NFT_DEPOSITS.has(deps.storage, token_id.to_string()),
-            InfinityError::InvalidInput("token_id is not owned by pair".to_string())
-        );
-        NFT_DEPOSITS.remove(deps.storage, token_id.to_string());
-
-        pair.internal.total_nfts -= 1u64;
+    let asset_recipient = address_or(asset_recipient.as_ref(), &pair.asset_recipient());
 
-        response =
-            transfer_nft(&pair.immutable.collection, token_id, &pair.asset_recipient(), response);
+    for token_id in &token_ids {
+        response = transfer_nft(&collection, token_id, &asset_recipient, response);
+
+        if pair.immutable.collection == collection
+            && NFT_DEPOSITS.has(deps.storage, token_id.to_string())
+        {
+            pair.internal.total_nfts -= 1u64;
+            NFT_DEPOSITS.remove(deps.storage, token_id.to_string());
+        }
     }
 
     Ok((pair, response))
 }
 
 pub fn execute_withdraw_any_nfts(
     deps: DepsMut,
+    env: Env,
     info: MessageInfo,
     pair: Pair,
+    collection: Addr,
     limit: u32,
+    asset_recipient: Option<Addr>,
 ) -> Result<(Pair, Response), ContractError> {
-    let token_ids = NFT_DEPOSITS
-        .range(deps.storage, None, None, Order::Ascending)
-        .take(limit as usize)
-        .map(|item| item.map(|(v, _)| v))
-        .collect::<StdResult<Vec<String>>>()?;
-
-    execute_withdraw_nfts(deps, info, pair, token_ids)
+    let token_ids = deps
+        .querier
+        .query_wasm_smart::<TokensResponse>(
+            &collection,
+            &Cw721QueryMsg::Tokens {
+                owner: env.contract.address.to_string(),
+                start_after: None,
+                limit: Some(limit),
+            },
+        )?
+        .tokens;
+
+    execute_withdraw_nfts(deps, info, pair, collection, token_ids, asset_recipient)
 }
 
 pub fn execute_deposit_tokens(
@@ -237,22 +270,20 @@ pub fn execute_withdraw_tokens(
     _info: MessageInfo,
     _env: Env,
     mut pair: Pair,
-    amount: Uint128,
+    funds: Vec<Coin>,
+    asset_recipient: Option<Addr>,
 ) -> Result<(Pair, Response), ContractError> {
-    ensure!(
-        amount <= pair.total_tokens,
-        InfinityError::InvalidInput("amount exceeds total tokens".to_string())
-    );
-
-    pair.total_tokens -= amount;
+    for fund in &funds {
+        if fund.denom == pair.immutable.denom {
+            pair.total_tokens -= fund.amount;
+        }
+    }
 
     let mut response = Response::new();
 
-    response = transfer_coin(
-        coin(amount.u128(), pair.immutable.denom.clone()),
-        &pair.asset_recipient(),
-        response,
-    );
+    let asset_recipient = address_or(asset_recipient.as_ref(), &pair.asset_recipient());
+
+    response = transfer_coins(funds, &asset_recipient, response);
 
     Ok((pair, response))
 }
@@ -262,9 +293,10 @@ pub fn execute_withdraw_all_tokens(
     info: MessageInfo,
     env: Env,
     pair: Pair,
+    asset_recipient: Option<Addr>,
 ) -> Result<(Pair, Response), ContractError> {
-    let total_tokens = pair.total_tokens;
-    execute_withdraw_tokens(deps, info, env, pair, total_tokens)
+    let all_tokens = deps.querier.query_all_balances(&env.contract.address)?;
+    execute_withdraw_tokens(deps, info, env, pair, all_tokens, asset_recipient)
 }
 
 #[allow(clippy::too_many_arguments)]

contracts/infinity-pair/src/msg.rs

@@ -32,20 +32,27 @@ pub enum ExecuteMsg {
     ReceiveNft(Cw721ReceiveMsg),
     /// Withdraw NFTs from the pair
     WithdrawNfts {
+        collection: String,
         token_ids: Vec<TokenId>,
+        asset_recipient: Option<String>,
     },
     /// Withdraw any NFTs, from the pair
     WithdrawAnyNfts {
+        collection: String,
         limit: u32,
+        asset_recipient: Option<String>,
     },
     /// Deposit tokens into the pair
     DepositTokens {},
     /// Withdraw tokens from the pair
     WithdrawTokens {
-        amount: Uint128,
+        funds: Vec<Coin>,
+        asset_recipient: Option<String>,
     },
     /// Withdraw all tokens from the pair
-    WithdrawAllTokens {},
+    WithdrawAllTokens {
+        asset_recipient: Option<String>,
+    },
     /// Update the parameters of a pair
     UpdatePairConfig {
         is_active: Option<bool>,
@@ -68,11 +75,6 @@ pub enum ExecuteMsg {
     SwapTokensForAnyNft {
         asset_recipient: Option<String>,
     },
-    // /// Remove a pair from contract storage and indexing
-    // RemovePair {
-    //     pair_id: u64,
-    //     asset_recipient: Option<String>,
-    // },
 }
 
 #[cw_serde]
@@ -84,10 +86,24 @@ pub enum QueryMsg {
     NftDeposits {
         query_options: Option<QueryOptions<String>>,
     },
+    #[returns(QuotesResponse)]
+    SellToPairQuotes {
+        limit: u32,
+    },
+    #[returns(QuotesResponse)]
+    BuyFromPairQuotes {
+        limit: u32,
+    },
 }
 
 #[cw_serde]
 pub struct NftDepositsResponse {
     pub collection: Addr,
     pub token_ids: Vec<TokenId>,
 }
+
+#[cw_serde]
+pub struct QuotesResponse {
+    pub denom: String,
+    pub quotes: Vec<Uint128>,
+}

contracts/infinity-pair/src/query.rs

@@ -1,11 +1,11 @@
 use crate::{
-    helpers::load_pair,
-    msg::{NftDepositsResponse, QueryMsg},
+    helpers::{load_pair, load_payout_context},
+    msg::{NftDepositsResponse, QueryMsg, QuotesResponse},
     pair::Pair,
-    state::{NFT_DEPOSITS, PAIR_IMMUTABLE},
+    state::{INFINITY_GLOBAL, NFT_DEPOSITS, PAIR_IMMUTABLE},
 };
 
-use cosmwasm_std::{to_binary, Binary, Deps, Env, StdError, StdResult};
+use cosmwasm_std::{to_binary, Binary, Deps, Env, StdError, StdResult, Uint128};
 use sg_index_query::{QueryOptions, QueryOptionsInternal};
 
 #[cfg(not(feature = "library"))]
@@ -18,6 +18,12 @@ pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult<Binary> {
         QueryMsg::NftDeposits {
             query_options,
         } => to_binary(&query_nft_deposits(deps, query_options.unwrap_or_default())?),
+        QueryMsg::SellToPairQuotes {
+            limit,
+        } => to_binary(&query_sell_to_pair_quotes(deps, env, limit)?),
+        QueryMsg::BuyFromPairQuotes {
+            limit,
+        } => to_binary(&query_buy_from_pair_quotes(deps, env, limit)?),
     }
 }
 
@@ -52,3 +58,71 @@ pub fn query_nft_deposits(
         token_ids,
     })
 }
+
+pub fn query_sell_to_pair_quotes(deps: Deps, env: Env, limit: u32) -> StdResult<QuotesResponse> {
+    let mut pair = load_pair(&env.contract.address, deps.storage, &deps.querier)
+        .map_err(|_| StdError::generic_err("failed to load pair".to_string()))?;
+
+    let infinity_global = INFINITY_GLOBAL.load(deps.storage)?;
+    let payout_context = load_payout_context(
+        deps,
+        &infinity_global,
+        &pair.immutable.collection,
+        &pair.immutable.denom,
+    )
+    .map_err(|_| StdError::generic_err("failed to load payout context".to_string()))?;
+
+    let mut quotes: Vec<Uint128> = vec![];
+
+    let mut idx = 0u32;
+    while idx < limit {
+        if let Some(quote_summary) = &pair.internal.sell_to_pair_quote_summary {
+            quotes.push(quote_summary.seller_amount);
+        } else {
+            break;
+        }
+
+        pair.sim_swap_nft_for_tokens(&payout_context);
+
+        idx += 1;
+    }
+
+    Ok(QuotesResponse {
+        denom: pair.immutable.denom,
+        quotes,
+    })
+}
+
+pub fn query_buy_from_pair_quotes(deps: Deps, env: Env, limit: u32) -> StdResult<QuotesResponse> {
+    let mut pair = load_pair(&env.contract.address, deps.storage, &deps.querier)
+        .map_err(|_| StdError::generic_err("failed to load pair".to_string()))?;
+
+    let infinity_global = INFINITY_GLOBAL.load(deps.storage)?;
+    let payout_context = load_payout_context(
+        deps,
+        &infinity_global,
+        &pair.immutable.collection,
+        &pair.immutable.denom,
+    )
+    .map_err(|_| StdError::generic_err("failed to load payout context".to_string()))?;
+
+    let mut quotes: Vec<Uint128> = vec![];
+
+    let mut idx = 0u32;
+    while idx < limit {
+        if let Some(quote_summary) = &pair.internal.buy_from_pair_quote_summary {
+            quotes.push(quote_summary.total());
+        } else {
+            break;
+        }
+
+        pair.sim_swap_tokens_for_nft(&payout_context);
+
+        idx += 1;
+    }
+
+    Ok(QuotesResponse {
+        denom: pair.immutable.denom,
+        quotes,
+    })
+}

unit-tests/src/helpers/nft_functions.rs

@@ -72,19 +72,19 @@ pub fn _approve_all(
     assert!(res.is_ok());
 }
 
-pub fn _transfer(
+pub fn transfer(
     router: &mut StargazeApp,
     creator: &Addr,
     recipient: &Addr,
     collection: &Addr,
-    token_id: u32,
+    token_id: &String,
 ) {
     let transfer_msg: Sg721ExecuteMsg<Empty, Empty> = Sg721ExecuteMsg::TransferNft {
         recipient: recipient.to_string(),
         token_id: token_id.to_string(),
     };
-    let res = router.execute_contract(creator.clone(), collection.clone(), &transfer_msg, &[]);
-    assert!(res.is_ok());
+    let response = router.execute_contract(creator.clone(), collection.clone(), &transfer_msg, &[]);
+    assert!(response.is_ok());
 }
 
 pub fn _burn(router: &mut StargazeApp, creator: &Addr, collection: &Addr, token_id: u32) {