feat(budgets): implement full ERC1155 support for SimpleBudget

src/budgets/Budget.sol

@@ -135,7 +135,7 @@ abstract contract Budget is Ownable, Cloneable, Receiver {
 
     /// @inheritdoc Receiver
     receive() external payable virtual override {
-        // Nothing to do here, but we need the function to be implemented
+        return;
     }
 
     /// @inheritdoc Receiver

src/budgets/SimpleBudget.sol

@@ -1,16 +1,21 @@
 // SPDX-License-Identifier: GPL-3.0
 pragma solidity ^0.8.24;
 
+import {IERC1155Receiver} from "lib/openzeppelin-contracts/contracts/token/ERC1155/IERC1155Receiver.sol";
+import {IERC1155} from "lib/openzeppelin-contracts/contracts/token/ERC1155/IERC1155.sol";
+import {IERC165} from "lib/openzeppelin-contracts/contracts/utils/introspection/IERC165.sol";
+
 import {LibZip} from "lib/solady/src/utils/LibZip.sol";
 import {SafeTransferLib} from "lib/solady/src/utils/SafeTransferLib.sol";
+import {ReentrancyGuard} from "lib/solady/src/utils/ReentrancyGuard.sol";
 
 import {Budget} from "src/budgets/Budget.sol";
 import {Cloneable} from "src/shared/Cloneable.sol";
 
 /// @title Simple Budget
 /// @notice A minimal budget implementation that simply holds and distributes tokens (ERC20-like and native)
 /// @dev This type of budget does NOT support NFTs (ERC-721, ERC-1155, etc.)
-contract SimpleBudget is Budget {
+contract SimpleBudget is Budget, IERC1155Receiver, ReentrancyGuard {
     using LibZip for bytes;
     using SafeTransferLib for address;
 
@@ -21,7 +26,10 @@ contract SimpleBudget is Budget {
     }
 
     /// @dev The total amount of each fungible asset distributed from the budget
-    mapping(address => uint256) private _distributed;
+    mapping(address => uint256) private _distributedFungible;
+
+    /// @dev The total amount of each ERC1155 asset and token ID distributed from the budget
+    mapping(address => mapping(uint256 => uint256)) private _distributedERC1155;
 
     /// @dev The mapping of authorized addresses
     mapping(address => bool) private _isAuthorized;
@@ -60,7 +68,9 @@ contract SimpleBudget is Budget {
             FungiblePayload memory payload = abi.decode(request.data, (FungiblePayload));
 
             // Ensure the value received is equal to the `payload.amount`
-            if (msg.value != payload.amount) revert InvalidAllocation(request.asset, payload.amount);
+            if (msg.value != payload.amount) {
+                revert InvalidAllocation(request.asset, payload.amount);
+            }
         } else if (request.assetType == AssetType.ERC20) {
             FungiblePayload memory payload = abi.decode(request.data, (FungiblePayload));
 
@@ -69,6 +79,16 @@ contract SimpleBudget is Budget {
             if (request.asset.balanceOf(address(this)) < payload.amount) {
                 revert InvalidAllocation(request.asset, payload.amount);
             }
+        } else if (request.assetType == AssetType.ERC1155) {
+            ERC1155Payload memory payload = abi.decode(request.data, (ERC1155Payload));
+
+            // Transfer `payload.amount` of `payload.tokenId` to this contract
+            IERC1155(request.asset).safeTransferFrom(
+                request.target, address(this), payload.tokenId, payload.amount, payload.data
+            );
+            if (IERC1155(request.asset).balanceOf(address(this), payload.tokenId) < payload.amount) {
+                revert InvalidAllocation(request.asset, payload.amount);
+            }
         } else {
             // Unsupported asset type
             return false;
@@ -79,8 +99,8 @@ contract SimpleBudget is Budget {
 
     /// @inheritdoc Budget
     /// @notice Reclaims assets from the budget
-    /// @param data_ The compressed data for the reclamation `(address asset, uint256 amount, address receiver)`
-    /// @return True if the reclamation was successful
+    /// @param data_ The compressed {Transfer} request
+    /// @return True if the request was successful
     /// @dev Only the owner can directly reclaim assets from the budget
     /// @dev If the amount is zero, the entire balance of the asset will be transferred to the receiver
     /// @dev If the asset transfer fails, the reclamation will revert
@@ -89,7 +109,14 @@ contract SimpleBudget is Budget {
         if (request.assetType == AssetType.ETH || request.assetType == AssetType.ERC20) {
             FungiblePayload memory payload = abi.decode(request.data, (FungiblePayload));
             uint256 amount = payload.amount == 0 ? available(request.asset) : payload.amount;
-            _transfer(request.asset, request.target, amount);
+            _transferFungible(request.asset, request.target, amount);
+        } else if (request.assetType == AssetType.ERC1155) {
+            ERC1155Payload memory payload = abi.decode(request.data, (ERC1155Payload));
+            uint256 amount =
+                payload.amount == 0 ? IERC1155(request.asset).balanceOf(address(this), payload.tokenId) : payload.amount;
+            IERC1155(request.asset).safeTransferFrom(
+                address(this), request.target, payload.tokenId, amount, payload.data
+            );
         } else {
             return false;
         }
@@ -106,9 +133,22 @@ contract SimpleBudget is Budget {
         Transfer memory request = abi.decode(data_.cdDecompress(), (Transfer));
         if (request.assetType == AssetType.ERC20 || request.assetType == AssetType.ETH) {
             FungiblePayload memory payload = abi.decode(request.data, (FungiblePayload));
+
             uint256 avail = available(request.asset);
-            if (payload.amount > avail) revert InsufficientFunds(request.asset, avail, payload.amount);
-            _transfer(request.asset, request.target, payload.amount);
+            if (payload.amount > avail) {
+                revert InsufficientFunds(request.asset, avail, payload.amount);
+            }
+
+            _transferFungible(request.asset, request.target, payload.amount);
+        } else if (request.assetType == AssetType.ERC1155) {
+            ERC1155Payload memory payload = abi.decode(request.data, (ERC1155Payload));
+
+            uint256 avail = IERC1155(request.asset).balanceOf(address(this), payload.tokenId);
+            if (payload.amount > avail) {
+                revert InsufficientFunds(request.asset, avail, payload.amount);
+            }
+
+            _transferERC1155(request.asset, request.target, payload.tokenId, payload.amount, payload.data);
         } else {
             return false;
         }
@@ -152,7 +192,15 @@ contract SimpleBudget is Budget {
     /// @return The total amount of assets
     /// @dev This is simply the sum of the current balance and the distributed amount
     function total(address asset_) external view virtual override returns (uint256) {
-        return available(asset_) + _distributed[asset_];
+        return available(asset_) + _distributedFungible[asset_];
+    }
+
+    /// @notice Get the total amount of ERC1155 assets allocated to the budget, including any that have been distributed
+    /// @param asset_ The address of the asset
+    /// @param tokenId_ The ID of the token
+    /// @return The total amount of assets
+    function total(address asset_, uint256 tokenId_) external view virtual returns (uint256) {
+        return IERC1155(asset_).balanceOf(address(this), tokenId_) + _distributedERC1155[asset_][tokenId_];
     }
 
     /// @inheritdoc Budget
@@ -165,12 +213,28 @@ contract SimpleBudget is Budget {
         return asset_ == address(0) ? address(this).balance : asset_.balanceOf(address(this));
     }
 
+    /// @notice Get the amount of ERC1155 assets available for distribution from the budget
+    /// @param asset_ The address of the asset
+    /// @param tokenId_ The ID of the token
+    /// @return The amount of assets available
+    function available(address asset_, uint256 tokenId_) public view virtual returns (uint256) {
+        return IERC1155(asset_).balanceOf(address(this), tokenId_);
+    }
+
     /// @inheritdoc Budget
     /// @notice Get the amount of assets that have been distributed from the budget
     /// @param asset_ The address of the asset
     /// @return The amount of assets distributed
     function distributed(address asset_) external view virtual override returns (uint256) {
-        return _distributed[asset_];
+        return _distributedFungible[asset_];
+    }
+
+    /// @notice Get the amount of ERC1155 assets that have been distributed from the budget
+    /// @param asset_ The address of the asset
+    /// @param tokenId_ The ID of the token
+    /// @return The amount of assets distributed
+    function distributed(address asset_, uint256 tokenId_) external view virtual returns (uint256) {
+        return _distributedERC1155[asset_][tokenId_];
     }
 
     /// @inheritdoc Budget
@@ -185,12 +249,14 @@ contract SimpleBudget is Budget {
     /// @param amount_ The amount of the asset to transfer
     /// @dev This function is used to transfer assets from the budget to a given recipient (typically an incentive contract)
     /// @dev If the destination address is the zero address, or the transfer fails for any reason, this function will revert
-    function _transfer(address asset_, address to_, uint256 amount_) internal virtual {
+    function _transferFungible(address asset_, address to_, uint256 amount_) internal virtual nonReentrant {
         // Increment the total amount of the asset distributed from the budget
         if (to_ == address(0)) revert TransferFailed(asset_, to_, amount_);
-        if (amount_ > available(asset_)) revert InsufficientFunds(asset_, available(asset_), amount_);
+        if (amount_ > available(asset_)) {
+            revert InsufficientFunds(asset_, available(asset_), amount_);
+        }
 
-        _distributed[asset_] += amount_;
+        _distributedFungible[asset_] += amount_;
 
         // Transfer the asset to the recipient
         if (asset_ == address(0)) {
@@ -201,4 +267,53 @@ contract SimpleBudget is Budget {
 
         emit Distributed(asset_, to_, amount_);
     }
+
+    function _transferERC1155(address asset_, address to_, uint256 tokenId_, uint256 amount_, bytes memory data_)
+        internal
+        virtual
+        nonReentrant
+    {
+        // Increment the total amount of the asset distributed from the budget
+        if (to_ == address(0)) revert TransferFailed(asset_, to_, amount_);
+        if (amount_ > available(asset_, tokenId_)) {
+            revert InsufficientFunds(asset_, available(asset_, tokenId_), amount_);
+        }
+
+        _distributedERC1155[asset_][tokenId_] += amount_;
+
+        // Transfer the asset to the recipient
+        // wake-disable-next-line reentrancy (`nonReentrant` modifier is applied to the function)
+        IERC1155(asset_).safeTransferFrom(address(this), to_, tokenId_, amount_, data_);
+
+        emit Distributed(asset_, to_, amount_);
+    }
+
+    /// @inheritdoc IERC1155Receiver
+    /// @dev This contract does not care about the specifics of the inbound token, so we simply return the magic value (i.e. the selector for `onERC1155Received`)
+    function onERC1155Received(address, address, uint256, uint256, bytes calldata)
+        external
+        pure
+        override
+        returns (bytes4)
+    {
+        // We don't need to do anything here
+        return IERC1155Receiver.onERC1155Received.selector;
+    }
+
+    /// @inheritdoc IERC1155Receiver
+    /// @dev This contract does not care about the specifics of the inbound token, so we simply return the magic value (i.e. the selector for `onERC1155Received`)
+    function onERC1155BatchReceived(address, address, uint256[] calldata, uint256[] calldata, bytes calldata)
+        external
+        pure
+        override
+        returns (bytes4)
+    {
+        // We don't need to do anything here
+        return IERC1155Receiver.onERC1155BatchReceived.selector;
+    }
+
+    /// @inheritdoc Cloneable
+    function supportsInterface(bytes4 interfaceId) public view virtual override(Budget, IERC165) returns (bool) {
+        return interfaceId == type(IERC1155Receiver).interfaceId || super.supportsInterface(interfaceId);
+    }
 }

src/shared/Mocks.sol

@@ -4,6 +4,7 @@ pragma solidity ^0.8.24;
 import {LibString} from "lib/solady/src/utils/LibString.sol";
 import {ERC20} from "lib/solady/src/tokens/ERC20.sol";
 import {ERC721} from "lib/solady/src/tokens/ERC721.sol";
+import {ERC1155} from "lib/solady/src/tokens/ERC1155.sol";
 
 /**
  * 🚨 WARNING: The mocks in this file are for testing purposes only. DO NOT use
@@ -12,7 +13,8 @@ import {ERC721} from "lib/solady/src/tokens/ERC721.sol";
  * associated with someone who makes such poor life choices.
  */
 
-/// @notice A mock ERC721 token
+/// @title MockERC721
+/// @notice A mock ERC721 token (FOR TESTING PURPOSES ONLY)
 contract MockERC721 is ERC721 {
     uint256 public totalSupply;
     uint256 public mintPrice = 0.1 ether;
@@ -36,6 +38,8 @@ contract MockERC721 is ERC721 {
     }
 }
 
+/// @title MockERC20
+/// @notice A mock ERC20 token (FOR TESTING PURPOSES ONLY)
 contract MockERC20 is ERC20 {
     function name() public pure override returns (string memory) {
         return "Mock ERC20";
@@ -49,3 +53,19 @@ contract MockERC20 is ERC20 {
         _mint(to, amount);
     }
 }
+
+/// @title MockERC1155
+/// @notice A mock ERC1155 token (FOR TESTING PURPOSES ONLY)
+contract MockERC1155 is ERC1155 {
+    function uri(uint256 id) public view virtual override returns (string memory) {
+        return string(abi.encodePacked("https://example.com/token/", LibString.toString(id)));
+    }
+
+    function mint(address to, uint256 id, uint256 amount) public {
+        _mint(to, id, amount, "");
+    }
+
+    function burn(address from, uint256 id, uint256 amount) public {
+        _burn(from, id, amount);
+    }
+}