util: introduce generic containers package

Package container provides three container types: heap, list, and ring.

The package is a direct copy of the Go standard library's `container` package,
but with interfaces replaced with generics. No other changes have been made,
even at the expense of ergonomics in a few places, and no other changes should
be made. The intention is to delete this library when the Go standard library
adds generics support to these packages.

The use of generics over (empty) interfaces for the contained value in these
containers provides three main benefits:

  - type safety: generics replace dynamic type assertions with static type
    checks. This ensures that code with type errors cannot compile, rather than
    panicking at runtime.

  - performance: generics allow values to be stored directly in the container,
    rather than boxed in an interface{}. This reduces allocations, eliminates
    runtime type assertions, and improves cache locality.

  - code documentation: generics require the container to be parameterized by
    the type of value it stores. This serves as a helpful form of documentation
    that describes the intended use of the container.

Container types are one of the quintessential application of generics, so much
of the benefit of generics that was described in the original design proposal[1]
is realized by this package.

[1]: https://go.googlesource.com/proposal/+/HEAD/design/43651-type-parameters.md

Epic: None
Release note: None

pkg/BUILD.bazel

@@ -629,6 +629,9 @@ ALL_TESTS = [
     "//pkg/util/cgroups:cgroups_test",
     "//pkg/util/circuit:circuit_test",
     "//pkg/util/cloudinfo:cloudinfo_test",
+    "//pkg/util/container/heap:heap_test",
+    "//pkg/util/container/list:list_test",
+    "//pkg/util/container/ring:ring_test",
     "//pkg/util/ctxgroup:ctxgroup_test",
     "//pkg/util/ctxlog:ctxlog_test",
     "//pkg/util/ctxutil:ctxutil_test",
@@ -2262,6 +2265,13 @@ GO_TARGETS = [
     "//pkg/util/cloudinfo:cloudinfo",
     "//pkg/util/cloudinfo:cloudinfo_test",
     "//pkg/util/collatedstring:collatedstring",
+    "//pkg/util/container/heap:heap",
+    "//pkg/util/container/heap:heap_test",
+    "//pkg/util/container/list:list",
+    "//pkg/util/container/list:list_test",
+    "//pkg/util/container/ring:ring",
+    "//pkg/util/container/ring:ring_test",
+    "//pkg/util/container:container",
     "//pkg/util/ctxgroup:ctxgroup",
     "//pkg/util/ctxgroup:ctxgroup_test",
     "//pkg/util/ctxlog:ctxlog",

pkg/util/container/BUILD.bazel

@@ -0,0 +1,8 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+
+go_library(
+    name = "container",
+    srcs = ["doc.go"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/util/container",
+    visibility = ["//visibility:public"],
+)

pkg/util/container/doc.go

@@ -0,0 +1,41 @@
+// Copyright 2023 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+/*
+Package container provides three container types: heap, list, and ring.
+
+The package is a direct copy of the Go standard library's container package,
+but with interfaces replaced with generics. No other changes have been made,
+even at the expense of ergonomics in a few places, and no other changes should
+be made. The intention is to delete this library when the Go standard library
+adds generics support to these packages.
+
+The use of generics over (empty) interfaces for the contained value in these
+containers provides three main benefits:
+
+  - type safety: generics replace dynamic type assertions with static type
+    checks. This ensures that code with type errors cannot compile, rather than
+    panicking at runtime.
+
+  - performance: generics allow values to be stored directly in the container,
+    rather than boxed in an interface{}. This reduces allocations, eliminates
+    runtime type assertions, and improves cache locality.
+
+  - code documentation: generics require the container to be parameterized by
+    the type of value it stores. This serves as a helpful form of documentation
+    that describes the intended use of the container.
+
+Container types are one of the quintessential application of generics, so much
+of the benefit of generics that was described in the original design proposal[1]
+is realized by this package.
+
+[1]: https://go.googlesource.com/proposal/+/HEAD/design/43651-type-parameters.md
+*/
+package container

pkg/util/container/heap/BUILD.bazel

@@ -0,0 +1,19 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "heap",
+    srcs = ["heap.go"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/util/container/heap",
+    visibility = ["//visibility:public"],
+)
+
+go_test(
+    name = "heap_test",
+    srcs = [
+        "example_intheap_test.go",
+        "example_pq_test.go",
+        "heap_test.go",
+    ],
+    args = ["-test.timeout=295s"],
+    embed = [":heap"],
+)

pkg/util/container/heap/example_intheap_test.go

@@ -0,0 +1,58 @@
+// Copyright 2023 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+// Copyright 2012 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// This example demonstrates an integer heap built using the heap interface.
+package heap_test
+
+import (
+	"fmt"
+
+	"github.com/cockroachdb/cockroach/pkg/util/container/heap"
+)
+
+// An IntHeap is a min-heap of ints.
+type IntHeap []int
+
+func (h IntHeap) Len() int           { return len(h) }
+func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] }
+func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
+
+func (h *IntHeap) Push(x int) {
+	// Push and Pop use pointer receivers because they modify the slice's length,
+	// not just its contents.
+	*h = append(*h, x)
+}
+
+func (h *IntHeap) Pop() int {
+	old := *h
+	n := len(old)
+	x := old[n-1]
+	*h = old[0 : n-1]
+	return x
+}
+
+// This example inserts several ints into an IntHeap, checks the minimum,
+// and removes them in order of priority.
+func Example_intHeap() {
+	h := &IntHeap{2, 1, 5}
+	heap.Init[int](h)
+	heap.Push[int](h, 3)
+	fmt.Printf("minimum: %d\n", (*h)[0])
+	for h.Len() > 0 {
+		fmt.Printf("%d ", heap.Pop[int](h))
+	}
+	// Output:
+	// minimum: 1
+	// 1 2 3 5
+}

pkg/util/container/heap/example_pq_test.go

@@ -0,0 +1,109 @@
+// Copyright 2023 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+// Copyright 2012 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// This example demonstrates a priority queue built using the heap interface.
+package heap_test
+
+import (
+	"fmt"
+
+	"github.com/cockroachdb/cockroach/pkg/util/container/heap"
+)
+
+// An Item is something we manage in a priority queue.
+type Item struct {
+	value    string // The value of the item; arbitrary.
+	priority int    // The priority of the item in the queue.
+	// The index is needed by update and is maintained by the heap.Interface methods.
+	index int // The index of the item in the heap.
+}
+
+// A PriorityQueue implements heap.Interface and holds Items.
+type PriorityQueue []*Item
+
+func (pq PriorityQueue) Len() int { return len(pq) }
+
+func (pq PriorityQueue) Less(i, j int) bool {
+	// We want Pop to give us the highest, not lowest, priority so we use greater than here.
+	return pq[i].priority > pq[j].priority
+}
+
+func (pq PriorityQueue) Swap(i, j int) {
+	pq[i], pq[j] = pq[j], pq[i]
+	pq[i].index = i
+	pq[j].index = j
+}
+
+func (pq *PriorityQueue) Push(x *Item) {
+	n := len(*pq)
+	item := x
+	item.index = n
+	*pq = append(*pq, item)
+}
+
+func (pq *PriorityQueue) Pop() *Item {
+	old := *pq
+	n := len(old)
+	item := old[n-1]
+	old[n-1] = nil  // avoid memory leak
+	item.index = -1 // for safety
+	*pq = old[0 : n-1]
+	return item
+}
+
+// update modifies the priority and value of an Item in the queue.
+func (pq *PriorityQueue) update(item *Item, value string, priority int) {
+	item.value = value
+	item.priority = priority
+	heap.Fix[*Item](pq, item.index)
+}
+
+// This example creates a PriorityQueue with some items, adds and manipulates an item,
+// and then removes the items in priority order.
+func Example_priorityQueue() {
+	// Some items and their priorities.
+	items := map[string]int{
+		"banana": 3, "apple": 2, "pear": 4,
+	}
+
+	// Create a priority queue, put the items in it, and
+	// establish the priority queue (heap) invariants.
+	pq := make(PriorityQueue, len(items))
+	i := 0
+	for value, priority := range items {
+		pq[i] = &Item{
+			value:    value,
+			priority: priority,
+			index:    i,
+		}
+		i++
+	}
+	heap.Init[*Item](&pq)
+
+	// Insert a new item and then modify its priority.
+	item := &Item{
+		value:    "orange",
+		priority: 1,
+	}
+	heap.Push[*Item](&pq, item)
+	pq.update(item, item.value, 5)
+
+	// Take the items out; they arrive in decreasing priority order.
+	for pq.Len() > 0 {
+		item := heap.Pop[*Item](&pq)
+		fmt.Printf("%.2d:%s ", item.priority, item.value)
+	}
+	// Output:
+	// 05:orange 04:pear 03:banana 02:apple
+}

pkg/util/container/heap/heap.go

@@ -0,0 +1,128 @@
+// Copyright 2023 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+// Copyright 2009 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package heap provides heap operations for any type that implements
+// heap.Interface. A heap is a tree with the property that each node is the
+// minimum-valued node in its subtree.
+//
+// The minimum element in the tree is the root, at index 0.
+//
+// A heap is a common way to implement a priority queue. To build a priority
+// queue, implement the Heap interface with the (negative) priority as the
+// ordering for the Less method, so Push adds items while Pop removes the
+// highest-priority item from the queue. The Examples include such an
+// implementation; the file example_pq_test.go has the complete source.
+package heap
+
+import "sort"
+
+// The Interface type describes the requirements
+// for a type using the routines in this package.
+// Any type that implements it may be used as a
+// min-heap with the following invariants (established after
+// Init has been called or if the data is empty or sorted):
+//
+//	!h.Less(j, i) for 0 <= i < h.Len() and 2*i+1 <= j <= 2*i+2 and j < h.Len()
+//
+// Note that Push and Pop in this interface are for package heap's
+// implementation to call. To add and remove things from the heap,
+// use heap.Push and heap.Pop.
+type Interface[T any] interface {
+	sort.Interface
+	Push(x T) // add x as element Len()
+	Pop() T   // remove and return element Len() - 1.
+}
+
+// Init establishes the heap invariants required by the other routines in this package.
+// Init is idempotent with respect to the heap invariants
+// and may be called whenever the heap invariants may have been invalidated.
+// The complexity is O(n) where n = h.Len().
+func Init[T any](h Interface[T]) {
+	// heapify
+	n := h.Len()
+	for i := n/2 - 1; i >= 0; i-- {
+		down(h, i, n)
+	}
+}
+
+// Push pushes the element x onto the heap.
+// The complexity is O(log n) where n = h.Len().
+func Push[T any](h Interface[T], x T) {
+	h.Push(x)
+	up(h, h.Len()-1)
+}
+
+// Pop removes and returns the minimum element (according to Less) from the heap.
+// The complexity is O(log n) where n = h.Len().
+// Pop is equivalent to Remove(h, 0).
+func Pop[T any](h Interface[T]) T {
+	n := h.Len() - 1
+	h.Swap(0, n)
+	down(h, 0, n)
+	return h.Pop()
+}
+
+// Remove removes and returns the element at index i from the heap.
+// The complexity is O(log n) where n = h.Len().
+func Remove[T any](h Interface[T], i int) T {
+	n := h.Len() - 1
+	if n != i {
+		h.Swap(i, n)
+		if !down(h, i, n) {
+			up(h, i)
+		}
+	}
+	return h.Pop()
+}
+
+// Fix re-establishes the heap ordering after the element at index i has changed its value.
+// Changing the value of the element at index i and then calling Fix is equivalent to,
+// but less expensive than, calling Remove(h, i) followed by a Push of the new value.
+// The complexity is O(log n) where n = h.Len().
+func Fix[T any](h Interface[T], i int) {
+	if !down(h, i, h.Len()) {
+		up(h, i)
+	}
+}
+
+func up[T any](h Interface[T], j int) {
+	for {
+		i := (j - 1) / 2 // parent
+		if i == j || !h.Less(j, i) {
+			break
+		}
+		h.Swap(i, j)
+		j = i
+	}
+}
+
+func down[T any](h Interface[T], i0, n int) bool {
+	i := i0
+	for {
+		j1 := 2*i + 1
+		if j1 >= n || j1 < 0 { // j1 < 0 after int overflow
+			break
+		}
+		j := j1 // left child
+		if j2 := j1 + 1; j2 < n && h.Less(j2, j1) {
+			j = j2 // = 2*i + 2  // right child
+		}
+		if !h.Less(j, i) {
+			break
+		}
+		h.Swap(i, j)
+		i = j
+	}
+	return i > i0
+}