Repository: thehubbleproject/hubble-contracts (ts)

Issue:
[Audit] Use mapping for babyTrees
> Storage allocated variable is used as in-cycle index. Also, lists are unoptimized data structure in solidity, each operation with list requires additional SLOAD to check index overflow.

> We propose replacing list to hashmap and using memory cached variables in cycles.

https://github.com/thehubbleproject/hubble-contracts/blob/32234da0ce3015b4bece04a264a136a6778191dd/contracts/DepositManager.sol#L68

Pull request:
Use mapping for baby tree
Fix #398 

Gas cost before:
```
Insert leaf 0 	 Operation cost: 63383
Insert leaf 1 	 Operation cost: 48232
Insert leaf 2 	 Operation cost: 33383
Insert leaf 3 	 Operation cost: 67278
Insert leaf 4 	 Operation cost: 33383
Insert leaf 5 	 Operation cost: 48232
Insert leaf 6 	 Operation cost: 33383
Insert leaf 7 	 Operation cost: 82124
Insert leaf 8 	 Operation cost: 33383
Insert leaf 9 	 Operation cost: 48232
Insert leaf 10 	 Operation cost: 33383
Insert leaf 11 	 Operation cost: 67278
Insert leaf 12 	 Operation cost: 33383
Insert leaf 13 	 Operation cost: 48232
Insert leaf 14 	 Operation cost: 33383
Insert leaf 15 	 Operation cost: 145649
```

