1//  Copyright (c) 2014 Couchbase, 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
15package collector
16
17import (
18	"container/heap"
19
20	"github.com/blevesearch/bleve/v2/search"
21)
22
23type collectStoreHeap struct {
24	heap    search.DocumentMatchCollection
25	compare collectorCompare
26}
27
28func newStoreHeap(capacity int, compare collectorCompare) *collectStoreHeap {
29	rv := &collectStoreHeap{
30		heap:    make(search.DocumentMatchCollection, 0, capacity),
31		compare: compare,
32	}
33	heap.Init(rv)
34	return rv
35}
36
37func (c *collectStoreHeap) AddNotExceedingSize(doc *search.DocumentMatch,
38	size int) *search.DocumentMatch {
39	c.add(doc)
40	if c.Len() > size {
41		return c.removeLast()
42	}
43	return nil
44}
45
46func (c *collectStoreHeap) add(doc *search.DocumentMatch) {
47	heap.Push(c, doc)
48}
49
50func (c *collectStoreHeap) removeLast() *search.DocumentMatch {
51	return heap.Pop(c).(*search.DocumentMatch)
52}
53
54func (c *collectStoreHeap) Final(skip int, fixup collectorFixup) (search.DocumentMatchCollection, error) {
55	count := c.Len()
56	size := count - skip
57	if size <= 0 {
58		return make(search.DocumentMatchCollection, 0), nil
59	}
60	rv := make(search.DocumentMatchCollection, size)
61	for i := size - 1; i >= 0; i-- {
62		doc := heap.Pop(c).(*search.DocumentMatch)
63		rv[i] = doc
64		err := fixup(doc)
65		if err != nil {
66			return nil, err
67		}
68	}
69	return rv, nil
70}
71
72// heap interface implementation
73
74func (c *collectStoreHeap) Len() int {
75	return len(c.heap)
76}
77
78func (c *collectStoreHeap) Less(i, j int) bool {
79	so := c.compare(c.heap[i], c.heap[j])
80	return -so < 0
81}
82
83func (c *collectStoreHeap) Swap(i, j int) {
84	c.heap[i], c.heap[j] = c.heap[j], c.heap[i]
85}
86
87func (c *collectStoreHeap) Push(x interface{}) {
88	c.heap = append(c.heap, x.(*search.DocumentMatch))
89}
90
91func (c *collectStoreHeap) Pop() interface{} {
92	var rv *search.DocumentMatch
93	rv, c.heap = c.heap[len(c.heap)-1], c.heap[:len(c.heap)-1]
94	return rv
95}
96