Skip to content

Commit

Permalink
code consolidation in util
Browse files Browse the repository at this point in the history
  • Loading branch information
@tbenr
tbenr committed Jul 16, 2024
1 parent ff3b40d commit 839dff5
Show file tree
Hide file tree
Showing 3 changed files with 176 additions and 140 deletions.
80 changes: 17 additions & 63 deletions ...ain/java/tech/pegasys/teku/infrastructure/ssz/schema/impl/AbstractSszContainerSchema.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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;
}
Expand All @@ -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<TreeNode> 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<TreeNode> fixedChildrenSubtrees = new ArrayDeque<>(childCount);
final IntList variableChildrenOffsets = new IntArrayList(childCount);

ArrayDeque<Integer> 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<Integer> variableChildrenSizes =
validateAndPrepareForVariableChildrenDeserialization(
reader, variableChildrenOffsets, endOffset);

List<TreeNode> 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);
Expand Down
106 changes: 29 additions & 77 deletions ...ech/pegasys/teku/infrastructure/ssz/schema/impl/AbstractSszStableContainerBaseSchema.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -458,57 +445,22 @@ 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<Integer> 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<TreeNode> childrenSubtrees = new ArrayList<>(childCount);
for (int i = 0; i < childCount; i++) {
if (!activeFields.getBit(i)) {
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);
Expand Down
Loading

0 comments on commit 839dff5

Please sign in to comment.