Support different display granularities in flame graph view. (#896)

* Support different display granularities in flame graph view.

Previously, flame-graph view could only display filefunctions
granularity. We can now support all the available granularities.

This allows the user to supply a granularity either via a command-line
flag (e.g., --files), or as a URL parameter (e.g., ?g=files).

Details
* Made the default initial granularity "" so that we can stick with
  the default flame-graph granularity of "filefunctions" while also
  allowing overrides.
* Use different computation of the list of shorter display names
  for file names (strip off leading path components instead of package
  prefixes).
* Similarly compute color based on directory name instead of package name.
* Include line number and column number in displayed names if the
  granularity includes line numbers.

* Improve granularity entry for options command

---------

Co-authored-by: Alexey Alexandrov <aalexand@users.noreply.github.com>

internal/driver/config.go

@@ -68,7 +68,7 @@ func defaultConfig() config {
 		Trim:         true,
 		DivideBy:     1.0,
 		Sort:         "flat",
-		Granularity:  "functions",
+		Granularity:  "", // Default depends on the display format
 	}
 }
 

internal/driver/driver.go

@@ -254,6 +254,8 @@ func aggregate(prof *profile.Profile, cfg config) error {
 	var function, filename, linenumber, address bool
 	inlines := !cfg.NoInlines
 	switch cfg.Granularity {
+	case "":
+		function = true // Default granularity is "functions"
 	case "addresses":
 		if inlines {
 			return nil

internal/driver/interactive.go

@@ -206,6 +206,9 @@ func printCurrentOptions(p *profile.Profile, ui plugin.UI) {
 			// Add quotes for empty values.
 			v = `""`
 		}
+		if n == "granularity" && v == "" {
+			v = "(default)"
+		}
 		if comment != "" {
 			comment = commentStart + " " + comment
 		}

internal/driver/interactive_test.go

@@ -186,40 +186,36 @@ func TestInteractiveCommands(t *testing.T) {
 		{
 			"top 10 --cum focus1 -ignore focus2",
 			map[string]string{
-				"granularity": "functions",
-				"nodecount":   "10",
-				"sort":        "cum",
-				"focus":       "focus1|focus2",
-				"ignore":      "ignore",
+				"nodecount": "10",
+				"sort":      "cum",
+				"focus":     "focus1|focus2",
+				"ignore":    "ignore",
 			},
 		},
 		{
 			"top10 --cum focus1 -ignore focus2",
 			map[string]string{
-				"granularity": "functions",
-				"nodecount":   "10",
-				"sort":        "cum",
-				"focus":       "focus1|focus2",
-				"ignore":      "ignore",
+				"nodecount": "10",
+				"sort":      "cum",
+				"focus":     "focus1|focus2",
+				"ignore":    "ignore",
 			},
 		},
 		{
 			"dot",
 			map[string]string{
-				"granularity": "functions",
-				"nodecount":   "80",
-				"sort":        "flat",
+				"nodecount": "80",
+				"sort":      "flat",
 			},
 		},
 		{
 			"tags   -ignore1 -ignore2 focus1 >out",
 			map[string]string{
-				"granularity": "functions",
-				"nodecount":   "80",
-				"sort":        "flat",
-				"output":      "out",
-				"tagfocus":    "focus1",
-				"tagignore":   "ignore1|ignore2",
+				"nodecount": "80",
+				"sort":      "flat",
+				"output":    "out",
+				"tagfocus":  "focus1",
+				"tagignore": "ignore1|ignore2",
 			},
 		},
 		{

internal/driver/stacks.go

@@ -28,7 +28,9 @@ func (ui *webInterface) stackView(w http.ResponseWriter, req *http.Request) {
 	rpt, errList := ui.makeReport(w, req, []string{"svg"}, func(cfg *config) {
 		cfg.CallTree = true
 		cfg.Trim = false
-		cfg.Granularity = "filefunctions"
+		if cfg.Granularity == "" {
+			cfg.Granularity = "filefunctions"
+		}
 	})
 	if rpt == nil {
 		return // error already reported

internal/report/shortnames.go

@@ -15,18 +15,34 @@
 package report
 
 import (
+	"path/filepath"
 	"regexp"
 
 	"github.com/google/pprof/internal/graph"
 )
 
-var sepRE = regexp.MustCompile(`::|\.`)
+var (
+	sepRE     = regexp.MustCompile(`::|\.`)
+	fileSepRE = regexp.MustCompile(`/`)
+)
+
+// fileNameSuffixes returns a non-empty sequence of shortened file names
+// (in decreasing preference) that can be used to represent name.
+func fileNameSuffixes(name string) []string {
+	return allSuffixes(filepath.ToSlash(filepath.Clean(name)), fileSepRE)
+}
 
 // shortNameList returns a non-empty sequence of shortened names
 // (in decreasing preference) that can be used to represent name.
 func shortNameList(name string) []string {
 	name = graph.ShortenFunctionName(name)
-	seps := sepRE.FindAllStringIndex(name, -1)
+	return allSuffixes(name, sepRE)
+}
+
+// allSuffixes returns a list of suffixes (in order of decreasing length)
+// found by splitting at re.
+func allSuffixes(name string, re *regexp.Regexp) []string {
+	seps := re.FindAllStringIndex(name, -1)
 	result := make([]string, 0, len(seps)+1)
 	result = append(result, name)
 	for _, sep := range seps {

internal/report/shortnames_test.go

@@ -32,3 +32,29 @@ func TestShortNames(t *testing.T) {
 		})
 	}
 }
+
+func TestFileNameSuffixes(t *testing.T) {
+	type testCase struct {
+		name string
+		in   string
+		out  []string
+	}
+	test := func(name, in string, out ...string) testCase {
+		return testCase{name, in, out}
+	}
+
+	for _, c := range []testCase{
+		test("empty", "", "."),
+		test("simple", "foo", "foo"),
+		test("manypaths", "a/b/c", "a/b/c", "b/c", "c"),
+		test("leading", "/a/b", "/a/b", "a/b", "b"),
+		test("trailing", "a/b", "a/b", "b"),
+	} {
+		t.Run(c.name, func(t *testing.T) {
+			got := fileNameSuffixes(c.in)
+			if !reflect.DeepEqual(c.out, got) {
+				t.Errorf("fileNameSuffixes(%q) = %#v, expecting %#v", c.in, got, c.out)
+			}
+		})
+	}
+}

internal/report/stacks.go

@@ -18,6 +18,7 @@ import (
 	"crypto/sha256"
 	"encoding/binary"
 	"fmt"
+	"path/filepath"
 
 	"github.com/google/pprof/internal/measurement"
 	"github.com/google/pprof/profile"
@@ -99,41 +100,57 @@ func (rpt *Report) Stacks() StackSet {
 	}
 	s.makeInitialStacks(rpt)
 	s.fillPlaces()
-	s.assignColors()
 	return *s
 }
 
 func (s *StackSet) makeInitialStacks(rpt *Report) {
 	type key struct {
-		line    profile.Line
-		inlined bool
+		funcName string
+		fileName string
+		line     int64
+		column   int64
+		inlined  bool
 	}
 	srcs := map[key]int{} // Sources identified so far.
 	seenFunctions := map[string]bool{}
 	unknownIndex := 1
+
 	getSrc := func(line profile.Line, inlined bool) int {
-		k := key{line, inlined}
+		fn := line.Function
+		if fn == nil {
+			fn = &profile.Function{Name: fmt.Sprintf("?%d?", unknownIndex)}
+			unknownIndex++
+		}
+
+		k := key{fn.Name, fn.Filename, line.Line, line.Column, inlined}
 		if i, ok := srcs[k]; ok {
 			return i
 		}
-		x := StackSource{Places: []StackSlot{}} // Ensure Places is non-nil
-		if fn := line.Function; fn != nil {
-			x.FullName = fn.Name
-			x.FileName = fn.Filename
-			if !seenFunctions[fn.Name] {
-				x.UniqueName = fn.Name
-				seenFunctions[fn.Name] = true
-			} else {
-				// Assign a different name so pivoting picks this function.
-				x.UniqueName = fmt.Sprint(fn.Name, "#", fn.ID)
-			}
-		} else {
-			x.FullName = fmt.Sprintf("?%d?", unknownIndex)
+
+		fileName := trimPath(fn.Filename, rpt.options.TrimPath, rpt.options.SourcePath)
+		x := StackSource{
+			FileName: fileName,
+			Inlined:  inlined,
+			Places:   []StackSlot{}, // Ensure Places is non-nil
+		}
+		if fn.Name != "" {
+			x.FullName = addLineInfo(fn.Name, line)
+			x.Display = shortNameList(x.FullName)
+			x.Color = pickColor(packageName(fn.Name))
+		} else { // Use file name, e.g., for file granularity display.
+			x.FullName = addLineInfo(fileName, line)
+			x.Display = fileNameSuffixes(x.FullName)
+			x.Color = pickColor(filepath.Dir(fileName))
+		}
+
+		if !seenFunctions[x.FullName] {
 			x.UniqueName = x.FullName
-			unknownIndex++
+			seenFunctions[x.FullName] = true
+		} else {
+			// Assign a different name so pivoting picks this function.
+			x.UniqueName = fmt.Sprint(x.FullName, "#", fn.ID)
 		}
-		x.Inlined = inlined
-		x.Display = shortNameList(x.FullName)
+
 		s.Sources = append(s.Sources, x)
 		srcs[k] = len(s.Sources) - 1
 		return len(s.Sources) - 1
@@ -179,18 +196,25 @@ func (s *StackSet) fillPlaces() {
 	}
 }
 
-func (s *StackSet) assignColors() {
-	// Assign different color indices to different packages.
+// pickColor picks a color for key.
+func pickColor(key string) int {
 	const numColors = 1048576
-	for i, src := range s.Sources {
-		pkg := packageName(src.FullName)
-		h := sha256.Sum256([]byte(pkg))
-		index := binary.LittleEndian.Uint32(h[:])
-		s.Sources[i].Color = int(index % numColors)
-	}
+	h := sha256.Sum256([]byte(key))
+	index := binary.LittleEndian.Uint32(h[:])
+	return int(index % numColors)
 }
 
 // Legend returns the list of lines to display as the legend.
 func (s *StackSet) Legend() []string {
 	return reportLabels(s.report, s.report.total, len(s.Sources), len(s.Sources), 0, 0, false)
 }
+
+func addLineInfo(str string, line profile.Line) string {
+	if line.Column != 0 {
+		return fmt.Sprint(str, ":", line.Line, ":", line.Column)
+	}
+	if line.Line != 0 {
+		return fmt.Sprint(str, ":", line.Line)
+	}
+	return str
+}

internal/report/stacks_test.go

@@ -3,6 +3,7 @@ package report
 import (
 	"fmt"
 	"reflect"
+	"slices"
 	"strings"
 	"testing"
 
@@ -18,7 +19,13 @@ func makeTestStacks(samples ...*profile.Sample) StackSet {
 
 func TestStacks(t *testing.T) {
 	// See report_test.go for the functions available to use in tests.
-	main, foo, bar, tee := testL[0], testL[1], testL[2], testL[3]
+	locs := clearLineAndColumn(testL)
+	main, foo, bar, tee := locs[0], locs[1], locs[2], locs[3]
+
+	// Also make some file-only locations to test file granularity.
+	fileMain := makeFileLocation(main)
+	fileFoo := makeFileLocation(foo)
+	fileBar := makeFileLocation(bar)
 
 	// stack holds an expected stack value found in StackSet.
 	type stack struct {
@@ -57,6 +64,17 @@ func TestStacks(t *testing.T) {
 				makeStack(200, "0:root", "1:main", "2:foo", "2:foo", "3:bar"),
 			},
 		},
+		{
+			"files",
+			makeTestStacks(
+				testSample(100, fileFoo, fileMain),
+				testSample(200, fileBar, fileMain),
+			),
+			[]stack{
+				makeStack(100, "0:root", "1:dir/main", "2:dir/foo"),
+				makeStack(200, "0:root", "1:dir/main", "3:dir/bar"),
+			},
+		},
 	} {
 		t.Run(c.name, func(t *testing.T) {
 			var got []stack
@@ -79,7 +97,8 @@ func TestStacks(t *testing.T) {
 
 func TestStackSources(t *testing.T) {
 	// See report_test.go for the functions available to use in tests.
-	main, foo, bar, tee, inl := testL[0], testL[1], testL[2], testL[3], testL[5]
+	locs := clearLineAndColumn(testL)
+	main, foo, bar, tee, inl := locs[0], locs[1], locs[2], locs[3], locs[5]
 
 	type srcInfo struct {
 		name    string
@@ -189,3 +208,33 @@ func findSource(stacks StackSet, name string) StackSource {
 	}
 	return StackSource{}
 }
+
+// clearLineAndColumn drops line and column numbers to simplify tests that
+// do not care about line and column numbers.
+func clearLineAndColumn(locs []*profile.Location) []*profile.Location {
+	result := make([]*profile.Location, len(locs))
+	for i, loc := range locs {
+		newLoc := *loc
+		newLoc.Line = slices.Clone(loc.Line)
+		for j := range newLoc.Line {
+			newLoc.Line[j].Line = 0
+			newLoc.Line[j].Column = 0
+		}
+		result[i] = &newLoc
+	}
+	return result
+}
+
+// makeFileLocation switches loc from function to file-granularity.
+func makeFileLocation(loc *profile.Location) *profile.Location {
+	result := *loc
+	result.ID += 1000
+	result.Line = slices.Clone(loc.Line)
+	for i := range result.Line {
+		fn := *result.Line[i].Function
+		fn.Filename = "dir/" + fn.Name
+		fn.Name = ""
+		result.Line[i].Function = &fn
+	}
+	return &result
+}