From fa485dfd5f5f316a9c662dd8b5e0b8c2be594683 Mon Sep 17 00:00:00 2001
From: Erik Dubovyk <erik.dubovyk@cyber-bee.co.uk>
Date: Tue, 6 Aug 2024 20:41:27 +0300
Subject: [PATCH] Added stack generators

---
 test/abstract/FlowUtilsAbstractTest.sol | 98 +++++++++++++++++++++++++
 1 file changed, 98 insertions(+)

diff --git a/test/abstract/FlowUtilsAbstractTest.sol b/test/abstract/FlowUtilsAbstractTest.sol
index 771df4ff..b3eb31b5 100644
--- a/test/abstract/FlowUtilsAbstractTest.sol
+++ b/test/abstract/FlowUtilsAbstractTest.sol
@@ -4,6 +4,9 @@ pragma solidity ^0.8.18;
 import {Test, Vm} from "forge-std/Test.sol";
 import {ERC20Transfer, ERC721Transfer, ERC1155Transfer, RAIN_FLOW_SENTINEL} from "src/interface/unstable/IFlowV5.sol";
 import {Sentinel} from "rain.solmem/lib/LibStackSentinel.sol";
+import {ERC1155SupplyChange} from "src/interface/unstable/IFlowERC1155V5.sol";
+import {ERC721SupplyChange} from "src/interface/unstable/IFlowERC721V5.sol";
+import {ERC20SupplyChange} from "src/interface/unstable/IFlowERC20V5.sol";
 
 abstract contract FlowUtilsAbstractTest is Test {
     function generateTokenTransferStack(
@@ -46,6 +49,101 @@ abstract contract FlowUtilsAbstractTest is Test {
         return stack;
     }
 
+    function generateFlowERC1155Stack(
+        ERC1155Transfer[] memory erc1155Transfers,
+        ERC721Transfer[] memory erc721Transfers,
+        ERC20Transfer[] memory erc20Transfers,
+        ERC1155SupplyChange[] memory erc1155Burns,
+        ERC1155SupplyChange[] memory erc1155Mints
+    ) internal pure returns (uint256[] memory stack) {
+        uint256[] memory transfersStack = generateTokenTransferStack(erc1155Transfers, erc721Transfers, erc20Transfers);
+        uint256 totalItems = transfersStack.length + 1 + (erc1155Burns.length * 3) + 1 + (erc1155Mints.length * 3);
+
+        stack = new uint256[](totalItems);
+        uint256 index = 0;
+        uint256 separator = Sentinel.unwrap(RAIN_FLOW_SENTINEL);
+
+        for (uint256 i = 0; i < transfersStack.length; i++) {
+            stack[index++] = transfersStack[i];
+        }
+
+        stack[index++] = separator;
+        for (uint256 i = 0; i < erc1155Burns.length; i++) {
+            stack[index++] = uint256(uint160(erc1155Burns[i].account));
+            stack[index++] = erc1155Burns[i].id;
+            stack[index++] = erc1155Burns[i].amount;
+        }
+
+        stack[index++] = separator;
+        for (uint256 i = 0; i < erc1155Mints.length; i++) {
+            stack[index++] = uint256(uint160(erc1155Mints[i].account));
+            stack[index++] = erc1155Mints[i].id;
+            stack[index++] = erc1155Mints[i].amount;
+        }
+    }
+
+    function generateFlowERC721Stack(
+        ERC1155Transfer[] memory erc1155Transfers,
+        ERC721Transfer[] memory erc721Transfers,
+        ERC20Transfer[] memory erc20Transfers,
+        ERC721SupplyChange[] memory erc721Burns,
+        ERC721SupplyChange[] memory erc721Mints
+    ) internal pure returns (uint256[] memory stack) {
+        uint256[] memory transfersStack = generateTokenTransferStack(erc1155Transfers, erc721Transfers, erc20Transfers);
+        uint256 totalItems = transfersStack.length + 1 + (erc721Burns.length * 2) + 1 + (erc721Mints.length * 2);
+
+        stack = new uint256[](totalItems);
+        uint256 index = 0;
+        uint256 separator = Sentinel.unwrap(RAIN_FLOW_SENTINEL);
+
+        for (uint256 i = 0; i < transfersStack.length; i++) {
+            stack[index++] = transfersStack[i];
+        }
+
+        stack[index++] = separator;
+        for (uint256 i = 0; i < erc721Burns.length; i++) {
+            stack[index++] = uint256(uint160(erc721Burns[i].account));
+            stack[index++] = erc721Burns[i].id;
+        }
+
+        stack[index++] = separator;
+        for (uint256 i = 0; i < erc721Mints.length; i++) {
+            stack[index++] = uint256(uint160(erc721Mints[i].account));
+            stack[index++] = erc721Mints[i].id;
+        }
+    }
+
+    function generateFlowERC20Stack(
+        ERC1155Transfer[] memory erc1155Transfers,
+        ERC721Transfer[] memory erc721Transfers,
+        ERC20Transfer[] memory erc20Transfers,
+        ERC20SupplyChange[] memory erc20Burns,
+        ERC20SupplyChange[] memory erc20Mints
+    ) internal pure returns (uint256[] memory stack) {
+        uint256[] memory transfersStack = generateTokenTransferStack(erc1155Transfers, erc721Transfers, erc20Transfers);
+        uint256 totalItems = transfersStack.length + 1 + (erc20Mints.length * 2) + 1 + (erc20Mints.length * 2);
+
+        stack = new uint256[](totalItems);
+        uint256 index = 0;
+        uint256 separator = Sentinel.unwrap(RAIN_FLOW_SENTINEL);
+
+        for (uint256 i = 0; i < transfersStack.length; i++) {
+            stack[index++] = transfersStack[i];
+        }
+
+        stack[index++] = separator;
+        for (uint256 i = 0; i < erc20Burns.length; i++) {
+            stack[index++] = uint256(uint160(erc20Burns[i].account));
+            stack[index++] = erc20Burns[i].amount;
+        }
+
+        stack[index++] = separator;
+        for (uint256 i = 0; i < erc20Mints.length; i++) {
+            stack[index++] = uint256(uint160(erc20Mints[i].account));
+            stack[index++] = erc20Mints[i].amount;
+        }
+    }
+
     function findEvent(Vm.Log[] memory logs, bytes32 eventSignature) internal pure returns (Vm.Log memory) {
         for (uint256 i = 0; i < logs.length; i++) {
             if (logs[i].topics[0] == eventSignature) {