Extract map helper functions

report/latest_map.go

@@ -1,9 +1,7 @@
 package report
 
 import (
-	"bytes"
 	"fmt"
-	"sort"
 	"time"
 
 	"github.com/ugorji/go/codec"
@@ -60,34 +58,12 @@ func (m LatestMap) Size() int {
 // both inputs. When both inputs contain the same key, the newer value
 // is used.
 func (m LatestMap) Merge(other LatestMap) LatestMap {
-	var (
-		mSize     = m.Size()
-		otherSize = other.Size()
-		output    = m.Map
-		iter      = other.Map
-	)
-	switch {
-	case mSize == 0:
-		return other
-	case otherSize == 0:
-		return m
-	case mSize < otherSize:
-		output, iter = iter, output
-	}
 	if m.decoder != other.decoder {
 		panic(fmt.Sprintf("Cannot merge maps with different entry value types, this has %#v, other has %#v", m.decoder, other.decoder))
 	}
-
-	iter.ForEach(func(key string, iterVal interface{}) {
-		if existingVal, ok := output.Lookup(key); ok {
-			if existingVal.(LatestEntry).Timestamp.Before(iterVal.(LatestEntry).Timestamp) {
-				output = output.Set(key, iterVal)
-			}
-		} else {
-			output = output.Set(key, iterVal)
-		}
+	output := mergeMaps(m.Map, other.Map, func(a, b interface{}) bool {
+		return a.(LatestEntry).Timestamp.Before(b.(LatestEntry).Timestamp)
 	})
-
 	return LatestMap{output, m.decoder}
 }
 
@@ -138,44 +114,17 @@ func (m LatestMap) ForEach(fn func(k string, ts time.Time, v interface{})) {
 
 // String returns the LatestMap's string representation.
 func (m LatestMap) String() string {
-	keys := []string{}
-	if m.Map == nil {
-		m = MakeLatestMapWithDecoder(m.decoder)
-	}
-	for _, k := range m.Map.Keys() {
-		keys = append(keys, k)
-	}
-	sort.Strings(keys)
-
-	buf := bytes.NewBufferString("{")
-	for _, key := range keys {
-		val, _ := m.Map.Lookup(key)
-		fmt.Fprintf(buf, "%s: %s,\n", key, val)
-	}
-	fmt.Fprintf(buf, "}")
-	return buf.String()
+	return mapToString(m.Map)
 }
 
 // DeepEqual tests equality with other LatestMap.
 func (m LatestMap) DeepEqual(n LatestMap) bool {
-	if m.Size() != n.Size() {
-		return false
-	}
-	if m.Size() == 0 {
-		return true
-	}
 	if m.decoder != n.decoder {
 		panic(fmt.Sprintf("Cannot check equality of maps with different entry value types, this has %#v, other has %#v", m.decoder, n.decoder))
 	}
-	equal := true
-	m.Map.ForEach(func(k string, val interface{}) {
-		if otherValue, ok := n.Map.Lookup(k); !ok {
-			equal = false
-		} else {
-			equal = equal && val.(LatestEntry).Equal(otherValue.(LatestEntry))
-		}
+	return mapEqual(m.Map, n.Map, func(val, otherValue interface{}) bool {
+		return val.(LatestEntry).Equal(otherValue.(LatestEntry))
 	})
-	return equal
 }
 
 func (m LatestMap) toIntermediate() map[string]LatestEntry {
@@ -197,47 +146,15 @@ func (m *LatestMap) CodecEncodeSelf(encoder *codec.Encoder) {
 	}
 }
 
-// constants from https://github.com/ugorji/go/blob/master/codec/helper.go#L207
-const (
-	containerMapKey   = 2
-	containerMapValue = 3
-	containerMapEnd   = 4
-)
-
 // CodecDecodeSelf implements codec.Selfer.
-// This implementation does not use the intermediate form as that was a
-// performance issue; skipping it saved almost 10% CPU.  Note this means
-// we are using undocumented, internal APIs, which could break in the future.
-// See https://github.com/weaveworks/scope/pull/1709 for more information.
 func (m *LatestMap) CodecDecodeSelf(decoder *codec.Decoder) {
-	z, r := codec.GenHelperDecoder(decoder)
-	if r.TryDecodeAsNil() {
-		*m = MakeLatestMapWithDecoder(m.decoder)
-		return
-	}
-
-	length := r.ReadMapStart()
-	out := ps.NewMap()
-	for i := 0; length < 0 || i < length; i++ {
-		if length < 0 && r.CheckBreak() {
-			break
-		}
-
-		var key string
-		z.DecSendContainerState(containerMapKey)
-		if !r.TryDecodeAsNil() {
-			key = r.DecodeString()
-		}
-
-		var value LatestEntry
-		z.DecSendContainerState(containerMapValue)
-		if !r.TryDecodeAsNil() {
+	out := mapRead(decoder, func(isNil bool) interface{} {
+		value := LatestEntry{}
+		if !isNil {
 			m.decoder.Decode(decoder, &value)
 		}
-
-		out = out.UnsafeMutableSet(key, value)
-	}
-	z.DecSendContainerState(containerMapEnd)
+		return value
+	})
 	*m = LatestMap{out, m.decoder}
 }
 

report/map_helpers.go

@@ -0,0 +1,120 @@
+package report
+
+import (
+	"bytes"
+	"fmt"
+	"sort"
+
+	"github.com/ugorji/go/codec"
+	"github.com/weaveworks/ps"
+)
+
+// Helper functions for ps.Map, without considering what is inside
+
+// Return a new map containing all the elements of the two input maps
+// and where the same key is in both, pick 'b' where prefer(a,b) is true
+func mergeMaps(m, n ps.Map, prefer func(a, b interface{}) bool) ps.Map {
+	switch {
+	case m == nil:
+		return n
+	case n == nil:
+		return m
+	case m.Size() < n.Size():
+		m, n = n, m
+	}
+
+	n.ForEach(func(key string, val interface{}) {
+		if existingVal, found := m.Lookup(key); found {
+			if prefer(existingVal, val) {
+				m = m.Set(key, val)
+			}
+		} else {
+			m = m.Set(key, val)
+		}
+	})
+
+	return m
+}
+
+func mapEqual(m, n ps.Map, equalf func(a, b interface{}) bool) bool {
+	var mSize, nSize int
+	if m != nil {
+		mSize = m.Size()
+	}
+	if n != nil {
+		nSize = n.Size()
+	}
+	if mSize != nSize {
+		return false
+	}
+	if mSize == 0 {
+		return true
+	}
+	equal := true
+	m.ForEach(func(k string, val interface{}) {
+		if otherValue, ok := n.Lookup(k); !ok {
+			equal = false
+		} else {
+			equal = equal && equalf(val, otherValue)
+		}
+	})
+	return equal
+}
+
+// very similar to ps.Map.String() but with keys sorted
+func mapToString(m ps.Map) string {
+	keys := []string{}
+	if m != nil {
+		for _, k := range m.Keys() {
+			keys = append(keys, k)
+		}
+		sort.Strings(keys)
+	}
+
+	buf := bytes.NewBufferString("{")
+	for _, key := range keys {
+		val, _ := m.Lookup(key)
+		fmt.Fprintf(buf, "%s: %s,\n", key, val)
+	}
+	fmt.Fprintf(buf, "}")
+	return buf.String()
+}
+
+// constants from https://github.com/ugorji/go/blob/master/codec/helper.go#L207
+const (
+	containerMapKey   = 2
+	containerMapValue = 3
+	containerMapEnd   = 4
+)
+
+// CodecDecodeSelf implements codec.Selfer.
+// This implementation does not use the intermediate form as that was a
+// performance issue; skipping it saved almost 10% CPU.  Note this means
+// we are using undocumented, internal APIs, which could break in the future.
+// See https://github.com/weaveworks/scope/pull/1709 for more information.
+func mapRead(decoder *codec.Decoder, decodeValue func(isNil bool) interface{}) ps.Map {
+	z, r := codec.GenHelperDecoder(decoder)
+	if r.TryDecodeAsNil() {
+		return ps.NewMap()
+	}
+
+	length := r.ReadMapStart()
+	out := ps.NewMap()
+	for i := 0; length < 0 || i < length; i++ {
+		if length < 0 && r.CheckBreak() {
+			break
+		}
+
+		var key string
+		z.DecSendContainerState(containerMapKey)
+		if !r.TryDecodeAsNil() {
+			key = r.DecodeString()
+		}
+
+		z.DecSendContainerState(containerMapValue)
+		value := decodeValue(r.TryDecodeAsNil())
+		out = out.UnsafeMutableSet(key, value)
+	}
+	z.DecSendContainerState(containerMapEnd)
+	return out
+}