cleanup

yihuang · Aug 27, 2024 · 4e2aa82 · 4e2aa82
1 parent 156e548
commit 4e2aa82
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 39 deletions.
diff --git a/types/mempool/priority_nonce.go b/types/mempool/priority_nonce.go
@@ -351,9 +351,13 @@ func (i *PriorityNonceIterator[C]) Tx() sdk.Tx {
 //
 // NOTE: It is not safe to use this iterator while removing transactions from
 // the underlying mempool.
-func (mp *PriorityNonceMempool[C]) Select(_ context.Context, _ [][]byte) Iterator {
+func (mp *PriorityNonceMempool[C]) Select(ctx context.Context, txs [][]byte) Iterator {
 	mp.mtx.Lock()
 	defer mp.mtx.Unlock()
+	return mp.doSelect(ctx, txs)
+}
+
+func (mp *PriorityNonceMempool[C]) doSelect(_ context.Context, _ [][]byte) Iterator {
 	if mp.priorityIndex.Len() == 0 {
 		return nil
 	}
@@ -369,21 +373,11 @@ func (mp *PriorityNonceMempool[C]) Select(_ context.Context, _ [][]byte) Iterato
 }
 
 // SelectBy will hold the mutex during the iteration, callback returns if continue.
-func (mp *PriorityNonceMempool[C]) SelectBy(_ context.Context, _ [][]byte, callback func(sdk.Tx) bool) {
+func (mp *PriorityNonceMempool[C]) SelectBy(ctx context.Context, txs [][]byte, callback func(sdk.Tx) bool) {
 	mp.mtx.Lock()
 	defer mp.mtx.Unlock()
-	if mp.priorityIndex.Len() == 0 {
-		return
-	}
-
-	mp.reorderPriorityTies()
-
-	iterator := &PriorityNonceIterator[C]{
-		mempool:       mp,
-		senderCursors: make(map[string]*skiplist.Element),
-	}
 
-	iter := iterator.iteratePriority()
+	iter := mp.doSelect(ctx, txs)
 	for iter != nil && callback(iter.Tx()) {
 		iter = iter.Next()
 	}

diff --git a/types/mempool/sender_nonce.go b/types/mempool/sender_nonce.go
@@ -159,9 +159,13 @@ func (snm *SenderNonceMempool) Insert(_ context.Context, tx sdk.Tx) error {
 //
 // NOTE: It is not safe to use this iterator while removing transactions from
 // the underlying mempool.
-func (snm *SenderNonceMempool) Select(_ context.Context, _ [][]byte) Iterator {
+func (snm *SenderNonceMempool) Select(ctx context.Context, txs [][]byte) Iterator {
 	snm.mtx.Lock()
 	defer snm.mtx.Unlock()
+	return snm.doSelect(ctx, txs)
+}
+
+func (snm *SenderNonceMempool) doSelect(_ context.Context, _ [][]byte) Iterator {
 	var senders []string
 
 	senderCursors := make(map[string]*skiplist.Element)
@@ -190,34 +194,11 @@ func (snm *SenderNonceMempool) Select(_ context.Context, _ [][]byte) Iterator {
 }
 
 // SelectBy will hold the mutex during the iteration, callback returns if continue.
-func (snm *SenderNonceMempool) SelectBy(_ context.Context, _ [][]byte, callback func(sdk.Tx) bool) {
+func (snm *SenderNonceMempool) SelectBy(ctx context.Context, txs [][]byte, callback func(sdk.Tx) bool) {
 	snm.mtx.Lock()
 	defer snm.mtx.Unlock()
-	var senders []string
-
-	senderCursors := make(map[string]*skiplist.Element)
-	orderedSenders := skiplist.New(skiplist.String)
-
-	// #nosec
-	for s := range snm.senders {
-		orderedSenders.Set(s, s)
-	}
-
-	s := orderedSenders.Front()
-	for s != nil {
-		sender := s.Value.(string)
-		senders = append(senders, sender)
-		senderCursors[sender] = snm.senders[sender].Front()
-		s = s.Next()
-	}
-
-	iterator := &senderNonceMempoolIterator{
-		senders:       senders,
-		rnd:           snm.rnd,
-		senderCursors: senderCursors,
-	}
 
-	iter := iterator.Next()
+	iter := snm.doSelect(ctx, txs)
 	for iter != nil && callback(iter.Tx()) {
 		iter = iter.Next()
 	}