Fix incorrect DepositTree finalization when 0 finalized deposits

Author: zilm13

ethereum/pow/merkletree/src/main/java/tech/pegasys/teku/ethereum/pow/merkletree/DepositTree.java

@@ -95,6 +95,9 @@ public void finalize(final Eth1Data eth1Data, final UInt64 blockHeight) {
     checkArgument(
         totalDepositCount >= eth1Data.getDepositCount().longValue(),
         "Merkle tree does not contain all deposits to be finalized");
+    if (eth1Data.getDepositCount().equals(UInt64.ZERO)) {
+      return;
+    }
     finalizedExecutionBlock =
         Optional.of(new BlockHashAndHeight(eth1Data.getBlockHash(), blockHeight));
     finalizedDepositCount = eth1Data.getDepositCount().longValue();

ethereum/pow/merkletree/src/test/java/tech/pegasys/teku/ethereum/pow/merkletree/DepositTreeTest.java

@@ -19,8 +19,10 @@
 import static tech.pegasys.teku.ethereum.pow.api.DepositConstants.DEPOSIT_CONTRACT_TREE_DEPTH;
 
 import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
+import java.util.function.Function;
 import org.apache.tuweni.bytes.Bytes32;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -54,13 +56,58 @@ void shouldGenerateSnapshots(final int nonFinalizedDepositCount) throws Exceptio
 
       if (i >= nonFinalizedDepositCount) {
         final DepositTestCase finalisingTestCase = testCases.get(i - nonFinalizedDepositCount);
-        depositTree.finalize(finalisingTestCase.getEth1Data(), finalisingTestCase.getBlockHeight());
+        final Eth1Data correctEth1Data = finalisingTestCase.getEth1Data();
+        depositTree.finalize(
+            new Eth1Data(
+                correctEth1Data.getDepositRoot(),
+                correctEth1Data.getDepositCount().minus(1),
+                correctEth1Data.getBlockHash()),
+            finalisingTestCase.getBlockHeight());
         final Optional<DepositTreeSnapshot> snapshotOptional = depositTree.getSnapshot();
         assertThat(snapshotOptional).contains(finalisingTestCase.getSnapshot());
       }
     }
   }
 
+  @Test
+  void snapshotShouldHaveCorrectDepositCount() {
+    final int depositsCount = 10;
+    final Bytes32 depositRoot = Bytes32.random();
+    final Bytes32 blockHash = Bytes32.random();
+    final Function<Integer, Eth1Data> eth1Generator =
+        deposits -> new Eth1Data(depositRoot, UInt64.valueOf(deposits), blockHash);
+
+    final List<Bytes32> leafs = new ArrayList<>();
+    for (int i = 0; i < depositsCount; ++i) {
+      leafs.add(Bytes32.random());
+    }
+    final DepositTree fullTree = new DepositTree();
+    leafs.forEach(fullTree::pushLeaf);
+    fullTree.finalize(eth1Generator.apply(depositsCount), UInt64.ZERO);
+    final DepositTreeSnapshot fullTreeSnapshot = fullTree.getSnapshot().orElseThrow();
+
+    for (int finalized = 0; finalized <= depositsCount; ++finalized) {
+      final DepositTree customTree = new DepositTree();
+      for (int j = 0; j < depositsCount; ++j) {
+        customTree.pushLeaf(leafs.get(j));
+      }
+      customTree.finalize(eth1Generator.apply(finalized), UInt64.ZERO);
+      final Optional<DepositTreeSnapshot> customTreeSnapshot = customTree.getSnapshot();
+      System.out.println(customTreeSnapshot);
+      switch (finalized) {
+        case 0 -> assertThat(customTreeSnapshot).isEmpty();
+        case depositsCount -> assertThat(customTreeSnapshot).contains(fullTreeSnapshot);
+        default -> {
+          assertThat(customTreeSnapshot).isPresent();
+          final DepositTreeSnapshot snapshot = customTreeSnapshot.get();
+          assertThat(snapshot.getDepositCount()).isEqualTo(finalized);
+          assertThat(snapshot.getDepositRoot()).isNotEqualTo(fullTreeSnapshot.getDepositRoot());
+          assertThat(snapshot.getFinalized()).isNotEqualTo(fullTreeSnapshot.getFinalized());
+        }
+      }
+    }
+  }
+
   @Test
   void shouldRestoreFromSnapshots() throws Exception {
     final List<DepositTestCase> testCases = loadEipTestCases();