From f0ff4fad7401afe1600c631391c647e29d176191 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Wed, 9 Jul 2014 16:43:36 -0700
Subject: [PATCH] flatmap: add Keys() and Merge()

---
 flatmap/map.go      | 53 +++++++++++++++++++++++++++++++++++++----
 flatmap/map_test.go | 58 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 106 insertions(+), 5 deletions(-)

diff --git a/flatmap/map.go b/flatmap/map.go
index cd9498074..3035baba7 100644
--- a/flatmap/map.go
+++ b/flatmap/map.go
@@ -16,13 +16,56 @@ type Map map[string]string
 func (m Map) Delete(prefix string) {
 	for k, _ := range m {
 		match := k == prefix
-		if !match && !strings.HasPrefix(k, prefix) {
-			continue
-		}
-		if k[len(prefix):len(prefix)+1] != "." {
-			continue
+		if !match {
+			if !strings.HasPrefix(k, prefix) {
+				continue
+			}
+
+			if k[len(prefix):len(prefix)+1] != "." {
+				continue
+			}
 		}
 
 		delete(m, k)
 	}
 }
+
+// Keys returns all of the top-level keys in this map
+func (m Map) Keys() []string {
+	ks := make(map[string]struct{})
+	for k, _ := range m {
+		idx := strings.Index(k, ".")
+		if idx == -1 {
+			idx = len(k)
+		}
+
+		ks[k[:idx]] = struct{}{}
+	}
+
+	result := make([]string, 0, len(ks))
+	for k, _ := range ks {
+		result = append(result, k)
+	}
+
+	return result
+}
+
+// Merge merges the contents of the other Map into this one.
+//
+// This merge is smarter than a simple map iteration because it
+// will fully replace arrays and other complex structures that
+// are present in this map with the other map's. For example, if
+// this map has a 3 element "foo" list, and m2 has a 2 element "foo"
+// list, then the result will be that m has a 2 element "foo"
+// list.
+func (m Map) Merge(m2 Map) {
+	for _, prefix := range m2.Keys() {
+		m.Delete(prefix)
+
+		for k, v := range m2 {
+			if strings.HasPrefix(k, prefix) {
+				m[k] = v
+			}
+		}
+	}
+}
diff --git a/flatmap/map_test.go b/flatmap/map_test.go
index 4df4ef410..bb0f44f29 100644
--- a/flatmap/map_test.go
+++ b/flatmap/map_test.go
@@ -22,3 +22,61 @@ func TestMapDelete(t *testing.T) {
 		t.Fatalf("bad: %#v", m)
 	}
 }
+
+func TestMapKeys(t *testing.T) {
+	cases := []struct {
+		Input  map[string]string
+		Output []string
+	}{
+		{
+			Input: map[string]string{
+				"foo":       "bar",
+				"bar.#":     "bar",
+				"bar.0.foo": "bar",
+				"bar.0.baz": "bar",
+			},
+			Output: []string{
+				"foo",
+				"bar",
+			},
+		},
+	}
+
+	for _, tc := range cases {
+		actual := Map(tc.Input).Keys()
+		if !reflect.DeepEqual(actual, tc.Output) {
+			t.Fatalf("input: %#v\n\nbad: %#v", tc.Input, actual)
+		}
+	}
+}
+
+func TestMapMerge(t *testing.T) {
+	cases := []struct {
+		One    map[string]string
+		Two    map[string]string
+		Result map[string]string
+	}{
+		{
+			One: map[string]string{
+				"foo": "bar",
+				"bar": "nope",
+			},
+			Two: map[string]string{
+				"bar": "baz",
+				"baz": "buz",
+			},
+			Result: map[string]string{
+				"foo": "bar",
+				"bar": "baz",
+				"baz": "buz",
+			},
+		},
+	}
+
+	for i, tc := range cases {
+		Map(tc.One).Merge(Map(tc.Two))
+		if !reflect.DeepEqual(tc.One, tc.Result) {
+			t.Fatalf("case %d bad: %#v", i, tc.One)
+		}
+	}
+}