diff --git a/tx-pool/src/component/pool_map.rs b/tx-pool/src/component/pool_map.rs index 76623c0fdd..a421f10e04 100644 --- a/tx-pool/src/component/pool_map.rs +++ b/tx-pool/src/component/pool_map.rs @@ -153,7 +153,7 @@ impl PoolMap { } pub(crate) fn sorted_proposed_iter(&self) -> impl Iterator { - self.score_sorted_iter_by(vec![Status::Proposed]) + self.score_sorted_iter_by_status(Status::Proposed) } pub(crate) fn get(&self, id: &ProposalShortId) -> Option<&TxEntry> { @@ -329,23 +329,19 @@ impl PoolMap { conflicts } - // fill proposal txs - pub(crate) fn fill_proposals( + // find the pending txs sorted by score, and return their proposal short ids + pub(crate) fn get_proposals( &self, limit: usize, exclusion: &HashSet, - proposals: &mut HashSet, - status: Status, - ) { - for entry in self.score_sorted_iter_by(vec![status]) { - if proposals.len() == limit { - break; - } - let id = entry.proposal_short_id(); - if !exclusion.contains(&id) { - proposals.insert(id); - } - } + ) -> HashSet { + self.score_sorted_iter_by_status(Status::Pending) + .filter_map(|entry| { + let id = entry.proposal_short_id(); + (!exclusion.contains(&id)).then_some(id) + }) + .take(limit) + .collect() } pub(crate) fn iter(&self) -> impl Iterator { @@ -370,15 +366,24 @@ impl PoolMap { self.proposed_count = 0; } - pub(crate) fn score_sorted_iter_by( + pub(crate) fn score_sorted_iter_by_status( + &self, + status: Status, + ) -> impl Iterator { + self.entries + .iter_by_score() + .rev() + .filter_map(move |entry| (entry.status == status).then_some(&entry.inner)) + } + + pub(crate) fn score_sorted_iter_by_statuses( &self, statuses: Vec, ) -> impl Iterator { self.entries .iter_by_score() .rev() - .filter(move |entry| statuses.contains(&entry.status)) - .map(|entry| &entry.inner) + .filter_map(move |entry| statuses.contains(&entry.status).then_some(&entry.inner)) } fn remove_entry_links(&mut self, id: &ProposalShortId) { diff --git a/tx-pool/src/component/tests/pending.rs b/tx-pool/src/component/tests/pending.rs index 33a8d5d373..e262248354 100644 --- a/tx-pool/src/component/tests/pending.rs +++ b/tx-pool/src/component/tests/pending.rs @@ -163,7 +163,7 @@ fn test_remove_entry() { } #[test] -fn test_fill_proposals() { +fn test_get_proposals() { let mut pool = PoolMap::new(1000); let tx1 = build_tx(vec![(&Byte32::zero(), 1), (&h256!("0x1").pack(), 1)], 1); let tx2 = build_tx( @@ -193,30 +193,26 @@ fn test_fill_proposals() { let id2 = tx2.proposal_short_id(); let id3 = tx3.proposal_short_id(); - let mut ret = HashSet::new(); - pool.fill_proposals(10, &HashSet::new(), &mut ret, Status::Pending); + let ret = pool.get_proposals(10, &HashSet::new()); assert_eq!( ret, HashSet::from_iter(vec![id1.clone(), id2.clone(), id3.clone()]) ); - let mut ret = HashSet::new(); - pool.fill_proposals(1, &HashSet::new(), &mut ret, Status::Pending); + let ret = pool.get_proposals(1, &HashSet::new()); assert_eq!(ret.len(), 1); - let mut ret = HashSet::new(); - pool.fill_proposals(2, &HashSet::new(), &mut ret, Status::Pending); + let ret = pool.get_proposals(2, &HashSet::new()); assert_eq!(ret.len(), 2); - let mut ret = HashSet::new(); let mut exclusion = HashSet::new(); exclusion.insert(id2); - pool.fill_proposals(2, &exclusion, &mut ret, Status::Pending); + let ret = pool.get_proposals(2, &exclusion); assert_eq!(ret, HashSet::from_iter(vec![id1, id3])); } #[test] -fn test_fill_proposals_with_high_score() { +fn test_get_proposals_with_high_score() { let mut pool = PoolMap::new(1000); let tx1 = build_tx(vec![(&Byte32::zero(), 1), (&h256!("0x1").pack(), 1)], 1); let tx2 = build_tx( @@ -242,25 +238,21 @@ fn test_fill_proposals_with_high_score() { let id2 = tx2.proposal_short_id(); let id3 = tx3.proposal_short_id(); - let mut ret = HashSet::new(); - pool.fill_proposals(10, &HashSet::new(), &mut ret, Status::Pending); + let ret = pool.get_proposals(10, &HashSet::new()); assert_eq!( ret, HashSet::from_iter(vec![id3.clone(), id2.clone(), id1.clone()]) ); - let mut ret = HashSet::new(); - pool.fill_proposals(1, &HashSet::new(), &mut ret, Status::Pending); + let ret = pool.get_proposals(1, &HashSet::new()); assert_eq!(ret, HashSet::from_iter(vec![id3.clone()])); - let mut ret = HashSet::new(); - pool.fill_proposals(2, &HashSet::new(), &mut ret, Status::Pending); + let ret = pool.get_proposals(2, &HashSet::new()); assert_eq!(ret, HashSet::from_iter(vec![id3.clone(), id2.clone()])); - let mut ret = HashSet::new(); let mut exclusion = HashSet::new(); exclusion.insert(id2); - pool.fill_proposals(2, &exclusion, &mut ret, Status::Pending); + let ret = pool.get_proposals(2, &exclusion); assert_eq!(ret, HashSet::from_iter(vec![id1, id3])); } diff --git a/tx-pool/src/pool.rs b/tx-pool/src/pool.rs index 6399f18602..54b2395e4c 100644 --- a/tx-pool/src/pool.rs +++ b/tx-pool/src/pool.rs @@ -402,10 +402,7 @@ impl TxPool { limit: usize, exclusion: &HashSet, ) -> HashSet { - let mut proposals = HashSet::with_capacity(limit); - self.pool_map - .fill_proposals(limit, exclusion, &mut proposals, Status::Pending); - proposals + self.pool_map.get_proposals(limit, exclusion) } /// Returns tx from tx-pool or storage corresponding to the id. @@ -426,7 +423,7 @@ impl TxPool { pub(crate) fn get_ids(&self) -> TxPoolIds { let pending = self .pool_map - .score_sorted_iter_by(vec![Status::Pending, Status::Gap]) + .score_sorted_iter_by_statuses(vec![Status::Pending, Status::Gap]) .map(|entry| entry.transaction().hash()) .collect(); @@ -442,7 +439,7 @@ impl TxPool { pub(crate) fn get_all_entry_info(&self) -> TxPoolEntryInfo { let pending = self .pool_map - .score_sorted_iter_by(vec![Status::Pending, Status::Gap]) + .score_sorted_iter_by_statuses(vec![Status::Pending, Status::Gap]) .map(|entry| (entry.transaction().hash(), entry.to_info())) .collect();