Implement tokensByIndex extension

- Remove restrictions from mock mint and burn calls

contracts/mocks/ERC721BasicTokenMock.sol

@@ -12,6 +12,6 @@ contract ERC721BasicTokenMock is ERC721BasicToken {
   }
 
   function burn(uint256 _tokenId) public {
-    super.doBurn(_tokenId);
+    super.doBurn(ownerOf(_tokenId), _tokenId);
   }
 }

contracts/mocks/ERC721TokenMock.sol

@@ -1,20 +1,26 @@
 pragma solidity ^0.4.18;
 
-import "./ERC721BasicTokenMock.sol";
 import "../token/ERC721/ERC721Token.sol";
 
 /**
  * @title ERC721TokenMock
  * This mock just provides a public mint and burn functions for testing purposes,
  * and a mock metadata URI implementation
  */
-contract ERC721TokenMock is ERC721Token, ERC721BasicTokenMock {
+contract ERC721TokenMock is ERC721Token {
   function ERC721TokenMock(string name, string symbol)
-  ERC721BasicTokenMock()
   ERC721Token(name, symbol)
   public
   { }
 
+  function mint(address _to, uint256 _tokenId) public {
+    doMint(_to, _tokenId);
+  }
+
+  function burn(uint256 _tokenId) public {
+    doBurn(ownerOf(_tokenId), _tokenId);
+  }
+
   // Mock implementation for testing.
   // Do not use this code in production!
   function tokenURI(uint256 _tokenId) public view returns (string) {

contracts/token/ERC721/ERC721.sol

@@ -7,7 +7,7 @@ import "./ERC721Basic.sol";
 contract ERC721Enumerable is ERC721Basic {
   function totalSupply() public view returns (uint256);
   function tokenOfOwnerByIndex(address _owner, uint256 _index) public view returns (uint256 _tokenId);
-  // function tokenByIndex(uint256 _index) public view returns (uint256);
+  function tokenByIndex(uint256 _index) public view returns (uint256);
 }
 
 /// @title ERC-721 Non-Fungible Token Standard, optional metadata extension

contracts/token/ERC721/ERC721BasicToken.sol

@@ -19,9 +19,6 @@ contract ERC721BasicToken is ERC721Basic {
   // Equals to bytes4(keccak256("onERC721Received(address,uint256,bytes)"))
   bytes4 ERC721_RECEIVED = 0xf0b9e5ba; 
 
-  // Total amount of tokens
-  uint256 internal totalTokens;
-
   // Mapping from token ID to owner
   mapping (uint256 => address) internal tokenOwner;
 
@@ -194,7 +191,6 @@ contract ERC721BasicToken is ERC721Basic {
   function doMint(address _to, uint256 _tokenId) internal {
     require(_to != address(0));
     addToken(_to, _tokenId);
-    totalTokens = totalTokens.add(1);
     Transfer(0x0, _to, _tokenId);
   }
 
@@ -203,11 +199,10 @@ contract ERC721BasicToken is ERC721Basic {
   * @dev Reverts if the token does not exist
   * @param _tokenId uint256 ID of the token being burned by the msg.sender
   */
-  function doBurn(uint256 _tokenId) onlyOwnerOf(_tokenId) internal {
-    clearApproval(msg.sender, _tokenId);
-    removeToken(msg.sender, _tokenId);
-    totalTokens = totalTokens.sub(1);
-    Transfer(msg.sender, 0x0, _tokenId);
+  function doBurn(address _owner, uint256 _tokenId) internal {
+    clearApproval(_owner, _tokenId);
+    removeToken(_owner, _tokenId);
+    Transfer(_owner, 0x0, _tokenId);
   }
 
   /**

contracts/token/ERC721/ERC721Token.sol

@@ -24,6 +24,12 @@ contract ERC721Token is ERC721, ERC721BasicToken {
   // Mapping from token ID to index of the owner tokens list
   mapping(uint256 => uint256) internal ownedTokensIndex;
 
+  // Array with all token ids, used for enumeration
+  uint256[] internal allTokens;
+
+  // Mapping from token id to position in the allTokens array
+  mapping(uint256 => uint256) internal allTokensIndex;
+
   /**
   * @dev Constructor function
   */
@@ -73,7 +79,18 @@ contract ERC721Token is ERC721, ERC721BasicToken {
   * @return uint256 representing the total amount of tokens
   */
   function totalSupply() public view returns (uint256) {
-    return totalTokens;
+    return allTokens.length;
+  }
+
+  /**
+  * @dev Gets the token ID at a given index of all the tokens in this contract
+  * @dev Reverts if the index is greater or equal to the total number of tokens
+  * @param _index uint256 representing the index to be accessed of the tokens list
+  * @return uint256 token ID at the given index of the tokens list
+  */
+  function tokenByIndex(uint256 _index) public view returns (uint256) {
+    require(_index < totalSupply());
+    return allTokens[_index];
   }
 
   /**
@@ -113,4 +130,38 @@ contract ERC721Token is ERC721, ERC721BasicToken {
     ownedTokensIndex[lastToken] = tokenIndex;
   }
 
+  /**
+  * @dev Internal function to mint a new token
+  * @dev Reverts if the given token ID already exists
+  * @param _to address the beneficiary that will own the minted token
+  * @param _tokenId uint256 ID of the token to be minted by the msg.sender
+  */
+  function doMint(address _to, uint256 _tokenId) internal {
+    super.doMint(_to, _tokenId);
+
+    allTokensIndex[_tokenId] = allTokens.length;
+    allTokens.push(_tokenId);
+  }
+
+  /**
+  * @dev Internal function to burn a specific token
+  * @dev Reverts if the token does not exist
+  * @param _owner owner of the token to burn
+  * @param _tokenId uint256 ID of the token being burned by the msg.sender
+  */
+  function doBurn(address _owner, uint256 _tokenId) internal {
+    super.doBurn(_owner, _tokenId);
+
+    uint256 tokenIndex = allTokensIndex[_tokenId];
+    uint256 lastTokenIndex = allTokens.length.sub(1);
+    uint256 lastToken = allTokens[lastTokenIndex];
+
+    allTokens[tokenIndex] = lastToken;
+    allTokens[lastTokenIndex] = 0;
+
+    allTokens.length--;
+    allTokensIndex[_tokenId] = 0;
+    allTokensIndex[lastToken] = tokenIndex;
+  }
+
 }

test/token/ERC721/ERC721BasicToken.behaviour.js

@@ -138,13 +138,13 @@ export default function shouldBehaveLikeERC721BasicToken (accounts) {
         });
 
         it('adjusts owners tokens by index', async function () {
-          if (!this.token.tokensOfOwnerByIndex) return;
+          if (!this.token.tokenOfOwnerByIndex) return;
 
-          const newOwnerToken = await this.tokensOfOwnerByIndex(this.to, 0);
-          newOwnerToken.should.be.equal(tokenId);
+          const newOwnerToken = await this.token.tokenOfOwnerByIndex(this.to, 0);
+          newOwnerToken.toNumber().should.be.equal(tokenId);
 
-          const previousOwnerToken = await this.tokensOfOwnerByIndex(owner, 0);
-          previousOwnerToken.should.not.be.equal(tokenId);
+          const previousOwnerToken = await this.token.tokenOfOwnerByIndex(owner, 0);
+          previousOwnerToken.toNumber().should.not.be.equal(tokenId);
         });
       };
 

test/token/ERC721/ERC721MintBurn.behaviour.js

@@ -40,13 +40,6 @@ export default function shouldMintAndBurnERC721Token (accounts) {
           balance.should.be.bignumber.equal(1);
         });
 
-        it('adjusts owner tokens by index', async function () {
-          if (!this.token.tokensOfOwnerByIndex) return;
-
-          const token = await this.tokensOfOwnerByIndex(to, 0);
-          token.should.be.equal(tokenId);
-        });
-
         it('emits a transfer event', async function () {
           logs.length.should.be.equal(1);
           logs[0].event.should.be.eq('Transfer');
@@ -86,13 +79,6 @@ export default function shouldMintAndBurnERC721Token (accounts) {
           balance.should.be.bignumber.equal(1);
         });
 
-        it('removes that token from the token list of the owner', async function () {
-          if (!this.token.tokensOfOwnerByIndex) return;
-
-          const token = await this.tokensOfOwnerByIndex(sender, 0);
-          token.should.not.be.equal(tokenId);
-        });
-
         it('emits a burn event', async function () {
           logs.length.should.be.equal(1);
           logs[0].event.should.be.eq('Transfer');
@@ -126,12 +112,6 @@ export default function shouldMintAndBurnERC721Token (accounts) {
         });
       });
 
-      describe('when the msg.sender does not own given token', function () {
-        it('reverts', async function () {
-          await assertRevert(this.token.burn(tokenId, { from: accounts[1] }));
-        });
-      });
-
       describe('when the given token ID was not tracked by this contract', function () {
         it('reverts', async function () {
           await assertRevert(this.token.burn(unknownTokenId, { from: creator }));

test/token/ERC721/ERC721Token.test.js

@@ -1,6 +1,7 @@
 import assertRevert from '../../helpers/assertRevert';
 import shouldBehaveLikeERC721BasicToken from './ERC721BasicToken.behaviour';
 import shouldMintAndBurnERC721Token from './ERC721MintBurn.behaviour';
+import _ from 'lodash';
 
 const BigNumber = web3.BigNumber;
 const ERC721Token = artifacts.require('ERC721TokenMock.sol');
@@ -13,8 +14,8 @@ require('chai')
 contract('ERC721Token', function (accounts) {
   const name = 'Non Fungible Token';
   const symbol = 'NFT';
-  const firstTokenId = 1;
-  const secondTokenId = 2;
+  const firstTokenId = 100;
+  const secondTokenId = 200;
   const creator = accounts[0];
 
   beforeEach(async function () {
@@ -30,6 +31,51 @@ contract('ERC721Token', function (accounts) {
       await this.token.mint(creator, secondTokenId, { from: creator });
     });
 
+    describe('mint', function () {
+      const to = accounts[1];
+      const tokenId = 3;
+
+      beforeEach(async function () {
+        await this.token.mint(to, tokenId);
+      });
+
+      it('adjusts owner tokens by index', async function () {
+        const token = await this.token.tokenOfOwnerByIndex(to, 0);
+        token.toNumber().should.be.equal(tokenId);
+      });
+
+      it('adjusts all tokens list', async function () {
+        const newToken = await this.token.tokenByIndex(2);
+        newToken.toNumber().should.be.equal(tokenId);
+      });
+    });
+
+    describe('burn', function () {
+      const tokenId = firstTokenId;
+      const sender = creator;
+
+      beforeEach(async function () {
+        await this.token.burn(tokenId, { from: sender });
+      });
+
+      it('removes that token from the token list of the owner', async function () {
+        const token = await this.token.tokenOfOwnerByIndex(sender, 0);
+        token.toNumber().should.be.equal(secondTokenId);
+      });
+
+      it('adjusts all tokens list', async function () {
+        const token = await this.token.tokenByIndex(0);
+        token.toNumber().should.be.equal(secondTokenId);
+      });
+
+      it('burns all tokens', async function () {
+        await this.token.burn(secondTokenId, { from: sender });
+        const total = await this.token.totalSupply();
+        total.toNumber().should.be.equal(0);
+        await assertRevert(this.token.tokenByIndex(0));
+      });
+    });
+
     describe('metadata', function () {
       it('has a name', async function () {
         const name = await this.token.name();
@@ -56,32 +102,78 @@ contract('ERC721Token', function (accounts) {
     });
 
     describe('tokenOfOwnerByIndex', function () {
-      describe('when the given address owns some tokens', function () {
-        const owner = creator;
+      const owner = creator;
+      const another = accounts[1];
+
+      describe('when the given index is lower than the amount of tokens owned by the given address', function () {
+        it('returns the token ID placed at the given index', async function () {
+          const tokenId = await this.token.tokenOfOwnerByIndex(owner, 0);
+          tokenId.should.be.bignumber.equal(firstTokenId);
+        });
+      });
+
+      describe('when the index is greater than or equal to the total tokens owned by the given address', function () {
+        it('reverts', async function () {
+          await assertRevert(this.token.tokenOfOwnerByIndex(owner, 2));
+        });
+      });
 
-        describe('when the given index is lower than the amount of tokens owned by the given address', function () {
-          const index = 0;
+      describe('when the given address does not own any token', function () {
+        it('reverts', async function () {
+          await assertRevert(this.token.tokenOfOwnerByIndex(another, 0));
+        });
+      });
 
-          it('returns the token ID placed at the given index', async function () {
-            const tokenId = await this.token.tokenOfOwnerByIndex(owner, index);
-            tokenId.should.be.bignumber.equal(firstTokenId);
-          });
+      describe('after transferring all tokens to another user', function () {
+        beforeEach(async function () {
+          await this.token.transferFrom(owner, another, firstTokenId, { from: owner });
+          await this.token.transferFrom(owner, another, secondTokenId, { from: owner });
         });
 
-        describe('when the index is greater than or equal to the total tokens owned by the given address', function () {
-          const index = 2;
+        it('returns correct token IDs for target', async function () {
+          const count = await this.token.balanceOf(another);
+          count.toNumber().should.be.equal(2);
+          const tokensListed = await Promise.all(_.range(2).map(i => this.token.tokenOfOwnerByIndex(another, i)));
+          tokensListed.map(t => t.toNumber()).should.have.members([firstTokenId, secondTokenId]);
+        });
 
-          it('reverts', async function () {
-            await assertRevert(this.token.tokenOfOwnerByIndex(owner, index));
-          });
+        it('returns empty collection for original owner', async function () {
+          const count = await this.token.balanceOf(owner);
+          count.toNumber().should.be.equal(0);
+          await assertRevert(this.token.tokenOfOwnerByIndex(owner, 0));
         });
       });
+    });
 
-      describe('when the given address does not own any token', function () {
-        const owner = accounts[1];
+    describe('tokenByIndex', function () {
+      it('should return all tokens', async function () {
+        const tokensListed = await Promise.all(_.range(2).map(i => this.token.tokenByIndex(i)));
+        tokensListed.map(t => t.toNumber()).should.have.members([firstTokenId, secondTokenId]);
+      });
 
-        it('reverts', async function () {
-          await assertRevert(this.token.tokenOfOwnerByIndex(owner, 0));
+      it('should revert if index is greater than supply', async function () {
+        await assertRevert(this.token.tokenByIndex(2));
+      });
+
+      [firstTokenId, secondTokenId].forEach(function (tokenId) {
+        it(`should return all tokens after burning token ${tokenId} and minting new tokens`, async function () {
+          const owner = accounts[0];
+          const newTokenId = 300;
+          const anotherNewTokenId = 400;
+
+          await this.token.burn(tokenId, { from: owner });
+          await this.token.mint(owner, newTokenId, { from: owner });
+          await this.token.mint(owner, anotherNewTokenId, { from: owner });
+
+          const count = await this.token.totalSupply();
+          count.toNumber().should.be.equal(3);
+
+          const tokensListed = await Promise.all(_.range(3).map(i => this.token.tokenByIndex(i)));
+          const expectedTokens = _.filter(
+            [firstTokenId, secondTokenId, newTokenId, anotherNewTokenId],
+            x => (x !== tokenId)
+          );
+          tokensListed.map(t => t.toNumber()).should.have.members(expectedTokens);
         });
       });
     });