From 69e7005a70d314996ea8395e47fb957109d7b448 Mon Sep 17 00:00:00 2001 From: Hadrien Croubois Date: Mon, 7 Jun 2021 16:26:22 +0200 Subject: [PATCH 1/9] Add a BitMap struct --- contracts/mocks/BitmapMock.sol | 23 ++++++++ contracts/utils/structs/BitMap.sol | 32 +++++++++++ test/utils/structs/BitMap.test.js | 90 ++++++++++++++++++++++++++++++ 3 files changed, 145 insertions(+) create mode 100644 contracts/mocks/BitmapMock.sol create mode 100644 contracts/utils/structs/BitMap.sol create mode 100644 test/utils/structs/BitMap.test.js diff --git a/contracts/mocks/BitmapMock.sol b/contracts/mocks/BitmapMock.sol new file mode 100644 index 00000000000..dc663f7f80a --- /dev/null +++ b/contracts/mocks/BitmapMock.sol @@ -0,0 +1,23 @@ +// SPDX-License-Identifier: MIT + +pragma solidity ^0.8.0; + +import "../utils/structs/BitMap.sol"; + +contract BitMapMock { + using BitMap for BitMap.UintBitMap; + + BitMap.UintBitMap private _bitmap; + + function get(uint256 index) public view returns (bool) { + return _bitmap.get(index); + } + + function set(uint256 index) public { + _bitmap.set(index); + } + + function unset(uint256 index) public { + _bitmap.unset(index); + } +} diff --git a/contracts/utils/structs/BitMap.sol b/contracts/utils/structs/BitMap.sol new file mode 100644 index 00000000000..e6564c69ea1 --- /dev/null +++ b/contracts/utils/structs/BitMap.sol @@ -0,0 +1,32 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.0; + +/** + * @dev Library for managing uint256 to bool mapping in a compact and efficient way, providing the keys are sequential. + * Largelly inspired by Uniswap's https://github.com/Uniswap/merkle-distributor/blob/master/contracts/MerkleDistributor.sol[merkle-distributor]. + */ +library BitMap { + struct UintBitMap { + mapping(uint256 => uint256) _data; + } + + function get(UintBitMap storage bitmap, uint256 index) internal view returns (bool) { + uint256 bucket = index / 256; + uint256 pos = index % 256; + uint256 word = bitmap._data[bucket]; + uint256 mask = (1 << pos); + return word & mask == mask; + } + + function set(UintBitMap storage bitmap, uint256 index) internal { + uint256 bucket = index / 256; + uint256 pos = index % 256; + bitmap._data[bucket] = bitmap._data[bucket] | (1 << pos); + } + + function unset(UintBitMap storage bitmap, uint256 index) internal { + uint256 bucket = index / 256; + uint256 pos = index % 256; + bitmap._data[bucket] = bitmap._data[bucket] & ~(1 << pos); + } +} diff --git a/test/utils/structs/BitMap.test.js b/test/utils/structs/BitMap.test.js new file mode 100644 index 00000000000..71c091505d5 --- /dev/null +++ b/test/utils/structs/BitMap.test.js @@ -0,0 +1,90 @@ +const { BN } = require('@openzeppelin/test-helpers'); +const { expect } = require('chai'); + +const BitMap = artifacts.require('BitMapMock'); + +contract('BitMap', function (accounts) { + const keyA = new BN('7891'); + const keyB = new BN('451'); + const keyC = new BN('9592328'); + + beforeEach(async function () { + this.bitmap = await BitMap.new(); + }); + + it('starts empty', async function () { + expect(await this.bitmap.get(keyA)).to.equal(false); + expect(await this.bitmap.get(keyB)).to.equal(false); + expect(await this.bitmap.get(keyC)).to.equal(false); + }); + + describe('set', function () { + it('adds a key', async function () { + await this.bitmap.set(keyA); + expect(await this.bitmap.get(keyA)).to.equal(true); + expect(await this.bitmap.get(keyB)).to.equal(false); + expect(await this.bitmap.get(keyC)).to.equal(false); + }); + + it('adds several keys', async function () { + await this.bitmap.set(keyA); + await this.bitmap.set(keyB); + expect(await this.bitmap.get(keyA)).to.equal(true); + expect(await this.bitmap.get(keyB)).to.equal(true); + expect(await this.bitmap.get(keyC)).to.equal(false); + }); + }); + + describe('unset', function () { + it('removes added keys', async function () { + await this.bitmap.set(keyA); + await this.bitmap.set(keyB); + await this.bitmap.unset(keyA); + expect(await this.bitmap.get(keyA)).to.equal(false); + expect(await this.bitmap.get(keyB)).to.equal(true); + expect(await this.bitmap.get(keyC)).to.equal(false); + }); + + it('adds and removes multiple keys', async function () { + // [] + + await this.bitmap.set(keyA); + await this.bitmap.set(keyC); + + // [A, C] + + await this.bitmap.unset(keyA); + await this.bitmap.unset(keyB); + + // [C] + + await this.bitmap.set(keyB); + + // [C, B] + + await this.bitmap.set(keyA); + await this.bitmap.unset(keyC); + + // [A, B] + + await this.bitmap.set(keyA); + await this.bitmap.set(keyB); + + // [A, B] + + await this.bitmap.set(keyC); + await this.bitmap.unset(keyA); + + // [B, C] + + await this.bitmap.set(keyA); + await this.bitmap.unset(keyB); + + // [A, C] + + expect(await this.bitmap.get(keyA)).to.equal(true); + expect(await this.bitmap.get(keyB)).to.equal(false); + expect(await this.bitmap.get(keyC)).to.equal(true); + }); + }); +}); From 4df1b0eb6e47d09bd3b01161cf81da55d57ec37f Mon Sep 17 00:00:00 2001 From: Hadrien Croubois Date: Mon, 7 Jun 2021 16:47:17 +0200 Subject: [PATCH 2/9] improve bitmap testing with consecutive keys --- test/utils/structs/BitMap.test.js | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/test/utils/structs/BitMap.test.js b/test/utils/structs/BitMap.test.js index 71c091505d5..22ed9b85b3a 100644 --- a/test/utils/structs/BitMap.test.js +++ b/test/utils/structs/BitMap.test.js @@ -33,6 +33,17 @@ contract('BitMap', function (accounts) { expect(await this.bitmap.get(keyB)).to.equal(true); expect(await this.bitmap.get(keyC)).to.equal(false); }); + + it('adds several consecutive keys', async function () { + await this.bitmap.set(keyA.addn(0)); + await this.bitmap.set(keyA.addn(1)); + await this.bitmap.set(keyA.addn(3)); + expect(await this.bitmap.get(keyA.addn(0))).to.equal(true); + expect(await this.bitmap.get(keyA.addn(1))).to.equal(true); + expect(await this.bitmap.get(keyA.addn(2))).to.equal(false); + expect(await this.bitmap.get(keyA.addn(3))).to.equal(true); + expect(await this.bitmap.get(keyA.addn(4))).to.equal(false); + }); }); describe('unset', function () { @@ -45,6 +56,18 @@ contract('BitMap', function (accounts) { expect(await this.bitmap.get(keyC)).to.equal(false); }); + it('removes consecutive added keys', async function () { + await this.bitmap.set(keyA.addn(0)); + await this.bitmap.set(keyA.addn(1)); + await this.bitmap.set(keyA.addn(3)); + await this.bitmap.unset(keyA.addn(1)); + expect(await this.bitmap.get(keyA.addn(0))).to.equal(true); + expect(await this.bitmap.get(keyA.addn(1))).to.equal(false); + expect(await this.bitmap.get(keyA.addn(2))).to.equal(false); + expect(await this.bitmap.get(keyA.addn(3))).to.equal(true); + expect(await this.bitmap.get(keyA.addn(4))).to.equal(false); + }); + it('adds and removes multiple keys', async function () { // [] From 34dca5809c78e8e905be4e510f88468f86dc67d8 Mon Sep 17 00:00:00 2001 From: Hadrien Croubois Date: Wed, 9 Jun 2021 12:16:29 +0200 Subject: [PATCH 3/9] minor clarity and gas optimization --- contracts/utils/structs/BitMap.sol | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/contracts/utils/structs/BitMap.sol b/contracts/utils/structs/BitMap.sol index e6564c69ea1..55133ffdd80 100644 --- a/contracts/utils/structs/BitMap.sol +++ b/contracts/utils/structs/BitMap.sol @@ -15,18 +15,18 @@ library BitMap { uint256 pos = index % 256; uint256 word = bitmap._data[bucket]; uint256 mask = (1 << pos); - return word & mask == mask; + return word & mask != 0; } function set(UintBitMap storage bitmap, uint256 index) internal { uint256 bucket = index / 256; uint256 pos = index % 256; - bitmap._data[bucket] = bitmap._data[bucket] | (1 << pos); + bitmap._data[bucket] |= (1 << pos); } function unset(UintBitMap storage bitmap, uint256 index) internal { uint256 bucket = index / 256; uint256 pos = index % 256; - bitmap._data[bucket] = bitmap._data[bucket] & ~(1 << pos); + bitmap._data[bucket] &= ~(1 << pos); } } From b43bbefbe11b7deb89b9619fccae3db6433461ea Mon Sep 17 00:00:00 2001 From: Hadrien Croubois Date: Wed, 9 Jun 2021 22:49:25 +0200 Subject: [PATCH 4/9] rename Bitmap.UintBitMap to Bitmaps.Bitmap --- contracts/mocks/BitmapMock.sol | 6 +++--- contracts/utils/structs/{BitMap.sol => BitMaps.sol} | 10 +++++----- 2 files changed, 8 insertions(+), 8 deletions(-) rename contracts/utils/structs/{BitMap.sol => BitMaps.sol} (75%) diff --git a/contracts/mocks/BitmapMock.sol b/contracts/mocks/BitmapMock.sol index dc663f7f80a..ab679795b20 100644 --- a/contracts/mocks/BitmapMock.sol +++ b/contracts/mocks/BitmapMock.sol @@ -2,12 +2,12 @@ pragma solidity ^0.8.0; -import "../utils/structs/BitMap.sol"; +import "../utils/structs/BitMaps.sol"; contract BitMapMock { - using BitMap for BitMap.UintBitMap; + using BitMaps for BitMaps.BitMap; - BitMap.UintBitMap private _bitmap; + BitMaps.BitMap private _bitmap; function get(uint256 index) public view returns (bool) { return _bitmap.get(index); diff --git a/contracts/utils/structs/BitMap.sol b/contracts/utils/structs/BitMaps.sol similarity index 75% rename from contracts/utils/structs/BitMap.sol rename to contracts/utils/structs/BitMaps.sol index 55133ffdd80..859bad85031 100644 --- a/contracts/utils/structs/BitMap.sol +++ b/contracts/utils/structs/BitMaps.sol @@ -5,12 +5,12 @@ pragma solidity ^0.8.0; * @dev Library for managing uint256 to bool mapping in a compact and efficient way, providing the keys are sequential. * Largelly inspired by Uniswap's https://github.com/Uniswap/merkle-distributor/blob/master/contracts/MerkleDistributor.sol[merkle-distributor]. */ -library BitMap { - struct UintBitMap { +library BitMaps { + struct BitMap { mapping(uint256 => uint256) _data; } - function get(UintBitMap storage bitmap, uint256 index) internal view returns (bool) { + function get(BitMap storage bitmap, uint256 index) internal view returns (bool) { uint256 bucket = index / 256; uint256 pos = index % 256; uint256 word = bitmap._data[bucket]; @@ -18,13 +18,13 @@ library BitMap { return word & mask != 0; } - function set(UintBitMap storage bitmap, uint256 index) internal { + function set(BitMap storage bitmap, uint256 index) internal { uint256 bucket = index / 256; uint256 pos = index % 256; bitmap._data[bucket] |= (1 << pos); } - function unset(UintBitMap storage bitmap, uint256 index) internal { + function unset(BitMap storage bitmap, uint256 index) internal { uint256 bucket = index / 256; uint256 pos = index % 256; bitmap._data[bucket] &= ~(1 << pos); From 65c6296c977bb81ddcbfb964d8dd5a26092541e6 Mon Sep 17 00:00:00 2001 From: Hadrien Croubois Date: Thu, 10 Jun 2021 13:46:09 +0200 Subject: [PATCH 5/9] add changelog and document entries --- CHANGELOG.md | 1 + contracts/utils/README.adoc | 2 ++ 2 files changed, 3 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index d5672289eb6..a6cd96da153 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ * Tokens: Wrap definitely safe subtractions in `unchecked` blocks. * `Math`: Add a `ceilDiv` method for performing ceiling division. * `ERC1155Supply`: add a new `ERC1155` extension that keeps track of the totalSupply of each tokenId. ([#2593](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2593)) + * `BitMaps`: add a new `BitMaps` library that provides a storage efficient datastructure for `uint256` to `bool` mapping with contiguous keys. ([#2710](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2710)) ### Breaking Changes diff --git a/contracts/utils/README.adoc b/contracts/utils/README.adoc index ae5c38a50a9..4edcf923bb1 100644 --- a/contracts/utils/README.adoc +++ b/contracts/utils/README.adoc @@ -80,6 +80,8 @@ Note that, in all cases, accounts simply _declare_ their interfaces, but they ar == Data Structures +{{BitMaps}} + {{EnumerableMap}} {{EnumerableSet}} From 22ea07c5570410b10a2decbd6c6ffe7e4d2cc18a Mon Sep 17 00:00:00 2001 From: Hadrien Croubois Date: Thu, 10 Jun 2021 13:49:50 +0200 Subject: [PATCH 6/9] cleanup BitMaps code --- contracts/utils/structs/BitMaps.sol | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/contracts/utils/structs/BitMaps.sol b/contracts/utils/structs/BitMaps.sol index 859bad85031..2bde5b132fd 100644 --- a/contracts/utils/structs/BitMaps.sol +++ b/contracts/utils/structs/BitMaps.sol @@ -12,21 +12,19 @@ library BitMaps { function get(BitMap storage bitmap, uint256 index) internal view returns (bool) { uint256 bucket = index / 256; - uint256 pos = index % 256; - uint256 word = bitmap._data[bucket]; - uint256 mask = (1 << pos); - return word & mask != 0; + uint256 mask = 1 << (index % 256); + return bitmap._data[bucket] & mask != 0; } function set(BitMap storage bitmap, uint256 index) internal { uint256 bucket = index / 256; - uint256 pos = index % 256; - bitmap._data[bucket] |= (1 << pos); + uint256 mask = 1 << (index % 256); + bitmap._data[bucket] |= mask; } function unset(BitMap storage bitmap, uint256 index) internal { uint256 bucket = index / 256; - uint256 pos = index % 256; - bitmap._data[bucket] &= ~(1 << pos); + uint256 mask = 1 << (index % 256); + bitmap._data[bucket] &= ~mask; } } From 76c36767bd2c1881c850d189f7c8ea4cad77cce9 Mon Sep 17 00:00:00 2001 From: Hadrien Croubois Date: Thu, 10 Jun 2021 22:54:03 +0200 Subject: [PATCH 7/9] Apply suggestions from code review Co-authored-by: Francisco Giordano --- contracts/utils/structs/BitMaps.sol | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/contracts/utils/structs/BitMaps.sol b/contracts/utils/structs/BitMaps.sol index 2bde5b132fd..0aa466505a5 100644 --- a/contracts/utils/structs/BitMaps.sol +++ b/contracts/utils/structs/BitMaps.sol @@ -10,18 +10,27 @@ library BitMaps { mapping(uint256 => uint256) _data; } + /** + * @dev Returns whether the bit at `index` is set. + */ function get(BitMap storage bitmap, uint256 index) internal view returns (bool) { uint256 bucket = index / 256; uint256 mask = 1 << (index % 256); return bitmap._data[bucket] & mask != 0; } + /** + * @dev Sets the bit at `index`. + */ function set(BitMap storage bitmap, uint256 index) internal { uint256 bucket = index / 256; uint256 mask = 1 << (index % 256); bitmap._data[bucket] |= mask; } + /** + * @dev Unsets the bit at `index`. + */ function unset(BitMap storage bitmap, uint256 index) internal { uint256 bucket = index / 256; uint256 mask = 1 << (index % 256); From e70b191d9d0b1c2e45bff406a109bb087cab388a Mon Sep 17 00:00:00 2001 From: Hadrien Croubois Date: Fri, 11 Jun 2021 17:45:49 +0200 Subject: [PATCH 8/9] add a BitMap.setTo function --- contracts/mocks/BitmapMock.sol | 4 ++++ contracts/utils/structs/BitMaps.sol | 11 ++++++++++ test/utils/structs/BitMap.test.js | 32 +++++++++++++++++++++++++++++ 3 files changed, 47 insertions(+) diff --git a/contracts/mocks/BitmapMock.sol b/contracts/mocks/BitmapMock.sol index ab679795b20..ccf8486f5da 100644 --- a/contracts/mocks/BitmapMock.sol +++ b/contracts/mocks/BitmapMock.sol @@ -13,6 +13,10 @@ contract BitMapMock { return _bitmap.get(index); } + function setTo(uint256 index, bool value) public { + _bitmap.setTo(index, value); + } + function set(uint256 index) public { _bitmap.set(index); } diff --git a/contracts/utils/structs/BitMaps.sol b/contracts/utils/structs/BitMaps.sol index 0aa466505a5..8235f18ba41 100644 --- a/contracts/utils/structs/BitMaps.sol +++ b/contracts/utils/structs/BitMaps.sol @@ -19,6 +19,17 @@ library BitMaps { return bitmap._data[bucket] & mask != 0; } + /** + * @dev Sets the bit at `index` to the boolean `value`. + */ + function setTo(BitMap storage bitmap, uint256 index, bool value) internal { + if (value) { + set(bitmap, index); + } else { + unset(bitmap, index); + } + } + /** * @dev Sets the bit at `index`. */ diff --git a/test/utils/structs/BitMap.test.js b/test/utils/structs/BitMap.test.js index 22ed9b85b3a..58d70ca8f11 100644 --- a/test/utils/structs/BitMap.test.js +++ b/test/utils/structs/BitMap.test.js @@ -18,6 +18,38 @@ contract('BitMap', function (accounts) { expect(await this.bitmap.get(keyC)).to.equal(false); }); + describe('setTo', function () { + it('set a key to true', async function () { + await this.bitmap.setTo(keyA, true); + expect(await this.bitmap.get(keyA)).to.equal(true); + expect(await this.bitmap.get(keyB)).to.equal(false); + expect(await this.bitmap.get(keyC)).to.equal(false); + }); + + it('set a key to false', async function () { + await this.bitmap.setTo(keyA, true); + await this.bitmap.setTo(keyA, false); + expect(await this.bitmap.get(keyA)).to.equal(false); + expect(await this.bitmap.get(keyB)).to.equal(false); + expect(await this.bitmap.get(keyC)).to.equal(false); + }); + + it('set several consecutive keys', async function () { + await this.bitmap.setTo(keyA.addn(0), true); + await this.bitmap.setTo(keyA.addn(1), true); + await this.bitmap.setTo(keyA.addn(2), true); + await this.bitmap.setTo(keyA.addn(3), true); + await this.bitmap.setTo(keyA.addn(4), true); + await this.bitmap.setTo(keyA.addn(2), false); + await this.bitmap.setTo(keyA.addn(4), false); + expect(await this.bitmap.get(keyA.addn(0))).to.equal(true); + expect(await this.bitmap.get(keyA.addn(1))).to.equal(true); + expect(await this.bitmap.get(keyA.addn(2))).to.equal(false); + expect(await this.bitmap.get(keyA.addn(3))).to.equal(true); + expect(await this.bitmap.get(keyA.addn(4))).to.equal(false); + }); + }); + describe('set', function () { it('adds a key', async function () { await this.bitmap.set(keyA); From 7ac884e11b12d6270731739e6dc9c7047ae26bf8 Mon Sep 17 00:00:00 2001 From: Hadrien Croubois Date: Fri, 11 Jun 2021 17:53:21 +0200 Subject: [PATCH 9/9] fix lint --- contracts/utils/structs/BitMaps.sol | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/contracts/utils/structs/BitMaps.sol b/contracts/utils/structs/BitMaps.sol index 8235f18ba41..f41cb8c9799 100644 --- a/contracts/utils/structs/BitMaps.sol +++ b/contracts/utils/structs/BitMaps.sol @@ -22,7 +22,11 @@ library BitMaps { /** * @dev Sets the bit at `index` to the boolean `value`. */ - function setTo(BitMap storage bitmap, uint256 index, bool value) internal { + function setTo( + BitMap storage bitmap, + uint256 index, + bool value + ) internal { if (value) { set(bitmap, index); } else {