Update VC and BN APIs for naive aggregation (#950)

* Refactor `Attestation` production

* Add constant

* Start refactor for aggregation

* Return early when no attesting validators

* Refactor into individual functions

* Tidy, add comments

* Add first draft of NaiveAggregationPool

* Further progress on naive aggregation pool

* Fix compile errors in VC

* Change locking logic for naive pool

* Introduce AttesationType

* Add pruning, comments

* Add MAX_ATTESTATIONS_PER_SLOT restriction

* Add pruning based on slot

* Update BN for new aggregation fns

* Fix test compile errors

* Fix failing rest_api test

* Move SignedAggregateAndProof into own file

* Update docs, fix warning

* Tidy some formatting in validator API

* Remove T::default_spec from signing

* Fix failing rest test

* Tidy

* Add test, fix bug

* Improve naive pool tests

* Add max attestations test

* Revert changes to the op_pool

* Refactor timer

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -8,6 +8,7 @@ use crate::events::{EventHandler, EventKind};
 use crate::fork_choice::{Error as ForkChoiceError, ForkChoice};
 use crate::head_tracker::HeadTracker;
 use crate::metrics;
+use crate::naive_aggregation_pool::{Error as NaiveAggregationError, NaiveAggregationPool};
 use crate::persisted_beacon_chain::PersistedBeaconChain;
 use crate::shuffling_cache::ShufflingCache;
 use crate::snapshot_cache::SnapshotCache;
@@ -62,6 +63,23 @@ pub const OP_POOL_DB_KEY: [u8; 32] = [0; 32];
 pub const ETH1_CACHE_DB_KEY: [u8; 32] = [0; 32];
 pub const FORK_CHOICE_DB_KEY: [u8; 32] = [0; 32];
 
+#[derive(Debug, PartialEq)]
+pub enum AttestationType {
+    /// An attestation with a single-signature that has been published in accordance with the naive
+    /// aggregation strategy.
+    ///
+    /// These attestations may have come from a `committee_index{subnet_id}_beacon_attestation`
+    /// gossip subnet or they have have come directly from a validator attached to our API.
+    ///
+    /// If `should_store == true`, the attestation will be added to the `NaiveAggregationPool`.
+    Unaggregated { should_store: bool },
+    /// An attestation with one more more signatures that has passed through the aggregation phase
+    /// of the naive aggregation scheme.
+    ///
+    /// These attestations must have come from the `beacon_aggregate_and_proof` gossip subnet.
+    Aggregated,
+}
+
 #[derive(Debug, PartialEq)]
 pub enum AttestationProcessingOutcome {
     Processed,
@@ -142,6 +160,12 @@ pub struct BeaconChain<T: BeaconChainTypes> {
     /// Stores all operations (e.g., `Attestation`, `Deposit`, etc) that are candidates for
     /// inclusion in a block.
     pub op_pool: OperationPool<T::EthSpec>,
+    /// A pool of attestations dedicated to the "naive aggregation strategy" defined in the eth2
+    /// specs.
+    ///
+    /// This pool accepts `Attestation` objects that only have one aggregation bit set and provides
+    /// a method to get an aggregated `Attestation` for some `AttestationData`.
+    pub naive_aggregation_pool: NaiveAggregationPool<T::EthSpec>,
     /// Provides information from the Ethereum 1 (PoW) chain.
     pub eth1_chain: Option<Eth1Chain<T::Eth1Chain, T::EthSpec, T::Store>>,
     /// Stores a "snapshot" of the chain at the time the head-of-the-chain block was received.
@@ -676,27 +700,14 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         }
     }
 
-    /// Produce an aggregate attestation that has been collected for this slot and committee.
-    // TODO: Check and optimize
-    pub fn return_aggregate_attestation(
+    /// Returns an aggregated `Attestation`, if any, that has a matching `attestation.data`.
+    ///
+    /// The attestation will be obtained from `self.naive_aggregation_pool`.
+    pub fn get_aggregated_attestation(
         &self,
-        slot: Slot,
-        index: CommitteeIndex,
-    ) -> Result<Attestation<T::EthSpec>, Error> {
-        let epoch = |slot: Slot| slot.epoch(T::EthSpec::slots_per_epoch());
-        let head_state = &self.head()?.beacon_state;
-
-        let state = if epoch(slot) == epoch(head_state.slot) {
-            self.head()?.beacon_state
-        } else {
-            // The block proposer shuffling is not affected by the state roots, so we don't need to
-            // calculate them.
-            self.state_at_slot(slot, StateSkipConfig::WithoutStateRoots)?
-        };
-
-        self.op_pool
-            .get_raw_aggregated_attestations(&slot, &index, &state, &self.spec)
-            .map_err(Error::from)
+        data: &AttestationData,
+    ) -> Result<Option<Attestation<T::EthSpec>>, Error> {
+        self.naive_aggregation_pool.get(data).map_err(Into::into)
     }
 
     /// Produce a raw unsigned `Attestation` that is valid for the given `slot` and `index`.
@@ -824,12 +835,12 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
     pub fn process_attestation(
         &self,
         attestation: Attestation<T::EthSpec>,
-        store_raw: Option<bool>,
+        attestation_type: AttestationType,
     ) -> Result<AttestationProcessingOutcome, Error> {
         metrics::inc_counter(&metrics::ATTESTATION_PROCESSING_REQUESTS);
         let timer = metrics::start_timer(&metrics::ATTESTATION_PROCESSING_TIMES);
 
-        let outcome = self.process_attestation_internal(attestation.clone(), store_raw);
+        let outcome = self.process_attestation_internal(attestation.clone(), attestation_type);
 
         match &outcome {
             Ok(outcome) => match outcome {
@@ -883,7 +894,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
     pub fn process_attestation_internal(
         &self,
         attestation: Attestation<T::EthSpec>,
-        store_raw: Option<bool>,
+        attestation_type: AttestationType,
     ) -> Result<AttestationProcessingOutcome, Error> {
         let initial_validation_timer =
             metrics::start_timer(&metrics::ATTESTATION_PROCESSING_INITIAL_VALIDATION_TIMES);
@@ -1120,20 +1131,61 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             // subnet without a validator responsible for aggregating it, we don't store it in the
             // op pool.
             if self.eth1_chain.is_some() {
-                if let Some(is_raw) = store_raw {
-                    if is_raw {
-                        // This is a raw un-aggregated attestation received from a subnet with a
-                        // connected validator required to aggregate and publish these attestations
-                        self.op_pool
-                            .insert_raw_attestation(attestation, &fork, &self.spec)?;
-                    } else {
-                        // This an aggregate attestation received from the aggregate attestation
-                        // channel
-                        self.op_pool.insert_aggregate_attestation(
-                            attestation,
-                            &fork,
-                            &self.spec,
-                        )?;
+                match attestation_type {
+                    AttestationType::Unaggregated { should_store } if should_store => {
+                        match self.naive_aggregation_pool.insert(&attestation) {
+                            Ok(outcome) => trace!(
+                                self.log,
+                                "Stored unaggregated attestation";
+                                "outcome" => format!("{:?}", outcome),
+                                "index" => attestation.data.index,
+                                "slot" => attestation.data.slot.as_u64(),
+                            ),
+                            Err(NaiveAggregationError::SlotTooLow {
+                                slot,
+                                lowest_permissible_slot,
+                            }) => {
+                                trace!(
+                                    self.log,
+                                    "Refused to store unaggregated attestation";
+                                    "lowest_permissible_slot" => lowest_permissible_slot.as_u64(),
+                                    "slot" => slot.as_u64(),
+                                );
+                            }
+                            Err(e) => error!(
+                                    self.log,
+                                    "Failed to store unaggregated attestation";
+                                    "error" => format!("{:?}", e),
+                                    "index" => attestation.data.index,
+                                    "slot" => attestation.data.slot.as_u64(),
+                            ),
+                        }
+                    }
+                    AttestationType::Unaggregated { .. } => trace!(
+                        self.log,
+                        "Did not store unaggregated attestation";
+                        "index" => attestation.data.index,
+                        "slot" => attestation.data.slot.as_u64(),
+                    ),
+                    AttestationType::Aggregated => {
+                        let index = attestation.data.index;
+                        let slot = attestation.data.slot;
+
+                        match self
+                            .op_pool
+                            .insert_attestation(attestation, &fork, &self.spec)
+                        {
+                            Ok(_) => {}
+                            Err(e) => {
+                                error!(
+                                    self.log,
+                                    "Failed to add attestation to op pool";
+                                    "error" => format!("{:?}", e),
+                                    "index" => index,
+                                    "slot" => slot.as_u64(),
+                                );
+                            }
+                        }
                     }
                 }
             }
@@ -1181,7 +1233,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         if let Ok(Some(pubkey)) =
             self.validator_pubkey(aggregate_and_proof.aggregator_index as usize)
         {
-            if !signed_aggregate_and_proof.is_valid(&pubkey, &state.fork) {
+            if !signed_aggregate_and_proof.is_valid(&pubkey, &state.fork, &self.spec) {
                 Err(AttestationDropReason::AggregatorSignatureInvalid)
             } else {
                 Ok(())
@@ -1915,18 +1967,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
     pub fn per_slot_task(&self) {
         trace!(self.log, "Running beacon chain per slot tasks");
         if let Some(slot) = self.slot_clock.now() {
-            self.op_pool.prune_committee_attestations(&slot)
-        }
-    }
-
-    /// Called by the timer on every epoch.
-    ///
-    /// Performs epoch-based pruning.
-    pub fn per_epoch_task(&self) {
-        trace!(self.log, "Running beacon chain per epoch tasks");
-        if let Some(slot) = self.slot_clock.now() {
-            let current_epoch = slot.epoch(T::EthSpec::slots_per_epoch());
-            self.op_pool.prune_attestations(&current_epoch);
+            self.naive_aggregation_pool.prune(slot);
         }
     }
 

beacon_node/beacon_chain/src/builder.rs

@@ -415,6 +415,8 @@ where
             op_pool: self
                 .op_pool
                 .ok_or_else(|| "Cannot build without op pool".to_string())?,
+            // TODO: allow for persisting and loading the pool from disk.
+            naive_aggregation_pool: <_>::default(),
             eth1_chain: self.eth1_chain,
             canonical_head: TimeoutRwLock::new(canonical_head.clone()),
             genesis_block_root: self

beacon_node/beacon_chain/src/errors.rs

@@ -1,5 +1,6 @@
 use crate::eth1_chain::Error as Eth1ChainError;
 use crate::fork_choice::Error as ForkChoiceError;
+use crate::naive_aggregation_pool::Error as NaiveAggregationError;
 use operation_pool::OpPoolError;
 use ssz::DecodeError;
 use ssz_types::Error as SszTypesError;
@@ -64,6 +65,7 @@ pub enum BeaconChainError {
     DuplicateValidatorPublicKey,
     ValidatorPubkeyCacheFileError(String),
     OpPoolError(OpPoolError),
+    NaiveAggregationError(NaiveAggregationError),
 }
 
 easy_from_to!(SlotProcessingError, BeaconChainError);
@@ -71,6 +73,7 @@ easy_from_to!(AttestationValidationError, BeaconChainError);
 easy_from_to!(SszTypesError, BeaconChainError);
 easy_from_to!(OpPoolError, BeaconChainError);
 easy_from_to!(BlockSignatureVerifierError, BeaconChainError);
+easy_from_to!(NaiveAggregationError, BeaconChainError);
 
 #[derive(Debug, PartialEq)]
 pub enum BlockProductionError {

beacon_node/beacon_chain/src/lib.rs

@@ -12,6 +12,7 @@ pub mod events;
 mod fork_choice;
 mod head_tracker;
 mod metrics;
+mod naive_aggregation_pool;
 mod persisted_beacon_chain;
 mod shuffling_cache;
 mod snapshot_cache;
@@ -20,7 +21,7 @@ mod timeout_rw_lock;
 mod validator_pubkey_cache;
 
 pub use self::beacon_chain::{
-    AttestationProcessingOutcome, BeaconChain, BeaconChainTypes, StateSkipConfig,
+    AttestationProcessingOutcome, AttestationType, BeaconChain, BeaconChainTypes, StateSkipConfig,
 };
 pub use self::beacon_snapshot::BeaconSnapshot;
 pub use self::errors::{BeaconChainError, BlockProductionError};

beacon_node/beacon_chain/src/naive_aggregation_pool.rs

@@ -0,0 +1,478 @@
+use parking_lot::RwLock;
+use std::collections::HashMap;
+use types::{Attestation, AttestationData, EthSpec, Slot};
+
+/// The number of slots that will be stored in the pool.
+///
+/// For example, if `SLOTS_RETAINED == 3` and the pool is pruned at slot `6`, then all attestations
+/// at slots less than `4` will be dropped and any future attestation with a slot less than `4`
+/// will be refused.
+const SLOTS_RETAINED: usize = 3;
+
+/// The maximum number of distinct `AttestationData` that will be stored in each slot.
+///
+/// This is a DoS protection measure.
+const MAX_ATTESTATIONS_PER_SLOT: usize = 16_384;
+
+/// Returned upon successfully inserting an attestation into the pool.
+#[derive(Debug, PartialEq)]
+pub enum InsertOutcome {
+    /// The `attestation.data` had not been seen before and was added to the pool.
+    NewAttestationData { committee_index: usize },
+    /// A validator signature for the given `attestation.data` was already known. No changes were
+    /// made.
+    SignatureAlreadyKnown { committee_index: usize },
+    /// The `attestation.data` was known, but a signature for the given validator was not yet
+    /// known. The signature was aggregated into the pool.
+    SignatureAggregated { committee_index: usize },
+}
+
+#[derive(Debug, PartialEq)]
+pub enum Error {
+    /// The given `attestation.data.slot` was too low to be stored. No changes were made.
+    SlotTooLow {
+        slot: Slot,
+        lowest_permissible_slot: Slot,
+    },
+    /// The given `attestation.aggregation_bits` field was empty.
+    NoAggregationBitsSet,
+    /// The given `attestation.aggregation_bits` field had more than one signature. The number of
+    /// signatures found is included.
+    MoreThanOneAggregationBitSet(usize),
+    /// We have reached the maximum number of unique `AttestationData` that can be stored in a
+    /// slot. This is a DoS protection function.
+    ReachedMaxAttestationsPerSlot(usize),
+    /// The given `attestation.aggregation_bits` field had a different length to the one currently
+    /// stored. This indicates a fairly serious error somewhere in the code that called this
+    /// function.
+    InconsistentBitfieldLengths,
+    /// The function to obtain a map index failed, this is an internal error.
+    InvalidMapIndex(usize),
+    /// The given `attestation` was for the incorrect slot. This is an internal error.
+    IncorrectSlot { expected: Slot, attestation: Slot },
+}
+
+/// A collection of `Attestation` objects, keyed by their `attestation.data`. Enforces that all
+/// `attestation` are from the same slot.
+struct AggregatedAttestationMap<E: EthSpec> {
+    map: HashMap<AttestationData, Attestation<E>>,
+    slot: Slot,
+}
+
+impl<E: EthSpec> AggregatedAttestationMap<E> {
+    /// Create an empty collection that will only contain attestation for the given `slot`.
+    pub fn new(slot: Slot) -> Self {
+        Self {
+            slot,
+            map: <_>::default(),
+        }
+    }
+
+    /// Insert an attestation into `self`, aggregating it into the pool.
+    ///
+    /// The given attestation (`a`) must only have one signature and be from the slot that `self`
+    /// was initialized with.
+    pub fn insert(&mut self, a: &Attestation<E>) -> Result<InsertOutcome, Error> {
+        if a.data.slot != self.slot {
+            return Err(Error::IncorrectSlot {
+                expected: self.slot,
+                attestation: a.data.slot,
+            });
+        }
+
+        let set_bits = a
+            .aggregation_bits
+            .iter()
+            .enumerate()
+            .filter(|(_i, bit)| *bit)
+            .map(|(i, _bit)| i)
+            .collect::<Vec<_>>();
+
+        let committee_index = set_bits
+            .first()
+            .copied()
+            .ok_or_else(|| Error::NoAggregationBitsSet)?;
+
+        if set_bits.len() > 1 {
+            return Err(Error::MoreThanOneAggregationBitSet(set_bits.len()));
+        }
+
+        if let Some(existing_attestation) = self.map.get_mut(&a.data) {
+            if existing_attestation
+                .aggregation_bits
+                .get(committee_index)
+                .map_err(|_| Error::InconsistentBitfieldLengths)?
+            {
+                Ok(InsertOutcome::SignatureAlreadyKnown { committee_index })
+            } else {
+                existing_attestation.aggregate(a);
+                Ok(InsertOutcome::SignatureAggregated { committee_index })
+            }
+        } else {
+            if self.map.len() >= MAX_ATTESTATIONS_PER_SLOT {
+                return Err(Error::ReachedMaxAttestationsPerSlot(
+                    MAX_ATTESTATIONS_PER_SLOT,
+                ));
+            }
+
+            self.map.insert(a.data.clone(), a.clone());
+            Ok(InsertOutcome::NewAttestationData { committee_index })
+        }
+    }
+
+    /// Returns an aggregated `Attestation` with the given `data`, if any.
+    ///
+    /// The given `a.data.slot` must match the slot that `self` was initialized with.
+    pub fn get(&self, data: &AttestationData) -> Result<Option<Attestation<E>>, Error> {
+        if data.slot != self.slot {
+            return Err(Error::IncorrectSlot {
+                expected: self.slot,
+                attestation: data.slot,
+            });
+        }
+
+        Ok(self.map.get(data).cloned())
+    }
+}
+
+/// A pool of `Attestation` that is specially designed to store "unaggregated" attestations from
+/// the native aggregation scheme.
+///
+/// **The `NaiveAggregationPool` does not do any signature or attestation verification. It assumes
+/// that all `Attestation` objects provided are valid.**
+///
+/// ## Details
+///
+/// The pool sorts the `Attestation` by `attestation.data.slot`, then by `attestation.data`.
+///
+/// As each unaggregated attestation is added it is aggregated with any existing `attestation` with
+/// the same `AttestationData`. Considering that the pool only accepts attestations with a single
+/// signature, there should only ever be a single aggregated `Attestation` for any given
+/// `AttestationData`.
+///
+/// The pool has a capacity for `SLOTS_RETAINED` slots, when a new `attestation.data.slot` is
+/// provided, the oldest slot is dropped and replaced with the new slot. The pool can also be
+/// pruned by supplying a `current_slot`; all existing attestations with a slot lower than
+/// `current_slot - SLOTS_RETAINED` will be removed and any future attestation with a slot lower
+/// than that will also be refused. Pruning is done automatically based upon the attestations it
+/// receives and it can be triggered manually.
+pub struct NaiveAggregationPool<E: EthSpec> {
+    lowest_permissible_slot: RwLock<Slot>,
+    maps: RwLock<Vec<AggregatedAttestationMap<E>>>,
+}
+
+impl<E: EthSpec> Default for NaiveAggregationPool<E> {
+    fn default() -> Self {
+        Self {
+            lowest_permissible_slot: RwLock::new(Slot::new(0)),
+            maps: RwLock::new(vec![]),
+        }
+    }
+}
+
+impl<E: EthSpec> NaiveAggregationPool<E> {
+    /// Insert an attestation into `self`, aggregating it into the pool.
+    ///
+    /// The given attestation (`a`) must only have one signature and have an
+    /// `attestation.data.slot` that is not lower than `self.lowest_permissible_slot`.
+    ///
+    /// The pool may be pruned if the given `attestation.data` has a slot higher than any
+    /// previously seen.
+    pub fn insert(&self, attestation: &Attestation<E>) -> Result<InsertOutcome, Error> {
+        let lowest_permissible_slot = *self.lowest_permissible_slot.read();
+
+        // Reject any attestations that are too old.
+        if attestation.data.slot < lowest_permissible_slot {
+            return Err(Error::SlotTooLow {
+                slot: attestation.data.slot,
+                lowest_permissible_slot,
+            });
+        }
+
+        // Prune the pool if this attestation indicates that the current slot has advanced.
+        if (lowest_permissible_slot + SLOTS_RETAINED as u64) < attestation.data.slot + 1 {
+            self.prune(attestation.data.slot)
+        }
+
+        let index = self.get_map_index(attestation.data.slot);
+
+        self.maps
+            .write()
+            .get_mut(index)
+            .ok_or_else(|| Error::InvalidMapIndex(index))?
+            .insert(attestation)
+    }
+
+    /// Returns an aggregated `Attestation` with the given `data`, if any.
+    pub fn get(&self, data: &AttestationData) -> Result<Option<Attestation<E>>, Error> {
+        self.maps
+            .read()
+            .iter()
+            .find(|map| map.slot == data.slot)
+            .map(|map| map.get(data))
+            .unwrap_or_else(|| Ok(None))
+    }
+
+    /// Removes any attestations with a slot lower than `current_slot` and bars any future
+    /// attestations with a slot lower than `current_slot - SLOTS_RETAINED`.
+    pub fn prune(&self, current_slot: Slot) {
+        // Taking advantage of saturating subtraction on `Slot`.
+        let lowest_permissible_slot = current_slot - Slot::from(SLOTS_RETAINED);
+
+        self.maps
+            .write()
+            .retain(|map| map.slot >= lowest_permissible_slot);
+
+        *self.lowest_permissible_slot.write() = lowest_permissible_slot;
+    }
+
+    /// Returns the index of `self.maps` that matches `slot`.
+    ///
+    /// If there is no existing map for this slot one will be created. If `self.maps.len() >=
+    /// SLOTS_RETAINED`, the map with the lowest slot will be replaced.
+    fn get_map_index(&self, slot: Slot) -> usize {
+        let mut maps = self.maps.write();
+
+        if let Some(index) = maps.iter().position(|map| map.slot == slot) {
+            return index;
+        }
+
+        if maps.len() < SLOTS_RETAINED || maps.is_empty() {
+            let index = maps.len();
+            maps.push(AggregatedAttestationMap::new(slot));
+            return index;
+        }
+
+        let index = maps
+            .iter()
+            .enumerate()
+            .min_by_key(|(_i, map)| map.slot)
+            .map(|(i, _map)| i)
+            .expect("maps cannot be empty due to previous .is_empty() check");
+
+        maps[index] = AggregatedAttestationMap::new(slot);
+
+        index
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use ssz_types::BitList;
+    use types::{
+        test_utils::{generate_deterministic_keypair, test_random_instance},
+        Fork, Hash256,
+    };
+
+    type E = types::MainnetEthSpec;
+
+    fn get_attestation(slot: Slot) -> Attestation<E> {
+        let mut a: Attestation<E> = test_random_instance();
+        a.data.slot = slot;
+        a.aggregation_bits = BitList::with_capacity(4).expect("should create bitlist");
+        a
+    }
+
+    fn sign(a: &mut Attestation<E>, i: usize) {
+        a.sign(
+            &generate_deterministic_keypair(i).sk,
+            i,
+            &Fork::default(),
+            &E::default_spec(),
+        )
+        .expect("should sign attestation");
+    }
+
+    fn unset_bit(a: &mut Attestation<E>, i: usize) {
+        a.aggregation_bits
+            .set(i, false)
+            .expect("should unset aggregation bit")
+    }
+
+    #[test]
+    fn single_attestation() {
+        let mut a = get_attestation(Slot::new(0));
+
+        let pool = NaiveAggregationPool::default();
+
+        assert_eq!(
+            pool.insert(&a),
+            Err(Error::NoAggregationBitsSet),
+            "should not accept attestation without any signatures"
+        );
+
+        sign(&mut a, 0);
+
+        assert_eq!(
+            pool.insert(&a),
+            Ok(InsertOutcome::NewAttestationData { committee_index: 0 }),
+            "should accept new attestation"
+        );
+        assert_eq!(
+            pool.insert(&a),
+            Ok(InsertOutcome::SignatureAlreadyKnown { committee_index: 0 }),
+            "should acknowledge duplicate signature"
+        );
+
+        let retrieved = pool
+            .get(&a.data)
+            .expect("should not error while getting attestation")
+            .expect("should get an attestation");
+        assert_eq!(
+            retrieved, a,
+            "retrieved attestation should equal the one inserted"
+        );
+
+        sign(&mut a, 1);
+
+        assert_eq!(
+            pool.insert(&a),
+            Err(Error::MoreThanOneAggregationBitSet(2)),
+            "should not accept attestation with multiple signatures"
+        );
+    }
+
+    #[test]
+    fn multiple_attestations() {
+        let mut a_0 = get_attestation(Slot::new(0));
+        let mut a_1 = a_0.clone();
+
+        sign(&mut a_0, 0);
+        sign(&mut a_1, 1);
+
+        let pool = NaiveAggregationPool::default();
+
+        assert_eq!(
+            pool.insert(&a_0),
+            Ok(InsertOutcome::NewAttestationData { committee_index: 0 }),
+            "should accept a_0"
+        );
+        assert_eq!(
+            pool.insert(&a_1),
+            Ok(InsertOutcome::SignatureAggregated { committee_index: 1 }),
+            "should accept a_1"
+        );
+
+        let retrieved = pool
+            .get(&a_0.data)
+            .expect("should not error while getting attestation")
+            .expect("should get an attestation");
+
+        let mut a_01 = a_0.clone();
+        a_01.aggregate(&a_1);
+
+        assert_eq!(
+            retrieved, a_01,
+            "retrieved attestation should be aggregated"
+        );
+
+        /*
+         * Throw a different attestation data in there and ensure it isn't aggregated
+         */
+
+        let mut a_different = a_0.clone();
+        let different_root = Hash256::from_low_u64_be(1337);
+        unset_bit(&mut a_different, 0);
+        sign(&mut a_different, 2);
+        assert!(a_different.data.beacon_block_root != different_root);
+        a_different.data.beacon_block_root = different_root;
+
+        assert_eq!(
+            pool.insert(&a_different),
+            Ok(InsertOutcome::NewAttestationData { committee_index: 2 }),
+            "should accept a_different"
+        );
+
+        assert_eq!(
+            pool.get(&a_0.data)
+                .expect("should not error while getting attestation")
+                .expect("should get an attestation"),
+            retrieved,
+            "should not have aggregated different attestation data"
+        );
+    }
+
+    #[test]
+    fn auto_pruning() {
+        let mut base = get_attestation(Slot::new(0));
+        sign(&mut base, 0);
+
+        let pool = NaiveAggregationPool::default();
+
+        for i in 0..SLOTS_RETAINED * 2 {
+            let slot = Slot::from(i);
+            let mut a = base.clone();
+            a.data.slot = slot;
+
+            assert_eq!(
+                pool.insert(&a),
+                Ok(InsertOutcome::NewAttestationData { committee_index: 0 }),
+                "should accept new attestation"
+            );
+
+            if i < SLOTS_RETAINED {
+                let len = i + 1;
+                assert_eq!(
+                    pool.maps.read().len(),
+                    len,
+                    "the pool should have length {}",
+                    len
+                );
+            } else {
+                assert_eq!(
+                    pool.maps.read().len(),
+                    SLOTS_RETAINED,
+                    "the pool should have length SLOTS_RETAINED"
+                );
+
+                let mut pool_slots = pool
+                    .maps
+                    .read()
+                    .iter()
+                    .map(|map| map.slot)
+                    .collect::<Vec<_>>();
+
+                pool_slots.sort_unstable();
+
+                for (j, pool_slot) in pool_slots.iter().enumerate() {
+                    let expected_slot = slot - (SLOTS_RETAINED - 1 - j) as u64;
+                    assert_eq!(
+                        *pool_slot, expected_slot,
+                        "the slot of the map should be {}",
+                        expected_slot
+                    )
+                }
+            }
+        }
+    }
+
+    #[test]
+    fn max_attestations() {
+        let mut base = get_attestation(Slot::new(0));
+        sign(&mut base, 0);
+
+        let pool = NaiveAggregationPool::default();
+
+        for i in 0..=MAX_ATTESTATIONS_PER_SLOT {
+            let mut a = base.clone();
+            a.data.beacon_block_root = Hash256::from_low_u64_be(i as u64);
+
+            if i < MAX_ATTESTATIONS_PER_SLOT {
+                assert_eq!(
+                    pool.insert(&a),
+                    Ok(InsertOutcome::NewAttestationData { committee_index: 0 }),
+                    "should accept attestation below limit"
+                );
+            } else {
+                assert_eq!(
+                    pool.insert(&a),
+                    Err(Error::ReachedMaxAttestationsPerSlot(
+                        MAX_ATTESTATIONS_PER_SLOT
+                    )),
+                    "should not accept attestation above limit"
+                );
+            }
+        }
+    }
+}

beacon_node/beacon_chain/src/test_utils.rs

@@ -6,7 +6,7 @@ use crate::{
     builder::{BeaconChainBuilder, Witness},
     eth1_chain::CachingEth1Backend,
     events::NullEventHandler,
-    AttestationProcessingOutcome, BeaconChain, BeaconChainTypes, StateSkipConfig,
+    AttestationProcessingOutcome, AttestationType, BeaconChain, BeaconChainTypes, StateSkipConfig,
 };
 use genesis::interop_genesis_state;
 use rayon::prelude::*;
@@ -342,7 +342,7 @@ where
         .for_each(|attestation| {
             match self
                 .chain
-                .process_attestation(attestation, Some(false))
+                .process_attestation(attestation, AttestationType::Aggregated)
                 .expect("should not error during attestation processing")
             {
                 AttestationProcessingOutcome::Processed => (),