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/mocks/BitmapMock.sol b/contracts/mocks/BitmapMock.sol new file mode 100644 index 00000000000..ccf8486f5da --- /dev/null +++ b/contracts/mocks/BitmapMock.sol @@ -0,0 +1,27 @@ +// SPDX-License-Identifier: MIT + +pragma solidity ^0.8.0; + +import "../utils/structs/BitMaps.sol"; + +contract BitMapMock { + using BitMaps for BitMaps.BitMap; + + BitMaps.BitMap private _bitmap; + + function get(uint256 index) public view returns (bool) { + return _bitmap.get(index); + } + + function setTo(uint256 index, bool value) public { + _bitmap.setTo(index, value); + } + + function set(uint256 index) public { + _bitmap.set(index); + } + + function unset(uint256 index) public { + _bitmap.unset(index); + } +} 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}} diff --git a/contracts/utils/structs/BitMaps.sol b/contracts/utils/structs/BitMaps.sol new file mode 100644 index 00000000000..f41cb8c9799 --- /dev/null +++ b/contracts/utils/structs/BitMaps.sol @@ -0,0 +1,54 @@ +// 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 BitMaps { + struct BitMap { + 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` 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`. + */ + 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); + bitmap._data[bucket] &= ~mask; + } +} diff --git a/test/utils/structs/BitMap.test.js b/test/utils/structs/BitMap.test.js new file mode 100644 index 00000000000..58d70ca8f11 --- /dev/null +++ b/test/utils/structs/BitMap.test.js @@ -0,0 +1,145 @@ +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('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); + 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); + }); + + 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 () { + 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('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 () { + // [] + + 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); + }); + }); +});