✨ STL: totalSupply (#1212)

src/utils/SafeTransferLib.sol

@@ -25,6 +25,9 @@ library SafeTransferLib {
     /// @dev The ERC20 `approve` has failed.
     error ApproveFailed();
 
+    /// @dev The ERC20 `totalSupply` query has failed.
+    error TotalSupplyQueryFailed();
+
     /// @dev The Permit2 operation has failed.
     error Permit2Failed();
 
@@ -396,6 +399,22 @@ library SafeTransferLib {
         }
     }
 
+    /// @dev Returns the total supply of the `token`.
+    /// Reverts if the token does not exist or does not implement `totalSupply()`.
+    function totalSupply(address token) internal view returns (uint256 result) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            mstore(0x00, 0x18160ddd) // `totalSupply()`.
+            if iszero(
+                and(gt(returndatasize(), 0x1f), staticcall(gas(), token, 0x1c, 0x04, 0x00, 0x20))
+            ) {
+                mstore(0x00, 0x54cd9435) // `TotalSupplyQueryFailed()`.
+                revert(0x1c, 0x04)
+            }
+            result := mload(0x00)
+        }
+    }
+
     /// @dev Sends `amount` of ERC20 `token` from `from` to `to`.
     /// If the initial attempt fails, try to use Permit2 to transfer the token.
     /// Reverts upon failure.

test/SafeTransferLib.t.sol

@@ -1091,4 +1091,16 @@ contract SafeTransferLibTest is SoladyTest {
             t.s
         );
     }
+
+    function testTotalSupplyQuery() public {
+        uint256 totalSupplyBefore = this.totalSupplyQuery(address(erc20));
+        erc20.burn(address(this), 123);
+        assertEq(this.totalSupplyQuery(address(erc20)), totalSupplyBefore - 123);
+        vm.expectRevert(SafeTransferLib.TotalSupplyQueryFailed.selector);
+        this.totalSupplyQuery(address(0));
+    }
+
+    function totalSupplyQuery(address token) public view returns (uint256) {
+        return SafeTransferLib.totalSupply(token);
+    }
 }