Add owner enumeration with packed token IDs

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<felt252, felt252>,
+    }
+
+    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<TContractState>,
+        impl ERC721: ERC721Component::HasComponent<TContractState>,
+        +ERC721Component::ERC721HooksTrait<TContractState>,
+        +SRC5Component::HasComponent<TContractState>,
+        +Drop<TContractState>,
+    > of IBeastsOwnerEnumerable<ComponentState<TContractState>> {
+        fn token_of_owner_by_index(
+            self: @ComponentState<TContractState>, 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<TContractState>,
+        impl ERC721: ERC721Component::HasComponent<TContractState>,
+        +ERC721Component::ERC721HooksTrait<TContractState>,
+        impl SRC5: SRC5Component::HasComponent<TContractState>,
+        +Drop<TContractState>,
+    > of InternalTrait<TContractState> {
+        fn initializer(ref self: ComponentState<TContractState>) {
+            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<TContractState>, 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<TContractState>, 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<TContractState>, 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);
+        }
+    }
+}

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<u256> {
+        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<u256>, 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);
+    }
+}