diff --git a/cis2-fractionalizer/src/contract.rs b/cis2-fractionalizer/src/contract.rs index 6fbcd43..bc67ca1 100644 --- a/cis2-fractionalizer/src/contract.rs +++ b/cis2-fractionalizer/src/contract.rs @@ -197,7 +197,8 @@ fn contract_transfer( if to.address().matches_contract(&ctx.self_address()) { // tokens are being transferred to self // burn the tokens - let remaining_amount: ContractTokenAmount = state.burn(&token_id, amount, &from)?; + state.burn(&token_id, amount, &from)?; + let remaining_amount: ContractTokenAmount = state.get_supply(&token_id); // log burn event logger.log(&Cis2Event::Burn(BurnEvent { diff --git a/cis2-fractionalizer/src/state.rs b/cis2-fractionalizer/src/state.rs index 7f669f9..27252fd 100644 --- a/cis2-fractionalizer/src/state.rs +++ b/cis2-fractionalizer/src/state.rs @@ -1,5 +1,5 @@ use concordium_cis2::*; -use concordium_std::{*}; +use concordium_std::*; use crate::{ error::{ContractError, CustomContractError}, @@ -59,6 +59,7 @@ pub struct State { pub(crate) state: StateMap, S>, /// All of the token IDs pub(crate) tokens: StateMap, + pub(crate) token_supply: StateMap, /// Map with contract addresses providing implementations of additional /// standards. pub(crate) implementors: StateMap, S>, @@ -73,6 +74,7 @@ impl State { tokens: state_builder.new_map(), implementors: state_builder.new_map(), collaterals: state_builder.new_map(), + token_supply: state_builder.new_map(), } } @@ -85,14 +87,18 @@ impl State { owner: &Address, state_builder: &mut StateBuilder, ) { - self.tokens - .insert(*token_id, token_metadata.to_metadata_url()); - let mut owner_state = self - .state - .entry(*owner) - .or_insert_with(|| AddressState::empty(state_builder)); - let mut owner_balance = owner_state.balances.entry(*token_id).or_insert(0.into()); - *owner_balance += amount; + { + self.tokens + .insert(*token_id, token_metadata.to_metadata_url()); + let mut owner_state = self + .state + .entry(*owner) + .or_insert_with(|| AddressState::empty(state_builder)); + let mut owner_balance = owner_state.balances.entry(*token_id).or_insert(0.into()); + *owner_balance += amount; + } + + self.increase_supply(*token_id, amount); } pub(crate) fn burn( @@ -101,21 +107,27 @@ impl State { amount: ContractTokenAmount, owner: &Address, ) -> ContractResult { - match self.state.get_mut(owner) { - Some(address_state) => match address_state.balances.get_mut(token_id) { - Some(mut b) => { - ensure!( - b.cmp(&amount).is_ge(), - Cis2Error::Custom(CustomContractError::NoBalanceToBurn) - ); - - *b -= amount; - Ok(*b) - } + let ret = { + match self.state.get_mut(owner) { + Some(address_state) => match address_state.balances.get_mut(token_id) { + Some(mut b) => { + ensure!( + b.cmp(&amount).is_ge(), + Cis2Error::Custom(CustomContractError::NoBalanceToBurn) + ); + + *b -= amount; + Ok(*b) + } + None => Err(Cis2Error::Custom(CustomContractError::NoBalanceToBurn)), + }, None => Err(Cis2Error::Custom(CustomContractError::NoBalanceToBurn)), - }, - None => Err(Cis2Error::Custom(CustomContractError::NoBalanceToBurn)), - } + } + }; + + self.decrease_supply(*token_id, amount); + + ret } /// Check that the token ID currently exists in this contract. @@ -141,6 +153,29 @@ impl State { Ok(balance) } + fn increase_supply(&mut self, token_id: ContractTokenId, amount: ContractTokenAmount) { + let curr_supply = self.get_supply(&token_id); + self.token_supply.insert(token_id, curr_supply + amount); + } + + fn decrease_supply(&mut self, token_id: ContractTokenId, amount: ContractTokenAmount) { + let curr_supply = self.get_supply(&token_id); + let remaining_supply = curr_supply - amount; + + if remaining_supply.cmp(&ContractTokenAmount::from(0)).is_eq() { + self.token_supply.remove(&token_id); + } else { + self.token_supply.insert(token_id, curr_supply - amount); + } + } + + pub(crate) fn get_supply(&self, token_id: &ContractTokenId) -> ContractTokenAmount { + match self.token_supply.get(token_id) { + Some(amount) => *amount, + None => ContractTokenAmount::from(0), + } + } + /// Check if an address is an operator of a given owner address. pub(crate) fn is_operator(&self, address: &Address, owner: &Address) -> bool { self.state