Add EnumerableMap, refactor ERC721 (#2160)

* Implement AddressSet in terms of a generic Set

* Add Uint256Set

* Add EnumerableMap

* Fix wording on EnumerableSet docs and tests

* Refactor ERC721 using EnumerableSet and EnumerableMap

* Fix tests

* Fix linter error

* Gas optimization for EnumerableMap

* Gas optimization for EnumerableSet

* Remove often not-taken if from Enumerable data structures

* Fix failing test

* Gas optimization for EnumerableMap

* Fix linter errors

* Add comment for clarification

* Improve test naming

* Rename EnumerableMap.add to set

* Add overload for EnumerableMap.get with custom error message

* Improve Enumerable docs

* Rename Uint256Set to UintSet

* Add changelog entry
This commit is contained in:
Nicolás Venturo
2020-04-02 15:43:06 -03:00
committed by GitHub
parent 0408e51ae6
commit bd0778461d
12 changed files with 606 additions and 252 deletions

View File

@ -176,27 +176,17 @@ describe('ERC721', function () {
expect(await this.token.ownerOf(tokenId)).to.be.equal(this.toWhom);
});
it('emits a Transfer event', async function () {
expectEvent.inLogs(logs, 'Transfer', { from: owner, to: this.toWhom, tokenId: tokenId });
});
it('clears the approval for the token ID', async function () {
expect(await this.token.getApproved(tokenId)).to.be.equal(ZERO_ADDRESS);
});
if (approved) {
it('emit only a transfer event', async function () {
expectEvent.inLogs(logs, 'Transfer', {
from: owner,
to: this.toWhom,
tokenId: tokenId,
});
});
} else {
it('emits only a transfer event', async function () {
expectEvent.inLogs(logs, 'Transfer', {
from: owner,
to: this.toWhom,
tokenId: tokenId,
});
});
}
it('emits an Approval event', async function () {
expectEvent.inLogs(logs, 'Approval', { owner, approved: ZERO_ADDRESS, tokenId: tokenId });
});
it('adjusts owners balances', async function () {
expect(await this.token.balanceOf(owner)).to.be.bignumber.equal('1');
@ -708,15 +698,6 @@ describe('ERC721', function () {
});
});
describe('tokensOfOwner', function () {
it('returns total tokens of owner', async function () {
const tokenIds = await this.token.tokensOfOwner(owner);
expect(tokenIds.length).to.equal(2);
expect(tokenIds[0]).to.be.bignumber.equal(firstTokenId);
expect(tokenIds[1]).to.be.bignumber.equal(secondTokenId);
});
});
describe('totalSupply', function () {
it('returns total token supply', async function () {
expect(await this.token.totalSupply()).to.be.bignumber.equal('2');
@ -733,7 +714,7 @@ describe('ERC721', function () {
describe('when the index is greater than or equal to the total tokens owned by the given address', function () {
it('reverts', async function () {
await expectRevert(
this.token.tokenOfOwnerByIndex(owner, 2), 'ERC721Enumerable: owner index out of bounds'
this.token.tokenOfOwnerByIndex(owner, 2), 'EnumerableSet: index out of bounds'
);
});
});
@ -741,7 +722,7 @@ describe('ERC721', function () {
describe('when the given address does not own any token', function () {
it('reverts', async function () {
await expectRevert(
this.token.tokenOfOwnerByIndex(other, 0), 'ERC721Enumerable: owner index out of bounds'
this.token.tokenOfOwnerByIndex(other, 0), 'EnumerableSet: index out of bounds'
);
});
});
@ -764,7 +745,7 @@ describe('ERC721', function () {
it('returns empty collection for original owner', async function () {
expect(await this.token.balanceOf(owner)).to.be.bignumber.equal('0');
await expectRevert(
this.token.tokenOfOwnerByIndex(owner, 0), 'ERC721Enumerable: owner index out of bounds'
this.token.tokenOfOwnerByIndex(owner, 0), 'EnumerableSet: index out of bounds'
);
});
});
@ -781,7 +762,7 @@ describe('ERC721', function () {
it('should revert if index is greater than supply', async function () {
await expectRevert(
this.token.tokenByIndex(2), 'ERC721Enumerable: global index out of bounds'
this.token.tokenByIndex(2), 'EnumerableMap: index out of bounds'
);
});
@ -790,7 +771,7 @@ describe('ERC721', function () {
const newTokenId = new BN(300);
const anotherNewTokenId = new BN(400);
await this.token.burn(tokenId, { from: owner });
await this.token.burn(tokenId);
await this.token.mint(newOwner, newTokenId);
await this.token.mint(newOwner, anotherNewTokenId);
@ -865,6 +846,10 @@ describe('ERC721', function () {
expectEvent.inLogs(this.logs, 'Transfer', { from: owner, to: ZERO_ADDRESS, tokenId: firstTokenId });
});
it('emits an Approval event', function () {
expectEvent.inLogs(this.logs, 'Approval', { owner, approved: ZERO_ADDRESS, tokenId: firstTokenId });
});
it('deletes the token', async function () {
expect(await this.token.balanceOf(owner)).to.be.bignumber.equal('1');
await expectRevert(
@ -884,7 +869,7 @@ describe('ERC721', function () {
await this.token.burn(secondTokenId, { from: owner });
expect(await this.token.totalSupply()).to.be.bignumber.equal('0');
await expectRevert(
this.token.tokenByIndex(0), 'ERC721Enumerable: global index out of bounds'
this.token.tokenByIndex(0), 'EnumerableMap: index out of bounds'
);
});

View File

@ -0,0 +1,139 @@
const { accounts, contract } = require('@openzeppelin/test-environment');
const { BN, expectEvent } = require('@openzeppelin/test-helpers');
const { expect } = require('chai');
const zip = require('lodash.zip');
const EnumerableMapMock = contract.fromArtifact('EnumerableMapMock');
describe('EnumerableMap', function () {
const [ accountA, accountB, accountC ] = accounts;
const keyA = new BN('7891');
const keyB = new BN('451');
const keyC = new BN('9592328');
beforeEach(async function () {
this.map = await EnumerableMapMock.new();
});
async function expectMembersMatch (map, keys, values) {
expect(keys.length).to.equal(values.length);
await Promise.all(keys.map(async key =>
expect(await map.contains(key)).to.equal(true)
));
expect(await map.length()).to.bignumber.equal(keys.length.toString());
expect(await Promise.all(keys.map(key =>
map.get(key)
))).to.have.same.members(values);
// To compare key-value pairs, we zip keys and values, and convert BNs to
// strings to workaround Chai limitations when dealing with nested arrays
expect(await Promise.all([...Array(keys.length).keys()].map(async (index) => {
const entry = await map.at(index);
return [entry.key.toString(), entry.value];
}))).to.have.same.deep.members(
zip(keys.map(k => k.toString()), values)
);
}
it('starts empty', async function () {
expect(await this.map.contains(keyA)).to.equal(false);
await expectMembersMatch(this.map, [], []);
});
it('adds a key', async function () {
const receipt = await this.map.set(keyA, accountA);
expectEvent(receipt, 'OperationResult', { result: true });
await expectMembersMatch(this.map, [keyA], [accountA]);
});
it('adds several keys', async function () {
await this.map.set(keyA, accountA);
await this.map.set(keyB, accountB);
await expectMembersMatch(this.map, [keyA, keyB], [accountA, accountB]);
expect(await this.map.contains(keyC)).to.equal(false);
});
it('returns false when adding keys already in the set', async function () {
await this.map.set(keyA, accountA);
const receipt = (await this.map.set(keyA, accountA));
expectEvent(receipt, 'OperationResult', { result: false });
await expectMembersMatch(this.map, [keyA], [accountA]);
});
it('updates values for keys already in the set', async function () {
await this.map.set(keyA, accountA);
await this.map.set(keyA, accountB);
await expectMembersMatch(this.map, [keyA], [accountB]);
});
it('removes added keys', async function () {
await this.map.set(keyA, accountA);
const receipt = await this.map.remove(keyA);
expectEvent(receipt, 'OperationResult', { result: true });
expect(await this.map.contains(keyA)).to.equal(false);
await expectMembersMatch(this.map, [], []);
});
it('returns false when removing keys not in the set', async function () {
const receipt = await this.map.remove(keyA);
expectEvent(receipt, 'OperationResult', { result: false });
expect(await this.map.contains(keyA)).to.equal(false);
});
it('adds and removes multiple keys', async function () {
// []
await this.map.set(keyA, accountA);
await this.map.set(keyC, accountC);
// [A, C]
await this.map.remove(keyA);
await this.map.remove(keyB);
// [C]
await this.map.set(keyB, accountB);
// [C, B]
await this.map.set(keyA, accountA);
await this.map.remove(keyC);
// [A, B]
await this.map.set(keyA, accountA);
await this.map.set(keyB, accountB);
// [A, B]
await this.map.set(keyC, accountC);
await this.map.remove(keyA);
// [B, C]
await this.map.set(keyA, accountA);
await this.map.remove(keyB);
// [A, C]
await expectMembersMatch(this.map, [keyA, keyC], [accountA, accountC]);
expect(await this.map.contains(keyB)).to.equal(false);
});
});

View File

@ -11,18 +11,16 @@ describe('EnumerableSet', function () {
this.set = await EnumerableSetMock.new();
});
async function expectMembersMatch (set, members) {
await Promise.all(members.map(async account =>
async function expectMembersMatch (set, values) {
await Promise.all(values.map(async account =>
expect(await set.contains(account)).to.equal(true)
));
expect(await set.enumerate()).to.have.same.members(members);
expect(await set.length()).to.bignumber.equal(values.length.toString());
expect(await set.length()).to.bignumber.equal(members.length.toString());
expect(await Promise.all([...Array(members.length).keys()].map(index =>
expect(await Promise.all([...Array(values.length).keys()].map(index =>
set.at(index)
))).to.have.same.members(members);
))).to.have.same.members(values);
}
it('starts empty', async function () {
@ -33,7 +31,7 @@ describe('EnumerableSet', function () {
it('adds a value', async function () {
const receipt = await this.set.add(accountA);
expectEvent(receipt, 'TransactionResult', { result: true });
expectEvent(receipt, 'OperationResult', { result: true });
await expectMembersMatch(this.set, [accountA]);
});
@ -46,11 +44,11 @@ describe('EnumerableSet', function () {
expect(await this.set.contains(accountC)).to.equal(false);
});
it('returns false when adding elements already in the set', async function () {
it('returns false when adding values already in the set', async function () {
await this.set.add(accountA);
const receipt = (await this.set.add(accountA));
expectEvent(receipt, 'TransactionResult', { result: false });
expectEvent(receipt, 'OperationResult', { result: false });
await expectMembersMatch(this.set, [accountA]);
});
@ -63,15 +61,15 @@ describe('EnumerableSet', function () {
await this.set.add(accountA);
const receipt = await this.set.remove(accountA);
expectEvent(receipt, 'TransactionResult', { result: true });
expectEvent(receipt, 'OperationResult', { result: true });
expect(await this.set.contains(accountA)).to.equal(false);
await expectMembersMatch(this.set, []);
});
it('returns false when removing elements not in the set', async function () {
it('returns false when removing values not in the set', async function () {
const receipt = await this.set.remove(accountA);
expectEvent(receipt, 'TransactionResult', { result: false });
expectEvent(receipt, 'OperationResult', { result: false });
expect(await this.set.contains(accountA)).to.equal(false);
});