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
5package trace
6
7import (
8	"fmt"
9	"strings"
10	"testing"
11
12	"slices"
13)
14
15func 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
95func makeBatchCursor(v int64) *batchCursor {
96	return &batchCursor{ev: baseEvent{time: Time(v)}}
97}
98
99func 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
112func 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