Track shares

contracts/transmuter/src/contract.rs

@@ -1,7 +1,10 @@
 use cosmwasm_schema::cw_serde;
-use cosmwasm_std::{ensure_eq, BankMsg, Coin, Deps, DepsMut, Env, MessageInfo, Response, StdError};
+use cosmwasm_std::{
+    ensure, ensure_eq, Addr, BankMsg, Coin, Deps, DepsMut, Env, MessageInfo, Response, StdError,
+    Uint128,
+};
 use cw_controllers::{Admin, AdminResponse};
-use cw_storage_plus::Item;
+use cw_storage_plus::{Item, Map};
 use sylvia::contract;
 
 use crate::{error::ContractError, transmuter_pool::TransmuterPool};
@@ -12,6 +15,7 @@ const CONTRACT_VERSION: &str = env!("CARGO_PKG_VERSION");
 
 pub struct Transmuter<'a> {
     pub(crate) pool: Item<'a, TransmuterPool>,
+    pub(crate) shares: Map<'a, &'a Addr, Uint128>,
     pub(crate) admin: Admin<'a>,
 }
 
@@ -22,6 +26,7 @@ impl Transmuter<'_> {
         Self {
             pool: Item::new("pool"),
             admin: Admin::new("admin"),
+            shares: Map::new("shares"),
         }
     }
 
@@ -66,10 +71,21 @@ impl Transmuter<'_> {
         // ensure funds length == 1
         ensure_eq!(info.funds.len(), 1, ContractError::SingleCoinExpected {});
 
+        let supplying_coin = info.funds[0].clone();
+
+        // update shares
+        self.shares.update(
+            deps.storage,
+            &info.sender,
+            |shares| -> Result<Uint128, StdError> {
+                Ok(shares.unwrap_or_default() + supplying_coin.amount)
+            },
+        )?;
+
         // update pool
         self.pool
             .update(deps.storage, |mut pool| -> Result<_, ContractError> {
-                pool.supply(&info.funds[0])?;
+                pool.supply(&supplying_coin)?;
                 Ok(pool)
             })?;
 
@@ -129,17 +145,39 @@ impl Transmuter<'_> {
     ) -> Result<Response, ContractError> {
         let (deps, _env, info) = ctx;
 
-        // allow only admin to withdraw
-        self.admin
-            .assert_admin(deps.as_ref(), &info.sender)
-            .map_err(|_| ContractError::Unauthorized {})?;
+        // check if sender's shares is enough
+        let sender_shares = self
+            .shares
+            .may_load(deps.storage, &info.sender)?
+            .unwrap_or_default();
+
+        let required_shares = coins
+            .iter()
+            .fold(Uint128::zero(), |acc, curr| acc + curr.amount);
+
+        ensure!(
+            sender_shares >= required_shares,
+            ContractError::InsufficientShares {
+                required: required_shares,
+                available: sender_shares
+            }
+        );
+
+        // update shares
+        self.shares.update(
+            deps.storage,
+            &info.sender,
+            |sender_shares| -> Result<Uint128, StdError> {
+                Ok(sender_shares.unwrap_or_default() - required_shares)
+            },
+        )?;
 
         // withdraw
-        let mut pool = self.pool.load(deps.storage)?;
-        pool.withdraw(&coins)?;
-
-        // save pool
-        self.pool.save(deps.storage, &pool)?;
+        self.pool
+            .update(deps.storage, |mut pool| -> Result<_, ContractError> {
+                pool.withdraw(&coins)?;
+                Ok(pool)
+            })?;
 
         let bank_send_msg = BankMsg::Send {
             to_address: info.sender.to_string(),
@@ -166,6 +204,22 @@ impl Transmuter<'_> {
             pool: self.pool.load(deps.storage)?,
         })
     }
+
+    #[msg(query)]
+    fn shares(&self, ctx: (Deps, Env), address: String) -> Result<SharesResponse, ContractError> {
+        let (deps, _env) = ctx;
+        Ok(SharesResponse {
+            shares: self
+                .shares
+                .may_load(deps.storage, &deps.api.addr_validate(&address)?)?
+                .unwrap_or_default(),
+        })
+    }
+}
+
+#[cw_serde]
+pub struct SharesResponse {
+    pub shares: Uint128,
 }
 
 #[cw_serde]

contracts/transmuter/src/error.rs

