From 07552ca1490eba846c85bf423aea606347135d18 Mon Sep 17 00:00:00 2001
From: Matthias Seitz <matthias.seitz@outlook.de>
Date: Mon, 30 Oct 2023 13:57:52 +0100
Subject: [PATCH] style: builder style segment set (#5209)

---
 bin/reth/src/node/mod.rs         | 51 ++++++++++++++++++--------------
 crates/prune/src/segments/set.rs | 10 ++++++-
 2 files changed, 38 insertions(+), 23 deletions(-)

diff --git a/bin/reth/src/node/mod.rs b/bin/reth/src/node/mod.rs
index 5aa011bce..eed1f3c52 100644
--- a/bin/reth/src/node/mod.rs
+++ b/bin/reth/src/node/mod.rs
@@ -914,34 +914,41 @@ impl<Ext: RethCliExt> NodeCommand<Ext> {
         Ok(pipeline)
     }
 
+    /// Builds a [Pruner] with the given config.
     fn build_pruner<DB: Database>(
         &self,
         config: &PruneConfig,
         db: DB,
         highest_snapshots_rx: HighestSnapshotsTracker,
     ) -> Pruner<DB> {
-        let mut segments = SegmentSet::new();
-
-        if let Some(mode) = config.segments.receipts {
-            segments = segments.add_segment(reth_prune::segments::Receipts::new(mode));
-        }
-        if !config.segments.receipts_log_filter.is_empty() {
-            segments = segments.add_segment(reth_prune::segments::ReceiptsByLogs::new(
-                config.segments.receipts_log_filter.clone(),
-            ));
-        }
-        if let Some(mode) = config.segments.transaction_lookup {
-            segments = segments.add_segment(reth_prune::segments::TransactionLookup::new(mode));
-        }
-        if let Some(mode) = config.segments.sender_recovery {
-            segments = segments.add_segment(reth_prune::segments::SenderRecovery::new(mode));
-        }
-        if let Some(mode) = config.segments.account_history {
-            segments = segments.add_segment(reth_prune::segments::AccountHistory::new(mode));
-        }
-        if let Some(mode) = config.segments.storage_history {
-            segments = segments.add_segment(reth_prune::segments::StorageHistory::new(mode));
-        }
+        let segments = SegmentSet::default()
+            // Receipts
+            .segment_opt(config.segments.receipts.map(reth_prune::segments::Receipts::new))
+            // Receipts by logs
+            .segment_opt((!config.segments.receipts_log_filter.is_empty()).then(|| {
+                reth_prune::segments::ReceiptsByLogs::new(
+                    config.segments.receipts_log_filter.clone(),
+                )
+            }))
+            // Transaction lookup
+            .segment_opt(
+                config
+                    .segments
+                    .transaction_lookup
+                    .map(reth_prune::segments::TransactionLookup::new),
+            )
+            // Sender recovery
+            .segment_opt(
+                config.segments.sender_recovery.map(reth_prune::segments::SenderRecovery::new),
+            )
+            // Account history
+            .segment_opt(
+                config.segments.account_history.map(reth_prune::segments::AccountHistory::new),
+            )
+            // Storage history
+            .segment_opt(
+                config.segments.storage_history.map(reth_prune::segments::StorageHistory::new),
+            );
 
         Pruner::new(
             db,
diff --git a/crates/prune/src/segments/set.rs b/crates/prune/src/segments/set.rs
index 92402825e..8593dd42e 100644
--- a/crates/prune/src/segments/set.rs
+++ b/crates/prune/src/segments/set.rs
@@ -15,11 +15,19 @@ impl<DB: Database> SegmentSet<DB> {
     }
 
     /// Adds new [Segment] to collection.
-    pub fn add_segment<S: Segment<DB> + 'static>(mut self, segment: S) -> Self {
+    pub fn segment<S: Segment<DB> + 'static>(mut self, segment: S) -> Self {
         self.inner.push(Arc::new(segment));
         self
     }
 
+    /// Adds new [Segment] to collection if it's [Some].
+    pub fn segment_opt<S: Segment<DB> + 'static>(self, segment: Option<S>) -> Self {
+        if let Some(segment) = segment {
+            return self.segment(segment)
+        }
+        self
+    }
+
     /// Consumes [SegmentSet] and returns a [Vec].
     pub fn into_vec(self) -> Vec<Arc<dyn Segment<DB>>> {
         self.inner