load-store backing nodes fix. address some feedbacks.

Consensys · Jun 26, 2024 · 2e1bdd0 · 2e1bdd0
1 parent c421005
commit 2e1bdd0
Show file tree

Hide file tree

Showing 8 changed files with 160 additions and 25 deletions.
diff --git a/...ure/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/impl/AbstractSszComposite.java b/...ure/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/impl/AbstractSszComposite.java
@@ -96,7 +96,7 @@ protected IntCache<SszChildT> createCache() {
   }
 
   @Override
-  public SszChildT get(final int index) {
+  public final SszChildT get(final int index) {
     return childrenViewCache.getInt(index, this::getImplWithIndexCheck);
   }
 

diff --git a/...cture/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/schema/SszProfileSchema.java b/...cture/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/schema/SszProfileSchema.java
@@ -36,6 +36,10 @@ default TreeNode createTreeFromOptionalFieldValues(
 
   SszBitvector getActiveFields();
 
+  default List<? extends SszSchema<?>> getActiveChildrenSchemas() {
+    return getActiveFields().streamAllSetBits().mapToObj(this::getChildSchema).toList();
+  }
+
   /**
    * This method resolves the index of the nth active field.
    *

diff --git a/...z/src/main/java/tech/pegasys/teku/infrastructure/ssz/schema/SszStableContainerSchema.java b/...z/src/main/java/tech/pegasys/teku/infrastructure/ssz/schema/SszStableContainerSchema.java
@@ -76,7 +76,6 @@ default C createFromOptionalFieldValues(final List<Optional<? extends SszData>>
   SszBitvector getActiveFieldsBitvectorFromBackingNode(TreeNode node);
 
   @Override
-  @SuppressWarnings("unchecked")
   default Optional<SszStableContainerSchema<?>> toStableContainerSchema() {
     return Optional.of(this);
   }

diff --git a/.../main/java/tech/pegasys/teku/infrastructure/ssz/schema/impl/AbstractSszProfileSchema.java b/.../main/java/tech/pegasys/teku/infrastructure/ssz/schema/impl/AbstractSszProfileSchema.java
@@ -14,13 +14,15 @@
 package tech.pegasys.teku.infrastructure.ssz.schema.impl;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkState;
 import static tech.pegasys.teku.infrastructure.ssz.schema.impl.AbstractSszStableContainerSchema.CONTAINER_G_INDEX;
 
 import it.unimi.dsi.fastutil.ints.IntList;
 import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.List;
 import java.util.Set;
+import org.apache.tuweni.bytes.Bytes32;
 import tech.pegasys.teku.infrastructure.json.types.DeserializableTypeDefinition;
 import tech.pegasys.teku.infrastructure.ssz.SszData;
 import tech.pegasys.teku.infrastructure.ssz.SszProfile;
@@ -37,6 +39,10 @@
 import tech.pegasys.teku.infrastructure.ssz.tree.BranchNode;
 import tech.pegasys.teku.infrastructure.ssz.tree.GIndexUtil;
 import tech.pegasys.teku.infrastructure.ssz.tree.TreeNode;
+import tech.pegasys.teku.infrastructure.ssz.tree.TreeNodeSource;
+import tech.pegasys.teku.infrastructure.ssz.tree.TreeNodeSource.CompressedBranchInfo;
+import tech.pegasys.teku.infrastructure.ssz.tree.TreeNodeStore;
+import tech.pegasys.teku.infrastructure.ssz.tree.TreeUtil;
 
 public abstract class AbstractSszProfileSchema<C extends SszProfile>
     extends AbstractSszContainerSchema<C> implements SszProfileSchema<C> {
@@ -172,6 +178,59 @@ public boolean isActiveField(final int index) {
     return activeFields.getBit(index);
   }
 
+  @Override
+  public void storeBackingNodes(
+      final TreeNodeStore nodeStore,
+      final int maxBranchLevelsSkipped,
+      final long rootGIndex,
+      final TreeNode node) {
+    final TreeNode containerSubtree = node.get(CONTAINER_G_INDEX);
+    super.storeBackingNodes(
+        nodeStore,
+        maxBranchLevelsSkipped,
+        GIndexUtil.gIdxLeftGIndex(rootGIndex),
+        node.get(CONTAINER_G_INDEX));
+
+    nodeStore.storeBranchNode(
+        node.hashTreeRoot(), rootGIndex, 1, new Bytes32[] {containerSubtree.hashTreeRoot()});
+  }
+
+  @Override
+  public TreeNode loadBackingNodes(
+      final TreeNodeSource nodeSource, final Bytes32 rootHash, final long rootGIndex) {
+    if (TreeUtil.ZERO_TREES_BY_ROOT.containsKey(rootHash) || rootHash.equals(Bytes32.ZERO)) {
+      return getDefaultTree();
+    }
+
+    final CompressedBranchInfo branchData = nodeSource.loadBranchNode(rootHash, rootGIndex);
+    checkState(
+        branchData.getChildren().length == 1, "Profile root node must have exactly 1 children");
+    checkState(branchData.getDepth() == 1, "Profile root node must have depth of 1");
+    final Bytes32 containerHash = branchData.getChildren()[0];
+
+    long containerRootGIndex = GIndexUtil.gIdxLeftGIndex(rootGIndex);
+
+    final long lastUsefulGIndex =
+        GIndexUtil.gIdxChildGIndex(containerRootGIndex, maxChunks() - 1, treeDepth());
+    TreeNode containerTree =
+        LoadingUtil.loadNodesToDepth(
+            nodeSource,
+            containerHash,
+            containerRootGIndex,
+            treeDepth(),
+            super.getDefaultTree(),
+            lastUsefulGIndex,
+            this::loadChildNode);
+
+    return BranchNode.create(containerTree, activeFields.getBackingNode());
+  }
+
+  private TreeNode loadChildNode(
+      final TreeNodeSource nodeSource, final Bytes32 childHash, final long childGIndex) {
+    final int childIndex = GIndexUtil.gIdxChildIndexFromGIndex(childGIndex, treeDepth());
+    return getChildSchema(childIndex).loadBackingNodes(nodeSource, childHash, childGIndex);
+  }
+
   @Override
   public SszBitvector getActiveFields() {
     return activeFields;

diff --git a/...va/tech/pegasys/teku/infrastructure/ssz/schema/impl/AbstractSszStableContainerSchema.java b/...va/tech/pegasys/teku/infrastructure/ssz/schema/impl/AbstractSszStableContainerSchema.java
@@ -128,9 +128,9 @@ public AbstractSszStableContainerSchema(
   }
 
   protected TreeNode createDefaultContainerTreeNode() {
-    final List<TreeNode> defaultChildren = new ArrayList<>((int) getMaxLength());
+    final List<TreeNode> defaultChildren = new ArrayList<>(getFieldsCount());
     for (int i = 0; i < getFieldsCount(); i++) {
-      defaultChildren.add(getChildSchema(i).getDefault().getBackingNode());
+      defaultChildren.add(SszNone.INSTANCE.getBackingNode());
     }
     return TreeUtil.createTree(defaultChildren);
   }
@@ -375,14 +375,6 @@ public SszLengthBounds getSszLengthBounds() {
     return super.getSszLengthBounds().add(activeFieldsSchema.getSszLengthBounds());
   }
 
-  public SszLengthBounds getSszLengthBoundsAsProfile() {
-    return super.getSszLengthBounds();
-  }
-
-  public int getSszSizeAsProfile(final TreeNode node) {
-    return super.getSszSize(node);
-  }
-
   @Override
   public void storeBackingNodes(
       final TreeNodeStore nodeStore,
@@ -392,8 +384,12 @@ public void storeBackingNodes(
     final TreeNode containerSubtree = node.get(CONTAINER_G_INDEX);
     final TreeNode activeFieldsBitvectorSubtree = node.get(BITVECTOR_G_INDEX);
 
-    super.storeBackingNodes(
-        nodeStore, maxBranchLevelsSkipped, getContainerGIndex(rootGIndex), containerSubtree);
+    storeContainerBackingNodes(
+        activeFieldsSchema.createFromBackingNode(activeFieldsBitvectorSubtree),
+        nodeStore,
+        maxBranchLevelsSkipped,
+        getContainerGIndex(rootGIndex),
+        containerSubtree);
 
     activeFieldsSchema.storeBackingNodes(
         nodeStore,
@@ -404,12 +400,58 @@ public void storeBackingNodes(
     nodeStore.storeBranchNode(
         node.hashTreeRoot(),
         rootGIndex,
-        2,
+        1,
         new Bytes32[] {
           containerSubtree.hashTreeRoot(), activeFieldsBitvectorSubtree.hashTreeRoot()
         });
   }
 
+  private void storeContainerBackingNodes(
+      final SszBitvector activeFields,
+      final TreeNodeStore nodeStore,
+      final int maxBranchLevelsSkipped,
+      final long rootGIndex,
+      final TreeNode node) {
+
+    final int childDepth = treeDepth();
+    if (childDepth == 0) {
+      // Only one child so wrapper is omitted
+      storeChildNode(nodeStore, maxBranchLevelsSkipped, rootGIndex, node);
+      return;
+    }
+    final long lastUsefulGIndex =
+        GIndexUtil.gIdxChildGIndex(rootGIndex, maxChunks() - 1, childDepth);
+    StoringUtil.storeNodesToDepth(
+        nodeStore,
+        maxBranchLevelsSkipped,
+        node,
+        rootGIndex,
+        childDepth,
+        lastUsefulGIndex,
+        (targetDepthNode, targetDepthGIndex) ->
+            storeChildNode(
+                nodeStore,
+                maxBranchLevelsSkipped,
+                targetDepthGIndex,
+                targetDepthNode,
+                activeFields));
+  }
+
+  public void storeChildNode(
+      final TreeNodeStore nodeStore,
+      final int maxBranchLevelsSkipped,
+      final long gIndex,
+      final TreeNode node,
+      final SszBitvector activeFields) {
+    final int childIndex = GIndexUtil.gIdxChildIndexFromGIndex(gIndex, treeDepth());
+    if (activeFields.getBit(childIndex)) {
+      final SszSchema<?> childSchema = getChildSchema(childIndex);
+      childSchema.storeBackingNodes(nodeStore, maxBranchLevelsSkipped, gIndex, node);
+    }
+    SszPrimitiveSchemas.NONE_SCHEMA.storeBackingNodes(
+        nodeStore, maxBranchLevelsSkipped, gIndex, node);
+  }
+
   @Override
   public TreeNode loadBackingNodes(
       final TreeNodeSource nodeSource, final Bytes32 rootHash, final long rootGIndex) {
@@ -424,6 +466,13 @@ public TreeNode loadBackingNodes(
     final Bytes32 containerHash = branchData.getChildren()[0];
     final Bytes32 activeFieldsBitvectorHash = branchData.getChildren()[1];
 
+    final TreeNode activeFieldsTreeNode =
+        activeFieldsSchema.loadBackingNodes(
+            nodeSource, activeFieldsBitvectorHash, getBitvectorGIndex(rootGIndex));
+
+    final SszBitvector activeFields =
+        activeFieldsSchema.createFromBackingNode(activeFieldsTreeNode);
+
     final long lastUsefulGIndex =
         GIndexUtil.gIdxChildGIndex(rootGIndex, maxChunks() - 1, treeDepth());
     final TreeNode containerTreeNode =
@@ -434,18 +483,22 @@ public TreeNode loadBackingNodes(
             treeDepth(),
             defaultTreeNode.get(CONTAINER_G_INDEX),
             lastUsefulGIndex,
-            this::loadChildNode);
+            (tns, childHash, childGIndex) ->
+                loadChildNode(activeFields, tns, childHash, childGIndex));
 
-    return BranchNode.create(
-        containerTreeNode,
-        activeFieldsSchema.loadBackingNodes(
-            nodeSource, activeFieldsBitvectorHash, getBitvectorGIndex(rootGIndex)));
+    return BranchNode.create(containerTreeNode, activeFieldsTreeNode);
   }
 
   private TreeNode loadChildNode(
-      final TreeNodeSource nodeSource, final Bytes32 childHash, final long childGIndex) {
+      final SszBitvector activeFields,
+      final TreeNodeSource nodeSource,
+      final Bytes32 childHash,
+      final long childGIndex) {
     final int childIndex = GIndexUtil.gIdxChildIndexFromGIndex(childGIndex, treeDepth());
-    return getChildSchema(childIndex).loadBackingNodes(nodeSource, childHash, childGIndex);
+    if (activeFields.getBit(childIndex)) {
+      return getChildSchema(childIndex).loadBackingNodes(nodeSource, childHash, childGIndex);
+    }
+    return SszPrimitiveSchemas.NONE_SCHEMA.loadBackingNodes(nodeSource, childHash, childGIndex);
   }
 
   private static List<? extends NamedIndexedSchema<?>> createAllSchemas(

diff --git a/...tructure/ssz/src/test/java/tech/pegasys/teku/infrastructure/ssz/TestStableContainers.java b/...tructure/ssz/src/test/java/tech/pegasys/teku/infrastructure/ssz/TestStableContainers.java
@@ -133,7 +133,7 @@ public ShapeStableContainer createFromBackingNode(final TreeNode node) {
 
   public static final SszStableContainerSchema<CircleProfile> CIRCLE_PROFILE_SCHEMA =
       new AbstractSszProfileSchema<>(
-          "Circle", SHAPE_STABLE_CONTAINER_SCHEMA, CIRCLE_SCHEMA_INDICES) {
+          "CircleProfile", SHAPE_STABLE_CONTAINER_SCHEMA, CIRCLE_SCHEMA_INDICES) {
         @Override
         public CircleProfile createFromBackingNode(final TreeNode node) {
           return new CircleProfile(this, node);
@@ -142,7 +142,7 @@ public CircleProfile createFromBackingNode(final TreeNode node) {
 
   public static final SszStableContainerSchema<SquareProfile> SQUARE_PROFILE_SCHEMA =
       new AbstractSszProfileSchema<>(
-          "Square", SHAPE_STABLE_CONTAINER_SCHEMA, SQUARE_SCHEMA_INDICES) {
+          "SquareProfile", SHAPE_STABLE_CONTAINER_SCHEMA, SQUARE_SCHEMA_INDICES) {
         @Override
         public SquareProfile createFromBackingNode(final TreeNode node) {
           return new SquareProfile(this, node);

diff --git a/...ture/ssz/src/test/java/tech/pegasys/teku/infrastructure/ssz/schema/SszSchemaTestBase.java b/...ture/ssz/src/test/java/tech/pegasys/teku/infrastructure/ssz/schema/SszSchemaTestBase.java
@@ -126,6 +126,7 @@ protected void assertTreeRoundtrip(
     assertThatTreeNode(result).isTreeEqual(node);
     final SszData rebuiltData = schema.createFromBackingNode(result);
     assertThat(rebuiltData).isEqualTo(data);
+    assertThat(rebuiltData.sszSerialize()).isEqualTo(data.sszSerialize());
   }
 
   @MethodSource("testSchemaArguments")

diff --git a/...sz/src/testFixtures/java/tech/pegasys/teku/infrastructure/ssz/RandomSszDataGenerator.java b/...sz/src/testFixtures/java/tech/pegasys/teku/infrastructure/ssz/RandomSszDataGenerator.java
@@ -38,6 +38,7 @@
 import tech.pegasys.teku.infrastructure.ssz.schema.SszUnionSchema;
 import tech.pegasys.teku.infrastructure.ssz.schema.impl.AbstractSszContainerSchema;
 import tech.pegasys.teku.infrastructure.ssz.schema.impl.AbstractSszPrimitiveSchema;
+import tech.pegasys.teku.infrastructure.ssz.schema.impl.AbstractSszProfileSchema;
 import tech.pegasys.teku.infrastructure.ssz.schema.impl.AbstractSszStableContainerSchema;
 import tech.pegasys.teku.infrastructure.ssz.schema.impl.AbstractSszStableContainerSchema.NamedIndexedSchema;
 import tech.pegasys.teku.infrastructure.unsigned.UInt64;
@@ -76,6 +77,13 @@ public <T extends SszData> T randomData(final SszSchema<T> schema) {
     return randomDataStream(schema).findFirst().orElseThrow();
   }
 
+  public <T extends SszData> Optional<T> randomOptionalData(final SszSchema<T> schema) {
+    if (random.nextBoolean()) {
+      return Optional.of(randomDataStream(schema).findFirst().orElseThrow());
+    }
+    return Optional.empty();
+  }
+
   @SuppressWarnings("unchecked")
   public <T extends SszData> Stream<T> randomDataStream(final SszSchema<T> schema) {
     if (schema instanceof AbstractSszPrimitiveSchema) {
@@ -104,11 +112,22 @@ public <T extends SszData> Stream<T> randomDataStream(final SszSchema<T> schema)
           containerSchema.toStableContainerSchema().orElseThrow().getDefinedChildrenSchemas();
       return Stream.generate(
           () -> {
-            List<SszData> children =
+            List<Optional<? extends SszData>> children =
                 definedFieldSchemas.stream()
                     .map(
                         definedFieldSchema ->
                             containerSchema.getChildSchema(definedFieldSchema.getIndex()))
+                    .map(this::randomOptionalData)
+                    .collect(Collectors.toList());
+            return (T) containerSchema.createFromOptionalFieldValues(children);
+          });
+    } else if (schema instanceof AbstractSszProfileSchema<?>) {
+      AbstractSszProfileSchema<SszProfile> containerSchema =
+          (AbstractSszProfileSchema<SszProfile>) schema;
+      return Stream.generate(
+          () -> {
+            List<SszData> children =
+                containerSchema.getActiveChildrenSchemas().stream()
                     .map(this::randomData)
                     .collect(Collectors.toList());
             return (T) containerSchema.createFromFieldValues(children);