From 000a57689d1e6552f09adee1647d2d6825babc7f Mon Sep 17 00:00:00 2001
From: Brian Picciano <mediocregopher@gmail.com>
Date: Fri, 8 Feb 2019 18:44:20 -0500
Subject: [PATCH] mctx: implement beginnings of Annotations functionality

---
 mctx/annotate.go      | 71 +++++++++++++++++++++++++++++++++++++++++++
 mctx/annotate_test.go | 55 +++++++++++++++++++++++++++++++++
 mctx/ctx.go           | 21 +++++++++----
 3 files changed, 141 insertions(+), 6 deletions(-)
 create mode 100644 mctx/annotate.go
 create mode 100644 mctx/annotate_test.go

diff --git a/mctx/annotate.go b/mctx/annotate.go
new file mode 100644
index 0000000..b2d72e9
--- /dev/null
+++ b/mctx/annotate.go
@@ -0,0 +1,71 @@
+package mctx
+
+import (
+	"context"
+	"fmt"
+)
+
+type annotateKey struct {
+	userKey interface{}
+}
+
+// Annotate takes in one or more key/value pairs (kvs' length must be even) and
+// returns a Context carrying them. Annotations only exist on the local level,
+// i.e. a child and parent share different annotation namespaces.
+func Annotate(ctx context.Context, kvs ...interface{}) context.Context {
+	for i := 0; i < len(kvs); i += 2 {
+		ctx = WithLocalValue(ctx, annotateKey{kvs[i]}, kvs[i+1])
+	}
+	return ctx
+}
+
+// Annotations describes a set of keys/values which were set on a Context (but
+// not its parents or children) using Annotate.
+type Annotations [][2]interface{}
+
+// LocalAnnotations returns all key/value pairs which have been set via Annotate
+// on this Context (but not its parent or children). If a key was set twice then
+// only the most recent value is included. The returned order is
+// non-deterministic.
+func LocalAnnotations(ctx context.Context) Annotations {
+	var annotations Annotations
+	localValuesIter(ctx, func(key, val interface{}) {
+		aKey, ok := key.(annotateKey)
+		if !ok {
+			return
+		}
+		annotations = append(annotations, [2]interface{}{aKey.userKey, val})
+	})
+	return annotations
+}
+
+// StringMap formats each of the key/value pairs into strings using fmt.Sprint.
+// If any two keys format to the same string, then type information will be
+// prefaced to each one.
+func (aa Annotations) StringMap() map[string]string {
+	keyTypes := make(map[string]interface{}, len(aa))
+	keyVals := map[string]string{}
+	setKey := func(k, v interface{}) {
+		kStr := fmt.Sprint(k)
+		oldType := keyTypes[kStr]
+		if oldType == nil {
+			keyTypes[kStr] = k
+			keyVals[kStr] = fmt.Sprint(v)
+			return
+		}
+
+		// check if oldKey is in kvm, if so it needs to be moved to account for
+		// its type info
+		if oldV, ok := keyVals[kStr]; ok {
+			delete(keyVals, kStr)
+			keyVals[fmt.Sprintf("%T(%s)", oldType, kStr)] = oldV
+		}
+
+		keyVals[fmt.Sprintf("%T(%s)", k, kStr)] = fmt.Sprint(v)
+	}
+
+	for _, kv := range aa {
+		setKey(kv[0], kv[1])
+	}
+	return keyVals
+}
diff --git a/mctx/annotate_test.go b/mctx/annotate_test.go
new file mode 100644
index 0000000..3db3c04
--- /dev/null
+++ b/mctx/annotate_test.go
@@ -0,0 +1,55 @@
+package mctx
+
+import (
+	"context"
+	. "testing"
+
+	"github.com/mediocregopher/mediocre-go-lib/mtest/massert"
+)
+
+func TestAnnotate(t *T) {
+	parent := context.Background()
+	parent = Annotate(parent, "a", "foo")
+	parent = Annotate(parent, "b", "bar")
+
+	child := NewChild(parent, "child")
+	child = Annotate(child, "a", "FOO")
+	child = Annotate(child, "c", "BAZ")
+	parent = WithChild(parent, child)
+
+	parentAnnotations := LocalAnnotations(parent)
+	childAnnotations := LocalAnnotations(child)
+	massert.Fatal(t, massert.All(
+		massert.Len(parentAnnotations, 2),
+		massert.Has(parentAnnotations, [2]interface{}{"a", "foo"}),
+		massert.Has(parentAnnotations, [2]interface{}{"b", "bar"}),
+		massert.Len(childAnnotations, 2),
+		massert.Has(childAnnotations, [2]interface{}{"a", "FOO"}),
+		massert.Has(childAnnotations, [2]interface{}{"c", "BAZ"}),
+	))
+}
+
+func TestAnnotationsStingMap(t *T) {
+	type A int
+	type B int
+	aa := Annotations{
+		{"foo", "bar"},
+		{"1", "one"},
+		{1, 1},
+		{0, 0},
+		{A(0), 0},
+		{B(0), 0},
+	}
+
+	err := massert.Equal(map[string]string{
+		"foo":       "bar",
+		"string(1)": "one",
+		"int(1)":    "1",
+		"int(0)":    "0",
+		"mctx.A(0)": "0",
+		"mctx.B(0)": "0",
+	}, aa.StringMap()).Assert()
+	if err != nil {
+		t.Fatal(err)
+	}
+}
diff --git a/mctx/ctx.go b/mctx/ctx.go
index de4e3a8..304163f 100644
--- a/mctx/ctx.go
+++ b/mctx/ctx.go
@@ -218,17 +218,26 @@ func LocalValue(ctx context.Context, key interface{}) interface{} {
 	}
 }
 
-// LocalValues returns all key/value pairs which have been set on the Context
-// via WithLocalValue.
-func LocalValues(ctx context.Context) map[interface{}]interface{} {
-	m := map[interface{}]interface{}{}
+func localValuesIter(ctx context.Context, callback func(key, val interface{})) {
+	m := map[interface{}]struct{}{}
 	lv, _ := ctx.Value(localValsKey(0)).(*localVal)
 	for {
 		if lv == nil {
-			return m
+			return
 		} else if _, ok := m[lv.key]; !ok {
-			m[lv.key] = lv.val
+			callback(lv.key, lv.val)
+			m[lv.key] = struct{}{}
 		}
 		lv = lv.prev
 	}
 }
+
+// LocalValues returns all key/value pairs which have been set on the Context
+// via WithLocalValue.
+func LocalValues(ctx context.Context) map[interface{}]interface{} {
+	m := map[interface{}]interface{}{}
+	localValuesIter(ctx, func(key, val interface{}) {
+		m[key] = val
+	})
+	return m
+}