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