types/lazy: add lazy.GMap: a map of lazily computed GValues (#16532)

Fixes tailscale/corp#30360

Signed-off-by: Simon Law <[email protected]>

cmd/stund/depaware.txt

@@ -76,7 +76,7 @@ tailscale.com/cmd/stund dependencies: (generated by github.com/tailscale/depawar
    L 💣 tailscale.com/util/dirwalk                                   from tailscale.com/metrics
         tailscale.com/util/dnsname                                   from tailscale.com/tailcfg
         tailscale.com/util/lineiter                                  from tailscale.com/version/distro
-        tailscale.com/util/mak                                       from tailscale.com/syncs
+        tailscale.com/util/mak                                       from tailscale.com/syncs+
         tailscale.com/util/nocasemaps                                from tailscale.com/types/ipproto
         tailscale.com/util/rands                                     from tailscale.com/tsweb
         tailscale.com/util/slicesx                                   from tailscale.com/tailcfg

types/lazy/map.go

@@ -0,0 +1,62 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package lazy
+
+import "tailscale.com/util/mak"
+
+// GMap is a map of lazily computed [GValue] pointers, keyed by a comparable
+// type.
+//
+// Use either Get or GetErr, depending on whether your fill function returns an
+// error.
+//
+// GMap is not safe for concurrent use.
+type GMap[K comparable, V any] struct {
+	store map[K]*GValue[V]
+}
+
+// Len returns the number of entries in the map.
+func (s *GMap[K, V]) Len() int {
+	return len(s.store)
+}
+
+// Set attempts to set the value of k to v, and reports whether it succeeded.
+// Set only succeeds if k has never been called with Get/GetErr/Set before.
+func (s *GMap[K, V]) Set(k K, v V) bool {
+	z, ok := s.store[k]
+	if !ok {
+		z = new(GValue[V])
+		mak.Set(&s.store, k, z)
+	}
+	return z.Set(v)
+}
+
+// MustSet sets the value of k to v, or panics if k already has a value.
+func (s *GMap[K, V]) MustSet(k K, v V) {
+	if !s.Set(k, v) {
+		panic("Set after already filled")
+	}
+}
+
+// Get returns the value for k, computing it with fill if it's not already
+// present.
+func (s *GMap[K, V]) Get(k K, fill func() V) V {
+	z, ok := s.store[k]
+	if !ok {
+		z = new(GValue[V])
+		mak.Set(&s.store, k, z)
+	}
+	return z.Get(fill)
+}
+
+// GetErr returns the value for k, computing it with fill if it's not already
+// present.
+func (s *GMap[K, V]) GetErr(k K, fill func() (V, error)) (V, error) {
+	z, ok := s.store[k]
+	if !ok {
+		z = new(GValue[V])
+		mak.Set(&s.store, k, z)
+	}
+	return z.GetErr(fill)
+}

types/lazy/map_test.go

@@ -0,0 +1,95 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package lazy
+
+import (
+	"errors"
+	"testing"
+)
+
+func TestGMap(t *testing.T) {
+	var gm GMap[string, int]
+	n := int(testing.AllocsPerRun(1000, func() {
+		got := gm.Get("42", fortyTwo)
+		if got != 42 {
+			t.Fatalf("got %v; want 42", got)
+		}
+	}))
+	if n != 0 {
+		t.Errorf("allocs = %v; want 0", n)
+	}
+}
+
+func TestGMapErr(t *testing.T) {
+	var gm GMap[string, int]
+	n := int(testing.AllocsPerRun(1000, func() {
+		got, err := gm.GetErr("42", func() (int, error) {
+			return 42, nil
+		})
+		if got != 42 || err != nil {
+			t.Fatalf("got %v, %v; want 42, nil", got, err)
+		}
+	}))
+	if n != 0 {
+		t.Errorf("allocs = %v; want 0", n)
+	}
+
+	var gmErr GMap[string, int]
+	wantErr := errors.New("test error")
+	n = int(testing.AllocsPerRun(1000, func() {
+		got, err := gmErr.GetErr("42", func() (int, error) {
+			return 0, wantErr
+		})
+		if got != 0 || err != wantErr {
+			t.Fatalf("got %v, %v; want 0, %v", got, err, wantErr)
+		}
+	}))
+	if n != 0 {
+		t.Errorf("allocs = %v; want 0", n)
+	}
+}
+
+func TestGMapSet(t *testing.T) {
+	var gm GMap[string, int]
+	if !gm.Set("42", 42) {
+		t.Fatalf("Set failed")
+	}
+	if gm.Set("42", 43) {
+		t.Fatalf("Set succeeded after first Set")
+	}
+	n := int(testing.AllocsPerRun(1000, func() {
+		got := gm.Get("42", fortyTwo)
+		if got != 42 {
+			t.Fatalf("got %v; want 42", got)
+		}
+	}))
+	if n != 0 {
+		t.Errorf("allocs = %v; want 0", n)
+	}
+}
+
+func TestGMapMustSet(t *testing.T) {
+	var gm GMap[string, int]
+	gm.MustSet("42", 42)
+	defer func() {
+		if e := recover(); e == nil {
+			t.Errorf("unexpected success; want panic")
+		}
+	}()
+	gm.MustSet("42", 43)
+}
+
+func TestGMapRecursivePanic(t *testing.T) {
+	defer func() {
+		if e := recover(); e != nil {
+			t.Logf("got panic, as expected")
+		} else {
+			t.Errorf("unexpected success; want panic")
+		}
+	}()
+	gm := GMap[string, int]{}
+	gm.Get("42", func() int {
+		return gm.Get("42", func() int { return 42 })
+	})
+}