From 39777a46e185916a71b920d0e2dcf550bd6f1c22 Mon Sep 17 00:00:00 2001
From: Alexey <peretsmobil@gmail.com>
Date: Tue, 16 Jul 2019 14:04:14 +0300
Subject: [PATCH] max elements tree count test

---
 contracts/MerkleTreeWithHistory.sol | 10 +++++-----
 test/MerkleTreeWithHistory.test.js  | 23 ++++++++++++++++++++---
 2 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/contracts/MerkleTreeWithHistory.sol b/contracts/MerkleTreeWithHistory.sol
index 28aedbe..5c055c3 100644
--- a/contracts/MerkleTreeWithHistory.sol
+++ b/contracts/MerkleTreeWithHistory.sol
@@ -5,7 +5,7 @@ library MiMC {
 }
 
 contract MerkleTreeWithHistory {
-  uint8 levels;
+  uint256 public levels;
 
   uint8 constant ROOT_HISTORY_SIZE = 100;
   uint256[] private _roots;
@@ -18,7 +18,7 @@ contract MerkleTreeWithHistory {
 
   event LeafAdded(uint256 leaf, uint32 leaf_index);
 
-  constructor(uint8 tree_levels, uint256 zero_value) public {
+  constructor(uint256 tree_levels, uint256 zero_value) public {
     levels = tree_levels;
 
     _zeros.push(zero_value);
@@ -50,13 +50,13 @@ contract MerkleTreeWithHistory {
   function _insert(uint256 leaf) internal {
     uint32 leaf_index = next_index;
     uint32 current_index = next_index;
+    require(current_index != 2**(levels - 1), "Merkle tree is full");
     next_index += 1;
-
     uint256 current_level_hash = leaf;
     uint256 left;
     uint256 right;
 
-    for (uint8 i = 0; i < levels; i++) {
+    for (uint256 i = 0; i < levels; i++) {
       if (current_index % 2 == 0) {
         left = current_level_hash;
         right = _zeros[i];
@@ -78,7 +78,7 @@ contract MerkleTreeWithHistory {
     emit LeafAdded(leaf, leaf_index);
   }
 
-  function isKnownRoot(uint root) public view returns(bool) {
+  function isKnownRoot(uint256 root) public view returns(bool) {
     if (root == 0) {
       return false;
     }
diff --git a/test/MerkleTreeWithHistory.test.js b/test/MerkleTreeWithHistory.test.js
index 5f2f633..452a38a 100644
--- a/test/MerkleTreeWithHistory.test.js
+++ b/test/MerkleTreeWithHistory.test.js
@@ -12,6 +12,8 @@ const MiMC = artifacts.require('./MiMC.sol')
 const MerkleTree = require('../lib/MerkleTree')
 const MimcHasher = require('../lib/MiMC')
 
+const { AMOUNT, MERKLE_TREE_HEIGHT, EMPTY_ELEMENT } = process.env
+
 function BNArrayToStringArray(array) {
   const arrayToPrint = []
   array.forEach(item => {
@@ -25,8 +27,9 @@ contract('MerkleTreeWithHistory', async accounts => {
   let miMC
   const sender = accounts[0]
   const emptyAddress = '0x0000000000000000000000000000000000000000'
-  const levels = 16
-  const zeroValue = 1337
+  const levels = MERKLE_TREE_HEIGHT || 16
+  const zeroValue = EMPTY_ELEMENT || 1337
+  const value = AMOUNT || '1000000000000000000'
   let snapshotId
   let prefix = 'test'
   let tree
@@ -145,7 +148,6 @@ contract('MerkleTreeWithHistory', async accounts => {
 
   describe('#insert', async () => {
     it('should insert', async () => {
-      let filled_subtrees
       let rootFromContract
 
       for (i = 1; i < 11; i++) {
@@ -161,6 +163,21 @@ contract('MerkleTreeWithHistory', async accounts => {
         // console.log('rootFromCon', root.toString())
       }
     })
+
+    it.skip('should reject if tree is full', async () => {
+      // consider change the merkle tree height to run the test
+      // it should be changed in .env and ./circuits/withdraw.circom files (default is 16)
+
+      for (i = 0; i < 2**(MERKLE_TREE_HEIGHT - 1); i++) {
+        await merkleTreeWithHistory.insert(i+42).should.be.fulfilled
+      }
+
+      let error = await merkleTreeWithHistory.insert(1337).should.be.rejected
+      error.reason.should.be.equal('Merkle tree is full')
+
+      error = await merkleTreeWithHistory.insert(1).should.be.rejected
+      error.reason.should.be.equal('Merkle tree is full')
+    })
   })
 
   describe('#MIMC', async () => {