github.com/ledgerwatch/erigon-lib@v1.0.0/etl/heap.go (about)

     1  /*
     2     Copyright 2021 Erigon contributors
     3  
     4     Licensed under the Apache License, Version 2.0 (the "License");
     5     you may not use this file except in compliance with the License.
     6     You may obtain a copy of the License at
     7  
     8         http://www.apache.org/licenses/LICENSE-2.0
     9  
    10     Unless required by applicable law or agreed to in writing, software
    11     distwributed under the License is distributed on an "AS IS" BASIS,
    12     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13     See the License for the specific language governing permissions and
    14     limitations under the License.
    15  */
    16  
    17  package etl
    18  
    19  import (
    20  	"bytes"
    21  )
    22  
    23  type HeapElem struct {
    24  	Key     []byte
    25  	Value   []byte
    26  	TimeIdx int
    27  }
    28  
    29  type Heap struct {
    30  	elems []*HeapElem
    31  }
    32  
    33  func (h *Heap) Len() int {
    34  	return len(h.elems)
    35  }
    36  
    37  func (h *Heap) Less(i, j int) bool {
    38  	if c := bytes.Compare(h.elems[i].Key, h.elems[j].Key); c != 0 {
    39  		return c < 0
    40  	}
    41  	return h.elems[i].TimeIdx < h.elems[j].TimeIdx
    42  }
    43  
    44  func (h *Heap) Swap(i, j int) {
    45  	h.elems[i], h.elems[j] = h.elems[j], h.elems[i]
    46  }
    47  
    48  func (h *Heap) Push(x *HeapElem) {
    49  	h.elems = append(h.elems, x)
    50  }
    51  
    52  func (h *Heap) Pop() *HeapElem {
    53  	old := h.elems
    54  	n := len(old) - 1
    55  	x := old[n]
    56  	//old[n].Key, old[n].Value, old[n].TimeIdx = nil, nil, 0
    57  	old[n] = nil
    58  	h.elems = old[0:n]
    59  	return x
    60  }
    61  
    62  // ------ Copy-Paste of `container/heap/heap.go` without interface conversion
    63  
    64  // Init establishes the heap invariants required by the other routines in this package.
    65  // Init is idempotent with respect to the heap invariants
    66  // and may be called whenever the heap invariants may have been invalidated.
    67  // The complexity is O(n) where n = h.Len().
    68  func heapInit(h *Heap) {
    69  	// heapify
    70  	n := h.Len()
    71  	for i := n/2 - 1; i >= 0; i-- {
    72  		down(h, i, n)
    73  	}
    74  }
    75  
    76  // Push pushes the element x onto the heap.
    77  // The complexity is O(log n) where n = h.Len().
    78  func heapPush(h *Heap, x *HeapElem) {
    79  	h.Push(x)
    80  	up(h, h.Len()-1)
    81  }
    82  
    83  // Pop removes and returns the minimum element (according to Less) from the heap.
    84  // The complexity is O(log n) where n = h.Len().
    85  // Pop is equivalent to Remove(h, 0).
    86  func heapPop(h *Heap) *HeapElem {
    87  	n := h.Len() - 1
    88  	h.Swap(0, n)
    89  	down(h, 0, n)
    90  	return h.Pop()
    91  }
    92  
    93  func up(h *Heap, j int) {
    94  	for {
    95  		i := (j - 1) / 2 // parent
    96  		if i == j || !h.Less(j, i) {
    97  			break
    98  		}
    99  		h.Swap(i, j)
   100  		j = i
   101  	}
   102  }
   103  
   104  func down(h *Heap, i0, n int) bool {
   105  	i := i0
   106  	for {
   107  		j1 := 2*i + 1
   108  		if j1 >= n || j1 < 0 { // j1 < 0 after int overflow
   109  			break
   110  		}
   111  		j := j1 // left child
   112  		if j2 := j1 + 1; j2 < n && h.Less(j2, j1) {
   113  			j = j2 // = 2*i + 2  // right child
   114  		}
   115  		if !h.Less(j, i) {
   116  			break
   117  		}
   118  		h.Swap(i, j)
   119  		i = j
   120  	}
   121  	return i > i0
   122  }