From 53426fd8cf628101ab02f246a305c37a2e4a3344 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Dulat?= <michal@dulat.pl>
Date: Mon, 23 Apr 2018 13:34:45 +0200
Subject: [PATCH] Fix clearing of ownedTokensIndex and allTokensIndex in
 ERC721Token.sol It fixes #839.

---
 contracts/mocks/ERC721TokenMock.sol    |  8 ++++++++
 contracts/token/ERC721/ERC721Token.sol | 15 ++++++++++-----
 test/token/ERC721/ERC721Token.test.js  | 25 +++++++++++++++++++------
 3 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/contracts/mocks/ERC721TokenMock.sol b/contracts/mocks/ERC721TokenMock.sol
index 5fe3668efa5..d47e775c77c 100644
--- a/contracts/mocks/ERC721TokenMock.sol
+++ b/contracts/mocks/ERC721TokenMock.sol
@@ -24,4 +24,12 @@ contract ERC721TokenMock is ERC721Token {
   function setTokenURI(uint256 _tokenId, string _uri) public {
     super._setTokenURI(_tokenId, _uri);
   }
+
+  function ownedTokenIndex(uint256 _tokenId) public view returns(uint256) {
+    return ownedTokensIndex[_tokenId];
+  }
+
+  function allTokenIndex(uint256 _tokenId) public view returns(uint256) {
+    return allTokensIndex[_tokenId];
+  }
 }
diff --git a/contracts/token/ERC721/ERC721Token.sol b/contracts/token/ERC721/ERC721Token.sol
index 7e43ec7e959..ca9f4335fec 100644
--- a/contracts/token/ERC721/ERC721Token.sol
+++ b/contracts/token/ERC721/ERC721Token.sol
@@ -138,8 +138,11 @@ contract ERC721Token is ERC721, ERC721BasicToken {
     // the lastToken to the first position, and then dropping the element placed in the last position of the list
 
     ownedTokens[_from].length--;
-    ownedTokensIndex[_tokenId] = 0;
-    ownedTokensIndex[lastToken] = tokenIndex;
+    delete ownedTokensIndex[_tokenId];
+
+    if (_tokenId != lastToken) {
+      ownedTokensIndex[lastToken] = tokenIndex;
+    }
   }
 
   /**
@@ -178,8 +181,10 @@ contract ERC721Token is ERC721, ERC721BasicToken {
     allTokens[lastTokenIndex] = 0;
 
     allTokens.length--;
-    allTokensIndex[_tokenId] = 0;
-    allTokensIndex[lastToken] = tokenIndex;
-  }
+    delete allTokensIndex[_tokenId];
 
+    if (_tokenId != lastToken) {
+      allTokensIndex[lastToken] = tokenIndex;
+    }
+  }
 }
diff --git a/test/token/ERC721/ERC721Token.test.js b/test/token/ERC721/ERC721Token.test.js
index 9d53a6cef09..73c10b9c003 100644
--- a/test/token/ERC721/ERC721Token.test.js
+++ b/test/token/ERC721/ERC721Token.test.js
@@ -51,7 +51,7 @@ contract('ERC721Token', function (accounts) {
     });
 
     describe('burn', function () {
-      const tokenId = firstTokenId;
+      const tokenId = secondTokenId;
       const sender = creator;
 
       beforeEach(async function () {
@@ -59,17 +59,30 @@ contract('ERC721Token', function (accounts) {
       });
 
       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);
+        const firstToken = await this.token.tokenOfOwnerByIndex(sender, 0);
+        firstToken.toNumber().should.be.equal(firstTokenId);
+        await assertRevert(this.token.tokenOfOwnerByIndex(sender, 1));
       });
 
       it('adjusts all tokens list', async function () {
-        const token = await this.token.tokenByIndex(0);
-        token.toNumber().should.be.equal(secondTokenId);
+        const firstToken = await this.token.tokenByIndex(0);
+        firstToken.toNumber().should.be.equal(firstTokenId);
+        await assertRevert(this.token.tokenByIndex(1));
+      });
+
+      ['ownedTokenIndex', 'allTokenIndex'].forEach(fnName => {
+        const mappingName = fnName.replace('Token', 'Tokens');
+
+        it(`adjusts ${mappingName}`, async function () {
+          const firstTokenIndex = await this.token[fnName](firstTokenId);
+          const secondTokenIndex = await this.token[fnName](secondTokenId);
+          firstTokenIndex.toNumber().should.be.equal(0);
+          secondTokenIndex.toNumber().should.be.equal(0);
+        });
       });
 
       it('burns all tokens', async function () {
-        await this.token.burn(secondTokenId, { from: sender });
+        await this.token.burn(firstTokenId, { from: sender });
         const total = await this.token.totalSupply();
         total.toNumber().should.be.equal(0);
         await assertRevert(this.token.tokenByIndex(0));