Merge pull request #109 from lknix/fix_id_to_index_bug

Fix "id to index" bug

contracts/tokens/NFTokenEnumerable.sol

@@ -17,7 +17,7 @@ contract NFTokenEnumerable is
   uint256[] internal tokens;
 
   /**
-   * @dev Mapping from owner address to a list of owned NFT IDs.
+   * @dev Mapping from token ID its index in global tokens array.
    */
   mapping(uint256 => uint256) internal idToIndex;
 
@@ -56,6 +56,7 @@ contract NFTokenEnumerable is
   {
     super._mint(_to, _tokenId);
     tokens.push(_tokenId);
+    idToIndex[_tokenId] = tokens.length.sub(1);
   }
 
   /**
@@ -72,15 +73,16 @@ contract NFTokenEnumerable is
   )
     internal
   {
-    assert(tokens.length > 0);
     super._burn(_owner, _tokenId);
+    assert(tokens.length > 0);
 
     uint256 tokenIndex = idToIndex[_tokenId];
+    // Sanity check. This could be removed in the future.
+    assert(tokens[tokenIndex] == _tokenId);
     uint256 lastTokenIndex = tokens.length.sub(1);
     uint256 lastToken = tokens[lastTokenIndex];
 
     tokens[tokenIndex] = lastToken;
-    tokens[lastTokenIndex] = 0;
 
     tokens.length--;
     idToIndex[_tokenId] = 0;
@@ -155,6 +157,8 @@ contract NFTokenEnumerable is
     returns (uint256)
   {
     require(_index < tokens.length);
+    // Sanity check. This could be removed in the future.
+    assert(idToIndex[tokens[_index]] == _index);
     return tokens[_index];
   }
 

test/tokens/NFTokenEnumerable.test.js

@@ -46,13 +46,19 @@ contract('NFTokenEnumerableMock', (accounts) => {
     await nftoken.mint(accounts[1], id2);
     await nftoken.mint(accounts[2], id3);
 
-    const tokenId = await nftoken.tokenByIndex(1);
+    let tokenId = await nftoken.tokenByIndex(0);
+    assert.equal(tokenId, id1);
+
+    tokenId = await nftoken.tokenByIndex(1);
     assert.equal(tokenId, id2);
+
+    tokenId = await nftoken.tokenByIndex(2);
+    assert.equal(tokenId, id3);
   });
 
-  it('throws when trying to get token by unexistant index', async () => {
+  it('throws when trying to get token by non-existing index', async () => {
     await nftoken.mint(accounts[1], id1);
-    await assertRevert(nftoken.tokenByIndex(1));
+    await assertRevert(nftoken.tokenByIndex(10));
   });
 
   it('returns the correct token of owner by index', async () => {
@@ -64,29 +70,36 @@ contract('NFTokenEnumerableMock', (accounts) => {
     assert.equal(tokenId, id2);
   });
 
-  it('throws when trying to get token of owner by unexistant index', async () => {
+  it('throws when trying to get token of owner by non-existing index', async () => {
     await nftoken.mint(accounts[1], id1);
     await nftoken.mint(accounts[2], id3);
 
-    await assertRevert(nftoken.tokenOfOwnerByIndex(accounts[1], 1));
+    await assertRevert(nftoken.tokenOfOwnerByIndex(accounts[1], 4));
   });
 
   it('corectly burns a NFT', async () => {
+    await nftoken.mint(accounts[1], id1);
     await nftoken.mint(accounts[1], id2);
+    await nftoken.mint(accounts[1], id3);
     const { logs } = await nftoken.burn(accounts[1], id2);
     const transferEvent = logs.find(e => e.event === 'Transfer');
     assert.notEqual(transferEvent, undefined);
 
     const balance = await nftoken.balanceOf(accounts[1]);
-    assert.equal(balance, 0);
+    assert.equal(balance, 2);
 
     await assertRevert(nftoken.ownerOf(id2));
 
     const totalSupply = await nftoken.totalSupply();
-    assert.equal(totalSupply, 0);
+    assert.equal(totalSupply, 2);
 
-    await assertRevert(nftoken.tokenByIndex(0));
-    await assertRevert(nftoken.tokenOfOwnerByIndex(accounts[1], 0));
-  });
+    await assertRevert(nftoken.tokenByIndex(2));
+    await assertRevert(nftoken.tokenOfOwnerByIndex(accounts[1], 2));
+
+    let tokenId = await nftoken.tokenByIndex(0);
+    assert.equal(tokenId, id1);
 
+    tokenId = await nftoken.tokenByIndex(1);
+    assert.equal(tokenId, id3);
+  });
 });