Refactor SafeMath to avoid memory leaks (#2462)

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

CHANGELOG.md

@@ -7,9 +7,11 @@
  * `ERC20Permit`: added an implementation of the ERC20 permit extension for gasless token approvals. ([#2237](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2237))
  * Presets: added token presets with preminted fixed supply `ERC20PresetFixedSupply` and `ERC777PresetFixedSupply`. ([#2399](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2399))
  * `Address`: added `functionDelegateCall`, similar to the existing `functionCall`. ([#2333](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2333))
- * `Context`: moved from `contracts/GSN` to `contracts/utils`. ([#2453](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2453)) 
+ * `Context`: moved from `contracts/GSN` to `contracts/utils`. ([#2453](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2453))
  * `PaymentSplitter`: replace usage of `.transfer()` with `Address.sendValue` for improved compatibility with smart wallets. ([#2455](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2455))
- * `UpgradeableProxy`: bubble revert reasons from initialization calls. ([#2454](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2454)) 
+ * `UpgradeableProxy`: bubble revert reasons from initialization calls. ([#2454](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2454))
+ * `SafeMath`: fix a memory allocation issue by adding new `SafeMath.tryOp(uint,uint)→(bool,uint)` functions. `SafeMath.op(uint,uint,string)→uint` are now deprecated. ([#2462](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2462))
+ * `EnumerableMap`: fix a memory allocation issue by adding new `EnumerableMap.tryGet(uint)→(bool,address)` functions. `EnumerableMap.get(uint)→string` is now deprecated. ([#2462](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2462))
 
 ## 3.3.0 (2020-11-26)
 

contracts/math/SafeMath.sol

@@ -17,6 +17,52 @@ pragma solidity >=0.6.0 <0.8.0;
  */
 library SafeMath {
     /**
+     * @dev Returns the addition of two unsigned integers, with an overflow flag.
+     */
+    function tryAdd(uint256 a, uint256 b) internal pure returns (bool, uint256) {
+        uint256 c = a + b;
+        if (c < a) return (false, 0);
+        return (true, c);
+    }
+
+    /**
+     * @dev Returns the substraction of two unsigned integers, with an overflow flag.
+     */
+    function trySub(uint256 a, uint256 b) internal pure returns (bool, uint256) {
+        if (b > a) return (false, 0);
+        return (true, a - b);
+    }
+
+    /**
+     * @dev Returns the multiplication of two unsigned integers, with an overflow flag.
+     */
+    function tryMul(uint256 a, uint256 b) internal pure returns (bool, uint256) {
+        // Gas optimization: this is cheaper than requiring 'a' not being zero, but the
+        // benefit is lost if 'b' is also tested.
+        // See: https://github.com/OpenZeppelin/openzeppelin-contracts/pull/522
+        if (a == 0) return (true, 0);
+        uint256 c = a * b;
+        if (c / a != b) return (false, 0);
+        return (true, c);
+    }
+
+    /**
+     * @dev Returns the division of two unsigned integers, with a division by zero flag.
+     */
+    function tryDiv(uint256 a, uint256 b) internal pure returns (bool, uint256) {
+        if (b == 0) return (false, 0);
+        return (true, a / b);
+    }
+
+    /**
+     * @dev Returns the remainder of dividing two unsigned integers, with a division by zero flag.
+     */
+    function tryMod(uint256 a, uint256 b) internal pure returns (bool, uint256) {
+        if (b == 0) return (false, 0);
+        return (true, a % b);
+    }
+
+    /**
      * @dev Returns the addition of two unsigned integers, reverting on
      * overflow.
      *
@@ -29,7 +75,6 @@ library SafeMath {
     function add(uint256 a, uint256 b) internal pure returns (uint256) {
         uint256 c = a + b;
         require(c >= a, "SafeMath: addition overflow");
-
         return c;
     }
 
@@ -44,24 +89,8 @@ library SafeMath {
      * - Subtraction cannot overflow.
      */
     function sub(uint256 a, uint256 b) internal pure returns (uint256) {
-        return sub(a, b, "SafeMath: subtraction overflow");
-    }
-
-    /**
-     * @dev Returns the subtraction of two unsigned integers, reverting with custom message on
-     * overflow (when the result is negative).
-     *
-     * Counterpart to Solidity's `-` operator.
-     *
-     * Requirements:
-     *
-     * - Subtraction cannot overflow.
-     */
-    function sub(uint256 a, uint256 b, string memory errorMessage) internal pure returns (uint256) {
-        require(b <= a, errorMessage);
-        uint256 c = a - b;
-
-        return c;
+        require(b <= a, "SafeMath: subtraction overflow");
+        return a - b;
     }
 
     /**
@@ -75,21 +104,14 @@ library SafeMath {
      * - Multiplication cannot overflow.
      */
     function mul(uint256 a, uint256 b) internal pure returns (uint256) {
-        // Gas optimization: this is cheaper than requiring 'a' not being zero, but the
-        // benefit is lost if 'b' is also tested.
-        // See: https://github.com/OpenZeppelin/openzeppelin-contracts/pull/522
-        if (a == 0) {
-            return 0;
-        }
-
+        if (a == 0) return 0;
         uint256 c = a * b;
         require(c / a == b, "SafeMath: multiplication overflow");
-
         return c;
     }
 
     /**
-     * @dev Returns the integer division of two unsigned integers. Reverts on
+     * @dev Returns the integer division of two unsigned integers, reverting on
      * division by zero. The result is rounded towards zero.
      *
      * Counterpart to Solidity's `/` operator. Note: this function uses a
@@ -101,48 +123,71 @@ library SafeMath {
      * - The divisor cannot be zero.
      */
     function div(uint256 a, uint256 b) internal pure returns (uint256) {
-        return div(a, b, "SafeMath: division by zero");
+        require(b > 0, "SafeMath: division by zero");
+        return a / b;
     }
 
     /**
-     * @dev Returns the integer division of two unsigned integers. Reverts with custom message on
-     * division by zero. The result is rounded towards zero.
+     * @dev Returns the remainder of dividing two unsigned integers. (unsigned integer modulo),
+     * reverting when dividing by zero.
      *
-     * Counterpart to Solidity's `/` operator. Note: this function uses a
-     * `revert` opcode (which leaves remaining gas untouched) while Solidity
-     * uses an invalid opcode to revert (consuming all remaining gas).
+     * Counterpart to Solidity's `%` operator. This function uses a `revert`
+     * opcode (which leaves remaining gas untouched) while Solidity uses an
+     * invalid opcode to revert (consuming all remaining gas).
      *
      * Requirements:
      *
      * - The divisor cannot be zero.
      */
-    function div(uint256 a, uint256 b, string memory errorMessage) internal pure returns (uint256) {
-        require(b > 0, errorMessage);
-        uint256 c = a / b;
-        // assert(a == b * c + a % b); // There is no case in which this doesn't hold
+    function mod(uint256 a, uint256 b) internal pure returns (uint256) {
+        require(b > 0, "SafeMath: modulo by zero");
+        return a % b;
+    }
 
-        return c;
+    /**
+     * @dev Returns the subtraction of two unsigned integers, reverting with custom message on
+     * overflow (when the result is negative).
+     *
+     * CAUTION: This function is deprecated because it requires allocating memory for the error
+     * message unnecessarily. For custom revert reasons use {trySub}.
+     *
+     * Counterpart to Solidity's `-` operator.
+     *
+     * Requirements:
+     *
+     * - Subtraction cannot overflow.
+     */
+    function sub(uint256 a, uint256 b, string memory errorMessage) internal pure returns (uint256) {
+        require(b <= a, errorMessage);
+        return a - b;
     }
 
     /**
-     * @dev Returns the remainder of dividing two unsigned integers. (unsigned integer modulo),
-     * Reverts when dividing by zero.
+     * @dev Returns the integer division of two unsigned integers, reverting with custom message on
+     * division by zero. The result is rounded towards zero.
      *
-     * Counterpart to Solidity's `%` operator. This function uses a `revert`
-     * opcode (which leaves remaining gas untouched) while Solidity uses an
-     * invalid opcode to revert (consuming all remaining gas).
+     * CAUTION: This function is deprecated because it requires allocating memory for the error
+     * message unnecessarily. For custom revert reasons use {tryDiv}.
+     *
+     * Counterpart to Solidity's `/` operator. Note: this function uses a
+     * `revert` opcode (which leaves remaining gas untouched) while Solidity
+     * uses an invalid opcode to revert (consuming all remaining gas).
      *
      * Requirements:
      *
      * - The divisor cannot be zero.
      */
-    function mod(uint256 a, uint256 b) internal pure returns (uint256) {
-        return mod(a, b, "SafeMath: modulo by zero");
+    function div(uint256 a, uint256 b, string memory errorMessage) internal pure returns (uint256) {
+        require(b > 0, errorMessage);
+        return a / b;
     }
 
     /**
      * @dev Returns the remainder of dividing two unsigned integers. (unsigned integer modulo),
-     * Reverts with custom message when dividing by zero.
+     * reverting with custom message when dividing by zero.
+     *
+     * CAUTION: This function is deprecated because it requires allocating memory for the error
+     * message unnecessarily. For custom revert reasons use {tryMod}.
      *
      * Counterpart to Solidity's `%` operator. This function uses a `revert`
      * opcode (which leaves remaining gas untouched) while Solidity uses an
@@ -153,7 +198,7 @@ library SafeMath {
      * - The divisor cannot be zero.
      */
     function mod(uint256 a, uint256 b, string memory errorMessage) internal pure returns (uint256) {
-        require(b != 0, errorMessage);
+        require(b > 0, errorMessage);
         return a % b;
     }
 }

contracts/mocks/EnumerableMapMock.sol

@@ -34,7 +34,15 @@ contract EnumerableMapMock {
     }
 
 
+    function tryGet(uint256 key) public view returns (bool, address) {
+        return _map.tryGet(key);
+    }
+
     function get(uint256 key) public view returns (address) {
         return _map.get(key);
     }
+
+    function getWithMessage(uint256 key, string calldata errorMessage) public view returns (address) {
+        return _map.get(key, errorMessage);
+    }
 }

contracts/mocks/SafeMathMock.sol

@@ -5,23 +5,101 @@ pragma solidity >=0.6.0 <0.8.0;
 import "../math/SafeMath.sol";
 
 contract SafeMathMock {
-    function mul(uint256 a, uint256 b) public pure returns (uint256) {
-        return SafeMath.mul(a, b);
+    function tryAdd(uint256 a, uint256 b) public pure returns (bool flag, uint256 value) {
+        return SafeMath.tryAdd(a, b);
     }
 
-    function div(uint256 a, uint256 b) public pure returns (uint256) {
-        return SafeMath.div(a, b);
+    function trySub(uint256 a, uint256 b) public pure returns (bool flag, uint256 value) {
+        return SafeMath.trySub(a, b);
     }
 
-    function sub(uint256 a, uint256 b) public pure returns (uint256) {
-        return SafeMath.sub(a, b);
+    function tryMul(uint256 a, uint256 b) public pure returns (bool flag, uint256 value) {
+        return SafeMath.tryMul(a, b);
+    }
+
+    function tryDiv(uint256 a, uint256 b) public pure returns (bool flag, uint256 value) {
+        return SafeMath.tryDiv(a, b);
+    }
+
+    function tryMod(uint256 a, uint256 b) public pure returns (bool flag, uint256 value) {
+        return SafeMath.tryMod(a, b);
     }
 
     function add(uint256 a, uint256 b) public pure returns (uint256) {
         return SafeMath.add(a, b);
     }
 
+    function sub(uint256 a, uint256 b) public pure returns (uint256) {
+        return SafeMath.sub(a, b);
+    }
+
+    function mul(uint256 a, uint256 b) public pure returns (uint256) {
+        return SafeMath.mul(a, b);
+    }
+
+    function div(uint256 a, uint256 b) public pure returns (uint256) {
+        return SafeMath.div(a, b);
+    }
+
     function mod(uint256 a, uint256 b) public pure returns (uint256) {
         return SafeMath.mod(a, b);
     }
+
+    function subWithMessage(uint256 a, uint256 b, string memory errorMessage) public pure returns (uint256) {
+        return SafeMath.sub(a, b, errorMessage);
+    }
+
+    function divWithMessage(uint256 a, uint256 b, string memory errorMessage) public pure returns (uint256) {
+        return SafeMath.div(a, b, errorMessage);
+    }
+
+    function modWithMessage(uint256 a, uint256 b, string memory errorMessage) public pure returns (uint256) {
+        return SafeMath.mod(a, b, errorMessage);
+    }
+
+    function addMemoryCheck() public pure returns (uint256 mem) {
+        uint256 length = 32;
+        // solhint-disable-next-line no-inline-assembly
+        assembly { mem := mload(0x40) }
+        for (uint256 i = 0; i < length; ++i) { SafeMath.add(1, 1); }
+        // solhint-disable-next-line no-inline-assembly
+        assembly { mem := sub(mload(0x40), mem) }
+    }
+
+    function subMemoryCheck() public pure returns (uint256 mem) {
+        uint256 length = 32;
+        // solhint-disable-next-line no-inline-assembly
+        assembly { mem := mload(0x40) }
+        for (uint256 i = 0; i < length; ++i) { SafeMath.sub(1, 1); }
+        // solhint-disable-next-line no-inline-assembly
+        assembly { mem := sub(mload(0x40), mem) }
+    }
+
+    function mulMemoryCheck() public pure returns (uint256 mem) {
+        uint256 length = 32;
+        // solhint-disable-next-line no-inline-assembly
+        assembly { mem := mload(0x40) }
+        for (uint256 i = 0; i < length; ++i) { SafeMath.mul(1, 1); }
+        // solhint-disable-next-line no-inline-assembly
+        assembly { mem := sub(mload(0x40), mem) }
+    }
+
+    function divMemoryCheck() public pure returns (uint256 mem) {
+        uint256 length = 32;
+        // solhint-disable-next-line no-inline-assembly
+        assembly { mem := mload(0x40) }
+        for (uint256 i = 0; i < length; ++i) { SafeMath.div(1, 1); }
+        // solhint-disable-next-line no-inline-assembly
+        assembly { mem := sub(mload(0x40), mem) }
+    }
+
+    function modMemoryCheck() public pure returns (uint256 mem) {
+        uint256 length = 32;
+        // solhint-disable-next-line no-inline-assembly
+        assembly { mem := mload(0x40) }
+        for (uint256 i = 0; i < length; ++i) { SafeMath.mod(1, 1); }
+        // solhint-disable-next-line no-inline-assembly
+        assembly { mem := sub(mload(0x40), mem) }
+    }
+
 }

contracts/utils/EnumerableMap.sol

@@ -144,6 +144,16 @@ library EnumerableMap {
     }
 
     /**
+     * @dev Tries to returns the value associated with `key`.  O(1).
+     * Does not revert if `key` is not in the map.
+     */
+    function _tryGet(Map storage map, bytes32 key) private view returns (bool, bytes32) {
+        uint256 keyIndex = map._indexes[key];
+        if (keyIndex == 0) return (false, 0); // Equivalent to contains(map, key)
+        return (true, map._entries[keyIndex - 1]._value); // All indexes are 1-based
+    }
+
+    /**
      * @dev Returns the value associated with `key`.  O(1).
      *
      * Requirements:
@@ -151,11 +161,16 @@ library EnumerableMap {
      * - `key` must be in the map.
      */
     function _get(Map storage map, bytes32 key) private view returns (bytes32) {
-        return _get(map, key, "EnumerableMap: nonexistent key");
+        uint256 keyIndex = map._indexes[key];
+        require(keyIndex != 0, "EnumerableMap: nonexistent key"); // Equivalent to contains(map, key)
+        return map._entries[keyIndex - 1]._value; // All indexes are 1-based
     }
 
     /**
      * @dev Same as {_get}, with a custom error message when `key` is not in the map.
+     *
+     * CAUTION: This function is deprecated because it requires allocating memory for the error
+     * message unnecessarily. For custom revert reasons use {_tryGet}.
      */
     function _get(Map storage map, bytes32 key, string memory errorMessage) private view returns (bytes32) {
         uint256 keyIndex = map._indexes[key];
@@ -218,6 +233,15 @@ library EnumerableMap {
     }
 
     /**
+     * @dev Tries to returns the value associated with `key`.  O(1).
+     * Does not revert if `key` is not in the map.
+     */
+    function tryGet(UintToAddressMap storage map, uint256 key) internal view returns (bool, address) {
+        (bool success, bytes32 value) = _tryGet(map._inner, bytes32(key));
+        return (success, address(uint160(uint256(value))));
+    }
+
+    /**
      * @dev Returns the value associated with `key`.  O(1).
      *
      * Requirements:
@@ -230,6 +254,9 @@ library EnumerableMap {
 
     /**
      * @dev Same as {get}, with a custom error message when `key` is not in the map.
+     *
+     * CAUTION: This function is deprecated because it requires allocating memory for the error
+     * message unnecessarily. For custom revert reasons use {tryGet}.
      */
     function get(UintToAddressMap storage map, uint256 key, string memory errorMessage) internal view returns (address) {
         return address(uint160(uint256(_get(map._inner, bytes32(key), errorMessage))));

test/utils/EnumerableMap.test.js

-const { BN, expectEvent } = require('@openzeppelin/test-helpers');
+const { BN, constants, expectEvent, expectRevert } = require('@openzeppelin/test-helpers');
 const { expect } = require('chai');
 
 const zip = require('lodash.zip');
@@ -139,4 +139,43 @@ contract('EnumerableMap', function (accounts) {
       expect(await this.map.contains(keyB)).to.equal(false);
     });
   });
+
+  describe('read', function () {
+    beforeEach(async function () {
+      await this.map.set(keyA, accountA);
+    });
+
+    describe('get', function () {
+      it('existing value', async function () {
+        expect(await this.map.get(keyA)).to.be.equal(accountA);
+      });
+      it('missing value', async function () {
+        await expectRevert(this.map.get(keyB), 'EnumerableMap: nonexistent key');
+      });
+    });
+
+    describe('get with message', function () {
+      it('existing value', async function () {
+        expect(await this.map.getWithMessage(keyA, 'custom error string')).to.be.equal(accountA);
+      });
+      it('missing value', async function () {
+        await expectRevert(this.map.getWithMessage(keyB, 'custom error string'), 'custom error string');
+      });
+    });
+
+    describe('tryGet', function () {
+      it('existing value', async function () {
+        expect(await this.map.tryGet(keyA)).to.be.deep.equal({
+          0: true,
+          1: accountA,
+        });
+      });
+      it('missing value', async function () {
+        expect(await this.map.tryGet(keyB)).to.be.deep.equal({
+          0: false,
+          1: constants.ZERO_ADDRESS,
+        });
+      });
+    });
+  });
 });