diff --git a/sputnik-staking/res/sputnik_staking.wasm b/sputnik-staking/res/sputnik_staking.wasm old mode 100755 new mode 100644 index 5d7052227..2bbfedb74 Binary files a/sputnik-staking/res/sputnik_staking.wasm and b/sputnik-staking/res/sputnik_staking.wasm differ diff --git a/sputnik-staking/src/lib.rs b/sputnik-staking/src/lib.rs index b02a72f1c..434a8be63 100644 --- a/sputnik-staking/src/lib.rs +++ b/sputnik-staking/src/lib.rs @@ -199,6 +199,16 @@ mod tests { testing_env!(context.attached_deposit(to_yocto("1")).build()); contract.storage_deposit(Some(delegate_from_user.clone()), None); + let user = contract.get_user(delegate_from_user.clone()); + assert_eq!(user.near_amount, U128(to_yocto("1"))); + + contract.storage_deposit(Some(delegate_from_user.clone()), None); + let user = contract.get_user(delegate_from_user.clone()); + assert_eq!(user.near_amount, U128(2 * to_yocto("1"))); + + contract.storage_deposit(Some(delegate_from_user.clone()), Some(true)); + let user = contract.get_user(delegate_from_user.clone()); + assert_eq!(user.near_amount, U128(2 * to_yocto("1"))); testing_env!(context.predecessor_account_id(voting_token.clone()).build()); contract.ft_on_transfer( @@ -228,8 +238,14 @@ mod tests { contract.delegate(delegate_to_user.clone(), U128(to_yocto("10"))); let user = contract.get_user(delegate_from_user.clone()); assert_eq!(user.delegated_amount(), to_yocto("10")); + assert_eq!(user.near_amount, 165); + + contract.delegate(delegate_to_user.clone(), U128(to_yocto("10"))); + let user = contract.get_user(delegate_from_user.clone()); + assert_eq!(user.delegated_amount(), to_yocto("20")); + assert_eq!(user.near_amount, 165); - contract.undelegate(delegate_to_user, U128(to_yocto("10"))); + contract.undelegate(delegate_to_user, U128(to_yocto("20"))); let user = contract.get_user(delegate_from_user); assert_eq!(user.delegated_amount(), 0); assert_eq!(user.next_action_timestamp, U64(UNSTAKE_PERIOD)); diff --git a/sputnik-staking/src/storage_impl.rs b/sputnik-staking/src/storage_impl.rs index 9c1f9dca9..545f50c50 100644 --- a/sputnik-staking/src/storage_impl.rs +++ b/sputnik-staking/src/storage_impl.rs @@ -20,9 +20,14 @@ impl StorageManagement for Contract { let account_id = account_id.unwrap_or_else(env::predecessor_account_id); if self.users.contains_key(&account_id) { - log!("ERR_ACC_REGISTERED"); - if deposit_amount > 0 { - Promise::new(env::predecessor_account_id()).transfer(deposit_amount); + if registration_only.unwrap_or(false) { + log!("ERR_ACC_REGISTERED"); + + if deposit_amount > 0 { + Promise::new(env::predecessor_account_id()).transfer(deposit_amount); + } + } else { + self.internal_update_user_storage(&account_id, deposit_amount); } } else { let min_balance = User::min_storage() as Balance * env::storage_byte_cost(); diff --git a/sputnik-staking/src/user.rs b/sputnik-staking/src/user.rs index 2d17fb2a3..d5af5ddc6 100644 --- a/sputnik-staking/src/user.rs +++ b/sputnik-staking/src/user.rs @@ -75,8 +75,29 @@ impl User { env::block_timestamp() >= self.next_action_timestamp.0, "ERR_NOT_ENOUGH_TIME_PASSED" ); - self.storage_used += delegate_id.as_bytes().len() as StorageUsage + U128_LEN; - self.delegated_amounts.push((delegate_id, U128(amount))); + + let pos = self + .delegated_amounts + .iter() + .enumerate() + .position(|(_, (account_id, _))| account_id == &delegate_id); + + match pos { + Some(position) => { + let delegation = self.delegated_amounts[position].clone(); + + let _ = std::mem::replace( + &mut self.delegated_amounts[position], + (delegate_id, U128((delegation.1).0 + amount)), + ); + } + None => { + self.storage_used += delegate_id.as_bytes().len() as StorageUsage + U128_LEN; + + self.delegated_amounts.push((delegate_id, U128(amount))) + } + }; + self.assert_storage(); } @@ -159,6 +180,27 @@ impl Contract { ); } + pub fn internal_update_user_storage(&mut self, account_id: &AccountId, amount: Balance) { + let User { + storage_used, + near_amount, + vote_amount, + next_action_timestamp, + delegated_amounts, + } = self.internal_get_user(&account_id); + + self.users.insert( + account_id, + &VersionedUser::Default(User { + storage_used, + near_amount: U128(near_amount.0 + amount), + vote_amount, + next_action_timestamp, + delegated_amounts, + }), + ); + } + /// Deposit voting token. pub fn internal_deposit(&mut self, sender_id: &AccountId, amount: Balance) { let mut sender = self.internal_get_user(&sender_id);