From e9108b33f779b788069a7b36ac81615968b2da61 Mon Sep 17 00:00:00 2001
From: Varun Gandhi <varun.gandhi@sourcegraph.com>
Date: Tue, 25 Jun 2024 21:13:14 +0800
Subject: [PATCH] chore: Add collection type - OrderedSet (#63469)

It is missing from the stdlib and the orderedmap library we're using.
---
 internal/collections/BUILD.bazel         |  3 ++
 internal/collections/ordered_set.go      | 45 ++++++++++++++++++++++++
 internal/collections/ordered_set_test.go | 29 +++++++++++++++
 3 files changed, 77 insertions(+)
 create mode 100644 internal/collections/ordered_set.go
 create mode 100644 internal/collections/ordered_set_test.go

diff --git a/internal/collections/BUILD.bazel b/internal/collections/BUILD.bazel
index 24e53d33ae3..6b1843891b4 100644
--- a/internal/collections/BUILD.bazel
+++ b/internal/collections/BUILD.bazel
@@ -4,6 +4,7 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library")
 go_library(
     name = "collections",
     srcs = [
+        "ordered_set.go",
         "set.go",
         "slice_utils.go",
     ],
@@ -11,6 +12,7 @@ go_library(
     visibility = ["//:__subpackages__"],
     deps = [
         "//lib/errors",
+        "@com_github_wk8_go_ordered_map_v2//:go-ordered-map",
         "@org_golang_x_exp//constraints",
         "@org_golang_x_exp//maps",
     ],
@@ -20,6 +22,7 @@ go_test(
     name = "collections_test",
     timeout = "short",
     srcs = [
+        "ordered_set_test.go",
         "set_test.go",
         "slice_utils_test.go",
     ],
diff --git a/internal/collections/ordered_set.go b/internal/collections/ordered_set.go
new file mode 100644
index 00000000000..7f8a5706784
--- /dev/null
+++ b/internal/collections/ordered_set.go
@@ -0,0 +1,45 @@
+package collections
+
+import orderedmap "github.com/wk8/go-ordered-map/v2"
+
+// OrderedSet keeps track of values in insertion order.
+type OrderedSet[T comparable] orderedmap.OrderedMap[T, struct{}]
+
+// NewOrderedSet creates a OrderedSet[T] with the given values.
+// T must be a comparable type (implementing sort.Interface or == operator).
+func NewOrderedSet[T comparable](values ...T) *OrderedSet[T] {
+	s := OrderedSet[T](*orderedmap.New[T, struct{}]())
+	s.Add(values...)
+	return &s
+}
+
+func (s *OrderedSet[T]) impl() *orderedmap.OrderedMap[T, struct{}] {
+	return (*orderedmap.OrderedMap[T, struct{}])(s)
+}
+
+func (s *OrderedSet[T]) Add(values ...T) {
+	for _, v := range values {
+		s.impl().Set(v, struct{}{})
+	}
+}
+
+func (s *OrderedSet[T]) Remove(values ...T) {
+	for _, v := range values {
+		s.impl().Delete(v)
+	}
+}
+
+func (s *OrderedSet[T]) Has(value T) bool {
+	_, found := s.impl().Get(value)
+	return found
+}
+
+// Values returns a slice with all the values in the set.
+// The values are returned in insertion order.
+func (s *OrderedSet[T]) Values() []T {
+	out := make([]T, 0, s.impl().Len())
+	for x := s.impl().Oldest(); x != nil; x = x.Next() {
+		out = append(out, x.Key)
+	}
+	return out
+}
diff --git a/internal/collections/ordered_set_test.go b/internal/collections/ordered_set_test.go
new file mode 100644
index 00000000000..4e4cc2e4f60
--- /dev/null
+++ b/internal/collections/ordered_set_test.go
@@ -0,0 +1,29 @@
+package collections
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"pgregory.net/rapid"
+)
+
+func TestOrderedSet(t *testing.T) {
+	rapid.Check(t, func(t *rapid.T) {
+		data := rapid.SliceOfN(rapid.IntRange(-3, 6), 0, 10).Draw(t, "data")
+		set := NewSet(data...)
+		uniquedData := set.Values()
+		ordset := NewOrderedSet(uniquedData...)
+		require.Equal(t, uniquedData, ordset.Values())
+
+		otherData := rapid.SliceOfN(rapid.IntRange(-5, 5), 10, 10).Draw(t, "data")
+		for _, x := range otherData {
+			require.Equal(t, set.Has(x), ordset.Has(x))
+		}
+
+		for _, x := range uniquedData {
+			require.True(t, ordset.Has(x))
+			ordset.Remove(x)
+			require.False(t, ordset.Has(x))
+		}
+	})
+}