From b59fbcf1695ab1819911b9a63af051f9b90429de Mon Sep 17 00:00:00 2001
From: Alexey <peretsmobil@gmail.com>
Date: Thu, 4 Feb 2021 20:47:47 +0300
Subject: [PATCH] move deposits length to a separate var

---
 contracts/TornadoTrees.sol           | 36 +++++++++++++++-------------
 contracts/mocks/TornadoTreesMock.sol |  7 ++++--
 test/tornadoTrees.test.js            |  6 +++--
 3 files changed, 29 insertions(+), 20 deletions(-)

diff --git a/contracts/TornadoTrees.sol b/contracts/TornadoTrees.sol
index e10935e..4c59545 100644
--- a/contracts/TornadoTrees.sol
+++ b/contracts/TornadoTrees.sol
@@ -25,10 +25,12 @@ contract TornadoTrees is EnsResolve {
   uint256 public constant BYTES_SIZE = 32 + 32 + 4 + CHUNK_SIZE * ITEM_SIZE;
   uint256 public constant SNARK_FIELD = 21888242871839275222246405745257275088548364400416034343698204186575808495617;
 
-  bytes32[] public deposits;
+  mapping(uint256 => bytes32) public deposits;
+  uint256 public depositsLength;
   uint256 public lastProcessedDepositLeaf;
 
-  bytes32[] public withdrawals;
+  mapping(uint256 => bytes32) public withdrawals;
+  uint256 public withdrawalsLength;
   uint256 public lastProcessedWithdrawalLeaf;
 
   bool public initialized;
@@ -86,8 +88,6 @@ contract TornadoTrees is EnsResolve {
     lastProcessedWithdrawalLeaf = withdrawalLeaf;
 
     uint256 i = depositLeaf;
-
-    // todo deposits.length = _tornadoTreesV1.deposits.length
     while (true) {
       (bool success, bytes memory data) = address(_tornadoTreesV1).staticcall{ gas: 3000 }( // todo define more specise gas value.
         abi.encodeWithSignature("deposits(uint256)", i)
@@ -96,9 +96,10 @@ contract TornadoTrees is EnsResolve {
         break;
       }
       bytes32 deposit = abi.decode(data, (bytes32));
+      deposits[i] = deposit;
       i++;
-      deposits.push(deposit);
     }
+    depositsLength = depositLeaf + i;
 
     i = withdrawalLeaf;
     while (true) {
@@ -110,25 +111,28 @@ contract TornadoTrees is EnsResolve {
         break;
       }
       bytes32 withdrawal = abi.decode(data, (bytes32));
+      withdrawals[i] = withdrawal;
       i++;
-      withdrawals.push(withdrawal);
     }
+    withdrawalsLength = withdrawalLeaf + i;
   }
 
   function registerDeposit(address _instance, bytes32 _commitment) external onlyTornadoProxy onlyInitialized {
-    deposits.push(keccak256(abi.encode(_instance, _commitment, blockNumber())));
-    emit DepositData(_instance, _commitment, blockNumber(), deposits.length - 1);
+    uint256 _depositsLength = depositsLength;
+    deposits[_depositsLength] = keccak256(abi.encode(_instance, _commitment, blockNumber()));
+    emit DepositData(_instance, _commitment, blockNumber(), _depositsLength - 1);
   }
 
   function registerWithdrawal(address _instance, bytes32 _nullifierHash) external onlyTornadoProxy onlyInitialized {
-    withdrawals.push(keccak256(abi.encode(_instance, _nullifierHash, blockNumber())));
-    emit WithdrawalData(_instance, _nullifierHash, blockNumber(), withdrawals.length - 1);
+    uint256 _withdrawalsLength = withdrawalsLength;
+    withdrawals[_withdrawalsLength] = keccak256(abi.encode(_instance, _nullifierHash, blockNumber()));
+    emit WithdrawalData(_instance, _nullifierHash, blockNumber(), _withdrawalsLength - 1);
   }
 
   function migrate(TreeLeaf[] calldata _depositEvents, TreeLeaf[] calldata _withdrawalEvents) external {
     require(!initialized, "Already migrated");
     uint256 _lastProcessedDepositLeaf = lastProcessedDepositLeaf;
-    uint256 _depositLength = deposits.length;
+    uint256 _depositLength = depositsLength;
     for (uint256 i = 0; i < _depositLength - _lastProcessedDepositLeaf; i++) {
       bytes32 leafHash = keccak256(abi.encode(_depositEvents[i].instance, _depositEvents[i].hash, _depositEvents[i].block));
       require(leafHash == deposits[_lastProcessedDepositLeaf + i], "Incorrect deposit");
@@ -140,7 +144,7 @@ contract TornadoTrees is EnsResolve {
       );
     }
 
-    uint256 _withdrawalLength = withdrawals.length;
+    uint256 _withdrawalLength = withdrawalsLength;
     uint256 _lastProcessedWithdrawalLeaf = lastProcessedWithdrawalLeaf;
     for (uint256 i = 0; i < _withdrawalLength - _lastProcessedWithdrawalLeaf; i++) {
       bytes32 leafHash = keccak256(
@@ -245,16 +249,16 @@ contract TornadoTrees is EnsResolve {
     require(_withdrawalRoot == withdrawalRoot || _withdrawalRoot == previousWithdrawalRoot, "Incorrect withdrawal tree root");
   }
 
-  function getRegisteredDeposits() external view returns (uint256 count, bytes32[] memory _deposits) {
-    count = deposits.length - lastProcessedDepositLeaf;
+  function getRegisteredDeposits() external view returns (bytes32[] memory _deposits) {
+    uint256 count = depositsLength - lastProcessedDepositLeaf;
     _deposits = new bytes32[](count);
     for (uint256 i = 0; i < count; i++) {
       _deposits[i] = deposits[lastProcessedDepositLeaf + i];
     }
   }
 
-  function getRegisteredWithdrawals() external view returns (uint256 count, bytes32[] memory _withdrawals) {
-    count = withdrawals.length - lastProcessedWithdrawalLeaf;
+  function getRegisteredWithdrawals() external view returns (bytes32[] memory _withdrawals) {
+    uint256 count = withdrawalsLength - lastProcessedWithdrawalLeaf;
     _withdrawals = new bytes32[](count);
     for (uint256 i = 0; i < count; i++) {
       _withdrawals[i] = withdrawals[lastProcessedWithdrawalLeaf + i];
diff --git a/contracts/mocks/TornadoTreesMock.sol b/contracts/mocks/TornadoTreesMock.sol
index b525d77..a185a65 100644
--- a/contracts/mocks/TornadoTreesMock.sol
+++ b/contracts/mocks/TornadoTreesMock.sol
@@ -37,9 +37,12 @@ contract TornadoTreesMock is TornadoTrees {
     uint256 _withdrawBlockNumber
   ) public {
     setBlockNumber(_depositBlockNumber);
-    deposits.push(keccak256(abi.encode(_instance, _commitment, blockNumber())));
+    deposits[depositsLength] = keccak256(abi.encode(_instance, _commitment, blockNumber()));
+    depositsLength++;
+
     setBlockNumber(_withdrawBlockNumber);
-    withdrawals.push(keccak256(abi.encode(_instance, _nullifier, blockNumber())));
+    withdrawals[withdrawalsLength] = keccak256(abi.encode(_instance, _nullifier, blockNumber()));
+    withdrawalsLength++;
   }
 
   function updateDepositTreeMock(
diff --git a/test/tornadoTrees.test.js b/test/tornadoTrees.test.js
index 6e8aa1a..7136b06 100644
--- a/test/tornadoTrees.test.js
+++ b/test/tornadoTrees.test.js
@@ -122,7 +122,8 @@ describe('TornadoTrees', function () {
   describe('#getRegisteredDeposits', () => {
     it('should work', async () => {
       const abi = new ethers.utils.AbiCoder()
-      let { count, _deposits } = await tornadoTrees.getRegisteredDeposits()
+      const count = await tornadoTrees.depositsLength()
+      const _deposits = await tornadoTrees.getRegisteredDeposits()
       expect(count).to.be.equal(notes.length)
       _deposits.forEach((hash, i) => {
         const encodedData = abi.encode(
@@ -149,7 +150,8 @@ describe('TornadoTrees', function () {
   describe('#getRegisteredWithdrawals', () => {
     it('should work', async () => {
       const abi = new ethers.utils.AbiCoder()
-      let { count, _withdrawals } = await tornadoTrees.getRegisteredWithdrawals()
+      const count = await tornadoTrees.withdrawalsLength()
+      const _withdrawals = await tornadoTrees.getRegisteredWithdrawals()
       expect(count).to.be.equal(notes.length)
       _withdrawals.forEach((hash, i) => {
         const encodedData = abi.encode(