Add a BitMap struct (#2710)

Co-authored-by: Francisco Giordano <frangio.1@gmail.com>

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
 

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);
+    }
+}

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}}

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;
+    }
+}

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);
+    });
+  });
+});