github.com/google/syzkaller@v0.0.0-20240517125934-c0f1611a36d6/pkg/fuzzer/queue/prio_queue.go (about)

     1  // Copyright 2024 syzkaller project authors. All rights reserved.
     2  // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
     3  
     4  package queue
     5  
     6  import (
     7  	"container/heap"
     8  )
     9  
    10  type priorityQueueOps[T any] struct {
    11  	impl priorityQueueImpl[T]
    12  }
    13  
    14  func (pq *priorityQueueOps[T]) Len() int {
    15  	return pq.impl.Len()
    16  }
    17  
    18  func (pq *priorityQueueOps[T]) Push(item T, prio int) {
    19  	heap.Push(&pq.impl, &priorityQueueItem[T]{item, prio})
    20  }
    21  
    22  func (pq *priorityQueueOps[T]) Pop() T {
    23  	if len(pq.impl) == 0 {
    24  		var def T
    25  		return def
    26  	}
    27  	return heap.Pop(&pq.impl).(*priorityQueueItem[T]).value
    28  }
    29  
    30  // The implementation below is based on the example provided
    31  // by https://pkg.go.dev/container/heap.
    32  
    33  type priorityQueueItem[T any] struct {
    34  	value T
    35  	prio  int
    36  }
    37  
    38  type priorityQueueImpl[T any] []*priorityQueueItem[T]
    39  
    40  func (pq priorityQueueImpl[T]) Len() int { return len(pq) }
    41  
    42  func (pq priorityQueueImpl[T]) Less(i, j int) bool {
    43  	// We want Pop to give us the lowest priority.
    44  	return pq[i].prio < pq[j].prio
    45  }
    46  
    47  func (pq priorityQueueImpl[T]) Swap(i, j int) {
    48  	pq[i], pq[j] = pq[j], pq[i]
    49  }
    50  
    51  func (pq *priorityQueueImpl[T]) Push(x any) {
    52  	*pq = append(*pq, x.(*priorityQueueItem[T]))
    53  }
    54  
    55  func (pq *priorityQueueImpl[T]) Pop() any {
    56  	n := len(*pq)
    57  	item := (*pq)[n-1]
    58  	(*pq)[n-1] = nil
    59  	*pq = (*pq)[:n-1]
    60  	return item
    61  }