@@ -1,4 +1,4 @@
-use cosmwasm_std::{Coin, StdError};
+use cosmwasm_std::{Coin, StdError, Uint128};
 use thiserror::Error;
 
 #[derive(Error, Debug, PartialEq)]
@@ -32,4 +32,10 @@ pub enum ContractError {
 
     #[error("Insufficient fund: required: {required}, available: {available}")]
     InsufficientFund { required: Coin, available: Coin },
+
+    #[error("Insufficient shares: required: {required}, available: {available}")]
+    InsufficientShares {
+        required: Uint128,
+        available: Uint128,
+    },
 }

contracts/transmuter/src/multitest/suite.rs

@@ -1,10 +1,10 @@
 use super::test_env::*;
 use crate::{
-    contract::{ExecMsg, InstantiateMsg, PoolResponse, QueryMsg},
+    contract::{ExecMsg, InstantiateMsg, PoolResponse, QueryMsg, SharesResponse},
     transmuter_pool::TransmuterPool,
     ContractError,
 };
-use cosmwasm_std::{Addr, Coin};
+use cosmwasm_std::{Addr, Coin, Uint128};
 use cw_controllers::AdminResponse;
 use cw_multi_test::Executor;
 
@@ -107,6 +107,20 @@ fn test_supply() {
             out_coin_reserve: supplied_amount[0].clone()
         }
     );
+
+    // check shares
+    let SharesResponse { shares } = t
+        .app
+        .wrap()
+        .query_wasm_smart(
+            t.contract.clone(),
+            &QueryMsg::Shares {
+                address: t.accounts["provider"].to_string(),
+            },
+        )
+        .unwrap();
+
+    assert_eq!(shares, supplied_amount[0].amount);
 }
 
 #[test]
