Gas optimization storage NFTs

[code423n4]

Handle

a_delamo

Vulnerability details

Impact

In Visor.sol, NFTs are being stored using an array

    Nft[] public nfts;

This seems an optimal structure, but when needing to remove an NFT or look for an NFT by contract and id, we need to do O(n) iterations.

    function _removeNft(address nftContract, uint256 tokenId) internal {
        uint256 len = nfts.length;
        for (uint256 i = 0; i < len; i++) {
            Nft memory nftInfo = nfts[i];
            if (
                nftContract == nftInfo.nftContract && tokenId == nftInfo.tokenId
            ) {
                if (i != len - 1) {
                    nfts[i] = nfts[len - 1];
                }
                nfts.pop();
                emit RemoveNftToken(nftContract, tokenId);
                break;
            }
        }
    }

    function getNftIdByTokenIdAndAddr(address nftContract, uint256 tokenId)
        external
        view
        returns (uint256)
    {
        uint256 len = nfts.length;
        for (uint256 i = 0; i < len; i++) {
            if (
                nftContract == nfts[i].nftContract && tokenId == nfts[i].tokenId
            ) {
                return i;
            }
        }
        require(false, "Token not found");
    }

So, seems like we could improve these functions by storing the NFTs like a mapping. This is all the functions using the NFTs array:

function _addNft(address nftContract, uint256 tokenId) internal {
        nfts.push(Nft({tokenId: tokenId, nftContract: nftContract}));
        emit AddNftToken(nftContract, tokenId);
    }

    function _removeNft(address nftContract, uint256 tokenId) internal {
        uint256 len = nfts.length;
        for (uint256 i = 0; i < len; i++) {
            Nft memory nftInfo = nfts[i];
            if (
                nftContract == nftInfo.nftContract && tokenId == nftInfo.tokenId
            ) {
                if (i != len - 1) {
                    nfts[i] = nfts[len - 1];
                }
                nfts.pop();
                emit RemoveNftToken(nftContract, tokenId);
                break;
            }
        }
    }

function getNftById(uint256 i)
        external
        view
        returns (address nftContract, uint256 tokenId)
    {
        require(i < nfts.length, "ID overflow");

        //FIXME: Storage keyword
        Nft memory ni = nfts[i];
        nftContract = ni.nftContract;
        tokenId = ni.tokenId;
    }

    // @notice Get index of ERC721 in nfts[]
    /// @param nftContract Address of ERC721
    /// @param tokenId tokenId for NFT in nftContract
    function getNftIdByTokenIdAndAddr(address nftContract, uint256 tokenId)
        external
        view
        returns (uint256)
    {
        uint256 len = nfts.length;
        for (uint256 i = 0; i < len; i++) {
            if (
                nftContract == nfts[i].nftContract && tokenId == nfts[i].tokenId
            ) {
                return i;
            }
        }
        require(false, "Token not found");
    }

Checking the function getNftById, seems like a lookup by id (mapping(uint=> NFT)) should be fine. In case we need more, we could do mapping(uint => mapping(address => NFT)), but doesn't seems necessary.

Proof of Concept

Provide direct links to all referenced code in GitHub. Add screenshots, logs, or any other relevant proof that illustrates the concept.

Tools Used

None

Recommended Mitigation Steps

[ghost]

sponsor confirmed We will be applying this optimization in our next version

[ztcrypto]

patch link