diff --git a/platform-sdk/swirlds-common/src/testFixtures/java/com/swirlds/common/test/fixtures/merkle/util/MerkleTestUtils.java b/platform-sdk/swirlds-common/src/testFixtures/java/com/swirlds/common/test/fixtures/merkle/util/MerkleTestUtils.java index 2f67a01b36d2..73f2c2ba840d 100644 --- a/platform-sdk/swirlds-common/src/testFixtures/java/com/swirlds/common/test/fixtures/merkle/util/MerkleTestUtils.java +++ b/platform-sdk/swirlds-common/src/testFixtures/java/com/swirlds/common/test/fixtures/merkle/util/MerkleTestUtils.java @@ -53,6 +53,9 @@ import com.swirlds.config.api.Configuration; import com.swirlds.config.extensions.test.fixtures.TestConfigBuilder; import com.swirlds.metrics.api.Metrics; +import com.swirlds.virtualmap.VirtualKey; +import com.swirlds.virtualmap.VirtualMap; +import com.swirlds.virtualmap.internal.merkle.VirtualLeafNode; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.util.ArrayList; @@ -909,6 +912,36 @@ public static boolean areTreesEqual(final MerkleNode rootA, final MerkleNode roo return !iteratorB.hasNext(); } + /** + * For every virtual map in the trees and for every virtual key in the given key set, make + * sure either the map in both trees contains the key, or the map in both trees doesn't + * contain the key. + */ + @SuppressWarnings({"rawtypes", "unchecked"}) + public static boolean checkVirtualMapKeys( + final MerkleNode rootA, final MerkleNode rootB, final Set virtualKeys) { + final Iterator iteratorA = new MerkleIterator<>(rootA); + final Iterator iteratorB = new MerkleIterator<>(rootB); + while (iteratorA.hasNext()) { + if (!iteratorB.hasNext()) { + return false; + } + final MerkleNode a = iteratorA.next(); + final MerkleNode b = iteratorB.next(); + if (a instanceof VirtualMap vmA) { + if (!(b instanceof VirtualMap vmB)) { + return false; + } + for (final VirtualKey key : virtualKeys) { + if (vmA.containsKey(key) != vmB.containsKey(key)) { + return false; + } + } + } + } + return true; + } + /** * Check if a tree has had initialize() called on each internal node. */ @@ -1148,6 +1181,18 @@ private static boolean isVirtual(final MerkleNode node) { return node != null && (node.getClassId() == 0xaf2482557cfdb6bfL || node.getClassId() == 0x499677a326fb04caL); } + private static Set getVirtualKeys(final MerkleNode node) { + final Set keys = new HashSet<>(); + final Iterator it = new MerkleIterator<>(node); + while (it.hasNext()) { + final MerkleNode n = it.next(); + if (n instanceof VirtualLeafNode leaf) { + keys.add(leaf.getKey()); + } + } + return keys; + } + /** * Make sure the reconnect was valid. * @@ -1161,8 +1206,15 @@ private static boolean isVirtual(final MerkleNode node) { private static void assertReconnectValidity( final MerkleNode startingTree, final MerkleNode desiredTree, final MerkleNode generatedTree) { + // Checks that the trees are equal as merkle structures assertTrue(areTreesEqual(generatedTree, desiredTree), "reconnect should produce identical tree"); + final Set allKeys = new HashSet<>(); + allKeys.addAll(getVirtualKeys(startingTree)); + allKeys.addAll(getVirtualKeys(desiredTree)); + // A deeper check at VirtualMap level + assertTrue(checkVirtualMapKeys(generatedTree, desiredTree, allKeys)); + if (desiredTree != null) { assertNotSame(startingTree, desiredTree, "trees should be distinct objects"); diff --git a/platform-sdk/swirlds-common/src/testFixtures/java/module-info.java b/platform-sdk/swirlds-common/src/testFixtures/java/module-info.java index 2b1b1c777859..3dd94e2072c1 100644 --- a/platform-sdk/swirlds-common/src/testFixtures/java/module-info.java +++ b/platform-sdk/swirlds-common/src/testFixtures/java/module-info.java @@ -18,6 +18,7 @@ requires transitive com.swirlds.config.api; requires transitive com.swirlds.metrics.api; requires transitive com.swirlds.platform.core; + requires transitive com.swirlds.virtualmap; requires transitive com.hedera.pbj.runtime; requires com.swirlds.config.extensions.test.fixtures; requires com.swirlds.logging; diff --git a/platform-sdk/swirlds-platform-core/src/test/java/com/swirlds/platform/reconnect/VirtualMapReconnectTestBase.java b/platform-sdk/swirlds-platform-core/src/test/java/com/swirlds/platform/reconnect/VirtualMapReconnectTestBase.java index 531aa9ad12ce..c8b50216f95e 100644 --- a/platform-sdk/swirlds-platform-core/src/test/java/com/swirlds/platform/reconnect/VirtualMapReconnectTestBase.java +++ b/platform-sdk/swirlds-platform-core/src/test/java/com/swirlds/platform/reconnect/VirtualMapReconnectTestBase.java @@ -269,7 +269,12 @@ public void saveRecords( } delegate.saveRecords( - firstLeafPath, lastLeafPath, pathHashRecordsToUpdate, leaves.stream(), leafRecordsToDelete); + firstLeafPath, + lastLeafPath, + pathHashRecordsToUpdate, + leaves.stream(), + leafRecordsToDelete, + isReconnectContext); } @Override