Adds correct trimming with heuristics for trees

internal/graph/graph.go

@@ -18,6 +18,7 @@ package graph
 import (
 	"fmt"
 	"math"
+	"os"
 	"path/filepath"
 	"sort"
 	"strconv"
@@ -142,14 +143,17 @@ func (i *NodeInfo) NameComponents() []string {
 type NodeMap map[NodeInfo]*Node
 
 // NodeSet maps is a collection of node info structs.
-type NodeSet map[NodeInfo]bool
+type NodeSet struct {
+	Info map[NodeInfo]bool
+	Ptr  map[*Node]bool
+}
 
 // FindOrInsertNode takes the info for a node and either returns a matching node
 // from the node map if one exists, or adds one to the map if one does not.
 // If kept is non-nil, nodes are only added if they can be located on it.
 func (nm NodeMap) FindOrInsertNode(info NodeInfo, kept NodeSet) *Node {
-	if kept != nil {
-		if _, ok := kept[info]; !ok {
+	if kept.Info != nil {
+		if _, ok := kept.Info[info]; !ok {
 			return nil
 		}
 	}
@@ -331,6 +335,61 @@ func newTree(prof *profile.Profile, o *Options) (g *Graph) {
 	return selectNodesForGraph(nodes, o.DropNegative)
 }
 
+// Trims a Graph that is in forest form to contain only the nodes in kept. This
+// will not work correctly in the case that a node has multiple parents.
+func (g *Graph) TrimTree(kept NodeSet) *Graph {
+	// Creates a new list of nodes
+	oldNodes := g.Nodes
+	g.Nodes = make(Nodes, 0, len(kept.Ptr))
+
+	for _, cur := range oldNodes {
+		// A node may not have multiple parents
+		if len(cur.In) > 1 {
+			fmt.Fprintf(os.Stderr, "ERROR: TrimTree only works on trees.\n")
+		}
+
+		// If a node should be kept, add it to the next list of nodes
+		if _, ok := kept.Ptr[cur]; ok {
+			g.Nodes = append(g.Nodes, cur)
+			continue
+		}
+
+		// Get the parent. Since cur.In may only be of size 0 or 1, parent will be
+		// equal to either nil or the only node in cur.In
+		var parent *Node
+		for _, edge := range cur.In {
+			parent = edge.Src
+		}
+
+		if parent != nil {
+			// Remove the edge from the parent to this node
+			delete(parent.Out, cur)
+
+			// Reconfigure every edge from the current node to now begin at the parent.
+			for _, outEdge := range cur.Out {
+				child := outEdge.Dest
+
+				delete(child.In, cur)
+				child.In[parent] = outEdge
+				parent.Out[child] = outEdge
+
+				outEdge.Src = parent
+				outEdge.Residual = true
+				// Any reconfigured edge can no longer be Inline.
+				outEdge.Inline = false
+			}
+		} else {
+			// If a node has no parents, delete all the in edges of the children to make them
+			// all roots of their own trees.
+			for _, outEdge := range cur.Out {
+				delete(outEdge.Dest.In, cur)
+			}
+		}
+	}
+	g.RemoveRedundantEdges()
+	return g
+}
+
 func joinLabels(s *profile.Sample) string {
 	if len(s.Label) == 0 {
 		return ""
@@ -538,12 +597,16 @@ func (g *Graph) DiscardLowFrequencyNodes(nodeCutoff int64) NodeSet {
 }
 
 func makeNodeSet(nodes Nodes, nodeCutoff int64) NodeSet {
-	kept := make(NodeSet, len(nodes))
+	kept := NodeSet{
+		Info: make(map[NodeInfo]bool, len(nodes)),
+		Ptr:  make(map[*Node]bool, len(nodes)),
+	}
 	for _, n := range nodes {
 		if abs64(n.Cum) < nodeCutoff {
 			continue
 		}
-		kept[n.Info] = true
+		kept.Info[n.Info] = true
+		kept.Ptr[n] = true
 	}
 	return kept
 }

internal/graph/graph_test.go

@@ -0,0 +1,298 @@
+package graph
+
+import (
+	"fmt"
+	"testing"
+)
+
+func edgeDebugString(edge *Edge) string {
+	debug := ""
+	debug += fmt.Sprintf("\t\tSrc: %p\n", edge.Src)
+	debug += fmt.Sprintf("\t\tDest: %p\n", edge.Dest)
+	debug += fmt.Sprintf("\t\tResidual: %t\n", edge.Residual)
+	debug += fmt.Sprintf("\t\tInline: %t\n", edge.Inline)
+	return debug
+}
+
+func edgeMapsDebugString(in, out EdgeMap) string {
+	debug := ""
+	debug += "In Edges:\n"
+	for parent, edge := range in {
+		debug += fmt.Sprintf("\tParent: %p\n", parent)
+		debug += edgeDebugString(edge)
+	}
+	debug += "Out Edges:\n"
+	for child, edge := range out {
+		debug += fmt.Sprintf("\tChild: %p\n", child)
+		debug += edgeDebugString(edge)
+	}
+	return debug
+}
+
+func graphDebugString(graph *Graph) string {
+	debug := ""
+	for i, node := range graph.Nodes {
+		debug += fmt.Sprintf("Node %d: %p\n", i, node)
+	}
+
+	for i, node := range graph.Nodes {
+		debug += "\n"
+		debug += fmt.Sprintf("===  Node %d: %p  ===\n", i, node)
+		debug += edgeMapsDebugString(node.In, node.Out)
+	}
+	return debug
+}
+
+func expectedNodesDebugString(Expected []ExpectedNode) string {
+	debug := ""
+	for i, node := range Expected {
+		debug += fmt.Sprintf("Node %d: %p\n", i, node.Node)
+	}
+
+	for i, node := range Expected {
+		debug += "\n"
+		debug += fmt.Sprintf("===  Node %d: %p  ===\n", i, node.Node)
+		debug += edgeMapsDebugString(node.In, node.Out)
+	}
+	return debug
+}
+
+// Checks if two edges are equal
+func edgesEqual(this, that *Edge) bool {
+	return this.Src == that.Src && this.Dest == that.Dest &&
+		this.Residual == that.Residual && this.Inline == that.Inline
+}
+
+// Checks if all the edges in this equal all the edges in that.
+func edgeMapsEqual(this, that EdgeMap) bool {
+	if len(this) != len(that) {
+		return false
+	}
+	for node, thisEdge := range this {
+		if !edgesEqual(thisEdge, that[node]) {
+			return false
+		}
+	}
+	return true
+}
+
+// Check if node is equal to Expected
+func nodesEqual(node *Node, Expected ExpectedNode) bool {
+	return node == Expected.Node && edgeMapsEqual(node.In, Expected.In) &&
+		edgeMapsEqual(node.Out, Expected.Out)
+}
+
+// Check if the graph equals the one templated by Expected.
+func graphsEqual(graph *Graph, Expected []ExpectedNode) bool {
+	if len(graph.Nodes) != len(Expected) {
+		return false
+	}
+	ExpectedSet := make(map[*Node]ExpectedNode)
+	for i := range Expected {
+		ExpectedSet[Expected[i].Node] = Expected[i]
+	}
+
+	for _, node := range graph.Nodes {
+		ExpectedNode, found := ExpectedSet[node]
+		if !found || !nodesEqual(node, ExpectedNode) {
+			return false
+		}
+	}
+	return true
+}
+
+type ExpectedNode struct {
+	Node    *Node
+	In, Out EdgeMap
+}
+
+type TrimTreeTestCase struct {
+	Initial  *Graph
+	Expected []ExpectedNode
+	Keep     NodeSet
+}
+
+// Makes the edge from parent to child residual
+func makeExpectedEdgeResidual(parent, child ExpectedNode) {
+	parent.Out[child.Node].Residual = true
+	child.In[parent.Node].Residual = true
+}
+
+// Creates a directed edges from the parent to each of the children
+func createEdges(parent *Node, children ...*Node) {
+	for _, child := range children {
+		edge := &Edge{
+			Src:  parent,
+			Dest: child,
+		}
+		parent.Out[child] = edge
+		child.In[parent] = edge
+	}
+}
+
+// Creates a node without any edges
+func createEmptyNode() *Node {
+	return &Node{
+		In:  make(EdgeMap),
+		Out: make(EdgeMap),
+	}
+}
+
+// Creates an array of ExpectedNodes from nodes.
+func createExpectedNodes(nodes ...*Node) ([]ExpectedNode, NodeSet) {
+	Expected := make([]ExpectedNode, len(nodes))
+	Keep := NodeSet{
+		Ptr: make(map[*Node]bool, len(nodes)),
+	}
+
+	for i, node := range nodes {
+		Expected[i] = ExpectedNode{
+			Node: node,
+			In:   make(EdgeMap),
+			Out:  make(EdgeMap),
+		}
+		Keep.Ptr[node] = true
+	}
+
+	return Expected, Keep
+}
+
+// Creates a directed edges from the parent to each of the children
+func createExpectedEdges(parent ExpectedNode, children ...ExpectedNode) {
+	for _, child := range children {
+		edge := &Edge{
+			Src:  parent.Node,
+			Dest: child.Node,
+		}
+		parent.Out[child.Node] = edge
+		child.In[parent.Node] = edge
+	}
+}
+
+// The first test case looks like:
+//     0
+//     |
+//     1
+//   /   \
+//  2     3
+//
+// After Keeping 0, 2, 3. We should see:
+//     0
+//   /   \
+//  2     3
+func createTestCase1() TrimTreeTestCase {
+	// Create Initial graph
+	graph := &Graph{make(Nodes, 4)}
+	nodes := graph.Nodes
+	for i := range nodes {
+		nodes[i] = createEmptyNode()
+	}
+	createEdges(nodes[0], nodes[1])
+	createEdges(nodes[1], nodes[2], nodes[3])
+
+	// Create Expected graph
+	Expected, Keep := createExpectedNodes(nodes[0], nodes[2], nodes[3])
+	createExpectedEdges(Expected[0], Expected[1], Expected[2])
+	makeExpectedEdgeResidual(Expected[0], Expected[1])
+	makeExpectedEdgeResidual(Expected[0], Expected[2])
+	return TrimTreeTestCase{
+		Initial:  graph,
+		Expected: Expected,
+		Keep:     Keep,
+	}
+}
+
+// This test case looks like:
+//   3
+//   |
+//   1
+//   |
+//   2
+//   |
+//   0
+//   |
+//   4
+//
+// After Keeping 3 and 4. We should see:
+//   3
+//   |
+//   4
+func createTestCase2() TrimTreeTestCase {
+	// Create Initial graph
+	graph := &Graph{make(Nodes, 5)}
+	nodes := graph.Nodes
+	for i := range nodes {
+		nodes[i] = createEmptyNode()
+	}
+	createEdges(nodes[3], nodes[1])
+	createEdges(nodes[1], nodes[2])
+	createEdges(nodes[2], nodes[0])
+	createEdges(nodes[0], nodes[4])
+
+	// Create Expected graph
+	Expected, Keep := createExpectedNodes(nodes[3], nodes[4])
+	createExpectedEdges(Expected[0], Expected[1])
+	makeExpectedEdgeResidual(Expected[0], Expected[1])
+	return TrimTreeTestCase{
+		Initial:  graph,
+		Expected: Expected,
+		Keep:     Keep,
+	}
+}
+
+// If we trim an empty graph it should still be empty afterwards
+func createTestCase3() TrimTreeTestCase {
+	graph := &Graph{make(Nodes, 0)}
+	Expected, Keep := createExpectedNodes()
+	return TrimTreeTestCase{
+		Initial:  graph,
+		Expected: Expected,
+		Keep:     Keep,
+	}
+}
+
+// This test case looks like:
+//   0
+//
+// After Keeping 0. We should see:
+//   0
+func createTestCase4() TrimTreeTestCase {
+	graph := &Graph{make(Nodes, 1)}
+	nodes := graph.Nodes
+	for i := range nodes {
+		nodes[i] = createEmptyNode()
+	}
+	Expected, Keep := createExpectedNodes(nodes[0])
+	return TrimTreeTestCase{
+		Initial:  graph,
+		Expected: Expected,
+		Keep:     Keep,
+	}
+}
+
+func createTrimTreeTestCases() []TrimTreeTestCase {
+	caseGenerators := []func() TrimTreeTestCase{
+		createTestCase1,
+		createTestCase2,
+		createTestCase3,
+		createTestCase4,
+	}
+	cases := make([]TrimTreeTestCase, len(caseGenerators))
+	for i, gen := range caseGenerators {
+		cases[i] = gen()
+	}
+	return cases
+}
+
+func TestTrimTree(t *testing.T) {
+	tests := createTrimTreeTestCases()
+	for _, test := range tests {
+		graph := test.Initial
+		graph.TrimTree(test.Keep)
+		if !graphsEqual(graph, test.Expected) {
+			t.Fatalf("Graphs do not match.\nExpected: %s\nFound: %s\n",
+				expectedNodesDebugString(test.Expected),
+				graphDebugString(graph))
+		}
+	}
+}

internal/report/report.go

@@ -79,26 +79,20 @@ func (rpt *Report) newTrimmedGraph() (g *graph.Graph, origCount, droppedNodes, d
 	cumSort := o.CumSort
 
 	// First step: Build complete graph to identify low frequency nodes, based on their cum weight.
-	g = rpt.newGraph(nil)
+	g = rpt.newGraph(graph.NodeSet{nil, nil})
 	totalValue, _ := g.Nodes.Sum()
 	nodeCutoff := abs64(int64(float64(totalValue) * o.NodeFraction))
 	edgeCutoff := abs64(int64(float64(totalValue) * o.EdgeFraction))
 
-	// Do not apply edge cutoff to preserve tree structure.
-	if o.CallTree {
-		if o.OutputFormat == Dot {
-			fmt.Println("WARNING: Trimming trees is unsupported.")
-			fmt.Printf("Tree will contain at least %d nodes\n", o.NodeCount)
-			cumSort = true
-		}
-		edgeCutoff = 0
-	}
-
 	// Filter out nodes with cum value below nodeCutoff.
 	if nodeCutoff > 0 {
-		if nodesKept := g.DiscardLowFrequencyNodes(nodeCutoff); len(g.Nodes) != len(nodesKept) {
-			droppedNodes = len(g.Nodes) - len(nodesKept)
-			g = rpt.newGraph(nodesKept)
+		if nodesKept := g.DiscardLowFrequencyNodes(nodeCutoff); len(g.Nodes) != len(nodesKept.Ptr) {
+			droppedNodes = len(g.Nodes) - len(nodesKept.Ptr)
+			if o.CallTree {
+				g = g.TrimTree(nodesKept)
+			} else {
+				g = rpt.newGraph(nodesKept)
+			}
 		}
 	}
 	origCount = len(g.Nodes)
@@ -110,8 +104,12 @@ func (rpt *Report) newTrimmedGraph() (g *graph.Graph, origCount, droppedNodes, d
 		// Remove low frequency tags and edges as they affect selection.
 		g.TrimLowFrequencyTags(nodeCutoff)
 		g.TrimLowFrequencyEdges(edgeCutoff)
-		if nodesKept := g.SelectTopNodes(nodeCount, visualMode); len(nodesKept) != len(g.Nodes) {
-			g = rpt.newGraph(nodesKept)
+		if nodesKept := g.SelectTopNodes(nodeCount, visualMode); len(nodesKept.Ptr) != len(g.Nodes) {
+			if o.CallTree {
+				g = g.TrimTree(nodesKept)
+			} else {
+				g = rpt.newGraph(nodesKept)
+			}
 			g.SortNodes(cumSort, visualMode)
 		}
 	}
@@ -272,7 +270,7 @@ func printAssembly(w io.Writer, rpt *Report, obj plugin.ObjTool) error {
 	o := rpt.options
 	prof := rpt.prof
 
-	g := rpt.newGraph(nil)
+	g := rpt.newGraph(graph.NodeSet{nil, nil})
 
 	// If the regexp source can be parsed as an address, also match
 	// functions that land on that address.
@@ -580,7 +578,7 @@ func printTraces(w io.Writer, rpt *Report) error {
 
 	const separator = "-----------+-------------------------------------------------------"
 
-	_, locations := graph.CreateNodes(prof, false, nil)
+	_, locations := graph.CreateNodes(prof, false, graph.NodeSet{nil, nil})
 	for _, sample := range prof.Sample {
 		var stack graph.Nodes
 		for _, loc := range sample.Location {

internal/report/source.go

@@ -37,7 +37,7 @@ import (
 // eliminate potential nondeterminism.
 func printSource(w io.Writer, rpt *Report) error {
 	o := rpt.options
-	g := rpt.newGraph(nil)
+	g := rpt.newGraph(graph.NodeSet{nil, nil})
 
 	// Identify all the functions that match the regexp provided.
 	// Group nodes for each matching function.
@@ -117,7 +117,7 @@ func printSource(w io.Writer, rpt *Report) error {
 // functions with samples that match the regexp rpt.options.symbol.
 func printWebSource(w io.Writer, rpt *Report, obj plugin.ObjTool) error {
 	o := rpt.options
-	g := rpt.newGraph(nil)
+	g := rpt.newGraph(graph.NodeSet{nil, nil})
 
 	// If the regexp source can be parsed as an address, also match
 	// functions that land on that address.

internal/report/testdata/source.dot

@@ -1,17 +1,17 @@
 digraph "unnamed" {
 node [style=filled fillcolor="#f8f8f8"]
 subgraph cluster_L { "Duration: 10s, Total samples = 11111 " [shape=box fontsize=16 label="Duration: 10s, Total samples = 11111 \lShowing nodes accounting for 11111, 100% of 11111 total\l"] }
-N1 [label="main\nsource1:2\n1 (0.009%)\nof 11111 (100%)" fontsize=9 shape=box tooltip="main testdata/source1:2 (11111)" color="#b20000" fillcolor="#edd5d5"]
-N2 [label="tee\nsource2:2\n1000 (9.00%)\nof 11000 (99.00%)" fontsize=14 shape=box tooltip="tee testdata/source2:2 (11000)" color="#b20000" fillcolor="#edd5d5"]
-N3 [label="tee\nsource2:8\n10000 (90.00%)" fontsize=24 shape=box tooltip="tee testdata/source2:8 (10000)" color="#b20500" fillcolor="#edd6d5"]
-N4 [label="bar\nsource1:10\n0 of 100 (0.9%)" fontsize=8 shape=box tooltip="bar testdata/source1:10 (100)" color="#b2b0aa" fillcolor="#edecec"]
-N5 [label="tee\nsource2:8\n100 (0.9%)" fontsize=10 shape=box tooltip="tee testdata/source2:8 (100)" color="#b2b0aa" fillcolor="#edecec"]
-N6 [label="bar\nsource1:10\n10 (0.09%)" fontsize=9 shape=box tooltip="bar testdata/source1:10 (10)" color="#b2b2b1" fillcolor="#ededed"]
+N1 [label="tee\nsource2:8\n10000 (90.00%)" fontsize=24 shape=box tooltip="tee testdata/source2:8 (10000)" color="#b20500" fillcolor="#edd6d5"]
+N2 [label="main\nsource1:2\n1 (0.009%)\nof 11111 (100%)" fontsize=9 shape=box tooltip="main testdata/source1:2 (11111)" color="#b20000" fillcolor="#edd5d5"]
+N3 [label="tee\nsource2:2\n1000 (9.00%)\nof 11000 (99.00%)" fontsize=14 shape=box tooltip="tee testdata/source2:2 (11000)" color="#b20000" fillcolor="#edd5d5"]
+N4 [label="tee\nsource2:8\n100 (0.9%)" fontsize=10 shape=box tooltip="tee testdata/source2:8 (100)" color="#b2b0aa" fillcolor="#edecec"]
+N5 [label="bar\nsource1:10\n10 (0.09%)" fontsize=9 shape=box tooltip="bar testdata/source1:10 (10)" color="#b2b2b1" fillcolor="#ededed"]
+N6 [label="bar\nsource1:10\n0 of 100 (0.9%)" fontsize=8 shape=box tooltip="bar testdata/source1:10 (100)" color="#b2b0aa" fillcolor="#edecec"]
 N7 [label="foo\nsource1:4\n0 of 10 (0.09%)" fontsize=8 shape=box tooltip="foo testdata/source1:4 (10)" color="#b2b2b1" fillcolor="#ededed"]
-N1 -> N2 [label=" 11000" weight=100 penwidth=5 color="#b20000" tooltip="main testdata/source1:2 -> tee testdata/source2:2 (11000)" labeltooltip="main testdata/source1:2 -> tee testdata/source2:2 (11000)"]
-N2 -> N3 [label=" 10000" weight=91 penwidth=5 color="#b20500" tooltip="tee testdata/source2:2 -> tee testdata/source2:8 (10000)" labeltooltip="tee testdata/source2:2 -> tee testdata/source2:8 (10000)"]
-N4 -> N5 [label=" 100" color="#b2b0aa" tooltip="bar testdata/source1:10 -> tee testdata/source2:8 (100)" labeltooltip="bar testdata/source1:10 -> tee testdata/source2:8 (100)"]
-N1 -> N4 [label=" 100" color="#b2b0aa" tooltip="main testdata/source1:2 -> bar testdata/source1:10 (100)" labeltooltip="main testdata/source1:2 -> bar testdata/source1:10 (100)"]
-N7 -> N6 [label=" 10" color="#b2b2b1" tooltip="foo testdata/source1:4 -> bar testdata/source1:10 (10)" labeltooltip="foo testdata/source1:4 -> bar testdata/source1:10 (10)"]
-N1 -> N7 [label=" 10" color="#b2b2b1" tooltip="main testdata/source1:2 -> foo testdata/source1:4 (10)" labeltooltip="main testdata/source1:2 -> foo testdata/source1:4 (10)"]
+N2 -> N3 [label=" 11000" weight=100 penwidth=5 color="#b20000" tooltip="main testdata/source1:2 -> tee testdata/source2:2 (11000)" labeltooltip="main testdata/source1:2 -> tee testdata/source2:2 (11000)"]
+N3 -> N1 [label=" 10000" weight=91 penwidth=5 color="#b20500" tooltip="tee testdata/source2:2 -> tee testdata/source2:8 (10000)" labeltooltip="tee testdata/source2:2 -> tee testdata/source2:8 (10000)"]
+N6 -> N4 [label=" 100" color="#b2b0aa" tooltip="bar testdata/source1:10 -> tee testdata/source2:8 (100)" labeltooltip="bar testdata/source1:10 -> tee testdata/source2:8 (100)"]
+N2 -> N6 [label=" 100" color="#b2b0aa" tooltip="main testdata/source1:2 -> bar testdata/source1:10 (100)" labeltooltip="main testdata/source1:2 -> bar testdata/source1:10 (100)"]
+N7 -> N5 [label=" 10" color="#b2b2b1" tooltip="foo testdata/source1:4 -> bar testdata/source1:10 (10)" labeltooltip="foo testdata/source1:4 -> bar testdata/source1:10 (10)"]
+N2 -> N7 [label=" 10" color="#b2b2b1" tooltip="main testdata/source1:2 -> foo testdata/source1:4 (10)" labeltooltip="main testdata/source1:2 -> foo testdata/source1:4 (10)"]
 }