add tests for _spendAllowance and _approve, make branchless

ProjectOpenSea · Nov 30, 2023 · 86bac02 · 86bac02
1 parent 389d287
commit 86bac02
Show file tree

Hide file tree

Showing 3 changed files with 115 additions and 10 deletions.
diff --git a/src/reference/tokens/erc20/ERC20Preapproved_Solady.sol b/src/reference/tokens/erc20/ERC20Preapproved_Solady.sol
@@ -8,6 +8,16 @@ contract ERC20_Solady is ERC20ConduitPreapproved_Solady {
         _mint(to, amount);
     }
 
+    /// @dev Exposed to test internal function
+    function spendAllowance(address owner, address spender, uint256 amount) public {
+        _spendAllowance(owner, spender, amount);
+    }
+
+    /// @dev Exposed to test internal function
+    function approve(address owner, address spender, uint256 amount) public {
+        _approve(owner, spender, amount);
+    }
+
     function name() public pure override returns (string memory) {
         return "Test";
     }

diff --git a/src/tokens/erc20/ERC20ConduitPreapproved_Solady.sol b/src/tokens/erc20/ERC20ConduitPreapproved_Solady.sol
@@ -66,8 +66,8 @@ abstract contract ERC20ConduitPreapproved_Solady is ERC20, IPreapprovalForAll {
             allowance_ :=
                 xor(allowance_, mul(and(eq(caller(), CONDUIT), or(iszero(allowance_), iszero(not(allowance_)))), not(0)))
 
-            // If the allowance is not the maximum uint256 value:
-            if not(allowance_) {
+            // If the allowance is not the maximum uint256 value.
+            if add(allowance_, 1) {
                 // Revert if the amount to be transferred exceeds the allowance.
                 if gt(amount, allowance_) {
                     mstore(0x00, 0x13be252b) // `InsufficientAllowance()`.
@@ -181,15 +181,52 @@ abstract contract ERC20ConduitPreapproved_Solady is ERC20, IPreapprovalForAll {
     }
 
     function _spendAllowance(address owner, address spender, uint256 amount) internal virtual override {
-        if (spender == CONDUIT) {
-            uint256 allowance_ = super.allowance(owner, spender);
-            if (allowance_ == type(uint256).max) {
-                // Max allowance, no need to spend.
-                return;
-            } else if (allowance_ == 0) {
-                revert InsufficientAllowance();
+        /// @solidity memory-safe-assembly
+        assembly {
+            // Compute the allowance slot and load its value.
+            mstore(0x20, spender)
+            mstore(0x0c, SOLADY_ERC20_ALLOWANCE_SLOT_SEED)
+            mstore(0x00, owner)
+            let allowanceSlot := keccak256(0x0c, 0x34)
+            let allowance_ := sload(allowanceSlot)
+
+            // "flip" allowance if spender is CONDUIT and if allowance_ is 0 or type(uint256).max.
+            allowance_ :=
+                xor(allowance_, mul(and(eq(spender, CONDUIT), or(iszero(allowance_), iszero(not(allowance_)))), not(0)))
+
+            // If the allowance is not the maximum uint256 value.
+            if add(allowance_, 1) {
+                // Revert if the amount to be transferred exceeds the allowance.
+                if gt(amount, allowance_) {
+                    mstore(0x00, 0x13be252b) // `InsufficientAllowance()`.
+                    revert(0x1c, 0x04)
+                }
+                // Subtract and store the updated allowance.
+                sstore(
+                    allowanceSlot,
+                    xor(
+                        sub(allowance_, amount), mul(and(eq(spender, CONDUIT), iszero(sub(allowance_, amount))), not(0))
+                    )
+                )
             }
         }
-        super._spendAllowance(owner, spender, amount);
+    }
+
+    function _approve(address owner, address spender, uint256 amount) internal virtual override {
+        /// @solidity memory-safe-assembly
+        assembly {
+            let owner_ := shl(96, owner)
+            // Compute the allowance slot and store the amount.
+            mstore(0x20, spender)
+            mstore(0x0c, or(owner_, SOLADY_ERC20_ALLOWANCE_SLOT_SEED))
+
+            // "flip" amount if spender is CONDUIT and if amount is 0 or type(uint256).max.
+            amount := xor(amount, mul(and(eq(spender, CONDUIT), or(iszero(amount), iszero(not(amount)))), not(0)))
+
+            sstore(keccak256(0x0c, 0x34), amount)
+            // Emit the {Approval} event.
+            mstore(0x00, amount)
+            log3(0x00, 0x20, SOLADY_ERC20_APPROVAL_EVENT_SIGNATURE, shr(96, owner_), shr(96, mload(0x2c)))
+        }
     }
 }
diff --git a/test/tokens/ERC20ConduitPreapproved_Solady.t.sol b/test/tokens/ERC20ConduitPreapproved_Solady.t.sol
@@ -162,6 +162,64 @@ contract ERC20ConduitPreapproved_SoladyTest is Test, TestPlus, IPreapprovalForAl
         _checkAllowanceAndNonce(t);
     }
 
+    function testInternalSpendAllowance(address acct, address operator) public {
+        if (acct == address(0)) {
+            acct = address(1);
+        }
+        if (operator == address(0)) {
+            operator = address(1);
+        }
+        test.mint(acct, 1 ether);
+        vm.prank(acct);
+        test.approve(operator, 1 ether);
+        vm.prank(operator);
+        test.spendAllowance(acct, operator, 1 ether);
+        assertEq(test.allowance(acct, operator), 0);
+        vm.expectRevert(ERC20.InsufficientAllowance.selector);
+        test.spendAllowance(acct, operator, 1 ether);
+
+        // Allowance of type(uint256).max should not be reduced
+        vm.prank(acct);
+        test.approve(operator, type(uint256).max);
+        test.spendAllowance(acct, operator, 1 ether);
+        assertEq(test.allowance(acct, operator), type(uint256).max);
+
+        // Allowance of conduit with default type(uint256).max should not be reduced
+        test.spendAllowance(acct, CONDUIT, 1 ether);
+        assertEq(test.allowance(acct, CONDUIT), type(uint256).max);
+
+        // Allowance of conduit with lower allowance should be reduced
+        vm.prank(acct);
+        test.approve(CONDUIT, 1 ether);
+        test.spendAllowance(acct, CONDUIT, 1 ether);
+        assertEq(test.allowance(acct, CONDUIT), 0);
+    }
+
+    function testInternalApprove(address acct, address operator) public {
+        if (acct == address(0)) {
+            acct = address(1);
+        }
+        if (operator == address(0)) {
+            operator = address(1);
+        }
+        test.mint(acct, 1 ether);
+        test.approve(acct, operator, 1 ether);
+        assertEq(test.allowance(acct, operator), 1 ether);
+
+        test.approve(acct, operator, type(uint256).max);
+        assertEq(test.allowance(acct, operator), type(uint256).max);
+
+        assertEq(test.allowance(acct, CONDUIT), type(uint256).max);
+        test.approve(acct, CONDUIT, 1 ether);
+        assertEq(test.allowance(acct, CONDUIT), 1 ether);
+
+        test.approve(acct, CONDUIT, 0);
+        assertEq(test.allowance(acct, CONDUIT), 0);
+
+        test.approve(acct, CONDUIT, type(uint256).max);
+        assertEq(test.allowance(acct, CONDUIT), type(uint256).max);
+    }
+
     function _signPermit(_TestTemps memory t) internal view {
         bytes32 innerHash =
             keccak256(abi.encode(SOLADY_ERC20_PERMIT_TYPEHASH, t.owner, t.to, t.amount, t.nonce, t.deadline));