aboutsummaryrefslogtreecommitdiffstats
path: root/pkg/fuzzer/queue/prio_queue.go
blob: 5e7c65fba64e6551ad761da45e5156e458b120b4 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
// Copyright 2024 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package queue

import (
	"container/heap"
)

type priorityQueueOps[T any] struct {
	impl priorityQueueImpl[T]
}

func (pq *priorityQueueOps[T]) Len() int {
	return pq.impl.Len()
}

func (pq *priorityQueueOps[T]) Push(item T, prio int) {
	heap.Push(&pq.impl, &priorityQueueItem[T]{item, prio})
}

func (pq *priorityQueueOps[T]) Pop() T {
	if len(pq.impl) == 0 {
		var def T
		return def
	}
	return heap.Pop(&pq.impl).(*priorityQueueItem[T]).value
}

// The implementation below is based on the example provided
// by https://pkg.go.dev/container/heap.

type priorityQueueItem[T any] struct {
	value T
	prio  int
}

type priorityQueueImpl[T any] []*priorityQueueItem[T]

func (pq priorityQueueImpl[T]) Len() int { return len(pq) }

func (pq priorityQueueImpl[T]) Less(i, j int) bool {
	// We want Pop to give us the lowest priority.
	return pq[i].prio < pq[j].prio
}

func (pq priorityQueueImpl[T]) Swap(i, j int) {
	pq[i], pq[j] = pq[j], pq[i]
}

func (pq *priorityQueueImpl[T]) Push(x any) {
	*pq = append(*pq, x.(*priorityQueueItem[T]))
}

func (pq *priorityQueueImpl[T]) Pop() any {
	n := len(*pq)
	item := (*pq)[n-1]
	(*pq)[n-1] = nil
	*pq = (*pq)[:n-1]
	return item
}