batchcursor_test.go

     1  // Copyright 2023 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package trace
     6  
     7  import (
     8  	"fmt"
     9  	"strings"
    10  	"testing"
    11  
    12  	"slices"
    13  )
    14  
    15  func TestHeap(t *testing.T) {
    16  	var heap []*batchCursor
    17  
    18  	// Insert a bunch of values into the heap.
    19  	checkHeap(t, heap)
    20  	heap = heapInsert(heap, makeBatchCursor(5))
    21  	checkHeap(t, heap)
    22  	for i := int64(-20); i < 20; i++ {
    23  		heap = heapInsert(heap, makeBatchCursor(i))
    24  		checkHeap(t, heap)
    25  	}
    26  
    27  	// Update an element in the middle to be the new minimum.
    28  	for i := range heap {
    29  		if heap[i].ev.time == 5 {
    30  			heap[i].ev.time = -21
    31  			heapUpdate(heap, i)
    32  			break
    33  		}
    34  	}
    35  	checkHeap(t, heap)
    36  	if heap[0].ev.time != -21 {
    37  		t.Fatalf("heap update failed, expected %d as heap min: %s", -21, heapDebugString(heap))
    38  	}
    39  
    40  	// Update the minimum element to be smaller. There should be no change.
    41  	heap[0].ev.time = -22
    42  	heapUpdate(heap, 0)
    43  	checkHeap(t, heap)
    44  	if heap[0].ev.time != -22 {
    45  		t.Fatalf("heap update failed, expected %d as heap min: %s", -22, heapDebugString(heap))
    46  	}
    47  
    48  	// Update the last element to be larger. There should be no change.
    49  	heap[len(heap)-1].ev.time = 21
    50  	heapUpdate(heap, len(heap)-1)
    51  	checkHeap(t, heap)
    52  	if heap[len(heap)-1].ev.time != 21 {
    53  		t.Fatalf("heap update failed, expected %d as heap min: %s", 21, heapDebugString(heap))
    54  	}
    55  
    56  	// Update the last element to be smaller.
    57  	heap[len(heap)-1].ev.time = 7
    58  	heapUpdate(heap, len(heap)-1)
    59  	checkHeap(t, heap)
    60  	if heap[len(heap)-1].ev.time == 21 {
    61  		t.Fatalf("heap update failed, unexpected %d as heap min: %s", 21, heapDebugString(heap))
    62  	}
    63  
    64  	// Remove an element in the middle.
    65  	for i := range heap {
    66  		if heap[i].ev.time == 5 {
    67  			heap = heapRemove(heap, i)
    68  			break
    69  		}
    70  	}
    71  	checkHeap(t, heap)
    72  	for i := range heap {
    73  		if heap[i].ev.time == 5 {
    74  			t.Fatalf("failed to remove heap elem with time %d: %s", 5, heapDebugString(heap))
    75  		}
    76  	}
    77  
    78  	// Remove tail.
    79  	heap = heapRemove(heap, len(heap)-1)
    80  	checkHeap(t, heap)
    81  
    82  	// Remove from the head, and make sure the result is sorted.
    83  	l := len(heap)
    84  	var removed []*batchCursor
    85  	for i := 0; i < l; i++ {
    86  		removed = append(removed, heap[0])
    87  		heap = heapRemove(heap, 0)
    88  		checkHeap(t, heap)
    89  	}
    90  	if !slices.IsSortedFunc(removed, (*batchCursor).compare) {
    91  		t.Fatalf("heap elements not removed in sorted order, got: %s", heapDebugString(removed))
    92  	}
    93  }
    94  
    95  func makeBatchCursor(v int64) *batchCursor {
    96  	return &batchCursor{ev: baseEvent{time: Time(v)}}
    97  }
    98  
    99  func heapDebugString(heap []*batchCursor) string {
   100  	var sb strings.Builder
   101  	fmt.Fprintf(&sb, "[")
   102  	for i := range heap {
   103  		if i != 0 {
   104  			fmt.Fprintf(&sb, ", ")
   105  		}
   106  		fmt.Fprintf(&sb, "%d", heap[i].ev.time)
   107  	}
   108  	fmt.Fprintf(&sb, "]")
   109  	return sb.String()
   110  }
   111  
   112  func checkHeap(t *testing.T, heap []*batchCursor) {
   113  	t.Helper()
   114  
   115  	for i := range heap {
   116  		if i == 0 {
   117  			continue
   118  		}
   119  		if heap[(i-1)/2].compare(heap[i]) > 0 {
   120  			t.Errorf("heap invariant not maintained between index %d and parent %d: %s", i, i/2, heapDebugString(heap))
   121  		}
   122  	}
   123  	if t.Failed() {
   124  		t.FailNow()
   125  	}
   126  }
   127
View as plain text