ERC721Enumerable helper function

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- `ERC721EnumerableExtended` impl for fetching all ERC721 tokens of an owner in a single call (#1196)
 - Embeddable impls for ERC2981 component (#1173)
   - `ERC2981Info` with read functions for discovering the component's state
   - `ERC2981AdminOwnable` providing admin functions for a token that implements Ownable component

docs/modules/ROOT/pages/api/erc721.adoc

@@ -830,6 +830,26 @@ Use along with xref:#IERC721Enumerable-total_supply[IERC721Enumerable::total_sup
 Returns the token id owned by `owner` at a given `index` of its token list.
 Use along with xref:#IERC721-balance_of[IERC721::balance_of] to enumerate all of ``owner``'s tokens.
 
+[.contract]
+[[IERC721EnumerableExtended]]
+=== `++IERC721EnumerableExtended++` link:https://github.com/OpenZeppelin/cairo-contracts/blob/release-v0.18.0/packages/token/src/erc721/extensions/erc721_enumerable/interface.cairo[{github-icon},role=heading-link]
+
+Interface for the helper enumerable functions in {eip721}.
+
+[.contract-index]
+.Functions
+--
+* xref:#IERC721EnumerableExtended-all_tokens_of_owner[`++all_tokens_of_owner(owner)++`]
+--
+
+==== Functions
+
+[.contract-item]
+[[IERC721EnumerableExtended-all_tokens_of_owner]]
+==== `[.contract-item-name]#++all_tokens_of_owner++#++(owner: ContractAddress) -> Span<u256>++` [.item-kind]#external#
+
+Returns a list of all token ids owned by the specified `owner`.
+
 [.contract]
 [[ERC721EnumerableComponent]]
 === `++ERC721EnumerableComponent++` link:https://github.com/OpenZeppelin/cairo-contracts/blob/release-v0.18.0/packages/token/src/erc721/extensions/erc721_enumerable.cairo[{github-icon},role=heading-link]
@@ -878,6 +898,10 @@ mod ERC721EnumerableContract {
 * xref:#ERC721EnumerableComponent-total_supply[`++total_supply(self)++`]
 * xref:#ERC721EnumerableComponent-token_by_index[`++token_by_index(self, index)++`]
 * xref:#ERC721EnumerableComponent-token_of_owner_by_index[`++token_of_owner_by_index(self, address, index)++`]
+
+[.sub-index#ERC721EnumerableComponent-Embeddable-Impls-ERC721EnumerableExtendedImpl]
+.ERC721EnumerableExtendedImpl
+* xref:#ERC721EnumerableComponent-all_tokens_of_owner[`++all_tokens_of_owner(self, owner)++`]
 --
 
 [.contract-index]
@@ -922,6 +946,17 @@ Requirements:
 - `index` is less than ``owner``'s token balance.
 - `owner` is not the zero address.
 
+[.contract-item]
+[[ERC721EnumerableComponent-all_tokens_of_owner]]
+==== `[.contract-item-name]#++all_tokens_of_owner++#++(self: @ContractState, owner: ContractAddress) → Span<u256>++` [.item-kind]#external#
+
+Returns a list of all token ids owned by the specified `owner`.
+This function provides a more efficient alternative to calling `ERC721::balance_of`
+and iterating through tokens with `ERC721Enumerable::token_of_owner_by_index`.
+
+Requirements:
+- `owner` is not the zero address.
+
 [#ERC721EnumerableComponent-Internal-functions]
 ==== Internal functions
 

packages/test_common/src/mocks.cairo

@@ -5,6 +5,7 @@ pub mod erc1155;
 pub mod erc20;
 pub mod erc2981;
 pub mod erc721;
+pub mod erc721_enumerable;
 pub mod non_implementing;
 pub mod nonces;
 pub mod security;

packages/test_common/src/mocks/erc721_enumerable.cairo

@@ -0,0 +1,154 @@
+#[starknet::contract]
+pub mod ERC721EnumerableMock {
+    use openzeppelin_introspection::src5::SRC5Component;
+    use openzeppelin_token::erc721::ERC721Component;
+    use openzeppelin_token::erc721::extensions::ERC721EnumerableComponent;
+    use starknet::ContractAddress;
+
+    component!(path: ERC721Component, storage: erc721, event: ERC721Event);
+    component!(
+        path: ERC721EnumerableComponent, storage: erc721_enumerable, event: ERC721EnumerableEvent
+    );
+    component!(path: SRC5Component, storage: src5, event: SRC5Event);
+
+    // ERC721
+    #[abi(embed_v0)]
+    impl ERC721MixinImpl = ERC721Component::ERC721Impl<ContractState>;
+    impl ERC721InternalImpl = ERC721Component::InternalImpl<ContractState>;
+
+    // ERC721Enumerable
+    #[abi(embed_v0)]
+    impl ERC721EnumerableImpl =
+        ERC721EnumerableComponent::ERC721EnumerableImpl<ContractState>;
+    impl ERC721EnumerableInternalImpl = ERC721EnumerableComponent::InternalImpl<ContractState>;
+
+    // ERC721EnumerableExtended
+    #[abi(embed_v0)]
+    impl ERC721EnumerableExtendedImpl =
+        ERC721EnumerableComponent::ERC721EnumerableExtendedImpl<ContractState>;
+
+    // SRC5
+    #[abi(embed_v0)]
+    impl SRC5Impl = SRC5Component::SRC5Impl<ContractState>;
+
+    #[storage]
+    pub struct Storage {
+        #[substorage(v0)]
+        pub erc721: ERC721Component::Storage,
+        #[substorage(v0)]
+        pub erc721_enumerable: ERC721EnumerableComponent::Storage,
+        #[substorage(v0)]
+        pub src5: SRC5Component::Storage
+    }
+
+    #[event]
+    #[derive(Drop, starknet::Event)]
+    enum Event {
+        #[flat]
+        ERC721Event: ERC721Component::Event,
+        #[flat]
+        ERC721EnumerableEvent: ERC721EnumerableComponent::Event,
+        #[flat]
+        SRC5Event: SRC5Component::Event
+    }
+
+    impl ERC721HooksImpl of ERC721Component::ERC721HooksTrait<ContractState> {
+        fn before_update(
+            ref self: ERC721Component::ComponentState<ContractState>,
+            to: ContractAddress,
+            token_id: u256,
+            auth: ContractAddress
+        ) {
+            let mut contract_state = ERC721Component::HasComponent::get_contract_mut(ref self);
+            contract_state.erc721_enumerable.before_update(to, token_id);
+        }
+    }
+
+    #[constructor]
+    fn constructor(
+        ref self: ContractState,
+        name: ByteArray,
+        symbol: ByteArray,
+        base_uri: ByteArray,
+        recipient: ContractAddress,
+        token_id: u256
+    ) {
+        self.erc721.initializer(name, symbol, base_uri);
+        self.erc721_enumerable.initializer();
+        self.erc721.mint(recipient, token_id);
+    }
+}
+
+#[starknet::contract]
+pub mod SnakeERC721EnumerableMock {
+    use openzeppelin_introspection::src5::SRC5Component;
+    use openzeppelin_token::erc721::ERC721Component;
+    use openzeppelin_token::erc721::extensions::ERC721EnumerableComponent;
+
+    use starknet::ContractAddress;
+
+    component!(path: ERC721Component, storage: erc721, event: ERC721Event);
+    component!(
+        path: ERC721EnumerableComponent, storage: erc721_enumerable, event: ERC721EnumerableEvent
+    );
+    component!(path: SRC5Component, storage: src5, event: SRC5Event);
+
+    // ERC721
+    impl ERC721InternalImpl = ERC721Component::InternalImpl<ContractState>;
+
+    // ERC721Enumerable
+    #[abi(embed_v0)]
+    impl ERC721EnumerableImpl =
+        ERC721EnumerableComponent::ERC721EnumerableImpl<ContractState>;
+    impl ERC721EnumerableInternalImpl = ERC721EnumerableComponent::InternalImpl<ContractState>;
+
+    // SRC5
+    impl SRC5Impl = SRC5Component::SRC5Impl<ContractState>;
+
+    #[storage]
+    pub struct Storage {
+        #[substorage(v0)]
+        pub erc721: ERC721Component::Storage,
+        #[substorage(v0)]
+        pub erc721_enumerable: ERC721EnumerableComponent::Storage,
+        #[substorage(v0)]
+        pub src5: SRC5Component::Storage
+    }
+
+    #[event]
+    #[derive(Drop, starknet::Event)]
+    enum Event {
+        #[flat]
+        ERC721Event: ERC721Component::Event,
+        #[flat]
+        ERC721EnumerableEvent: ERC721EnumerableComponent::Event,
+        #[flat]
+        SRC5Event: SRC5Component::Event
+    }
+
+    impl ERC721HooksImpl of ERC721Component::ERC721HooksTrait<ContractState> {
+        fn before_update(
+            ref self: ERC721Component::ComponentState<ContractState>,
+            to: ContractAddress,
+            token_id: u256,
+            auth: ContractAddress
+        ) {
+            let mut contract_state = ERC721Component::HasComponent::get_contract_mut(ref self);
+            contract_state.erc721_enumerable.before_update(to, token_id);
+        }
+    }
+
+    #[constructor]
+    fn constructor(
+        ref self: ContractState,
+        name: ByteArray,
+        symbol: ByteArray,
+        base_uri: ByteArray,
+        recipient: ContractAddress,
+        token_id: u256
+    ) {
+        self.erc721.initializer(name, symbol, base_uri);
+        self.erc721_enumerable.initializer();
+        self.erc721.mint(recipient, token_id);
+    }
+}

packages/token/src/erc721/extensions/erc721_enumerable/erc721_enumerable.cairo

@@ -20,7 +20,6 @@ pub mod ERC721EnumerableComponent {
     use crate::erc721::ERC721Component::ERC721Impl;
     use crate::erc721::ERC721Component::InternalImpl as ERC721InternalImpl;
     use crate::erc721::ERC721Component;
-    use crate::erc721::extensions::erc721_enumerable::interface::IERC721Enumerable;
     use crate::erc721::extensions::erc721_enumerable::interface;
     use openzeppelin_introspection::src5::SRC5Component::InternalTrait as SRC5InternalTrait;
     use openzeppelin_introspection::src5::SRC5Component;
@@ -51,7 +50,7 @@ pub mod ERC721EnumerableComponent {
         +ERC721Component::ERC721HooksTrait<TContractState>,
         +SRC5Component::HasComponent<TContractState>,
         +Drop<TContractState>
-    > of IERC721Enumerable<ComponentState<TContractState>> {
+    > of interface::IERC721Enumerable<ComponentState<TContractState>> {
         /// Returns the total amount of tokens stored by the contract.
         fn total_supply(self: @ComponentState<TContractState>) -> u256 {
             self.ERC721Enumerable_all_tokens_len.read()
@@ -84,6 +83,36 @@ pub mod ERC721EnumerableComponent {
         }
     }
 
+    #[embeddable_as(ERC721EnumerableExtendedImpl)]
+    impl ERC721EnumerableExtended<
+        TContractState,
+        +HasComponent<TContractState>,
+        impl ERC721: ERC721Component::HasComponent<TContractState>,
+        +ERC721Component::ERC721HooksTrait<TContractState>,
+        +SRC5Component::HasComponent<TContractState>,
+        +Drop<TContractState>
+    > of interface::IERC721EnumerableExtended<ComponentState<TContractState>> {
+        /// Returns a list of all token ids owned by the specified `owner`.
+        /// This function provides a more efficient alternative to calling `ERC721::balance_of`
+        /// and iterating through tokens with `ERC721Enumerable::token_of_owner_by_index`.
+        ///
+        /// Requirements:
+        ///
+        /// - `owner` is not the zero address.
+        fn all_tokens_of_owner(
+            self: @ComponentState<TContractState>, owner: ContractAddress
+        ) -> Span<u256> {
+            let erc721_component = get_dep_component!(self, ERC721);
+            let balance = erc721_component.balance_of(owner);
+            let mut result = array![];
+            for index in 0
+                ..balance {
+                    result.append(self.ERC721Enumerable_owned_tokens.read((owner, index)));
+                };
+            result.span()
+        }
+    }
+
     //
     // Internal
     //

packages/token/src/erc721/extensions/erc721_enumerable/interface.cairo

@@ -13,3 +13,8 @@ pub trait IERC721Enumerable<TState> {
     fn token_by_index(self: @TState, index: u256) -> u256;
     fn token_of_owner_by_index(self: @TState, owner: ContractAddress, index: u256) -> u256;
 }
+
+#[starknet::interface]
+pub trait IERC721EnumerableExtended<TState> {
+    fn all_tokens_of_owner(self: @TState, owner: ContractAddress) -> Span<u256>;
+}

packages/token/src/tests/erc721/test_erc721_enumerable.cairo

@@ -1,6 +1,6 @@
 use crate::erc721::ERC721Component::{ERC721Impl, InternalImpl as ERC721InternalImpl};
 use crate::erc721::extensions::erc721_enumerable::ERC721EnumerableComponent::{
-    ERC721EnumerableImpl, InternalImpl
+    ERC721EnumerableImpl, ERC721EnumerableExtendedImpl, InternalImpl
 };
 use crate::erc721::extensions::erc721_enumerable::ERC721EnumerableComponent;
 use crate::erc721::extensions::erc721_enumerable::interface;
@@ -503,6 +503,27 @@ fn test__remove_token_from_all_tokens_enumeration_with_first_token() {
     assert_eq!(initial_supply - 1, new_supply);
 }
 
+//
+// all_tokens_of_owner
+//
+
+#[test]
+fn test_all_tokens_of_owner() {
+    let (_, tokens_list) = setup();
+    assert_all_tokens_of_owner(OWNER(), tokens_list);
+}
+
+#[test]
+fn test_all_tokens_of_owner_after_transfer() {
+    let _ = setup();
+    let mut contract_state = CONTRACT_STATE();
+
+    contract_state.erc721.transfer(OWNER(), RECIPIENT(), TOKEN_2);
+
+    assert_all_tokens_of_owner(OWNER(), array![TOKEN_1, TOKEN_3].span());
+    assert_all_tokens_of_owner(RECIPIENT(), array![TOKEN_2].span());
+}
+
 //
 // Helpers
 //
@@ -568,21 +589,24 @@ fn assert_all_tokens_index_to_id(index: u256, exp_token_id: u256) {
 
 fn assert_all_tokens_id_to_index(token_id: u256, exp_index: u256) {
     let state = @COMPONENT_STATE();
-
     let id_to_index = state.ERC721Enumerable_all_tokens_index.read(token_id);
     assert_eq!(id_to_index, exp_index);
 }
 
 fn assert_owner_tokens_index_to_id(owner: ContractAddress, index: u256, exp_token_id: u256) {
     let state = @COMPONENT_STATE();
-
     let index_to_id = state.ERC721Enumerable_owned_tokens.read((owner, index));
     assert_eq!(index_to_id, exp_token_id);
 }
 
 fn assert_owner_tokens_id_to_index(token_id: u256, exp_index: u256) {
     let state = @COMPONENT_STATE();
-
     let id_to_index = state.ERC721Enumerable_owned_tokens_index.read(token_id);
     assert_eq!(id_to_index, exp_index);
 }
+
+fn assert_all_tokens_of_owner(owner: ContractAddress, exp_tokens: Span<u256>) {
+    let state = @COMPONENT_STATE();
+    let tokens = state.all_tokens_of_owner(owner);
+    assert_eq!(tokens, exp_tokens);
+}