Simplifcation to get_data_columns.

Author: jimmygchen

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -1111,46 +1111,16 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         block_root: Hash256,
         indices: &[ColumnIndex],
     ) -> Result<Option<DataColumnSidecarList<T::EthSpec>>, Error> {
-        if indices.is_empty() {
-            return Ok(None);
-        }
-
-        let mut columns = vec![];
-        let mut indices: HashSet<ColumnIndex> = indices.iter().copied().collect();
-
-        if let Some(all_cols) = self
+        let all_cached_columns_opt = self
             .data_availability_checker
-            .get_data_columns(block_root)?
-        {
-            let filtered = all_cols
-                .into_iter()
-                .filter(|col| indices.remove(&col.index))
-                .collect::<Vec<_>>();
-            columns.extend(filtered);
-
-            if indices.is_empty() {
-                return Ok(Some(columns));
-            }
-        }
-        if let Some(all_cols) = self.early_attester_cache.get_data_columns(block_root) {
-            let filtered = all_cols
-                .into_iter()
-                .filter(|col| indices.remove(&col.index))
-                .collect::<Vec<_>>();
-            columns.extend(filtered);
-
-            if indices.is_empty() {
-                return Ok(Some(columns));
-            }
-        }
-        if let Some(filtered) = self.get_data_columns(&block_root, Some(&mut indices))? {
-            columns.extend(filtered)
-        };
+            .get_data_columns(block_root)
+            .or_else(|| self.early_attester_cache.get_data_columns(block_root));
 
-        if columns.is_empty() {
-            Ok(None)
+        if let Some(mut all_cached_columns) = all_cached_columns_opt {
+            all_cached_columns.retain(|col| indices.contains(&col.index));
+            Ok(Some(all_cached_columns))
         } else {
-            Ok(Some(columns))
+            self.get_data_columns(&block_root, Some(&indices.iter().cloned().collect()))
         }
     }
 
@@ -1235,7 +1205,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
     pub fn get_data_columns(
         &self,
         block_root: &Hash256,
-        indices: Option<&mut HashSet<ColumnIndex>>,
+        indices: Option<&HashSet<ColumnIndex>>,
     ) -> Result<Option<DataColumnSidecarList<T::EthSpec>>, Error> {
         self.store
             .get_data_columns(block_root, indices)

beacon_node/beacon_chain/src/data_availability_checker.rs

@@ -167,7 +167,7 @@ impl<T: BeaconChainTypes> DataAvailabilityChecker<T> {
     pub fn get_data_columns(
         &self,
         block_root: Hash256,
-    ) -> Result<Option<DataColumnSidecarList<T::EthSpec>>, AvailabilityCheckError> {
+    ) -> Option<DataColumnSidecarList<T::EthSpec>> {
         self.availability_cache.peek_data_columns(block_root)
     }
 

beacon_node/beacon_chain/src/data_availability_checker/overflow_lru_cache.rs

@@ -408,18 +408,12 @@ impl<T: BeaconChainTypes> DataAvailabilityCheckerInner<T> {
     pub fn peek_data_columns(
         &self,
         block_root: Hash256,
-    ) -> Result<Option<DataColumnSidecarList<T::EthSpec>>, AvailabilityCheckError> {
-        if let Some(pending_components) = self.critical.read().peek(&block_root) {
-            Ok(Some(
-                pending_components
+    ) -> Option<DataColumnSidecarList<T::EthSpec>> {
+        self.critical.read().peek(&block_root).map(|pending_components| pending_components
                     .verified_data_columns
                     .iter()
                     .map(|col| col.clone_arc())
-                    .collect(),
-            ))
-        } else {
-            Ok(None)
-        }
+                    .collect())
     }
 
     pub fn peek_pending_components<R, F: FnOnce(Option<&PendingComponents<T::EthSpec>>) -> R>(

beacon_node/store/src/hot_cold_store.rs

@@ -117,23 +117,10 @@ impl<E: EthSpec> BlockCache<E> {
     pub fn get_blobs<'a>(&'a mut self, block_root: &Hash256) -> Option<&'a BlobSidecarList<E>> {
         self.blob_cache.get(block_root)
     }
-    pub fn get_data_columns(
-        &mut self,
-        block_root: &Hash256,
-        indices: &mut Option<&mut HashSet<ColumnIndex>>,
-    ) -> Option<DataColumnSidecarList<E>> {
-        self.data_column_cache.get(block_root).map(|map| {
-            map.values()
-                .filter(|col| {
-                    if let Some(set) = indices.as_mut() {
-                        set.remove(&col.index)
-                    } else {
-                        true
-                    }
-                })
-                .cloned()
-                .collect::<Vec<_>>()
-        })
+    pub fn get_data_columns(&mut self, block_root: &Hash256) -> Option<DataColumnSidecarList<E>> {
+        self.data_column_cache
+            .get(block_root)
+            .map(|map| map.values().cloned().collect::<Vec<_>>())
     }
     pub fn get_data_column<'a>(
         &'a mut self,
@@ -2048,19 +2035,18 @@ impl<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>> HotColdDB<E, Hot, Cold>
     pub fn get_data_columns(
         &self,
         block_root: &Hash256,
-        mut indices: Option<&mut HashSet<ColumnIndex>>,
+        indices: Option<&HashSet<ColumnIndex>>,
     ) -> Result<Option<DataColumnSidecarList<E>>, Error> {
-        if let Some(columns) = self
-            .block_cache
-            .lock()
-            .get_data_columns(block_root, &mut indices)
-        {
+        // TODO: The block cache may not have all columns.
+        if let Some(columns) = self.block_cache.lock().get_data_columns(block_root) {
             metrics::inc_counter(&metrics::BEACON_DATA_COLUMNS_CACHE_HIT_COUNT);
             return Ok(Some(columns));
         }
+        // FIXME: refactor this
+        let indices = indices.cloned();
 
         let byte_columns: Result<Vec<Vec<u8>>, Error> = match indices {
-            Some(indices) => {
+            Some(mut indices) => {
                 let indices_snapshot: Vec<_> = indices.iter().copied().collect();
                 indices_snapshot
                     .iter()

testing/ef_tests/tests/tests.rs

@@ -670,10 +670,10 @@ mod ssz_static {
     #[ignore]
     // TODO(das): enable once EF tests are updated to latest release.
     fn data_column_by_root_identifier() {
-        SszStaticHandler::<DataColumnsByRootIdentifier, MinimalEthSpec>::default()
-            .run_for_feature(FeatureName::Fulu);
-        SszStaticHandler::<DataColumnsByRootIdentifier, MainnetEthSpec>::default()
-            .run_for_feature(FeatureName::Fulu);
+        // SszStaticHandler::<DataColumnsByRootIdentifier, MinimalEthSpec>::default()
+        //     .run_for_feature(FeatureName::Fulu);
+        // SszStaticHandler::<DataColumnsByRootIdentifier, MainnetEthSpec>::default()
+        //     .run_for_feature(FeatureName::Fulu);
     }
 
     #[test]