After (if we delete the right node)
```
Insert leaf 0    Operation cost: 63433
Insert leaf 1    Operation cost: 37320
Insert leaf 2    Operation cost: 33433
Insert leaf 3    Operation cost: 49604
Insert leaf 4    Operation cost: 33433
Insert leaf 5    Operation cost: 37320
Insert leaf 6    Operation cost: 33433
Insert leaf 7    Operation cost: 57688
Insert leaf 8    Operation cost: 33433
Insert leaf 9    Operation cost: 37320
Insert leaf 10   Operation cost: 33433
Insert leaf 11   Operation cost: 49604
Insert leaf 12   Operation cost: 33433
Insert leaf 13   Operation cost: 37320
Insert leaf 14   Operation cost: 33433
Insert leaf 15   Operation cost: 110908
```
(if we don't delete the right node)
```
Insert leaf 0    Operation cost: 63433
Insert leaf 1    Operation cost: 36454
Insert leaf 2    Operation cost: 18433
Insert leaf 3    Operation cost: 47872
Insert leaf 4    Operation cost: 18433
Insert leaf 5    Operation cost: 21454
Insert leaf 6    Operation cost: 18433
Insert leaf 7    Operation cost: 55090
Insert leaf 8    Operation cost: 18433
Insert leaf 9    Operation cost: 21454
Insert leaf 10   Operation cost: 18433
Insert leaf 11   Operation cost: 32872
Insert leaf 12   Operation cost: 18433
Insert leaf 13   Operation cost: 21454
Insert leaf 14   Operation cost: 18433
Insert leaf 15   Operation cost: 105671
```


Patch:
diff --git a/contracts/DepositManager.sol b/contracts/DepositManager.sol
index 36d60ee..9cdf345 100644
--- a/contracts/DepositManager.sol
+++ b/contracts/DepositManager.sol
@@ -36,10 +36,12 @@ contract SubtreeQueue {
 }
 
 contract DepositCore is SubtreeQueue {
-    // An element in this array is a deposit tree root of any depth.
+    // An element is a deposit tree root of any depth.
     // It could be just a leaf of a new deposit or
     // a root of a full grown subtree.
-    bytes32[] public babyTrees;
+    mapping(uint256 => bytes32) public babyTrees;
+    uint256 public babyTreesLength = 0;
+
     uint256 public depositCount = 0;
 
     uint256 public paramMaxSubtreeSize = 2;
@@ -48,47 +50,34 @@ contract DepositCore is SubtreeQueue {
         internal
         returns (bytes32 readySubtree)
     {
-        babyTrees.push(depositLeaf);
         depositCount++;
         uint256 i = depositCount;
+
+        uint256 len = babyTreesLength;
+        babyTrees[len] = depositLeaf;
+        len++;
         // As long as we have a pair to merge, we merge
         // the number of iteration is bounded by maxSubtreeDepth
         while (i & 1 == 0) {
-            // Override the left node with the merged result
-            babyTrees[babyTrees.length - 2] = keccak256(
-                abi.encode(
-                    // left node
-                    babyTrees[babyTrees.length - 2],
-                    // right node
-                    babyTrees[babyTrees.length - 1]
-                )
+            // Override the left node with the merged left and right nodes
+            babyTrees[len - 2] = keccak256(
+                abi.encode(babyTrees[len - 2], babyTrees[len - 1])
             );
-            // Discard the right node
-            delete babyTrees[babyTrees.length - 1];
-            babyTrees.length--;
-
+            len--;
             i >>= 1;
         }
+        babyTreesLength = len;
         // Subtree is ready, send to SubtreeQueue
         if (depositCount == paramMaxSubtreeSize) {
             readySubtree = babyTrees[0];
             enqueue(readySubtree);
-            reset();
+            // reset
+            babyTreesLength = 0;
+            depositCount = 0;
         } else {
             readySubtree = bytes32(0);
         }
     }
-
-    function reset() internal {
-        // Reset babyTrees
-        uint256 numberOfDeposits = babyTrees.length;
-        for (uint256 i = 0; i < numberOfDeposits; i++) {
-            delete babyTrees[i];
-        }
-        babyTrees.length = 0;
-        // Reset depositCount
-        depositCount = 0;
-    }
 }
 
 contract DepositManager is DepositCore, IDepositManager {


Test patch:
diff --git a/test/fast/deposit.test.ts b/test/fast/deposit.test.ts
index 2840d0c..b9013d4 100644
--- a/test/fast/deposit.test.ts
+++ b/test/fast/deposit.test.ts
@@ -23,39 +23,41 @@ describe("Deposit Core", async function() {
         );
     });
     it("insert and merge many deposits", async function() {
-        const maxSubtreeSize = 2 ** maxSubtreeDepth;
-        const leaves = randomLeaves(maxSubtreeSize);
-        const tree = Tree.new(maxSubtreeDepth);
-        for (let i = 0; i < maxSubtreeSize; i++) {
-            const {
-                gasCost,
-                readySubtree
-            } = await contract.callStatic.testInsertAndMerge(leaves[i]);
-            console.log(
-                `Insert leaf ${i} \t Operation cost: ${gasCost.toNumber()}`
-            );
-            await contract.testInsertAndMerge(leaves[i]);
-            tree.updateSingle(i, leaves[i]);
-            if (i !== maxSubtreeSize - 1) {
-                assert.equal(
-                    readySubtree,
-                    constants.HashZero,
-                    "Not a ready subtree yet"
-                );
-            } else {
-                assert.equal(
-                    readySubtree,
-                    tree.root,
-                    "Should be the merkle root of all leaves"
+        for (let j = 0; j < 5; j++) {
+            const maxSubtreeSize = 2 ** maxSubtreeDepth;
+            const leaves = randomLeaves(maxSubtreeSize);
+            const tree = Tree.new(maxSubtreeDepth);
+            for (let i = 0; i < maxSubtreeSize; i++) {
+                const {
+                    gasCost,
+                    readySubtree
+                } = await contract.callStatic.testInsertAndMerge(leaves[i]);
+                console.log(
+                    `Insert leaf ${i} \t Operation cost: ${gasCost.toNumber()}`
                 );
+                await contract.testInsertAndMerge(leaves[i]);
+                tree.updateSingle(i, leaves[i]);
+                if (i !== maxSubtreeSize - 1) {
+                    assert.equal(
+                        readySubtree,
+                        constants.HashZero,
+                        "Not a ready subtree yet"
+                    );
+                } else {
+                    assert.equal(
+                        readySubtree,
+                        tree.root,
+                        "Should be the merkle root of all leaves"
+                    );
+                }
             }
+            assert.equal((await contract.back()).toNumber(), j + 1);
+            assert.equal(
+                await contract.getQueue(j + 1),
+                tree.root,
+                "subtree root should be in the subtree queue now"
+            );
         }
-        assert.equal((await contract.back()).toNumber(), 1);
-        assert.equal(
-            await contract.getQueue(1),
-            tree.root,
-            "subtree root should be in the subtree queue now"
-        );
     });
 });