xref: /aosp_15_r20/external/licenseclassifier/stringclassifier/internal/pq/priority.go (revision 46c4c49da23cae783fa41bf46525a6505638499a)
1// Copyright 2017 Google Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15// Package pq provides a priority queue.
16package pq
17
18import "container/heap"
19
20// NewQueue returns an unbounded priority queue that compares elements using
21// less; the minimal element is at the top of the queue.
22//
23// If setIndex is not nil, the queue calls setIndex to inform each element of
24// its position in the queue.  If an element's priority changes, its position in
25// the queue may be incorrect.  Call Fix on the element's index to update the
26// queue.  Call Remove on the element's index to remove it from the queue.
27func NewQueue(less func(x, y interface{}) bool, setIndex func(x interface{}, idx int)) *Queue {
28	return &Queue{
29		heap: pqHeap{
30			less:     less,
31			setIndex: setIndex,
32		},
33	}
34}
35
36// Queue is a priority queue that supports updating the priority of an element.
37// A Queue must be created with NewQueue.
38type Queue struct {
39	heap pqHeap
40}
41
42// Len returns the number of elements in the queue.
43func (pq *Queue) Len() int {
44	return pq.heap.Len()
45}
46
47// Push adds x to the queue.
48func (pq *Queue) Push(x interface{}) {
49	heap.Push(&pq.heap, x)
50}
51
52// Min returns the minimal element.
53// Min panics if the queue is empty.
54func (pq *Queue) Min() interface{} {
55	return pq.heap.a[0]
56}
57
58// Pop removes and returns the minimal element.
59// Pop panics if the queue is empty.
60func (pq *Queue) Pop() interface{} {
61	return heap.Pop(&pq.heap)
62}
63
64// Fix adjusts the heap to reflect that the element at index has changed priority.
65func (pq *Queue) Fix(index int) {
66	heap.Fix(&pq.heap, index)
67}
68
69// Remove removes the element at index i from the heap.
70func (pq *Queue) Remove(index int) {
71	heap.Remove(&pq.heap, index)
72}
73
74// pqHeap implements heap.Interface.
75type pqHeap struct {
76	a        []interface{}
77	less     func(x, y interface{}) bool
78	setIndex func(x interface{}, idx int)
79}
80
81func (h pqHeap) Len() int {
82	return len(h.a)
83}
84
85func (h pqHeap) Less(i, j int) bool {
86	return h.less(h.a[i], h.a[j])
87}
88
89func (h pqHeap) Swap(i, j int) {
90	h.a[i], h.a[j] = h.a[j], h.a[i]
91	if h.setIndex != nil {
92		h.setIndex(h.a[i], i)
93		h.setIndex(h.a[j], j)
94	}
95}
96
97func (h *pqHeap) Push(x interface{}) {
98	n := len(h.a)
99	if h.setIndex != nil {
100		h.setIndex(x, n)
101	}
102	h.a = append(h.a, x)
103}
104
105func (h *pqHeap) Pop() interface{} {
106	old := h.a
107	n := len(old)
108	x := old[n-1]
109	h.a = old[:n-1]
110	return x
111}
112