treesteps: allow passthrough

This change improves tresteps to allow eliding nodes like
`invalidating.Iter` or `keyspan.assertIter` from the hierarchy.

If we simply pass through the `NodeInfo` from the wrapped node, we
still register the elided node instead of the wrapped node (and thus
the wrapped node ops are not recorded). To fix this, we the `Node` to
`NodeInfo`.

Note that passthrough happens implicitly when a node embeds another
node and relies on the embedded `TreeSpepsNode()` implementation (like
`colblk.KeyspanIter` embeds `colblk.keyspanIter`).

Author: RaduBerinde

internal/treesteps/testdata/passthrough

@@ -0,0 +1,18 @@
+# Here B is wrapped inside an invisible node; verify that we see the B op.
+run
+----
+step 1/5:
+  A
+   └── B
+step 2/5:
+  A  ← Op
+   └── B
+step 3/5:
+  A  ← Op
+   └── B  ← Op
+step 4/5:
+  A  ← Op
+   └── B  ← Op done
+step 5/5:
+  A  ← Op done
+   └── B

internal/treesteps/tree_steps_off.go

@@ -23,7 +23,7 @@ type Node interface {
 }
 type NodeInfo struct{}
 
-func NodeInfof(format string, args ...any) NodeInfo                  { return NodeInfo{} }
+func NodeInfof(node Node, format string, args ...any) NodeInfo       { return NodeInfo{} }
 func (ni *NodeInfo) AddPropf(key string, format string, args ...any) {}
 func (ni *NodeInfo) AddChildren(nodes ...Node)                       {}
 

internal/treesteps/tree_steps_on.go

@@ -36,21 +36,21 @@ func StartRecording(root Node, name string, opts ...RecordingOption) *Recording
 	mu.Lock()
 	defer mu.Unlock()
 	mu.recordingInProgress.Store(true)
-	w := &Recording{
+	rec := &Recording{
 		name:         name,
 		maxTreeDepth: 20,
 		maxOpDepth:   10,
+		root:         root,
 	}
 	for _, o := range opts {
-		o(w)
+		o(rec)
 	}
 
 	if mu.nodeMap == nil {
 		mu.nodeMap = make(map[Node]*nodeState)
 	}
-	w.root = nodeStateLocked(w, root)
-	w.stepLockedf("initial")
-	return w
+	rec.stepLockedf("initial")
+	return rec
 }
 
 // RecordingOption is an optional argument to StartRecording.
