diff --git a/contracts/registry/src/lib.rs b/contracts/registry/src/lib.rs index c58649d..ab3ff9b 100644 --- a/contracts/registry/src/lib.rs +++ b/contracts/registry/src/lib.rs @@ -855,6 +855,26 @@ impl Contract { } true } + + /// Helper function for `sbt_revoke_by_owner` + fn sbt_token_ids_by_owner( + &self, + account: AccountId, + issuer_id: u32, + limit: Option, + ) -> Vec<(TokenId, ClassId)> { + let first_key = balance_key(account.clone(), issuer_id, 0); + + let limit = limit.unwrap_or(1000); + assert!(limit > 0, "limit must be bigger than 0"); + + self.balances + .iter_from(first_key) + .take_while(|(key, _)| key.owner == account && key.issuer_id == issuer_id) + .map(|(key, token_id)| (token_id, key.class_id)) + .take(limit as usize) + .collect() + } } #[cfg(test)] @@ -2536,6 +2556,37 @@ mod tests { ); } + #[test] + fn sbt_token_ids_by_owner() { + let (mut ctx, mut ctr) = setup(&issuer1(), 20 * MINT_DEPOSIT); + let batch_metadata = mk_batch_metadata(20); + ctr.sbt_mint(vec![(alice(), batch_metadata[..10].to_vec())]); + + ctx.predecessor_account_id = issuer3(); + testing_env!(ctx.clone()); + ctr.sbt_mint(vec![(alice(), batch_metadata[..10].to_vec())]); + + ctx.predecessor_account_id = issuer2(); + testing_env!(ctx.clone()); + ctr.sbt_mint(vec![(alice(), batch_metadata[..10].to_vec())]); + + let alice_tokens: Vec<(u64, u64)> = (1..=20).map(|i| (i, i)).collect(); + + let res = ctr.sbt_token_ids_by_owner(alice(), ctr.assert_issuer(&issuer1()), None); + assert_eq!(res, &alice_tokens[0..10]); + let res = ctr.sbt_token_ids_by_owner(alice(), ctr.assert_issuer(&issuer2()), None); + assert_eq!(res, &alice_tokens[0..10]); + let res = ctr.sbt_token_ids_by_owner(alice(), ctr.assert_issuer(&issuer2()), None); + assert_eq!(res, &alice_tokens[0..10]); + + // mint more tokens for issuer1() + ctx.predecessor_account_id = issuer1(); + testing_env!(ctx); + ctr.sbt_mint(vec![(alice(), batch_metadata[10..20].to_vec())]); + let res = ctr.sbt_token_ids_by_owner(alice(), ctr.assert_issuer(&issuer1()), None); + assert_eq!(res, alice_tokens); + } + #[test] fn is_human_multiple_classes_with_expired_tokens() { let (mut ctx, mut ctr) = setup(&fractal_mainnet(), 150 * MINT_DEPOSIT); diff --git a/contracts/registry/src/registry.rs b/contracts/registry/src/registry.rs index 6509d68..1a5eb10 100644 --- a/contracts/registry/src/registry.rs +++ b/contracts/registry/src/registry.rs @@ -362,95 +362,99 @@ impl SBTRegistry for Contract { fn sbt_revoke_by_owner(&mut self, owner: AccountId, burn: bool) -> bool { let issuer = env::predecessor_account_id(); let issuer_id = self.assert_issuer(&issuer); - let mut tokens_by_owner = self.sbt_tokens_by_owner( - owner.clone(), - Some(issuer.clone()), - None, - Some(MAX_REVOKE_PER_CALL), - Some(true), - ); - let remaining_supply = self.sbt_supply_by_owner(owner.clone(), issuer.clone(), None); - let tokens_len = tokens_by_owner[0].1.len(); + let tokens_by_owner = + self.sbt_token_ids_by_owner(owner.clone(), issuer_id, Some(MAX_REVOKE_PER_CALL)); + if tokens_by_owner.is_empty() { return true; - }; - let (_, tokens) = tokens_by_owner.pop().unwrap(); - - let mut token_ids = Vec::new(); + } if burn { let mut burned_per_class: HashMap = HashMap::new(); - let tokens_burned = tokens.len() as u64; - for t in tokens { - token_ids.push(t.token); - let class_id = t.metadata.class; - self.balances.remove(&BalanceKey { + + // Batch updates for balances and issuer_tokens + for (token_id, class_id) in &tokens_by_owner { + let balance_key = BalanceKey { issuer_id, owner: owner.clone(), - class_id, - }); + class_id: *class_id, + }; - // collect the info about the tokens revoked per class - // to update the balance accordingly + self.balances.remove(&balance_key); + + // Collect info about tokens revoked per class to update the balance accordingly burned_per_class - .entry(class_id) + .entry(*class_id) .and_modify(|key_value| *key_value += 1) .or_insert(1); self.issuer_tokens.remove(&IssuerTokenId { issuer_id, - token: t.token, + token: *token_id, }); } - let key = &(owner, issuer_id); - let old_supply = self.supply_by_owner.get(key).unwrap(); - self.supply_by_owner - .insert(key, &(old_supply - tokens_burned)); + // Batch updates for supply values + let supply_update = tokens_by_owner.len() as u64; - let supply_by_issuer = self.supply_by_issuer.get(&issuer_id).unwrap_or(0); - self.supply_by_issuer - .insert(&issuer_id, &(supply_by_issuer - tokens_burned)); + // Update supply_by_owner + let owner_key = &(owner.clone(), issuer_id); + let supply_owner = self.supply_by_owner.get(owner_key).unwrap_or(0); + let new_supply_owner = supply_owner - supply_update; + self.supply_by_owner.insert(owner_key, &new_supply_owner); - // update supply by class + // Update supply_by_issuer + let supply_issuer = self.supply_by_issuer.get(&issuer_id).unwrap_or(0); + let new_supply_issuer = supply_issuer - supply_update; + self.supply_by_issuer.insert(&issuer_id, &new_supply_issuer); + + // Update supply_by_class for (class_id, tokens_revoked) in burned_per_class { - let key = &(issuer_id, class_id); - let old_supply = self.supply_by_class.get(key).unwrap(); - self.supply_by_class - .insert(key, &(old_supply - tokens_revoked)); + let class_key = &(issuer_id, class_id); + let supply_class = self.supply_by_class.get(class_key).unwrap_or(0); + let new_supply_class = supply_class - tokens_revoked; + self.supply_by_class.insert(class_key, &new_supply_class); } SbtTokensEvent { issuer: issuer.clone(), - tokens: token_ids.clone(), + tokens: tokens_by_owner + .iter() + .map(|(token_id, _)| *token_id) + .collect(), } .emit_burn(); } else { - // revoke - // update expire date for all tokens to current_timestamp + // Revoke: Update expire date for all tokens to current_timestamp let now = env::block_timestamp_ms(); - for mut t in tokens { - token_ids.push(t.token); - t.metadata.expires_at = Some(now); + for (token_id, _) in &tokens_by_owner { + let mut token = self.get_token(issuer_id, *token_id).to_token(*token_id); + token.metadata.expires_at = Some(now); let token_data = TokenData { owner: owner.clone(), - metadata: t.metadata.into(), + metadata: token.metadata.into(), }; self.issuer_tokens.insert( &IssuerTokenId { issuer_id, - token: t.token, + token: *token_id, }, &token_data, ); } } + SbtTokensEvent { - issuer, - tokens: token_ids, + issuer: issuer.clone(), + tokens: tokens_by_owner + .iter() + .map(|(token_id, _)| *token_id) + .collect(), } .emit_revoke(); - tokens_len >= remaining_supply as usize + + // Check if all tokens were revoked + tokens_by_owner.len() >= self.sbt_supply_by_owner(owner.clone(), issuer, None) as usize } /// Allows issuer to update token metadata reference and reference_hash.