1 files changed, 55 insertions, 62 deletions
diff --git a/pkg/state/stats.go b/pkg/state/stats.go
index eb51cda47..eaec664a1 100644
--- a/pkg/state/stats.go
+++ b/pkg/state/stats.go
@@ -17,7 +17,6 @@ package state
 import (
 	"bytes"
 	"fmt"
-	"reflect"
 	"sort"
 	"time"
 )
@@ -35,92 +34,81 @@ type statEntry struct {
 // All exported receivers accept nil.
 type Stats struct {
 	// byType contains a breakdown of time spent by type.
-	byType map[reflect.Type]*statEntry
+	//
+	// This is indexed *directly* by typeID, including zero.
+	byType []statEntry
 
 	// stack contains objects in progress.
-	stack []reflect.Type
+	stack []typeID
+
+	// names contains type names.
+	//
+	// This is also indexed *directly* by typeID, including zero, which we
+	// hard-code as "state.default". This is only resolved by calling fini
+	// on the stats object.
+	names []string
 
 	// last is the last start time.
 	last time.Time
 }
 
-// sample adds the samples to the given object.
-func (s *Stats) sample(typ reflect.Type) {
-	now := time.Now()
-	s.byType[typ].total += now.Sub(s.last)
-	s.last = now
+// init initializes statistics.
+func (s *Stats) init() {
+	s.last = time.Now()
+	s.stack = append(s.stack, 0)
 }
 
-// Add adds a sample count.
-func (s *Stats) Add(obj reflect.Value) {
-	if s == nil {
-		return
-	}
-	if s.byType == nil {
-		s.byType = make(map[reflect.Type]*statEntry)
-	}
-	typ := obj.Type()
-	entry, ok := s.byType[typ]
-	if !ok {
-		entry = new(statEntry)
-		s.byType[typ] = entry
+// fini finalizes statistics.
+func (s *Stats) fini(resolve func(id typeID) string) {
+	s.done()
+
+	// Resolve all type names.
+	s.names = make([]string, len(s.byType))
+	s.names[0] = "state.default" // See above.
+	for id := typeID(1); int(id) < len(s.names); id++ {
+		s.names[id] = resolve(id)
 	}
-	entry.count++
 }
 
-// Remove removes a sample count. It should only be called after a previous
-// Add().
-func (s *Stats) Remove(obj reflect.Value) {
-	if s == nil {
-		return
+// sample adds the samples to the given object.
+func (s *Stats) sample(id typeID) {
+	now := time.Now()
+	if len(s.byType) <= int(id) {
+		// Allocate all the missing entries in one fell swoop.
+		s.byType = append(s.byType, make([]statEntry, 1+int(id)-len(s.byType))...)
 	}
-	typ := obj.Type()
-	entry := s.byType[typ]
-	entry.count--
+	s.byType[id].total += now.Sub(s.last)
+	s.last = now
 }
 
-// Start starts a sample.
-func (s *Stats) Start(obj reflect.Value) {
-	if s == nil {
-		return
-	}
-	if len(s.stack) > 0 {
-		last := s.stack[len(s.stack)-1]
-		s.sample(last)
-	} else {
-		// First time sample.
-		s.last = time.Now()
-	}
-	s.stack = append(s.stack, obj.Type())
+// start starts a sample.
+func (s *Stats) start(id typeID) {
+	last := s.stack[len(s.stack)-1]
+	s.sample(last)
+	s.stack = append(s.stack, id)
 }
 
-// Done finishes the current sample.
-func (s *Stats) Done() {
-	if s == nil {
-		return
-	}
+// done finishes the current sample.
+func (s *Stats) done() {
 	last := s.stack[len(s.stack)-1]
 	s.sample(last)
+	s.byType[last].count++
 	s.stack = s.stack[:len(s.stack)-1]
 }
 
 type sliceEntry struct {
-	typ   reflect.Type
+	name  string
 	entry *statEntry
 }
 
 // String returns a table representation of the stats.
 func (s *Stats) String() string {
-	if s == nil || len(s.byType) == 0 {
-		return "(no data)"
-	}
-
 	// Build a list of stat entries.
 	ss := make([]sliceEntry, 0, len(s.byType))
-	for typ, entry := range s.byType {
+	for id := 0; id < len(s.names); id++ {
 		ss = append(ss, sliceEntry{
-			typ:   typ,
-			entry: entry,
+			name:  s.names[id],
+			entry: &s.byType[id],
 		})
 	}
 
@@ -136,17 +124,22 @@ func (s *Stats) String() string {
 		total time.Duration
 	)
 	buf.WriteString("\n")
-	buf.WriteString(fmt.Sprintf("%12s | %8s | %8s | %s\n", "total", "count", "per", "type"))
-	buf.WriteString("-------------+----------+----------+-------------\n")
+	buf.WriteString(fmt.Sprintf("% 16s | % 8s | % 16s | %s\n", "total", "count", "per", "type"))
+	buf.WriteString("-----------------+----------+------------------+----------------\n")
 	for _, se := range ss {
+		if se.entry.count == 0 {
+			// Since we store all types linearly, we are not
+			// guaranteed that any entry actually has time.
+			continue
+		}
 		count += se.entry.count
 		total += se.entry.total
 		per := se.entry.total / time.Duration(se.entry.count)
-		buf.WriteString(fmt.Sprintf("%12s | %8d | %8s | %s\n",
-			se.entry.total, se.entry.count, per, se.typ.String()))
+		buf.WriteString(fmt.Sprintf("% 16s | %8d | % 16s | %s\n",
+			se.entry.total, se.entry.count, per, se.name))
 	}
-	buf.WriteString("-------------+----------+----------+-------------\n")
-	buf.WriteString(fmt.Sprintf("%12s | %8d | %8s | [all]",
+	buf.WriteString("-----------------+----------+------------------+----------------\n")
+	buf.WriteString(fmt.Sprintf("% 16s | % 8d | % 16s | [all]",
 		total, count, total/time.Duration(count)))
 	return string(buf.Bytes())
 }