@@ -106,20 +106,27 @@ func NodeUpdated(n Node, reason string) {
 
 // Node must be implemented by every node in the hierarchy.
 type Node interface {
+	// TreeStepsNode returns the information about the current state of the node.
+	// The NodeInfo normally has the receiver as the Node; but it can be a
+	// descendant in the case of passthrough nodes.
 	TreeStepsNode() NodeInfo
 }
 
 // NodeInfo contains the information that we present for each node.
 type NodeInfo struct {
+	node       Node
 	name       string
 	properties [][2]string
 	children   []Node
 }
 
 // NodeInfof returns a NodeInfo with the name initialized with a formatted
 // string.
-func NodeInfof(format string, args ...any) NodeInfo {
-	return NodeInfo{name: fmt.Sprintf(format, args...)}
+func NodeInfof(node Node, format string, args ...any) NodeInfo {
+	return NodeInfo{
+		node: node,
+		name: fmt.Sprintf(format, args...),
+	}
 }
 
 // AddPropf adds a property to the NodeInfo.
@@ -148,7 +155,7 @@ type Recording struct {
 	name         string
 	maxTreeDepth int
 	maxOpDepth   int
-	root         *nodeState
+	root         Node
 	steps        []Step
 }
 
@@ -173,7 +180,7 @@ func (r *Recording) Finish() Steps {
 func (r *Recording) stepLockedf(format string, args ...any) {
 	r.steps = append(r.steps, Step{
 		Name: fmt.Sprintf(format, args...),
-		Root: buildTree(r.root),
+		Root: buildTree(r, r.root, 0),
 	})
 }
 
@@ -328,39 +335,43 @@ var mu struct {
 	nodeMap             map[Node]*nodeState
 }
 
-func buildTree(n *nodeState) TreeNode {
-	var t TreeNode
-	info := n.node.TreeStepsNode()
-	n.name = info.name
-	t.Name = info.name
-	t.Properties = info.properties
-	if len(n.ops) > 0 {
-		t.Ops = make([]string, len(n.ops))
+func buildTree(rec *Recording, n Node, depth int) TreeNode {
+	info := n.TreeStepsNode()
+	// Normally these are the same; they differ n is a "passthrough" node.
+	n = info.node
+	ns := nodeStateLocked(rec, n)
+	ns.name = info.name
+	ns.depth = depth
+
+	t := TreeNode{
+		Name:       info.name,
+		Properties: info.properties,
+	}
+	if len(ns.ops) > 0 {
+		t.Ops = make([]string, len(ns.ops))
 		for i := range t.Ops {
-			t.Ops[i] = n.ops[i].details
-			if n.ops[i].state != "" {
-				t.Ops[i] += " " + n.ops[i].state
+			t.Ops[i] = ns.ops[i].details
+			if ns.ops[i].state != "" {
+				t.Ops[i] += " " + ns.ops[i].state
 			}
 		}
 	}
-	if n.depth < n.recording.maxTreeDepth {
-		for i := range info.children {
-			c := nodeStateLocked(n.recording, info.children[i])
-			c.depth = n.depth + 1
-			t.Children = append(t.Children, buildTree(c))
+	if depth < rec.maxTreeDepth {
+		for _, c := range info.children {
+			t.Children = append(t.Children, buildTree(rec, c, depth+1))
 		}
 	} else if len(info.children) > 0 {
 		t.HasHiddenChildren = true
 	}
 	return t
 }
 
-func nodeStateLocked(w *Recording, n Node) *nodeState {
+func nodeStateLocked(rec *Recording, n Node) *nodeState {
 	ns, ok := mu.nodeMap[n]
 	if !ok {
-		ns = &nodeState{recording: w, node: n}
+		ns = &nodeState{recording: rec, node: n}
 		mu.nodeMap[n] = ns
-	} else if w != ns.recording {
+	} else if rec != ns.recording {
 		panic(fmt.Sprintf("node %v part of multiple recordings", n))
 	}
 	return ns

internal/treesteps/tree_steps_test.go

@@ -70,7 +70,7 @@ var _ Node = (*SegmentNode)(nil)
 
 // TreeStepsNode implements the Node interface.
 func (n *SegmentNode) TreeStepsNode() NodeInfo {
-	info := NodeInfof("n%d", n.Index)
+	info := NodeInfof(n, "n%d", n.Index)
 	info.AddPropf("range", "[%d, %d]", n.Start, n.End)
 	info.AddPropf("sum", "%d", n.Sum)
 	if n.Start < n.End {
@@ -178,3 +178,59 @@ func TestSegmentTree(t *testing.T) {
 		return ""
 	})
 }
+
+// TestSegmentTreePassthrough verifies that operations work as expected when the
+// hierarchy has passthrough nodes.
+func TestSegmentTreePassthrough(t *testing.T) {
+	if !Enabled {
+		t.Skip("treesteps not available in this build")
+	}
+	datadriven.RunTest(t, "testdata/passthrough", func(t *testing.T, td *datadriven.TestData) string {
+		a := &nodeA{child: &nodeBWrapper{}}
+		r := StartRecording(a, "foo")
+		a.Op()
+		steps := r.Finish()
+		return steps.String()
+	})
+}
+
+type nodeA struct {
+	child interface {
+		Node
+		Op()
+	}
+}
+
+func (n *nodeA) TreeStepsNode() NodeInfo {
+	ni := NodeInfof(n, "A")
+	ni.AddChildren(n.child)
+	return ni
+}
+
+func (n *nodeA) Op() {
+	if IsRecording(n) {
+		op := StartOpf(n, "Op")
+		defer op.Finishf("done")
+	}
+	n.child.Op()
+}
+
+type nodeB struct{}
+
+// nodeBWrapper is a passthrough node. When `(*nodeBWrapper).TreeStepsNode()` is called,
+// it is executed on `(*nodeB)`. Operations in `nodeB` would be associated with
+// the `*nodeB` node, not the `(*nodeBWrapper)` node.
+type nodeBWrapper struct {
+	nodeB
+}
+
+func (n *nodeB) Op() {
+	if IsRecording(n) {
+		op := StartOpf(n, "Op")
+		defer op.Finishf("done")
+	}
+}
+
+func (n *nodeB) TreeStepsNode() NodeInfo {
+	return NodeInfof(n, "B")
+}