diff --git a/src/enumerable.cairo b/src/enumerable.cairo new file mode 100644 index 0000000..6686229 --- /dev/null +++ b/src/enumerable.cairo @@ -0,0 +1,141 @@ +/// Owner enumeration for the Beasts collection. +/// +/// Answers "which tokens does this address hold" — the one thing standard +/// ERC721 cannot do, and the missing link for any client that wants to show a +/// wallet its Beasts. Because a token ID *is* the Beast (species, affixes, +/// tier, type and variant flags are all encoded), one enumeration call per +/// token is enough to reconstruct a wallet's whole collection with no further +/// chain reads. +/// +/// This is deliberately narrower than OpenZeppelin's `ERC721Enumerable`. That +/// extension also maintains a global token list (`token_by_index`, +/// `total_supply`), costing five storage writes per mint instead of two. A +/// global index buys little here: token IDs are derived rather than +/// sequential, the collection is unbounded by design, and `total_supply` is +/// already tracked directly by the contract. +/// +/// Token IDs are stored as `felt252` rather than `u256`. They occupy at most +/// 116 bits (see `pack.cairo`), so a felt holds one whole and halves the +/// storage each index entry costs. +#[starknet::component] +pub mod EnumerableComponent { + use core::num::traits::Zero; + use openzeppelin_introspection::src5::SRC5Component; + use openzeppelin_introspection::src5::SRC5Component::InternalTrait as SRC5InternalTrait; + use openzeppelin_token::erc721::ERC721Component; + use openzeppelin_token::erc721::ERC721Component::{ + ERC721Impl, InternalImpl as ERC721InternalImpl, + }; + use starknet::ContractAddress; + use starknet::storage::{Map, StorageMapReadAccess, StorageMapWriteAccess}; + use super::super::interfaces::{IBEASTS_OWNER_ENUMERABLE_ID, IBeastsOwnerEnumerable}; + + #[storage] + pub struct Storage { + /// (owner, index) -> token ID. Dense: indices always run 0..balance-1. + pub Enumerable_owned_tokens: Map<(ContractAddress, felt252), felt252>, + /// token ID -> its index in its owner's list. + pub Enumerable_owned_tokens_index: Map, + } + + pub mod Errors { + pub const OUT_OF_BOUNDS_INDEX: felt252 = 'ERC721Enum: out of bounds index'; + pub const BURN_NOT_SUPPORTED: felt252 = 'ERC721Enum: burn not supported'; + } + + #[embeddable_as(EnumerableImpl)] + pub impl Enumerable< + TContractState, + +HasComponent, + impl ERC721: ERC721Component::HasComponent, + +ERC721Component::ERC721HooksTrait, + +SRC5Component::HasComponent, + +Drop, + > of IBeastsOwnerEnumerable> { + fn token_of_owner_by_index( + self: @ComponentState, owner: ContractAddress, index: u256, + ) -> u256 { + let erc721_component = get_dep_component!(self, ERC721); + assert(index < erc721_component.balance_of(owner), Errors::OUT_OF_BOUNDS_INDEX); + let index_felt: felt252 = index.try_into().unwrap(); + self.Enumerable_owned_tokens.read((owner, index_felt)).into() + } + } + + #[generate_trait] + pub impl InternalImpl< + TContractState, + +HasComponent, + impl ERC721: ERC721Component::HasComponent, + +ERC721Component::ERC721HooksTrait, + impl SRC5: SRC5Component::HasComponent, + +Drop, + > of InternalTrait { + fn initializer(ref self: ComponentState) { + let mut src5_component = get_dep_component_mut!(ref self, SRC5); + src5_component.register_interface(IBEASTS_OWNER_ENUMERABLE_ID); + } + + /// Must be called from the contract's `ERC721HooksTrait::before_update`. + /// + /// Runs *before* the ERC721 component mutates ownership and balances, + /// so every `balance_of` read here is the pre-transfer value. Both + /// helpers depend on that: the new index is the recipient's current + /// length, and the last index is the sender's current length minus + /// one. + fn before_update( + ref self: ComponentState, to: ContractAddress, token_id: u256, + ) { + // Burning would leave a hole in the dense index. Beasts are not + // burnable, so this is an invariant guard rather than a + // limitation — and it is what lets the index stay dense without a + // tombstone scheme. + assert(!to.is_zero(), Errors::BURN_NOT_SUPPORTED); + + let erc721_component = get_dep_component!(@self, ERC721); + let previous_owner = erc721_component._owner_of(token_id); + let token_id_felt: felt252 = token_id.try_into().unwrap(); + + // A self-transfer changes nothing, and must not: re-adding would + // double-count the token in its owner's list. + if previous_owner == to { + return; + } + + if !previous_owner.is_zero() { + self._remove_token_from_owner_enumeration(previous_owner, token_id_felt); + } + self._add_token_to_owner_enumeration(to, token_id_felt); + } + + fn _add_token_to_owner_enumeration( + ref self: ComponentState, to: ContractAddress, token_id: felt252, + ) { + let erc721_component = get_dep_component!(@self, ERC721); + let len: felt252 = erc721_component.balance_of(to).try_into().unwrap(); + self.Enumerable_owned_tokens.write((to, len), token_id); + self.Enumerable_owned_tokens_index.write(token_id, len); + } + + /// Swap-and-pop: move the last token into the vacated slot so the + /// list stays dense and removal stays O(1). + fn _remove_token_from_owner_enumeration( + ref self: ComponentState, from: ContractAddress, token_id: felt252, + ) { + let erc721_component = get_dep_component!(@self, ERC721); + let last_token_index: felt252 = (erc721_component.balance_of(from) - 1) + .try_into() + .unwrap(); + let this_token_index = self.Enumerable_owned_tokens_index.read(token_id); + + if this_token_index != last_token_index { + let last_token_id = self.Enumerable_owned_tokens.read((from, last_token_index)); + self.Enumerable_owned_tokens.write((from, this_token_index), last_token_id); + self.Enumerable_owned_tokens_index.write(last_token_id, this_token_index); + } + + self.Enumerable_owned_tokens.write((from, last_token_index), 0); + self.Enumerable_owned_tokens_index.write(token_id, 0); + } + } +} diff --git a/src/enumerable_tests.cairo b/src/enumerable_tests.cairo new file mode 100644 index 0000000..44c1f19 --- /dev/null +++ b/src/enumerable_tests.cairo @@ -0,0 +1,252 @@ +#[cfg(test)] +mod enumerable_tests { + use beasts_nft::interfaces::{ + IBeastsDispatcher, IBeastsDispatcherTrait, IBeastsOwnerEnumerableDispatcher, + IBeastsOwnerEnumerableDispatcherTrait, + }; + use openzeppelin_interfaces::erc721::{IERC721Dispatcher, IERC721DispatcherTrait}; + use snforge_std::{ + ContractClassTrait, DeclareResultTrait, declare, start_cheat_caller_address, + start_mock_call, stop_cheat_caller_address, + }; + use starknet::ContractAddress; + + fn test_address(address: felt252) -> ContractAddress { + address.try_into().unwrap() + } + + fn deploy() -> ( + IBeastsDispatcher, IERC721Dispatcher, IBeastsOwnerEnumerableDispatcher, ContractAddress, + ) { + let owner = test_address('owner'); + let provider = test_address('provider'); + let art: ByteArray = "data:image/png;base64,iVBORw0KGgoAAAA1"; + start_mock_call(provider, selector!("get_data_uri"), art); + + let contract = declare("beasts_nft").unwrap().contract_class(); + let mut calldata = array![]; + let name: ByteArray = "Beasts"; + let symbol: ByteArray = "BEAST"; + name.serialize(ref calldata); + symbol.serialize(ref calldata); + owner.serialize(ref calldata); + owner.serialize(ref calldata); + 500_u128.serialize(ref calldata); + provider.serialize(ref calldata); + provider.serialize(ref calldata); + provider.serialize(ref calldata); + provider.serialize(ref calldata); + 0.serialize(ref calldata); + + let (contract_address, _) = contract.deploy(@calldata).unwrap(); + ( + IBeastsDispatcher { contract_address }, + IERC721Dispatcher { contract_address }, + IBeastsOwnerEnumerableDispatcher { contract_address }, + owner, + ) + } + + /// Every token an address holds, read back through enumeration. + fn tokens_of( + enumerable: IBeastsOwnerEnumerableDispatcher, + erc721: IERC721Dispatcher, + owner: ContractAddress, + ) -> Array { + let mut out = array![]; + let balance = erc721.balance_of(owner); + let mut i: u256 = 0; + while i < balance { + out.append(enumerable.token_of_owner_by_index(owner, i)); + i += 1; + } + out + } + + fn contains(haystack: @Array, needle: u256) -> bool { + let mut i = 0; + let mut found = false; + while i < haystack.len() { + if *haystack.at(i) == needle { + found = true; + break; + } + i += 1; + } + found + } + + #[test] + fn test_genesis_mints_are_all_enumerable() { + let (beasts, erc721, enumerable, owner) = deploy(); + + assert(erc721.balance_of(owner) == 75, 'Owner holds 75 genesis'); + let tokens = tokens_of(enumerable, erc721, owner); + assert(tokens.len() == 75, 'Enumerates 75 tokens'); + + // Every entry must be a real token owned by this address, and the + // list must be free of duplicates — a broken index shows up as the + // same token appearing twice. + let mut i = 0; + while i < tokens.len() { + let token_id = *tokens.at(i); + assert(erc721.owner_of(token_id) == owner, 'Enumerated token not owned'); + assert(beasts.get_beast(token_id).id != 0, 'Enumerated token invalid'); + + let mut j = i + 1; + while j < tokens.len() { + assert(*tokens.at(j) != token_id, 'Duplicate in enumeration'); + j += 1; + } + i += 1; + } + } + + #[test] + fn test_mint_appends_to_recipient() { + let (beasts, erc721, enumerable, owner) = deploy(); + let minter = test_address('minter'); + let player = test_address('player'); + + start_cheat_caller_address(beasts.contract_address, owner); + beasts.set_dungeon_address(minter); + stop_cheat_caller_address(beasts.contract_address); + + start_cheat_caller_address(beasts.contract_address, minter); + let (first, _, _) = beasts.mint(player, 3, 1, 1, 10, 100, 0, 0); + let (second, _, _) = beasts.mint(player, 3, 2, 2, 10, 100, 0, 0); + stop_cheat_caller_address(beasts.contract_address); + + let tokens = tokens_of(enumerable, erc721, player); + assert(tokens.len() == 2, 'Player enumerates 2'); + assert(*tokens.at(0) == first, 'First at index 0'); + assert(*tokens.at(1) == second, 'Second at index 1'); + } + + #[test] + fn test_transfer_moves_between_owners() { + let (beasts, erc721, enumerable, owner) = deploy(); + let minter = test_address('minter'); + let alice = test_address('alice'); + let bob = test_address('bob'); + + start_cheat_caller_address(beasts.contract_address, owner); + beasts.set_dungeon_address(minter); + stop_cheat_caller_address(beasts.contract_address); + + start_cheat_caller_address(beasts.contract_address, minter); + let (token, _, _) = beasts.mint(alice, 3, 1, 1, 10, 100, 0, 0); + stop_cheat_caller_address(beasts.contract_address); + + start_cheat_caller_address(erc721.contract_address, alice); + erc721.transfer_from(alice, bob, token); + stop_cheat_caller_address(erc721.contract_address); + + assert(tokens_of(enumerable, erc721, alice).len() == 0, 'Alice enumerates none'); + let bobs = tokens_of(enumerable, erc721, bob); + assert(bobs.len() == 1, 'Bob enumerates one'); + assert(*bobs.at(0) == token, 'Bob holds the token'); + } + + #[test] + fn test_removing_a_middle_token_keeps_the_list_dense() { + // The swap-and-pop path. Removing from the middle must move the last + // entry into the hole, or enumeration returns zero for that index and + // silently loses a token. + let (beasts, erc721, enumerable, owner) = deploy(); + let minter = test_address('minter'); + let alice = test_address('alice'); + let bob = test_address('bob'); + + start_cheat_caller_address(beasts.contract_address, owner); + beasts.set_dungeon_address(minter); + stop_cheat_caller_address(beasts.contract_address); + + start_cheat_caller_address(beasts.contract_address, minter); + let (a, _, _) = beasts.mint(alice, 3, 1, 1, 10, 100, 0, 0); + let (b, _, _) = beasts.mint(alice, 3, 2, 2, 10, 100, 0, 0); + let (c, _, _) = beasts.mint(alice, 3, 3, 3, 10, 100, 0, 0); + stop_cheat_caller_address(beasts.contract_address); + + // Drop the middle one. + start_cheat_caller_address(erc721.contract_address, alice); + erc721.transfer_from(alice, bob, b); + stop_cheat_caller_address(erc721.contract_address); + + let tokens = tokens_of(enumerable, erc721, alice); + assert(tokens.len() == 2, 'Alice enumerates 2'); + assert(contains(@tokens, a), 'Kept token a'); + assert(contains(@tokens, c), 'Kept token c'); + assert(!contains(@tokens, b), 'Dropped token b'); + assert(*tokens.at(0) != 0 && *tokens.at(1) != 0, 'No hole left behind'); + } + + #[test] + fn test_transfer_back_and_forth_stays_consistent() { + let (beasts, erc721, enumerable, owner) = deploy(); + let minter = test_address('minter'); + let alice = test_address('alice'); + let bob = test_address('bob'); + + start_cheat_caller_address(beasts.contract_address, owner); + beasts.set_dungeon_address(minter); + stop_cheat_caller_address(beasts.contract_address); + + start_cheat_caller_address(beasts.contract_address, minter); + let (a, _, _) = beasts.mint(alice, 3, 1, 1, 10, 100, 0, 0); + let (b, _, _) = beasts.mint(alice, 3, 2, 2, 10, 100, 0, 0); + stop_cheat_caller_address(beasts.contract_address); + + start_cheat_caller_address(erc721.contract_address, alice); + erc721.transfer_from(alice, bob, a); + stop_cheat_caller_address(erc721.contract_address); + start_cheat_caller_address(erc721.contract_address, bob); + erc721.transfer_from(bob, alice, a); + stop_cheat_caller_address(erc721.contract_address); + + let tokens = tokens_of(enumerable, erc721, alice); + assert(tokens.len() == 2, 'Alice enumerates 2 again'); + assert(contains(@tokens, a), 'Token a returned'); + assert(contains(@tokens, b), 'Token b retained'); + assert(tokens_of(enumerable, erc721, bob).len() == 0, 'Bob enumerates none'); + } + + #[test] + fn test_self_transfer_does_not_duplicate() { + // A self-transfer must be a no-op. Re-adding would list the token + // twice and inflate the index past balance_of. + let (beasts, erc721, enumerable, owner) = deploy(); + let minter = test_address('minter'); + let alice = test_address('alice'); + + start_cheat_caller_address(beasts.contract_address, owner); + beasts.set_dungeon_address(minter); + stop_cheat_caller_address(beasts.contract_address); + + start_cheat_caller_address(beasts.contract_address, minter); + let (token, _, _) = beasts.mint(alice, 3, 1, 1, 10, 100, 0, 0); + stop_cheat_caller_address(beasts.contract_address); + + start_cheat_caller_address(erc721.contract_address, alice); + erc721.transfer_from(alice, alice, token); + stop_cheat_caller_address(erc721.contract_address); + + let tokens = tokens_of(enumerable, erc721, alice); + assert(tokens.len() == 1, 'Still exactly one token'); + assert(*tokens.at(0) == token, 'Still the same token'); + } + + #[test] + #[should_panic(expected: ('ERC721Enum: out of bounds index',))] + fn test_index_past_balance_reverts() { + let (_, erc721, enumerable, owner) = deploy(); + enumerable.token_of_owner_by_index(owner, erc721.balance_of(owner)); + } + + #[test] + #[should_panic(expected: ('ERC721Enum: out of bounds index',))] + fn test_enumerating_a_holder_of_nothing_reverts() { + let (_, _, enumerable, _) = deploy(); + enumerable.token_of_owner_by_index(test_address('nobody'), 0); + } +} diff --git a/src/interfaces.cairo b/src/interfaces.cairo index 0f76fd3..7548449 100644 --- a/src/interfaces.cairo +++ b/src/interfaces.cairo @@ -1,6 +1,23 @@ use starknet::ContractAddress; use super::pack::PackableBeast; +/// Beasts-specific owner-only enumeration. +/// +/// Deliberately not the standard `IERC721Enumerable` ID: this implements only +/// the owner half of that interface, and claiming the full one would tell +/// callers `token_by_index` exists when it does not. +/// EFS: token_of_owner_by_index(ContractAddress,(u128,u128))->(u128,u128) +pub const IBEASTS_OWNER_ENUMERABLE_ID: felt252 = + 0x312c74a3a4f7aaf9aa3e80ddea171f958139ef0c3dbea524e0763682b7d57dd; + +#[starknet::interface] +pub trait IBeastsOwnerEnumerable { + /// Token held by `owner` at `index`, where index runs 0..balance_of(owner). + /// Order is not stable across transfers: removal swaps the last entry into + /// the vacated slot. + fn token_of_owner_by_index(self: @TContractState, owner: ContractAddress, index: u256) -> u256; +} + /// Interface for the Beasts NFT contract #[starknet::interface] pub trait IBeasts { diff --git a/src/lib.cairo b/src/lib.cairo index 17356f0..a8dcde7 100644 --- a/src/lib.cairo +++ b/src/lib.cairo @@ -8,6 +8,9 @@ pub mod beast_png_shiny_data; pub mod beast_ranking; pub mod beast_svg; pub mod encoding; +pub mod enumerable; +#[cfg(test)] +mod enumerable_tests; pub mod interfaces; pub mod metadata_generator; #[cfg(test)] @@ -32,13 +35,14 @@ pub mod beasts_nft { use openzeppelin_interfaces::erc721::IERC721Metadata; use openzeppelin_introspection::src5::SRC5Component; use openzeppelin_token::common::erc2981::ERC2981Component; - use openzeppelin_token::erc721::{ERC721Component, ERC721HooksEmptyImpl}; + use openzeppelin_token::erc721::ERC721Component; use starknet::ContractAddress; use starknet::storage::{ Map, StoragePathEntry, StoragePointerReadAccess, StoragePointerWriteAccess, }; use super::beast_manager::{BeastManagerTrait, BeastResult}; use super::beast_ranking::BeastRankingManagerTrait; + use super::enumerable::EnumerableComponent; use super::interfaces::{ IBeastImageDataProviderDispatcher, IBeastSystemsDispatcher, IBeastSystemsDispatcherTrait, IBeasts, IBeastsAnimation, @@ -51,6 +55,7 @@ pub mod beasts_nft { component!(path: ERC721Component, storage: erc721, event: ERC721Event); component!(path: SRC5Component, storage: src5, event: SRC5Event); component!(path: ERC2981Component, storage: erc2981, event: ERC2981Event); + component!(path: EnumerableComponent, storage: erc721_enumerable, event: EnumerableEvent); // Ownable Mixin #[abi(embed_v0)] @@ -64,6 +69,26 @@ pub mod beasts_nft { impl ERC721CamelOnlyImpl = ERC721Component::ERC721CamelOnlyImpl; impl ERC721InternalImpl = ERC721Component::InternalImpl; + // Owner enumeration + #[abi(embed_v0)] + impl EnumerableImpl = EnumerableComponent::EnumerableImpl; + impl EnumerableInternalImpl = EnumerableComponent::InternalImpl; + + /// Keeps the owner index in step with every mint and transfer. The + /// component reads pre-update balances, so it must run in `before_update` + /// and not `after_update`. + impl ERC721EnumerableHooks of ERC721Component::ERC721HooksTrait { + fn before_update( + ref self: ERC721Component::ComponentState, + to: ContractAddress, + token_id: u256, + auth: ContractAddress, + ) { + let mut contract_state = self.get_contract_mut(); + contract_state.erc721_enumerable.before_update(to, token_id); + } + } + // SRC5 Implementation #[abi(embed_v0)] impl SRC5Impl = SRC5Component::SRC5Impl; @@ -91,6 +116,8 @@ pub mod beasts_nft { pub src5: SRC5Component::Storage, #[substorage(v0)] pub erc2981: ERC2981Component::Storage, + #[substorage(v0)] + pub erc721_enumerable: EnumerableComponent::Storage, // Beast-specific storage pub beast_token_ranks: Map, // token_id -> current rank (for tokenURI) pub beast_species_lists: Map< @@ -127,6 +154,8 @@ pub mod beasts_nft { SRC5Event: SRC5Component::Event, #[flat] ERC2981Event: ERC2981Component::Event, + #[flat] + EnumerableEvent: EnumerableComponent::Event, MetadataUpdate: MetadataUpdate, } @@ -150,6 +179,7 @@ pub mod beasts_nft { ) { self.ownable.initializer(owner); self.erc721.initializer(name, symbol, ""); + self.erc721_enumerable.initializer(); self.erc2981.initializer(royalty_receiver, royalty_fraction); // Store external image data dispatchers