Improve attestation bits aggregator electra

ethereum/statetransition/src/test/java/tech/pegasys/teku/statetransition/attestation/utils/AttestationBitsAggregatorElectraTest.java

@@ -232,6 +232,34 @@ void aggregateOnMultipleOverlappingCommitteeBitsButWithSomeOfAggregationOverlapp
     assertThat(aggregator.getAggregationBits().streamAllSetBits()).containsExactly(0, 1, 2, 3, 8);
   }
 
+  @Test
+  void aggregateSingleAttestationFillUp() {
+    /*
+     01|234 <- committee 0 and 1 indices
+     10|100 <- bits
+    */
+    final ValidatableAttestation initialAttestation = createAttestation(List.of(0, 1), 0, 2);
+
+    /*
+     123 <- committee 1 indices
+     001 <- bits
+    */
+    final ValidatableAttestation otherAttestation = createAttestation(List.of(1), 2);
+
+    final AttestationBitsAggregator aggregator = AttestationBitsAggregator.of(initialAttestation);
+
+    // calculate the or
+    aggregator.or(otherAttestation.getAttestation());
+
+    /*
+     01|234 <- committee 0 and 1 indices
+     10|101 <- bits
+    */
+
+    assertThat(aggregator.getCommitteeBits().streamAllSetBits()).containsExactly(0, 1);
+    assertThat(aggregator.getAggregationBits().streamAllSetBits()).containsExactly(0, 2, 4);
+  }
+
   @Test
   void
       aggregateOnMultipleOverlappingCommitteeBitsButWithSomeOfAggregationOverlappingWhenNoCheck2() {
@@ -393,6 +421,107 @@ void bigAggregation() {
         .isEqualTo(result.getAttestation().getAggregationBits());
   }
 
+  @Test
+  void orIntoEmptyAggregator() {
+    final AttestationBitsAggregator aggregator =
+        AttestationBitsAggregator.fromEmptyFromAttestationSchema(
+            attestationSchema, Optional.of(committeeSizes));
+
+    final ValidatableAttestation otherAttestation = createAttestation(List.of(1), 2);
+
+    aggregator.or(otherAttestation.getAttestation());
+
+    assertThat(aggregator.getCommitteeBits().streamAllSetBits()).containsExactly(1);
+    assertThat(aggregator.getAggregationBits().streamAllSetBits()).containsExactly(2);
+  }
+
+  @Test
+  void orWithNewCommittee() {
+    final ValidatableAttestation initialAttestation = createAttestation(List.of(0), 1);
+    final AttestationBitsAggregator aggregator = AttestationBitsAggregator.of(initialAttestation);
+
+    final ValidatableAttestation otherAttestation = createAttestation(List.of(1), 2);
+
+    aggregator.or(otherAttestation.getAttestation());
+
+    assertThat(aggregator.getCommitteeBits().streamAllSetBits()).containsExactlyInAnyOrder(0, 1);
+    assertThat(aggregator.getAggregationBits().streamAllSetBits()).containsExactlyInAnyOrder(1, 4);
+  }
+
+  @Test
+  void orWithExistingCommitteeAddNewBits() {
+
+    final ValidatableAttestation initialAttestation = createAttestation(List.of(1), 0);
+    final AttestationBitsAggregator aggregator = AttestationBitsAggregator.of(initialAttestation);
+
+    final ValidatableAttestation otherAttestation = createAttestation(List.of(1), 2);
+
+    aggregator.or(otherAttestation.getAttestation());
+
+    assertThat(aggregator.getCommitteeBits().streamAllSetBits()).containsExactly(1);
+    assertThat(aggregator.getAggregationBits().streamAllSetBits()).containsExactlyInAnyOrder(0, 2);
+  }
+
+  @Test
+  void orWithExistingCommitteeOverlapAndNewBits() {
+    final ValidatableAttestation initialAttestation = createAttestation(List.of(1), 0, 1);
+    final AttestationBitsAggregator aggregator = AttestationBitsAggregator.of(initialAttestation);
+
+    final ValidatableAttestation otherAttestation = createAttestation(List.of(1), 1, 2);
+
+    aggregator.or(otherAttestation.getAttestation());
+
+    assertThat(aggregator.getCommitteeBits().streamAllSetBits()).containsExactly(1);
+    assertThat(aggregator.getAggregationBits().streamAllSetBits())
+        .containsExactlyInAnyOrder(0, 1, 2);
+  }
+
+  @Test
+  void orWithStrictSubsetAttestation() {
+    final ValidatableAttestation initialAttestation = createAttestation(List.of(1), 0, 2);
+    final AttestationBitsAggregator aggregator = AttestationBitsAggregator.of(initialAttestation);
+
+    final ValidatableAttestation otherAttestation = createAttestation(List.of(1), 0);
+
+    aggregator.or(otherAttestation.getAttestation());
+
+    assertThat(aggregator.getCommitteeBits().streamAllSetBits()).containsExactly(1);
+    assertThat(aggregator.getAggregationBits().streamAllSetBits()).containsExactlyInAnyOrder(0, 2);
+  }
+
+  @Test
+  void orWithMultipleCommitteesMixedNewAndExisting() {
+    final ValidatableAttestation initialAttestation = createAttestation(List.of(0, 1), 0, 3);
+    final AttestationBitsAggregator aggregator = AttestationBitsAggregator.of(initialAttestation);
+
+    final ValidatableAttestation otherAttestation = createAttestation(List.of(1, 2), 2, 5);
+
+    aggregator.or(otherAttestation.getAttestation());
+
+    assertThat(aggregator.getCommitteeBits().streamAllSetBits()).containsExactlyInAnyOrder(0, 1, 2);
+    assertThat(aggregator.getAggregationBits().streamAllSetBits())
+        .containsExactlyInAnyOrder(0, 3, 4, 7);
+  }
+
+  @Test
+  void orAggregatorWithAggregator() {
+    final ValidatableAttestation att1Data = createAttestation(List.of(0), 0);
+    final AttestationBitsAggregator aggregator1 = AttestationBitsAggregator.of(att1Data);
+
+    final ValidatableAttestation att2Data = createAttestation(List.of(0, 1), 1, 2);
+    final AttestationBitsAggregator aggregator2 = AttestationBitsAggregator.of(att2Data);
+
+    aggregator1.or(aggregator2);
+
+    assertThat(aggregator1.getCommitteeBits().streamAllSetBits()).containsExactlyInAnyOrder(0, 1);
+    assertThat(aggregator1.getAggregationBits().streamAllSetBits())
+        .containsExactlyInAnyOrder(0, 1, 2);
+
+    // aggregator2 should remain unchanged
+    assertThat(aggregator2.getCommitteeBits().streamAllSetBits()).containsExactlyInAnyOrder(0, 1);
+    assertThat(aggregator2.getAggregationBits().streamAllSetBits()).containsExactlyInAnyOrder(1, 2);
+  }
+
   @Test
   void isSuperSetOf1() {
 

infrastructure/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/collections/SszBitlist.java

@@ -14,6 +14,7 @@
 package tech.pegasys.teku.infrastructure.ssz.collections;
 
 import it.unimi.dsi.fastutil.ints.IntList;
+import java.util.BitSet;
 import java.util.stream.IntStream;
 import javax.annotation.Nullable;
 import tech.pegasys.teku.infrastructure.ssz.collections.impl.SszBitlistImpl;
@@ -43,6 +44,12 @@ default boolean isWritableSupported() {
 
   // Bitlist methods
 
+  BitSet getAsBitSet();
+
+  BitSet getAsBitSet(int start, int end);
+
+  int getLastSetBitIndex();
+
   /**
    * Performs a logical OR of this bit list with the bit list argument.
    *

infrastructure/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/collections/SszBitvector.java

@@ -14,6 +14,7 @@
 package tech.pegasys.teku.infrastructure.ssz.collections;
 
 import it.unimi.dsi.fastutil.ints.IntList;
+import java.util.BitSet;
 import java.util.stream.IntStream;
 import tech.pegasys.teku.infrastructure.ssz.primitive.SszBit;
 import tech.pegasys.teku.infrastructure.ssz.schema.collections.SszBitvectorSchema;
@@ -48,6 +49,8 @@ default boolean isWritableSupported() {
   /** Returns the number of bits set to {@code true} in this {@code SszBitlist}. */
   int getBitCount();
 
+  BitSet getAsBitSet();
+
   @Override
   default boolean isSet(final int i) {
     return i < size() && getBit(i);

infrastructure/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/collections/impl/BitlistImpl.java

@@ -51,12 +51,10 @@ public BitlistImpl(final int size, final long maxSize, final int... bitIndices)
     }
   }
 
-  public BitlistImpl(final int size, final long maxSize, final BitSet bitSet) {
+  public static BitlistImpl wrapBitSet(final int size, final long maxSize, final BitSet bitSet) {
     checkArgument(size >= 0, "Negative size");
     checkArgument(maxSize >= size, "maxSize should be >= size");
-    this.size = size;
-    this.data = bitSet;
-    this.maxSize = maxSize;
+    return new BitlistImpl(size, bitSet, maxSize);
   }
 
   private BitlistImpl(final int size, final BitSet data, final long maxSize) {
@@ -65,6 +63,18 @@ private BitlistImpl(final int size, final BitSet data, final long maxSize) {
     this.maxSize = maxSize;
   }
 
+  public BitSet getAsBitSet() {
+    return (BitSet) data.clone();
+  }
+
+  public BitSet getAsBitSet(final int start, final int end) {
+    return data.get(start, end);
+  }
+
+  public int getLastSetBitIndex() {
+    return data.length() - 1;
+  }
+
   /**
    * Returns new instance of this BitlistImpl with set bits from the other BitlistImpl
    *

infrastructure/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/collections/impl/BitvectorImpl.java

@@ -43,6 +43,12 @@ public static BitvectorImpl fromBytes(final Bytes bytes, final int size) {
     return new BitvectorImpl(bitset, size);
   }
 
+  public static BitvectorImpl wrapBitSet(final BitSet bitSet, final int size) {
+    final int length = bitSet.length();
+    checkArgument(length <= size, "BitSet length (%s) exceeds the size (%s)", length, size);
+    return new BitvectorImpl(bitSet, size);
+  }
+
   public static int sszSerializationLength(final int size) {
     return bitsCeilToBytes(size);
   }
@@ -112,6 +118,10 @@ public boolean getBit(final int i) {
     return data.get(i);
   }
 
+  public BitSet getAsBitSet() {
+    return (BitSet) data.clone();
+  }
+
   public int getSize() {
     return size;
   }

infrastructure/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/collections/impl/SszBitlistImpl.java

@@ -75,7 +75,7 @@ public static SszBitlistImpl ofBits(
 
   public static SszBitlistImpl wrapBitSet(
       final SszBitlistSchema<?> schema, final int size, final BitSet bitSet) {
-    return new SszBitlistImpl(schema, new BitlistImpl(size, schema.getMaxLength(), bitSet));
+    return new SszBitlistImpl(schema, BitlistImpl.wrapBitSet(size, schema.getMaxLength(), bitSet));
   }
 
   private final BitlistImpl value;
@@ -90,6 +90,21 @@ public SszBitlistImpl(final SszListSchema<SszBit, ?> schema, final BitlistImpl v
     this.value = value;
   }
 
+  @Override
+  public BitSet getAsBitSet() {
+    return value.getAsBitSet();
+  }
+
+  @Override
+  public BitSet getAsBitSet(final int start, final int end) {
+    return value.getAsBitSet(start, end);
+  }
+
+  @Override
+  public int getLastSetBitIndex() {
+    return value.getLastSetBitIndex();
+  }
+
   @SuppressWarnings("unchecked")
   @Override
   public SszBitlistSchema<SszBitlist> getSchema() {

infrastructure/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/collections/impl/SszBitvectorImpl.java

@@ -17,6 +17,7 @@
 
 import it.unimi.dsi.fastutil.ints.IntArrayList;
 import it.unimi.dsi.fastutil.ints.IntList;
+import java.util.BitSet;
 import java.util.stream.IntStream;
 import org.apache.tuweni.bytes.Bytes;
 import tech.pegasys.teku.infrastructure.ssz.cache.IntCache;
@@ -35,6 +36,11 @@ public static SszBitvectorImpl ofBits(final SszBitvectorSchema<?> schema, final
     return new SszBitvectorImpl(schema, new BitvectorImpl(schema.getLength(), bits));
   }
 
+  public static SszBitvectorImpl wrapBitSet(
+      final SszBitvectorSchema<?> schema, final int size, final BitSet bitSet) {
+    return new SszBitvectorImpl(schema, BitvectorImpl.wrapBitSet(bitSet, size));
+  }
+
   public static SszBitvector fromBytes(
       final SszBitvectorSchema<?> schema, final Bytes value, final int size) {
     return new SszBitvectorImpl(schema, BitvectorImpl.fromBytes(value, size));
@@ -58,6 +64,11 @@ public SszBitvectorImpl(final SszBitvectorSchema<?> schema, final BitvectorImpl
     this.value = value;
   }
 
+  @Override
+  public BitSet getAsBitSet() {
+    return value.getAsBitSet();
+  }
+
   @SuppressWarnings("unchecked")
   @Override
   public SszBitvectorSchema<SszBitvector> getSchema() {

infrastructure/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/schema/collections/SszBitlistSchema.java

@@ -33,13 +33,9 @@ default SszBitlistT empty() {
   SszBitlistT ofBits(int size, int... setBitIndices);
 
   /**
-   * Creates a SszBitlist by wrapping a given bitSet. This is an optimized constructor that DOES NOT
-   * clone the bitSet. It used in aggregating attestation pool. DO NOT MUTATE the after the
+   * Creates an SszBitlist by wrapping a given bitSet. This is an optimized constructor that DOES
+   * NOT clone the bitSet. It is used in aggregating attestation pool. DO NOT MUTATE after the
    * wrapping!! SszBitlist is supposed to be immutable.
-   *
-   * @param size size of the SszBitlist
-   * @param bitSet data backing the ssz
-   * @return SszBitlist instance
    */
   SszBitlistT wrapBitSet(int size, BitSet bitSet);
 

infrastructure/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/schema/collections/SszBitvectorSchema.java

@@ -13,6 +13,7 @@
 
 package tech.pegasys.teku.infrastructure.ssz.schema.collections;
 
+import java.util.BitSet;
 import java.util.stream.StreamSupport;
 import org.apache.tuweni.bytes.Bytes;
 import tech.pegasys.teku.infrastructure.ssz.collections.SszBitvector;
@@ -28,6 +29,13 @@ static SszBitvectorSchema<SszBitvector> create(final long length) {
 
   SszBitvectorT ofBits(int... setBitIndices);
 
+  /**
+   * Creates an SszBitvector by wrapping a given bitSet. This is an optimized constructor that DOES
+   * NOT clone the bitSet. It is used in aggregating attestation pool. DO NOT MUTATE after the
+   * wrapping!! SszBitvector is supposed to be immutable.
+   */
+  SszBitvectorT wrapBitSet(int size, BitSet bitSet);
+
   default SszBitvectorT fromBytes(final Bytes bitvectorBytes) {
     return sszDeserialize(bitvectorBytes);
   }

infrastructure/ssz/src/main/java/tech/pegasys/teku/infrastructure/ssz/schema/collections/impl/SszBitvectorSchemaImpl.java

@@ -16,6 +16,7 @@
 import static com.google.common.base.Preconditions.checkArgument;
 import static tech.pegasys.teku.infrastructure.ssz.schema.json.SszPrimitiveTypeDefinitions.sszSerializedType;
 
+import java.util.BitSet;
 import java.util.List;
 import java.util.stream.IntStream;
 import tech.pegasys.teku.infrastructure.json.types.DeserializableTypeDefinition;
@@ -54,6 +55,11 @@ public SszBitvector ofBits(final int... setBitIndices) {
     return SszBitvectorImpl.ofBits(this, setBitIndices);
   }
 
+  @Override
+  public SszBitvector wrapBitSet(final int size, final BitSet bitSet) {
+    return SszBitvectorImpl.wrapBitSet(this, size, bitSet);
+  }
+
   @Override
   public SszBitvector createFromElements(final List<? extends SszBit> elements) {
     return ofBits(IntStream.range(0, elements.size()).filter(i -> elements.get(i).get()).toArray());