remove share state

osmosis-labs · Jun 6, 2023 · eeaf5af · eeaf5af
1 parent 6c991a9
commit eeaf5af
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 81 deletions.
diff --git a/contracts/transmuter/src/contract.rs b/contracts/transmuter/src/contract.rs
@@ -106,10 +106,6 @@ impl Transmuter<'_> {
 
         let new_shares = Shares::calc_shares(&info.funds)?;
 
-        // update shares
-        self.shares
-            .add_share(deps.storage, &info.sender, new_shares)?;
-
         // update pool
         self.pool
             .update(deps.storage, |mut pool| -> Result<_, ContractError> {
@@ -142,7 +138,7 @@ impl Transmuter<'_> {
         let (deps, env, info) = ctx;
 
         // check if sender's shares is enough
-        let sender_shares = self.shares.get_share(deps.as_ref().storage, &info.sender)?;
+        let sender_shares = self.shares.get_share(deps.as_ref(), &info.sender)?;
 
         let required_shares = Shares::calc_shares(&tokens_out)?;
 
@@ -154,10 +150,6 @@ impl Transmuter<'_> {
             }
         );
 
-        // update shares
-        self.shares
-            .sub_share(deps.storage, &info.sender, required_shares)?;
-
         // exit pool
         self.pool
             .update(deps.storage, |mut pool| -> Result<_, ContractError> {
@@ -339,7 +331,7 @@ impl Transmuter<'_> {
         Ok(GetSharesResponse {
             shares: self
                 .shares
-                .get_share(deps.storage, &deps.api.addr_validate(&address)?)?,
+                .get_share(deps, &deps.api.addr_validate(&address)?)?,
         })
     }
 
@@ -376,7 +368,7 @@ impl Transmuter<'_> {
         ctx: (Deps, Env),
     ) -> Result<GetTotalSharesResponse, ContractError> {
         let (deps, _env) = ctx;
-        let total_shares = self.shares.get_total_shares(deps.storage)?;
+        let total_shares = self.shares.get_total_shares(deps)?;
         Ok(GetTotalSharesResponse { total_shares })
     }
 

diff --git a/contracts/transmuter/src/migration.rs b/contracts/transmuter/src/migration.rs
@@ -1 +1,4 @@
+use cosmwasm_schema::cw_serde;
+
+#[cw_serde]
 pub struct MigrateMsg {}
diff --git a/contracts/transmuter/src/shares.rs b/contracts/transmuter/src/shares.rs
@@ -1,17 +1,14 @@
-use cosmwasm_std::{Addr, Coin, StdError, StdResult, Storage, Uint128};
-use cw_storage_plus::{Item, Map};
+use cosmwasm_std::{Addr, Coin, Deps, StdError, StdResult, Storage, Uint128};
+use cw_storage_plus::Item;
+use osmosis_std::types::cosmos::bank::v1beta1::BankQuerier;
 
 pub struct Shares<'a> {
-    shares: Map<'a, &'a Addr, Uint128>,
-    total_shares: Item<'a, Uint128>,
     share_denom: Item<'a, String>,
 }
 
 impl Shares<'_> {
     pub const fn new() -> Self {
         Self {
-            shares: Map::new("shares"),
-            total_shares: Item::new("total_shares"),
             share_denom: Item::new("share_denom"),
         }
     }
@@ -27,73 +24,29 @@ impl Shares<'_> {
     }
 
     /// get the total shares
-    pub fn get_total_shares(&self, store: &dyn Storage) -> StdResult<Uint128> {
-        self.total_shares
-            .may_load(store)
-            .map(|opt| opt.unwrap_or_default())
-    }
+    pub fn get_total_shares(&self, deps: Deps) -> StdResult<Uint128> {
+        let share_denom = self.get_share_denom(deps.storage)?;
+        let bank_querier = BankQuerier::new(&deps.querier);
 
-    /// get shares for a given address
-    pub fn get_share(&self, store: &dyn Storage, address: &Addr) -> StdResult<Uint128> {
-        self.shares
-            .may_load(store, address)
-            .map(|opt| opt.unwrap_or_default())
+        bank_querier
+            .supply_of(share_denom)?
+            .amount
+            .ok_or_else(|| StdError::generic_err("No shares"))?
+            .amount
+            .parse()
     }
 
-    /// add shares for a given address
-    pub fn add_share(
-        &self,
-        store: &mut dyn Storage,
-        address: &Addr,
-        amount: Uint128,
-    ) -> StdResult<()> {
-        // add the amount to the current shares
-        self.shares
-            .update(store, address, |current| -> StdResult<_> {
-                current
-                    .unwrap_or_default()
-                    .checked_add(amount)
-                    .map_err(StdError::overflow)
-            })?;
-
-        // add the amount to the total shares
-        let total_shares = self.get_total_shares(store)?;
-        self.total_shares.save(
-            store,
-            &total_shares
-                .checked_add(amount)
-                .map_err(StdError::overflow)?,
-        )?;
-
-        Ok(())
-    }
-
-    /// subtract shares for a given address
-    pub fn sub_share(
-        &self,
-        store: &mut dyn Storage,
-        address: &Addr,
-        amount: Uint128,
-    ) -> StdResult<()> {
-        // subtract the amount from the current shares
-        self.shares
-            .update(store, address, |current| -> StdResult<_> {
-                current
-                    .unwrap_or_default()
-                    .checked_sub(amount)
-                    .map_err(StdError::overflow)
-            })?;
-
-        // subtract the amount from the total shares
-        let total_shares = self.get_total_shares(store)?;
-        self.total_shares.save(
-            store,
-            &total_shares
-                .checked_sub(amount)
-                .map_err(StdError::overflow)?,
-        )?;
-
-        Ok(())
+    /// get shares for a given address
+    pub fn get_share(&self, deps: Deps, address: &Addr) -> StdResult<Uint128> {
+        let share_denom = self.get_share_denom(deps.storage)?;
+        let bank_querier = BankQuerier::new(&deps.querier);
+
+        bank_querier
+            .balance(address.to_string(), share_denom)?
+            .balance
+            .ok_or_else(|| StdError::generic_err("No shares"))?
+            .amount
+            .parse()
     }
 
     /// calculate the amount of shares for a given amount of tokens