feat: Support other reward token types

contracts/reward-distributor/src/contract.rs

@@ -7,7 +7,7 @@ use cw_dex::astroport::AstroportPool;
 use cw_vault_standard::{VaultContract, VaultContractUnchecked};
 use neutron_astroport_reward_distributor::{
     Config, ConfigUnchecked, ContractError, ExecuteMsg, InstantiateMsg, InternalMsg, QueryMsg,
-    StateResponse, CONFIG, LAST_DISTRIBUTED, REWARD_POOL, REWARD_VAULT,
+    RewardInfo, RewardType, StateResponse, CONFIG, LAST_DISTRIBUTED, REWARD_TOKEN,
 };
 
 use crate::execute;
@@ -25,24 +25,41 @@ pub fn instantiate(
     cw2::set_contract_version(deps.storage, CONTRACT_NAME, CONTRACT_VERSION)?;
     cw_ownable::initialize_owner(deps.storage, deps.api, Some(&msg.owner))?;
 
-    let reward_vault: VaultContract =
-        VaultContractUnchecked::new(&msg.reward_vault_addr).check(deps.as_ref())?;
+    let reward_token = match msg.reward_token_info {
+        RewardInfo::VaultAddr(reward_vault_addr) => {
+            let reward_vault: VaultContract =
+                VaultContractUnchecked::new(&reward_vault_addr).check(deps.as_ref())?;
 
-    // Validate reward vault base token as CW20 Astroport LP token
-    let reward_lp_token = deps
-        .api
-        .addr_validate(&reward_vault.base_token)
-        .map_err(|_| StdError::generic_err("Invalid base token of reward vault"))?;
+            // Validate reward vault base token as CW20 Astroport LP token
+            let reward_lp_token = deps
+                .api
+                .addr_validate(&reward_vault.base_token)
+                .map_err(|_| StdError::generic_err("Invalid base token of reward vault"))?;
 
-    // Query minter of LP token to get reward pool address
-    let minter_res: MinterResponse = deps.querier.query(&QueryRequest::Wasm(WasmQuery::Smart {
-        contract_addr: reward_lp_token.to_string(),
-        msg: to_binary(&Cw20QueryMsg::Minter {})?,
-    }))?;
-    let reward_pool_addr = deps.api.addr_validate(&minter_res.minter)?;
+            // Query minter of LP token to get reward pool address
+            let minter_res: MinterResponse =
+                deps.querier.query(&QueryRequest::Wasm(WasmQuery::Smart {
+                    contract_addr: reward_lp_token.to_string(),
+                    msg: to_binary(&Cw20QueryMsg::Minter {})?,
+                }))?;
+            let reward_pool_addr = deps.api.addr_validate(&minter_res.minter)?;
 
-    // Query reward pool for pool info to create pool object
-    let reward_pool = AstroportPool::new(deps.as_ref(), reward_pool_addr)?;
+            // Query reward pool for pool info to create pool object
+            let reward_pool = AstroportPool::new(deps.as_ref(), reward_pool_addr)?;
+
+            RewardType::Vault {
+                vault: reward_vault,
+                pool: reward_pool,
+            }
+        }
+        RewardInfo::AstroportPoolAddr(pool_addr) => {
+            let reward_pool =
+                AstroportPool::new(deps.as_ref(), deps.api.addr_validate(&pool_addr)?)?;
+
+            RewardType::LP(reward_pool)
+        }
+        RewardInfo::NativeCoin(reward_coin_denom) => RewardType::Coin(reward_coin_denom),
+    };
 
     // Create config
     let config: Config = ConfigUnchecked {
@@ -54,8 +71,7 @@ pub fn instantiate(
 
     CONFIG.save(deps.storage, &config)?;
     LAST_DISTRIBUTED.save(deps.storage, &env.block.time.seconds())?;
-    REWARD_POOL.save(deps.storage, &reward_pool)?;
-    REWARD_VAULT.save(deps.storage, &reward_vault)?;
+    REWARD_TOKEN.save(deps.storage, &reward_token)?;
 
     Ok(Response::default())
 }
@@ -103,14 +119,11 @@ pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult<Binary> {
         }
         QueryMsg::State {} => {
             let config = CONFIG.load(deps.storage)?;
-            let reward_pool = REWARD_POOL.load(deps.storage)?;
-            let reward_vault = REWARD_VAULT.load(deps.storage)?;
             let last_distributed = LAST_DISTRIBUTED.load(deps.storage)?;
 
             to_binary(&StateResponse {
                 config,
-                reward_pool,
-                reward_vault,
+                reward_token: REWARD_TOKEN.load(deps.storage)?,
                 last_distributed,
             })
         }

contracts/reward-distributor/src/execute.rs

@@ -1,8 +1,10 @@
 use apollo_cw_asset::{Asset, AssetInfo, AssetList};
-use cosmwasm_std::{Deps, DepsMut, Env, Event, MessageInfo, Response, Uint128};
+use cosmwasm_std::{
+    coins, BankMsg, CosmosMsg, Deps, DepsMut, Env, Event, MessageInfo, Response, Uint128,
+};
 use cw_dex::traits::Pool as PoolTrait;
 use neutron_astroport_reward_distributor::{
-    ConfigUpdates, ContractError, InternalMsg, CONFIG, LAST_DISTRIBUTED, REWARD_POOL, REWARD_VAULT,
+    ConfigUpdates, ContractError, InternalMsg, RewardType, CONFIG, LAST_DISTRIBUTED, REWARD_TOKEN,
 };
 
 pub fn execute_distribute(deps: DepsMut, env: Env) -> Result<Response, ContractError> {
@@ -22,55 +24,77 @@ pub fn execute_distribute(deps: DepsMut, env: Env) -> Result<Response, ContractE
 
     // Calculate amount of rewards to be distributed
     let time_elapsed = current_time.saturating_sub(last_distributed.max(config.rewards_start_time));
-    let redeem_amount = config.emission_per_second * Uint128::from(time_elapsed);
-
-    // Query the vault to see how many base tokens would be returned after
-    // redeeming. If zero we return Ok, so that update_config does not fail when
-    // trying to distribute.
-    let base_token_amount = REWARD_VAULT
-        .load(deps.storage)?
-        .query_convert_to_assets(&deps.querier, redeem_amount)?;
-    if base_token_amount.is_zero() {
-        return Ok(Response::new());
-    }
-
-    // Check contract's balance of vault tokens and error if not enough. This is
-    // just so we get a clearer error message rather than the confusing "cannot
-    // sub 0 with x".
-    let reward_vault = REWARD_VAULT.load(deps.storage)?;
-    let vault_token_balance = deps
-        .querier
-        .query_balance(&env.contract.address, &reward_vault.vault_token)?;
-    if vault_token_balance.amount < redeem_amount {
-        return Err(ContractError::InsufficientVaultTokenBalance {
-            vault_token_balance: vault_token_balance.amount,
-            redeem_amount,
-        });
+    let reward_amount = config.emission_per_second * Uint128::from(time_elapsed);
+
+    let reward_token = REWARD_TOKEN.load(deps.storage)?;
+
+    let mut res = Response::new();
+
+    match reward_token {
+        RewardType::Vault { vault, pool: _ } => {
+            // Query the vault to see how many base tokens would be returned after
+            // redeeming. If zero we return Ok, so that update_config does not fail when
+            // trying to distribute.
+            let base_token_amount = vault.query_convert_to_assets(&deps.querier, reward_amount)?;
+            if base_token_amount.is_zero() {
+                return Ok(Response::new());
+            }
+
+            // Check contract's balance of vault tokens and error if not enough. This is
+            // just so we get a clearer error message rather than the confusing "cannot
+            // sub 0 with x".
+            let vault_token_balance = deps
+                .querier
+                .query_balance(&env.contract.address, &vault.vault_token)?;
+            if vault_token_balance.amount < reward_amount {
+                return Err(ContractError::InsufficientVaultTokenBalance {
+                    vault_token_balance: vault_token_balance.amount,
+                    redeem_amount: reward_amount,
+                });
+            }
+
+            // Redeem rewards from the vault
+            let redeem_msg = vault.redeem(reward_amount, None)?;
+
+            // Create internal callback msg
+            let callback_msg = InternalMsg::VaultTokensRedeemed {}.into_cosmos_msg(&env)?;
+
+            res = res.add_message(redeem_msg).add_message(callback_msg);
+        }
+        RewardType::LP(pool) => {
+            // Create message to withdraw liquidity from pool
+            let lp_tokens = Asset::new(AssetInfo::Cw20(pool.lp_token_addr.clone()), reward_amount);
+            res = pool.withdraw_liquidity(deps.as_ref(), &env, lp_tokens, AssetList::new())?;
+
+            // Create internal callback msg
+            let callback_msg = InternalMsg::LpRedeemed {}.into_cosmos_msg(&env)?;
+            res = res.add_message(callback_msg);
+        }
+        RewardType::Coin(reward_coin_denom) => {
+            // Create message to send coins to distribution address
+            let send_msg: CosmosMsg = BankMsg::Send {
+                to_address: config.distribution_addr.to_string(),
+                amount: coins(reward_amount.u128(), reward_coin_denom),
+            }
+            .into();
+            res = res.add_message(send_msg);
+        }
     }
 
     // Set last distributed time to current time
     LAST_DISTRIBUTED.save(deps.storage, &current_time)?;
 
-    // Redeem rewards from the vault
-    let redeem_msg = reward_vault.redeem(redeem_amount, None)?;
-
-    // Create internal callback msg
-    let callback_msg = InternalMsg::VaultTokensRedeemed {}.into_cosmos_msg(&env)?;
-
     let event = Event::new("apollo/neutron-astroport-reward-distributor/execute_distribute")
-        .add_attribute("vault_tokens_redeemed", redeem_amount);
+        .add_attribute("vault_tokens_redeemed", reward_amount);
 
-    Ok(Response::default()
-        .add_message(redeem_msg)
-        .add_message(callback_msg)
-        .add_event(event))
+    Ok(res.add_event(event))
 }
 
 pub fn execute_internal_vault_tokens_redeemed(
     deps: Deps,
     env: Env,
 ) -> Result<Response, ContractError> {
-    let reward_pool = REWARD_POOL.load(deps.storage)?;
+    let reward_pool = REWARD_TOKEN.load(deps.storage)?.into_pool()?;
 
     // Query lp token balance
     let reward_lp_token = AssetInfo::Cw20(reward_pool.lp_token_addr.clone());
@@ -93,7 +117,7 @@ pub fn execute_internal_vault_tokens_redeemed(
 
 pub fn execute_internal_lp_redeemed(deps: Deps, env: Env) -> Result<Response, ContractError> {
     let config = CONFIG.load(deps.storage)?;
-    let reward_pool = REWARD_POOL.load(deps.storage)?;
+    let reward_pool = REWARD_TOKEN.load(deps.storage)?.into_pool()?;
 
     // Query contracts balances of pool assets
     let pool_asset_balances: AssetList = AssetList::query_asset_info_balances(

contracts/reward-distributor/tests/integration_test.rs

@@ -3,10 +3,10 @@ use cosmwasm_std::{coin, Uint128};
 use cw_it::helpers::Unwrap;
 use cw_it::test_tube::Account;
 use cw_it::traits::CwItRunner;
-
 use locked_astroport_vault::helpers::INITIAL_VAULT_TOKENS_PER_BASE_TOKEN;
 use locked_astroport_vault_test_helpers::cw_vault_standard_test_helpers::traits::CwVaultStandardRobot;
 use locked_astroport_vault_test_helpers::robot::LockedAstroportVaultRobot;
+use neutron_astroport_reward_distributor::RewardType;
 use neutron_astroport_reward_distributor_test_helpers as test_helpers;
 
 use test_helpers::robot::RewardDistributorRobot;
@@ -39,16 +39,10 @@ fn test_initialization() {
     let config = state.config;
     assert_eq!(config.emission_per_second, Uint128::from(1000000u128));
     assert_eq!(config.distribution_addr, robot.distribution_acc.address());
-    assert_eq!(state.reward_pool, robot.reward_pool);
-    assert_eq!(
-        state.reward_pool.lp_token_addr.to_string(),
-        robot.reward_vault_robot.base_token()
-    );
-    assert_eq!(state.reward_vault.addr, robot.reward_vault_robot.vault_addr);
-    assert_eq!(
-        state.reward_vault.vault_token,
-        robot.reward_vault_robot.vault_token()
-    );
+    assert!(matches!(
+        state.reward_token,
+        RewardType::Vault { vault, pool } if vault.addr == robot.reward_vault_robot.vault_addr && pool.lp_token_addr == robot.reward_vault_robot.base_token()
+    ));
 
     // Query ownership
     let ownership = robot.query_ownership();

packages/reward-distributor/src/error.rs

@@ -21,3 +21,9 @@ pub enum ContractError {
         redeem_amount: Uint128,
     },
 }
+
+impl ContractError {
+    pub fn generic_err(msg: &str) -> Self {
+        StdError::generic_err(msg).into()
+    }
+}

packages/reward-distributor/src/msg.rs

@@ -1,19 +1,31 @@
 use cosmwasm_schema::{cw_serde, QueryResponses};
 use cosmwasm_std::{to_binary, CosmosMsg, Env, StdResult, Uint128, WasmMsg};
-use cw_dex::astroport::AstroportPool;
 use cw_ownable::{cw_ownable_execute, cw_ownable_query};
-use cw_vault_standard::VaultContract;
 
-use crate::{Config, ConfigUpdates};
+use crate::{Config, ConfigUpdates, RewardType};
+
+/// An enum for the information needed to instantiate the contract depending on
+/// the type of reward token used.
+#[cw_serde]
+pub enum RewardInfo {
+    /// The address of the vault if the reward token is a vault token
+    VaultAddr(String),
+    /// The address of the Astroport pool if the reward token is an Astroport LP
+    /// token
+    AstroportPoolAddr(String),
+    /// The denom of the native coin if the reward token is a native coin
+    NativeCoin(String),
+}
 
 #[cw_serde]
 pub struct InstantiateMsg {
     /// The account to be appointed the contract owner
     pub owner: String,
     /// The emission rate per second
     pub emission_per_second: Uint128,
-    /// The address of the vault contract in which rewards are being held
-    pub reward_vault_addr: String,
+    /// The info needed to instantiate the contract depending on the type of
+    /// reward token used
+    pub reward_token_info: RewardInfo,
     /// The address that rewards are being distributed to
     pub distribution_addr: String,
     /// The unix timestamp at which rewards start being distributed
@@ -67,7 +79,6 @@ pub enum QueryMsg {
 /// The response to a config query
 pub struct StateResponse {
     pub config: Config,
-    pub reward_pool: AstroportPool,
-    pub reward_vault: VaultContract,
+    pub reward_token: RewardType,
     pub last_distributed: u64,
 }

packages/reward-distributor/src/state.rs

@@ -1,17 +1,52 @@
+use cosmwasm_schema::cw_serde;
 use cw_dex::astroport::AstroportPool;
 use cw_storage_plus::Item;
 use cw_vault_standard::VaultContract;
 
 use crate::config::Config;
+use crate::ContractError;
+
+/// An enum representing different types of reward tokens
+#[cw_serde]
+pub enum RewardType {
+    /// The reward token is a vault token
+    Vault {
+        /// The vault contract
+        vault: VaultContract,
+        /// The Astroport pool that the vault holds liquidity in
+        pool: AstroportPool,
+    },
+    /// The reward token is an Astroport LP token
+    LP(AstroportPool),
+    /// The reward token is a native coin
+    Coin(String),
+}
+
+impl RewardType {
+    pub fn into_pool(self) -> Result<AstroportPool, ContractError> {
+        match self {
+            RewardType::Vault { vault: _, pool } => Ok(pool),
+            RewardType::LP(pool) => Ok(pool),
+            RewardType::Coin(_) => {
+                Err(ContractError::generic_err(
+                    "Cannot redeem vault tokens from coin reward",
+                ))
+            }
+        }
+    }
+}
 
 /// Stores the contract's config
 pub const CONFIG: Item<Config> = Item::new("config");
 
-/// Stores the Astroport pool in which rewards are being held
-pub const REWARD_POOL: Item<AstroportPool> = Item::new("reward_pool");
+/// Stores the reward token that this contract is distributing
+pub const REWARD_TOKEN: Item<RewardType> = Item::new("reward_token");
+
+// /// Stores the Astroport pool in which rewards are being held
+// pub const REWARD_POOL: Item<AstroportPool> = Item::new("reward_pool");
 
-/// Stores the vault contract in which rewards are being held
-pub const REWARD_VAULT: Item<VaultContract> = Item::new("reward_vault");
+// /// Stores the vault contract in which rewards are being held
+// pub const REWARD_VAULT: Item<VaultContract> = Item::new("reward_vault");
 
 /// Stores the last timestamp that rewards were distributed
 pub const LAST_DISTRIBUTED: Item<u64> = Item::new("last_distributed");

packages/test-helpers/src/robot.rs

@@ -22,7 +22,7 @@ use neutron_astroport_reward_distributor::InstantiateMsg;
 
 #[cfg(feature = "osmosis-test-tube")]
 use cw_it::Artifact;
-use reward_distributor::{Config, ConfigUpdates, QueryMsg};
+use reward_distributor::{Config, ConfigUpdates, QueryMsg, RewardInfo};
 
 pub const REWARD_DISTRIBUTOR_WASM_NAME: &str = "neutron_astroport_reward_distributor_contract.wasm";
 
@@ -109,7 +109,7 @@ impl<'a> RewardDistributorRobot<'a> {
             distribution_addr: distribution_acc.address(),
             emission_per_second: emission_per_second.into(),
             owner: admin.address(),
-            reward_vault_addr: reward_vault_robot.vault_addr.clone(),
+            reward_token_info: RewardInfo::VaultAddr(reward_vault_robot.vault_addr.clone()),
             rewards_start_time,
         };
         let contract_addr = Wasm::new(runner)