From 03ddb5f11871fbc6417476390d69d8cfe41f05ca Mon Sep 17 00:00:00 2001
From: ardocrat <ardocrat@gri.mw>
Date: Thu, 23 Apr 2026 00:52:16 +0300
Subject: [PATCH] pibd: check if segment not exists at request when selecting
 next required

---
 chain/src/txhashset/desegmenter.rs | 33 +++++++++++++++---------------
 1 file changed, 17 insertions(+), 16 deletions(-)

diff --git a/chain/src/txhashset/desegmenter.rs b/chain/src/txhashset/desegmenter.rs
index 5ea17f63..d8c41fe2 100644
--- a/chain/src/txhashset/desegmenter.rs
+++ b/chain/src/txhashset/desegmenter.rs
@@ -583,6 +583,16 @@ impl Desegmenter {
 			}
 		}
 
+		// Explicitly add segment identifier to request if not exists.
+		let mut maybe_add_to_request = |seg_id: SegmentTypeIdentifier| {
+			if return_vec.iter().any(|i| i == &seg_id) {
+				if return_vec.len() >= max_elements {
+					return_vec.pop();
+				}
+				return_vec.push(seg_id);
+			}
+		};
+
 		// Ensure we explicitly ask for the next output segment.
 		if let Some(next_output_idx) = self.next_required_output_segment_index() {
 			let seg_id = SegmentIdentifier {
@@ -590,10 +600,8 @@ impl Desegmenter {
 				idx: next_output_idx,
 			};
 			if !self.has_output_segment_with_id(seg_id) {
-				if return_vec.len() >= max_elements {
-					return_vec.pop();
-				}
-				return_vec.push(SegmentTypeIdentifier::new(SegmentType::Output, seg_id));
+				let next_output_seg_id = SegmentTypeIdentifier::new(SegmentType::Output, seg_id);
+				maybe_add_to_request(next_output_seg_id);
 			}
 		}
 
@@ -604,10 +612,8 @@ impl Desegmenter {
 				idx: next_rp_idx,
 			};
 			if !self.has_rangeproof_segment_with_id(seg_id) {
-				if return_vec.len() >= max_elements {
-					return_vec.pop();
-				}
-				return_vec.push(SegmentTypeIdentifier::new(SegmentType::RangeProof, seg_id));
+				let next_proof_seg_id = SegmentTypeIdentifier::new(SegmentType::RangeProof, seg_id);
+				maybe_add_to_request(next_proof_seg_id);
 			}
 		}
 
@@ -619,14 +625,9 @@ impl Desegmenter {
 				height: self.default_kernel_segment_height,
 				idx: next_kernel_idx,
 			};
-			let next_kernel_seg_id = SegmentTypeIdentifier::new(SegmentType::Kernel, seg_id);
-			if !self.has_kernel_segment_with_id(seg_id)
-				&& !return_vec.iter().any(|x| x == &next_kernel_seg_id)
-			{
-				if return_vec.len() >= max_elements {
-					return_vec.pop();
-				}
-				return_vec.push(next_kernel_seg_id);
+			if !self.has_kernel_segment_with_id(seg_id) {
+				let next_kernel_seg_id = SegmentTypeIdentifier::new(SegmentType::Kernel, seg_id);
+				maybe_add_to_request(next_kernel_seg_id);
 			}
 		}
 		if return_vec.is_empty() && self.bitmap_cache.is_some() {