@@ -342,7 +356,8 @@ fn test_admin() {
 fn test_withdraw() {
     let mut t = TestEnvBuilder::new()
         .with_account("user", vec![Coin::new(1_500, ETH_USDC)])
-        .with_account("provider", vec![Coin::new(200_000, COSMOS_USDC)])
+        .with_account("provider_1", vec![Coin::new(100_000, COSMOS_USDC)])
+        .with_account("provider_2", vec![Coin::new(100_000, COSMOS_USDC)])
         .with_instantiate_msg(InstantiateMsg {
             in_denom: ETH_USDC.to_string(),
             out_denom: COSMOS_USDC.to_string(),
@@ -353,7 +368,16 @@ fn test_withdraw() {
     // supply
     t.app
         .execute_contract(
-            t.accounts["provider"].clone(),
+            t.accounts["provider_1"].clone(),
+            t.contract.clone(),
+            &ExecMsg::Supply {},
+            &[Coin::new(100_000, COSMOS_USDC)],
+        )
+        .unwrap();
+
+    t.app
+        .execute_contract(
+            t.accounts["provider_2"].clone(),
             t.contract.clone(),
             &ExecMsg::Supply {},
             &[Coin::new(100_000, COSMOS_USDC)],
@@ -370,41 +394,59 @@ fn test_withdraw() {
         )
         .unwrap();
 
-    // non-admin cannot withdraw
+    // non-provider cannot withdraw
     let err = t
         .app
         .execute_contract(
             t.accounts["user"].clone(),
             t.contract.clone(),
             &ExecMsg::Withdraw {
-                coins: vec![Coin::new(1_000, COSMOS_USDC)],
+                coins: vec![Coin::new(1_500, ETH_USDC)],
             },
             &[],
         )
         .unwrap_err();
+
     assert_eq!(
         err.downcast_ref::<ContractError>().unwrap(),
-        &ContractError::Unauthorized {}
+        &ContractError::InsufficientShares {
+            required: 1500u128.into(),
+            available: Uint128::zero()
+        }
     );
 
-    // admin can withdraw
+    // provider can withdraw
     t.app
         .execute_contract(
-            Addr::unchecked("admin"),
+            t.accounts["provider_1"].clone(),
             t.contract.clone(),
             &ExecMsg::Withdraw {
-                coins: vec![Coin::new(1_000, COSMOS_USDC), Coin::new(1_000, ETH_USDC)],
+                coins: vec![Coin::new(500, ETH_USDC)],
             },
             &[],
         )
         .unwrap();
 
+    // check shares
+    let SharesResponse { shares } = t
+        .app
+        .wrap()
+        .query_wasm_smart(
+            t.contract.clone(),
+            &QueryMsg::Shares {
+                address: t.accounts["provider_1"].to_string(),
+            },
+        )
+        .unwrap();
+
+    assert_eq!(shares, Uint128::new(100_000 - 500));
+
     // check balances
     assert_eq!(
         t.app.wrap().query_all_balances(&t.contract).unwrap(),
         vec![
-            Coin::new(500, ETH_USDC),
-            Coin::new(100_000 - 1500 - 1000, COSMOS_USDC)
+            Coin::new(1500 - 500, ETH_USDC),
+            Coin::new(200_000 - 1500, COSMOS_USDC)
         ]
     );
 
@@ -417,37 +459,96 @@ fn test_withdraw() {
     assert_eq!(
         pool,
         TransmuterPool {
-            in_coin: Coin::new(500, ETH_USDC),
-            out_coin_reserve: Coin::new(100_000 - 1500 - 1000, COSMOS_USDC)
+            in_coin: Coin::new(1500 - 500, ETH_USDC),
+            out_coin_reserve: Coin::new(200_000 - 1500, COSMOS_USDC)
         }
     );
 
-    // check admin balance
+    // provider can withdraw both sides
+    t.app
+        .execute_contract(
+            t.accounts["provider_2"].clone(),
+            t.contract.clone(),
+            &ExecMsg::Withdraw {
+                coins: vec![Coin::new(1_000, ETH_USDC), Coin::new(99_000, COSMOS_USDC)],
+            },
+            &[],
+        )
+        .unwrap();
+
+    // check shares
+    let SharesResponse { shares } = t
+        .app
+        .wrap()
+        .query_wasm_smart(
+            t.contract.clone(),
+            &QueryMsg::Shares {
+                address: t.accounts["provider_2"].to_string(),
+            },
+        )
+        .unwrap();
+
+    assert_eq!(shares, Uint128::new(0));
+
+    // check balances
+    assert_eq!(
+        t.app.wrap().query_all_balances(&t.contract).unwrap(),
+        vec![Coin::new(200_000 - 1500 - 99_000, COSMOS_USDC)]
+    );
+
+    let PoolResponse { pool } = t
+        .app
+        .wrap()
+        .query_wasm_smart(t.contract.clone(), &QueryMsg::Pool {})
+        .unwrap();
+
+    assert_eq!(
+        pool,
+        TransmuterPool {
+            in_coin: Coin::new(0, ETH_USDC),
+            out_coin_reserve: Coin::new(200_000 - 1500 - 99_000, COSMOS_USDC)
+        }
+    );
+
+    // withdrawing excess shares fails
+    let err = t
+        .app
+        .execute_contract(
+            Addr::unchecked("provider_2"),
+            t.contract.clone(),
+            &ExecMsg::Withdraw {
+                coins: vec![Coin::new(1, ETH_USDC)],
+            },
+            &[],
+        )
+        .unwrap_err();
+
     assert_eq!(
-        t.app
-            .wrap()
-            .query_all_balances(&Addr::unchecked("admin"))
-            .unwrap(),
-        vec![Coin::new(1_000, ETH_USDC), Coin::new(1_000, COSMOS_USDC)]
+        err.downcast_ref::<ContractError>().unwrap(),
+        &ContractError::InsufficientShares {
+            required: Uint128::one(),
+            available: Uint128::zero()
+        }
     );
 
-    // withdrawing excess coins fails
+    // has remaining shares but no coins on the requested side
     let err = t
         .app
         .execute_contract(
-            Addr::unchecked("admin"),
+            Addr::unchecked("provider_1"),
             t.contract.clone(),
             &ExecMsg::Withdraw {
-                coins: vec![Coin::new(100_000, COSMOS_USDC)],
+                coins: vec![Coin::new(1, ETH_USDC), Coin::new(1, COSMOS_USDC)],
             },
             &[],
         )
         .unwrap_err();
+
     assert_eq!(
         err.downcast_ref::<ContractError>().unwrap(),
         &ContractError::InsufficientFund {
-            required: Coin::new(100_000, COSMOS_USDC),
-            available: Coin::new(100_000 - 1500 - 1000, COSMOS_USDC)
+            required: Coin::new(1, ETH_USDC),
+            available: Coin::new(0, ETH_USDC)
         }
     );
 }