Expose helper string->index logic for profile.Profile

internal/driver/commands.go

@@ -28,7 +28,6 @@ import (
 
 	"github.com/google/pprof/internal/plugin"
 	"github.com/google/pprof/internal/report"
-	"github.com/google/pprof/profile"
 	"github.com/google/pprof/third_party/svg"
 )
 
@@ -456,49 +455,6 @@ func invokeVisualizer(format PostProcessor, suffix string, visualizers []string)
 	}
 }
 
-// locateSampleIndex returns the appropriate index for a value of sample index.
-// If numeric, it returns the number, otherwise it looks up the text in the
-// profile sample types.
-func locateSampleIndex(p *profile.Profile, sampleIndex string) (int, error) {
-	if sampleIndex == "" {
-		if dst := p.DefaultSampleType; dst != "" {
-			for i, t := range sampleTypes(p) {
-				if t == dst {
-					return i, nil
-				}
-			}
-		}
-		// By default select the last sample value
-		return len(p.SampleType) - 1, nil
-	}
-	if i, err := strconv.Atoi(sampleIndex); err == nil {
-		if i < 0 || i >= len(p.SampleType) {
-			return 0, fmt.Errorf("sample_index %s is outside the range [0..%d]", sampleIndex, len(p.SampleType)-1)
-		}
-		return i, nil
-	}
-
-	// Remove the inuse_ prefix to support legacy pprof options
-	// "inuse_space" and "inuse_objects" for profiles containing types
-	// "space" and "objects".
-	noInuse := strings.TrimPrefix(sampleIndex, "inuse_")
-	for i, t := range p.SampleType {
-		if t.Type == sampleIndex || t.Type == noInuse {
-			return i, nil
-		}
-	}
-
-	return 0, fmt.Errorf("sample_index %q must be one of: %v", sampleIndex, sampleTypes(p))
-}
-
-func sampleTypes(p *profile.Profile) []string {
-	types := make([]string, len(p.SampleType))
-	for i, t := range p.SampleType {
-		types[i] = t.Type
-	}
-	return types
-}
-
 // variables describe the configuration parameters recognized by pprof.
 type variables map[string]*variable
 

internal/driver/driver.go

