diff --git a/infrastructure/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/schema/impl/AbstractSszContainerSchema.java b/infrastructure/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/schema/impl/AbstractSszContainerSchema.java index f092d157a4d..41e2dc7545e 100644 --- a/infrastructure/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/schema/impl/AbstractSszContainerSchema.java +++ b/infrastructure/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/schema/impl/AbstractSszContainerSchema.java @@ -14,6 +14,11 @@ package tech.pegasys.teku.infrastructure.ssz.schema.impl; import static com.google.common.base.Preconditions.checkArgument; +import static tech.pegasys.teku.infrastructure.ssz.schema.impl.ContainerSchemaUtil.deserializeFixedChild; +import static tech.pegasys.teku.infrastructure.ssz.schema.impl.ContainerSchemaUtil.deserializeVariableChild; +import static tech.pegasys.teku.infrastructure.ssz.schema.impl.ContainerSchemaUtil.serializeFixedChild; +import static tech.pegasys.teku.infrastructure.ssz.schema.impl.ContainerSchemaUtil.serializeVariableChild; +import static tech.pegasys.teku.infrastructure.ssz.schema.impl.ContainerSchemaUtil.validateAndPrepareForVariableChildrenDeserialization; import com.google.common.base.Suppliers; import it.unimi.dsi.fastutil.ints.IntArrayList; @@ -34,9 +39,7 @@ import tech.pegasys.teku.infrastructure.ssz.schema.SszContainerSchema; import tech.pegasys.teku.infrastructure.ssz.schema.SszFieldName; import tech.pegasys.teku.infrastructure.ssz.schema.SszSchema; -import tech.pegasys.teku.infrastructure.ssz.schema.SszType; import tech.pegasys.teku.infrastructure.ssz.schema.json.SszContainerTypeDefinition; -import tech.pegasys.teku.infrastructure.ssz.sos.SszDeserializeException; import tech.pegasys.teku.infrastructure.ssz.sos.SszLengthBounds; import tech.pegasys.teku.infrastructure.ssz.sos.SszReader; import tech.pegasys.teku.infrastructure.ssz.sos.SszWriter; @@ -277,25 +280,11 @@ public int sszSerializeTree(final TreeNode node, final SszWriter writer) { int variableChildOffset = getSszFixedPartSize(); int[] variableSizes = new int[getFieldsCount()]; for (int i = 0; i < getFieldsCount(); i++) { - TreeNode childSubtree = node.get(getChildGeneralizedIndex(i)); - SszSchema childType = getChildSchema(i); - if (childType.isFixedSize()) { - int size = childType.sszSerializeTree(childSubtree, writer); - assert size == childType.getSszFixedPartSize(); - } else { - writer.write(SszType.sszLengthToBytes(variableChildOffset)); - int childSize = childType.getSszSize(childSubtree); - variableSizes[i] = childSize; - variableChildOffset += childSize; - } + variableChildOffset += + serializeFixedChild(writer, this, i, node, variableSizes, variableChildOffset); } for (int i = 0; i < childrenSchemas.size(); i++) { - SszSchema childType = getChildSchema(i); - if (!childType.isFixedSize()) { - TreeNode childSubtree = node.get(getChildGeneralizedIndex(i)); - int size = childType.sszSerializeTree(childSubtree, writer); - assert size == variableSizes[i]; - } + serializeVariableChild(writer, this, i, variableSizes, node); } return variableChildOffset; } @@ -304,56 +293,21 @@ public int sszSerializeTree(final TreeNode node, final SszWriter writer) { public TreeNode sszDeserializeTree(final SszReader reader) { int endOffset = reader.getAvailableBytes(); int childCount = getFieldsCount(); - Queue fixedChildrenSubtrees = new ArrayDeque<>(childCount); - IntList variableChildrenOffsets = new IntArrayList(childCount); - for (int i = 0; i < childCount; i++) { - SszSchema childType = getChildSchema(i); - if (childType.isFixedSize()) { - try (SszReader sszReader = reader.slice(childType.getSszFixedPartSize())) { - TreeNode childNode = childType.sszDeserializeTree(sszReader); - fixedChildrenSubtrees.add(childNode); - } - } else { - int childOffset = SszType.sszBytesToLength(reader.read(SSZ_LENGTH_SIZE)); - variableChildrenOffsets.add(childOffset); - } - } - - if (variableChildrenOffsets.isEmpty()) { - if (reader.getAvailableBytes() > 0) { - throw new SszDeserializeException("Invalid SSZ: unread bytes for fixed size container"); - } - } else { - if (variableChildrenOffsets.getInt(0) != endOffset - reader.getAvailableBytes()) { - throw new SszDeserializeException( - "First variable element offset doesn't match the end of fixed part"); - } - } - - variableChildrenOffsets.add(endOffset); + final Queue fixedChildrenSubtrees = new ArrayDeque<>(childCount); + final IntList variableChildrenOffsets = new IntArrayList(childCount); - ArrayDeque variableChildrenSizes = - new ArrayDeque<>(variableChildrenOffsets.size() - 1); - for (int i = 0; i < variableChildrenOffsets.size() - 1; i++) { - variableChildrenSizes.add( - variableChildrenOffsets.getInt(i + 1) - variableChildrenOffsets.getInt(i)); + for (int i = 0; i < childCount; i++) { + deserializeFixedChild(reader, fixedChildrenSubtrees, variableChildrenOffsets, this, i); } - if (variableChildrenSizes.stream().anyMatch(s -> s < 0)) { - throw new SszDeserializeException("Invalid SSZ: wrong child offsets"); - } + final ArrayDeque variableChildrenSizes = + validateAndPrepareForVariableChildrenDeserialization( + reader, variableChildrenOffsets, endOffset); List childrenSubtrees = new ArrayList<>(childCount); for (int i = 0; i < childCount; i++) { - SszSchema childType = getChildSchema(i); - if (childType.isFixedSize()) { - childrenSubtrees.add(fixedChildrenSubtrees.remove()); - } else { - try (SszReader sszReader = reader.slice(variableChildrenSizes.remove())) { - TreeNode childNode = childType.sszDeserializeTree(sszReader); - childrenSubtrees.add(childNode); - } - } + deserializeVariableChild( + reader, childrenSubtrees, fixedChildrenSubtrees, variableChildrenSizes, this, i); } return TreeUtil.createTree(childrenSubtrees); diff --git a/infrastructure/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/schema/impl/AbstractSszStableContainerBaseSchema.java b/infrastructure/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/schema/impl/AbstractSszStableContainerBaseSchema.java index cecb7541715..825374e2717 100644 --- a/infrastructure/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/schema/impl/AbstractSszStableContainerBaseSchema.java +++ b/infrastructure/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/schema/impl/AbstractSszStableContainerBaseSchema.java @@ -15,6 +15,11 @@ 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.ContainerSchemaUtil.deserializeFixedChild; +import static tech.pegasys.teku.infrastructure.ssz.schema.impl.ContainerSchemaUtil.deserializeVariableChild; +import static tech.pegasys.teku.infrastructure.ssz.schema.impl.ContainerSchemaUtil.serializeFixedChild; +import static tech.pegasys.teku.infrastructure.ssz.schema.impl.ContainerSchemaUtil.serializeVariableChild; +import static tech.pegasys.teku.infrastructure.ssz.schema.impl.ContainerSchemaUtil.validateAndPrepareForVariableChildrenDeserialization; import com.google.common.base.Suppliers; import it.unimi.dsi.fastutil.ints.IntArrayList; @@ -26,7 +31,6 @@ import java.util.List; import java.util.Objects; import java.util.Optional; -import java.util.PrimitiveIterator.OfInt; import java.util.Queue; import java.util.Set; import java.util.function.Supplier; @@ -43,7 +47,6 @@ import tech.pegasys.teku.infrastructure.ssz.schema.collections.SszBitvectorSchema; import tech.pegasys.teku.infrastructure.ssz.schema.impl.AbstractSszContainerSchema.NamedSchema; import tech.pegasys.teku.infrastructure.ssz.schema.json.SszStableContainerBaseTypeDefinition; -import tech.pegasys.teku.infrastructure.ssz.sos.SszDeserializeException; import tech.pegasys.teku.infrastructure.ssz.sos.SszLengthBounds; import tech.pegasys.teku.infrastructure.ssz.sos.SszReader; import tech.pegasys.teku.infrastructure.ssz.sos.SszWriter; @@ -385,44 +388,28 @@ public int sszSerializeTree(final TreeNode node, final SszWriter writer) { // we won't write active field when no optional fields are permitted final int activeFieldsWroteBytes = sszSerializeActiveFields(activeFieldsBitvector, writer); - final TreeNode containerTree = node.get(CONTAINER_G_INDEX); - - int variableChildOffset = calcSszFixedPartSize(activeFieldsBitvector); - final int[] variableSizes = new int[activeFieldsBitvector.size()]; - for (final OfInt activeIndicesIterator = activeFieldsBitvector.streamAllSetBits().iterator(); - activeIndicesIterator.hasNext(); ) { - final int activeFieldIndex = activeIndicesIterator.next(); - - final TreeNode childSubtree = - containerTree.get( - SszStableContainerBaseSchema.super.getChildGeneralizedIndex(activeFieldIndex)); - final SszSchema childType = getChildSchema(activeFieldIndex); - if (childType.isFixedSize()) { - final int size = childType.sszSerializeTree(childSubtree, writer); - assert size == childType.getSszFixedPartSize(); - } else { - writer.write(SszType.sszLengthToBytes(variableChildOffset)); - final int childSize = childType.getSszSize(childSubtree); - variableSizes[activeFieldIndex] = childSize; - variableChildOffset += childSize; - } - } + + final int variableChildOffset = + activeFieldsBitvector + .streamAllSetBits() + .reduce( + calcSszFixedPartSize(activeFieldsBitvector), + (accumulatedOffset, activeFieldIndex) -> + accumulatedOffset + + serializeFixedChild( + writer, + this, + activeFieldIndex, + node, + variableSizes, + accumulatedOffset)); activeFieldsBitvector .streamAllSetBits() .forEach( - activeFieldIndex -> { - final SszSchema childType = getChildSchema(activeFieldIndex); - if (!childType.isFixedSize()) { - final TreeNode childSubtree = - containerTree.get( - SszStableContainerBaseSchema.super.getChildGeneralizedIndex( - activeFieldIndex)); - final int size = childType.sszSerializeTree(childSubtree, writer); - assert size == variableSizes[activeFieldIndex]; - } - }); + activeFieldIndex -> + serializeVariableChild(writer, this, activeFieldIndex, variableSizes, node)); return activeFieldsWroteBytes + variableChildOffset; } @@ -458,42 +445,13 @@ private TreeNode deserializeContainer(final SszReader reader, final SszBitvector activeFields .streamAllSetBits() .forEach( - i -> { - final SszSchema childType = getChildSchema(i); - if (childType.isFixedSize()) { - try (SszReader sszReader = reader.slice(childType.getSszFixedPartSize())) { - TreeNode childNode = childType.sszDeserializeTree(sszReader); - fixedChildrenSubtrees.add(childNode); - } - } else { - int childOffset = SszType.sszBytesToLength(reader.read(SSZ_LENGTH_SIZE)); - variableChildrenOffsets.add(childOffset); - } - }); - - if (variableChildrenOffsets.isEmpty()) { - if (reader.getAvailableBytes() > 0) { - throw new SszDeserializeException("Invalid SSZ: unread bytes for fixed size container"); - } - } else { - if (variableChildrenOffsets.getInt(0) != endOffset - reader.getAvailableBytes()) { - throw new SszDeserializeException( - "First variable element offset doesn't match the end of fixed part"); - } - } - - variableChildrenOffsets.add(endOffset); + i -> + deserializeFixedChild( + reader, fixedChildrenSubtrees, variableChildrenOffsets, this, i)); final ArrayDeque variableChildrenSizes = - new ArrayDeque<>(variableChildrenOffsets.size() - 1); - for (int i = 0; i < variableChildrenOffsets.size() - 1; i++) { - variableChildrenSizes.add( - variableChildrenOffsets.getInt(i + 1) - variableChildrenOffsets.getInt(i)); - } - - if (variableChildrenSizes.stream().anyMatch(s -> s < 0)) { - throw new SszDeserializeException("Invalid SSZ: wrong child offsets"); - } + validateAndPrepareForVariableChildrenDeserialization( + reader, variableChildrenOffsets, endOffset); final List childrenSubtrees = new ArrayList<>(childCount); for (int i = 0; i < childCount; i++) { @@ -501,14 +459,8 @@ private TreeNode deserializeContainer(final SszReader reader, final SszBitvector childrenSubtrees.add(SszPrimitiveSchemas.NONE_SCHEMA.getDefaultTree()); continue; } - final SszSchema childType = getChildSchema(i); - if (childType.isFixedSize()) { - childrenSubtrees.add(fixedChildrenSubtrees.remove()); - continue; - } - try (SszReader sszReader = reader.slice(variableChildrenSizes.remove())) { - childrenSubtrees.add(childType.sszDeserializeTree(sszReader)); - } + deserializeVariableChild( + reader, childrenSubtrees, fixedChildrenSubtrees, variableChildrenSizes, this, i); } return TreeUtil.createTree(childrenSubtrees); diff --git a/infrastructure/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/schema/impl/ContainerSchemaUtil.java b/infrastructure/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/schema/impl/ContainerSchemaUtil.java new file mode 100644 index 00000000000..61559db0c69 --- /dev/null +++ b/infrastructure/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/schema/impl/ContainerSchemaUtil.java @@ -0,0 +1,130 @@ +/* + * Copyright Consensys Software Inc., 2024 + * + * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on + * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the + * specific language governing permissions and limitations under the License. + */ + +package tech.pegasys.teku.infrastructure.ssz.schema.impl; + +import static tech.pegasys.teku.infrastructure.ssz.schema.SszType.SSZ_LENGTH_SIZE; + +import it.unimi.dsi.fastutil.ints.IntList; +import java.util.ArrayDeque; +import java.util.List; +import java.util.Queue; +import tech.pegasys.teku.infrastructure.ssz.schema.SszCompositeSchema; +import tech.pegasys.teku.infrastructure.ssz.schema.SszSchema; +import tech.pegasys.teku.infrastructure.ssz.schema.SszType; +import tech.pegasys.teku.infrastructure.ssz.sos.SszDeserializeException; +import tech.pegasys.teku.infrastructure.ssz.sos.SszReader; +import tech.pegasys.teku.infrastructure.ssz.sos.SszWriter; +import tech.pegasys.teku.infrastructure.ssz.tree.TreeNode; + +public class ContainerSchemaUtil { + static void deserializeFixedChild( + final SszReader reader, + final Queue fixedChildrenSubtrees, + final IntList variableChildrenOffsets, + final SszCompositeSchema containerSchema, + final int index) { + final SszSchema childType = containerSchema.getChildSchema(index); + if (childType.isFixedSize()) { + try (SszReader sszReader = reader.slice(childType.getSszFixedPartSize())) { + final TreeNode childNode = childType.sszDeserializeTree(sszReader); + fixedChildrenSubtrees.add(childNode); + } + } else { + final int childOffset = SszType.sszBytesToLength(reader.read(SSZ_LENGTH_SIZE)); + variableChildrenOffsets.add(childOffset); + } + } + + static ArrayDeque validateAndPrepareForVariableChildrenDeserialization( + final SszReader reader, final IntList variableChildrenOffsets, final int endOffset) { + + if (variableChildrenOffsets.isEmpty()) { + if (reader.getAvailableBytes() > 0) { + throw new SszDeserializeException("Invalid SSZ: unread bytes for fixed size container"); + } + } else { + if (variableChildrenOffsets.getInt(0) != endOffset - reader.getAvailableBytes()) { + throw new SszDeserializeException( + "First variable element offset doesn't match the end of fixed part"); + } + } + + variableChildrenOffsets.add(endOffset); + + ArrayDeque variableChildrenSizes = + new ArrayDeque<>(variableChildrenOffsets.size() - 1); + for (int i = 0; i < variableChildrenOffsets.size() - 1; i++) { + variableChildrenSizes.add( + variableChildrenOffsets.getInt(i + 1) - variableChildrenOffsets.getInt(i)); + } + + if (variableChildrenSizes.stream().anyMatch(s -> s < 0)) { + throw new SszDeserializeException("Invalid SSZ: wrong child offsets"); + } + + return variableChildrenSizes; + } + + static void deserializeVariableChild( + final SszReader reader, + final List childrenSubtrees, + final Queue fixedChildrenSubtrees, + final ArrayDeque variableChildrenSizes, + final SszCompositeSchema containerSchema, + final int index) { + final SszSchema childType = containerSchema.getChildSchema(index); + if (childType.isFixedSize()) { + childrenSubtrees.add(fixedChildrenSubtrees.remove()); + } else { + try (SszReader sszReader = reader.slice(variableChildrenSizes.remove())) { + TreeNode childNode = childType.sszDeserializeTree(sszReader); + childrenSubtrees.add(childNode); + } + } + } + + static int serializeFixedChild( + final SszWriter writer, + final SszCompositeSchema containerSchema, + final int index, + final TreeNode node, + final int[] variableSizes, + final int variableChildOffset) { + final TreeNode childSubtree = node.get(containerSchema.getChildGeneralizedIndex(index)); + final SszSchema childType = containerSchema.getChildSchema(index); + if (childType.isFixedSize()) { + int size = childType.sszSerializeTree(childSubtree, writer); + assert size == childType.getSszFixedPartSize(); + return 0; + } + writer.write(SszType.sszLengthToBytes(variableChildOffset)); + int childSize = childType.getSszSize(childSubtree); + variableSizes[index] = childSize; + return childSize; + } + + static void serializeVariableChild( + final SszWriter writer, + final SszCompositeSchema containerSchema, + final int index, + final int[] variableSizes, + final TreeNode node) { + SszSchema childType = containerSchema.getChildSchema(index); + if (!childType.isFixedSize()) { + TreeNode childSubtree = node.get(containerSchema.getChildGeneralizedIndex(index)); + int size = childType.sszSerializeTree(childSubtree, writer); + assert size == variableSizes[index]; + } + } +}