From b752bde280863f2ab1936e7cf626884d7ab44da3 Mon Sep 17 00:00:00 2001
From: Maisem Ali <maisem@tailscale.com>
Date: Tue, 6 Feb 2024 09:57:40 -0800
Subject: [PATCH] types/views: add SliceMapKey[T]

views.Slice are meant to be immutable, and if used as such it
is at times desirable to use them as a key in a map. For non-viewed
slices it was kinda doable by creating a custom key struct but views.Slice
didn't allow for the same so add a method to create that struct here.

Updates tailscale/corp#17122

Signed-off-by: Maisem Ali <maisem@tailscale.com>
---
 types/views/views.go      | 33 +++++++++++++++++++++++++++++++++
 types/views/views_test.go | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+)

diff --git a/types/views/views.go b/types/views/views.go
index 5a1671250..be267c275 100644
--- a/types/views/views.go
+++ b/types/views/views.go
@@ -39,6 +39,9 @@ func ByteSliceOf[T ~[]byte](x T) ByteSlice[T] {
 	return ByteSlice[T]{x}
 }
 
+// MapKey returns a unique key for a slice, based on its address and length.
+func (v ByteSlice[T]) MapKey() SliceMapKey[byte] { return mapKey(v.ж) }
+
 // Len returns the length of the slice.
 func (v ByteSlice[T]) Len() int {
 	return len(v.ж)
@@ -168,6 +171,22 @@ func (v SliceView[T, V]) SliceTo(i int) SliceView[T, V] { return SliceView[T, V]
 // Slice returns v[i:j]
 func (v SliceView[T, V]) Slice(i, j int) SliceView[T, V] { return SliceView[T, V]{v.ж[i:j]} }
 
+// SliceMapKey represents a comparable unique key for a slice, based on its
+// address and length. It can be used to key maps by slices but should only be
+// used when the underlying slice is immutable.
+//
+// Empty and nil slices have different keys.
+type SliceMapKey[T any] struct {
+	// t is the address of the first element, or nil if the slice is nil or
+	// empty.
+	t *T
+	// n is the length of the slice, or -1 if the slice is nil.
+	n int
+}
+
+// MapKey returns a unique key for a slice, based on its address and length.
+func (v SliceView[T, V]) MapKey() SliceMapKey[T] { return mapKey(v.ж) }
+
 // AppendTo appends the underlying slice values to dst.
 func (v SliceView[T, V]) AppendTo(dst []V) []V {
 	for _, x := range v.ж {
@@ -190,6 +209,20 @@ type Slice[T any] struct {
 	ж []T
 }
 
+// MapKey returns a unique key for a slice, based on its address and length.
+func (v Slice[T]) MapKey() SliceMapKey[T] { return mapKey(v.ж) }
+
+// mapKey returns a unique key for a slice, based on its address and length.
+func mapKey[T any](x []T) SliceMapKey[T] {
+	if x == nil {
+		return SliceMapKey[T]{nil, -1}
+	}
+	if len(x) == 0 {
+		return SliceMapKey[T]{nil, 0}
+	}
+	return SliceMapKey[T]{&x[0], len(x)}
+}
+
 // SliceOf returns a Slice for the provided slice for immutable values.
 // It is the caller's responsibility to make sure V is immutable.
 func SliceOf[T any](x []T) Slice[T] {
diff --git a/types/views/views_test.go b/types/views/views_test.go
index 2098b9f47..5b3f07733 100644
--- a/types/views/views_test.go
+++ b/types/views/views_test.go
@@ -166,3 +166,40 @@ func TestSliceEqual(t *testing.T) {
 		t.Error("got a[:2] == a[:1]")
 	}
 }
+
+// TestSliceMapKey tests that the MapKey method returns the same key for slices
+// with the same underlying slice and different keys for different slices or
+// with same underlying slice but different bounds.
+func TestSliceMapKey(t *testing.T) {
+	underlying := []string{"foo", "bar"}
+	nilSlice := SliceOf[string](nil)
+	empty := SliceOf([]string{})
+	u1 := SliceOf(underlying)
+	u2 := SliceOf(underlying)
+	u3 := SliceOf([]string{"foo", "bar"}) // different underlying slice
+
+	sub1 := u1.Slice(0, 1)
+	sub2 := u1.Slice(1, 2)
+	sub3 := u1.Slice(0, 2)
+
+	wantSame := []Slice[string]{u1, u2, sub3}
+	for i := 1; i < len(wantSame); i++ {
+		s0, si := wantSame[0], wantSame[i]
+		k0 := s0.MapKey()
+		ki := si.MapKey()
+		if ki != k0 {
+			t.Fatalf("wantSame[%d, %+v, %q) != wantSame[0, %+v, %q)", i, ki, si.AsSlice(), k0, s0.AsSlice())
+		}
+	}
+
+	wantDiff := []Slice[string]{nilSlice, empty, sub1, sub2, sub3, u3}
+	for i := 0; i < len(wantDiff); i++ {
+		for j := i + 1; j < len(wantDiff); j++ {
+			si, sj := wantDiff[i], wantDiff[j]
+			ki, kj := wantDiff[i].MapKey(), wantDiff[j].MapKey()
+			if ki == kj {
+				t.Fatalf("wantDiff[%d, %+v, %q] == wantDiff[%d, %+v, %q] ", i, ki, si.AsSlice(), j, kj, sj.AsSlice())
+			}
+		}
+	}
+}