@@ -258,7 +258,7 @@ func sampleFormat(p *profile.Profile, sampleIndex string, mean bool) (value, mea
 	if len(p.SampleType) == 0 {
 		return nil, nil, nil, fmt.Errorf("profile has no samples")
 	}
-	index, err := locateSampleIndex(p, sampleIndex)
+	index, err := p.SampleIndexByName(sampleIndex)
 	if err != nil {
 		return nil, nil, nil, err
 	}

internal/driver/interactive.go

@@ -69,7 +69,7 @@ func interactive(p *profile.Profile, o *plugin.Options) error {
 				if v := pprofVariables[name]; v != nil {
 					if name == "sample_index" {
 						// Error check sample_index=xxx to ensure xxx is a valid sample type.
-						index, err := locateSampleIndex(p, value)
+						index, err := p.SampleIndexByName(value)
 						if err != nil {
 							o.UI.PrintErr(err)
 							continue
@@ -161,6 +161,14 @@ func profileShortcuts(p *profile.Profile) shortcuts {
 	return s
 }
 
+func sampleTypes(p *profile.Profile) []string {
+	types := make([]string, len(p.SampleType))
+	for i, t := range p.SampleType {
+		types[i] = t.Type
+	}
+	return types
+}
+
 func printCurrentOptions(p *profile.Profile, ui plugin.UI) {
 	var args []string
 	type groupInfo struct {

profile/index.go

@@ -0,0 +1,64 @@
+// Copyright 2016 Google Inc. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package profile
+
+import (
+	"fmt"
+	"strconv"
+	"strings"
+)
+
+// SampleIndexByName returns the appropriate index for a value of sample index.
+// If numeric, it returns the number, otherwise it looks up the text in the
+// profile sample types.
+func (p *Profile) SampleIndexByName(sampleIndex string) (int, error) {
+	if sampleIndex == "" {
+		if dst := p.DefaultSampleType; dst != "" {
+			for i, t := range sampleTypes(p) {
+				if t == dst {
+					return i, nil
+				}
+			}
+		}
+		// By default select the last sample value
+		return len(p.SampleType) - 1, nil
+	}
+	if i, err := strconv.Atoi(sampleIndex); err == nil {
+		if i < 0 || i >= len(p.SampleType) {
+			return 0, fmt.Errorf("sample_index %s is outside the range [0..%d]", sampleIndex, len(p.SampleType)-1)
+		}
+		return i, nil
+	}
+
+	// Remove the inuse_ prefix to support legacy pprof options
+	// "inuse_space" and "inuse_objects" for profiles containing types
+	// "space" and "objects".
+	noInuse := strings.TrimPrefix(sampleIndex, "inuse_")
+	for i, t := range p.SampleType {
+		if t.Type == sampleIndex || t.Type == noInuse {
+			return i, nil
+		}
+	}
+
+	return 0, fmt.Errorf("sample_index %q must be one of: %v", sampleIndex, sampleTypes(p))
+}
+
+func sampleTypes(p *Profile) []string {
+	types := make([]string, len(p.SampleType))
+	for i, t := range p.SampleType {
+		types[i] = t.Type
+	}
+	return types
+}

profile/index_test.go

@@ -0,0 +1,100 @@
+package profile
+
+import (
+	"testing"
+)
+
+func TestSampleIndexByName(t *testing.T) {
+	for _, c := range []struct {
+		desc              string
+		sampleTypes       []string
+		defaultSampleType string
+		index             string
+		want              int
+		wantError         bool
+	}{
+		{
+			desc:        "use last by default",
+			index:       "",
+			want:        1,
+			sampleTypes: []string{"zero", "default"},
+		},
+		{
+			desc:              "honour specified default",
+			index:             "",
+			want:              1,
+			defaultSampleType: "default",
+			sampleTypes:       []string{"zero", "default", "two"},
+		},
+		{
+			desc:              "invalid default is ignored",
+			index:             "",
+			want:              2,
+			defaultSampleType: "non-existent",
+			sampleTypes:       []string{"zero", "one", "default"},
+		},
+		{
+			desc:        "index by int",
+			index:       "0",
+			want:        0,
+			sampleTypes: []string{"zero", "one", "two"},
+		},
+		{
+			desc:              "index by int ignores default",
+			index:             "0",
+			want:              0,
+			defaultSampleType: "default",
+			sampleTypes:       []string{"zero", "default", "two"},
+		},
+		{
+			desc:        "index by name",
+			index:       "two",
+			want:        2,
+			sampleTypes: []string{"zero", "one", "two", "three"},
+		},
+		{
+			desc:              "index by name ignores default",
+			index:             "zero",
+			want:              0,
+			defaultSampleType: "default",
+			sampleTypes:       []string{"zero", "default", "two"},
+		},
+		{
+			desc:        "out of bound int causes error",
+			index:       "100",
+			wantError:   true,
+			sampleTypes: []string{"zero", "default"},
+		},
+		{
+			desc:        "unknown name causes error",
+			index:       "does not exist",
+			wantError:   true,
+			sampleTypes: []string{"zero", "default"},
+		},
+		{
+			desc:        "'inused_{x}' recognized for legacy '{x}'",
+			index:       "inuse_zero",
+			want:        0,
+			sampleTypes: []string{"zero", "default"},
+		},
+	} {
+		p := &Profile{
+			DefaultSampleType: c.defaultSampleType,
+			SampleType:        []*ValueType{},
+		}
+		for _, st := range c.sampleTypes {
+			p.SampleType = append(p.SampleType, &ValueType{Type: st, Unit: "milliseconds"})
+		}
+
+		got, err := p.SampleIndexByName(c.index)
+
+		switch {
+		case c.wantError && err == nil:
+			t.Errorf("%s: error should have been returned not index=%d, err=%v", c.desc, got, err)
+		case !c.wantError && err != nil:
+			t.Errorf("%s: unexpected got index=%d, err=%v; wanted index=%d, err=nil", c.desc, got, err, c.want)
+		case !c.wantError && got != c.want:
+			t.Errorf("%s: got index=%d, want index=%d", c.desc, got, c.want)
+		}
+	}
+}