go/ast/inspector: add explicit parent pointer

This CL optimizes the Cursor.Parent and Cursor.Stack operations
by shrinking the event.index field to 32 bits (which is plenty)
and using the liberated space (in push nodes) to hold an explicit
parent pointer. This makes the CursorStack benchmark 4x faster,
bringing it to parity with WithStack. The cost is 8% slower NewInspector.

Before:
BenchmarkNewInspector-8    	    3435	   1695466 ns/op
BenchmarkInspectCalls/WithStack-8        	   44512	    140151 ns/op
BenchmarkInspectCalls/CursorStack-8      	   10000	    570207 ns/op

After:
BenchmarkNewInspector-8    	    3282	   1836795 ns/op
BenchmarkInspectCalls/WithStack-8        	   36932	    154218 ns/op
BenchmarkInspectCalls/CursorStack-8      	   37948	    155916 ns/op

Change-Id: I0607f5241de2db347cccb983e1ca4193572e6be1
Reviewed-on: https://go-review.googlesource.com/c/tools/+/638336
Auto-Submit: Alan Donovan <[email protected]>
Reviewed-by: Robert Findley <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>

Author: adonovan

go/ast/inspector/inspector.go

@@ -56,9 +56,10 @@ func New(files []*ast.File) *Inspector {
 // An event represents a push or a pop
 // of an ast.Node during a traversal.
 type event struct {
-	node  ast.Node
-	typ   uint64 // typeOf(node) on push event, or union of typ strictly between push and pop events on pop events
-	index int    // index of corresponding push or pop event
+	node   ast.Node
+	typ    uint64 // typeOf(node) on push event, or union of typ strictly between push and pop events on pop events
+	index  int32  // index of corresponding push or pop event
+	parent int32  // index of parent's push node (defined for push nodes only)
 }
 
 // TODO: Experiment with storing only the second word of event.node (unsafe.Pointer).
@@ -87,7 +88,7 @@ func (in *Inspector) Preorder(types []ast.Node, f func(ast.Node)) {
 	// })
 
 	mask := maskOf(types)
-	for i := 0; i < len(in.events); {
+	for i := int32(0); i < int32(len(in.events)); {
 		ev := in.events[i]
 		if ev.index > i {
 			// push
@@ -117,7 +118,7 @@ func (in *Inspector) Preorder(types []ast.Node, f func(ast.Node)) {
 // matches an element of the types slice.
 func (in *Inspector) Nodes(types []ast.Node, f func(n ast.Node, push bool) (proceed bool)) {
 	mask := maskOf(types)
-	for i := 0; i < len(in.events); {
+	for i := int32(0); i < int32(len(in.events)); {
 		ev := in.events[i]
 		if ev.index > i {
 			// push
@@ -151,7 +152,7 @@ func (in *Inspector) Nodes(types []ast.Node, f func(n ast.Node, push bool) (proc
 func (in *Inspector) WithStack(types []ast.Node, f func(n ast.Node, push bool, stack []ast.Node) (proceed bool)) {
 	mask := maskOf(types)
 	var stack []ast.Node
-	for i := 0; i < len(in.events); {
+	for i := int32(0); i < int32(len(in.events)); {
 		ev := in.events[i]
 		if ev.index > i {
 			// push
@@ -200,18 +201,24 @@ func traverse(files []*ast.File) []event {
 	events := make([]event, 0, capacity)
 
 	var stack []event
-	stack = append(stack, event{}) // include an extra event so file nodes have a parent
+	stack = append(stack, event{index: -1}) // include an extra event so file nodes have a parent
 	for _, f := range files {
 		ast.Inspect(f, func(n ast.Node) bool {
 			if n != nil {
 				// push
 				ev := event{
-					node:  n,
-					typ:   0,           // temporarily used to accumulate type bits of subtree
-					index: len(events), // push event temporarily holds own index
+					node:   n,
+					typ:    0,                  // temporarily used to accumulate type bits of subtree
+					index:  int32(len(events)), // push event temporarily holds own index
+					parent: stack[len(stack)-1].index,
 				}
 				stack = append(stack, ev)
 				events = append(events, ev)
+
+				// 2B nodes ought to be enough for anyone!
+				if int32(len(events)) < 0 {
+					panic("event index exceeded int32")
+				}
 			} else {
 				// pop
 				top := len(stack) - 1
@@ -220,9 +227,9 @@ func traverse(files []*ast.File) []event {
 				push := ev.index
 				parent := top - 1
 
-				events[push].typ = typ            // set type of push
-				stack[parent].typ |= typ | ev.typ // parent's typ contains push and pop's typs.
-				events[push].index = len(events)  // make push refer to pop
+				events[push].typ = typ                  // set type of push
+				stack[parent].typ |= typ | ev.typ       // parent's typ contains push and pop's typs.
+				events[push].index = int32(len(events)) // make push refer to pop
 
 				stack = stack[:top]
 				events = append(events, ev)

go/ast/inspector/iter.go

@@ -26,7 +26,7 @@ func (in *Inspector) PreorderSeq(types ...ast.Node) iter.Seq[ast.Node] {
 
 	return func(yield func(ast.Node) bool) {
 		mask := maskOf(types)
-		for i := 0; i < len(in.events); {
+		for i := int32(0); i < int32(len(in.events)); {
 			ev := in.events[i]
 			if ev.index > i {
 				// push
@@ -63,7 +63,7 @@ func All[N interface {
 
 	mask := typeOf((N)(nil))
 	return func(yield func(N) bool) {
-		for i := 0; i < len(in.events); {
+		for i := int32(0); i < int32(len(in.events)); {
 			ev := in.events[i]
 			if ev.index > i {
 				// push

internal/astutil/cursor/cursor.go

@@ -31,7 +31,7 @@ import (
 // Call [Root] to obtain a valid cursor.
 type Cursor struct {
 	in    *inspector.Inspector
-	index int // index of push node; -1 for virtual root node
+	index int32 // index of push node; -1 for virtual root node
 }
 
 // Root returns a cursor for the virtual root node,
@@ -63,9 +63,9 @@ func (c Cursor) String() string {
 }
 
 // indices return the [start, end) half-open interval of event indices.
-func (c Cursor) indices() (int, int) {
+func (c Cursor) indices() (int32, int32) {
 	if c.index < 0 {
-		return 0, len(c.events()) // root: all events
+		return 0, int32(len(c.events())) // root: all events
 	} else {
 		return c.index, c.events()[c.index].index + 1 // just one subtree
 	}
@@ -154,10 +154,7 @@ func (c Cursor) Inspect(types []ast.Node, f func(c Cursor, push bool) (descend b
 // from the [ast.File] down to the current cursor's node.
 //
 // To amortize allocation, it appends to the provided slice, which
-// must be empty. Even so, a traversal that frequently calls Stack on
-// a Cursor may be slower than using [Inspector.WithStack].
-// For best performance, call Stack selectively, or use
-// [Cursor.Parent] instead where practical.
+// must be empty.
 //
 // Stack must not be called on the Root node.
 func (c Cursor) Stack(stack []Cursor) []Cursor {
@@ -168,30 +165,9 @@ func (c Cursor) Stack(stack []Cursor) []Cursor {
 		panic("Cursor.Stack called on Root node")
 	}
 
-	// The events[i] operation is a hotspot due to the bounds
-	// check, unpredictable load, and high cache miss rate.
-	// This could be mitigated two ways:
-	//
-	// 1. Reducing the size of the event struct by 25% by changing
-	//    event.node to just the value portion (unsafe.Pointer) of
-	//    the interface. (The ast.Node itab can be reconstructed
-	//    from the event.typ bitmask.) This was measured in CL 637740
-	//    but found to deliver disappointing gains (<10%).
-	//
-	// 2. Change event.index to uint32 (which is plenty). Though this
-	//    wouldn't change the size of the table, it would create
-	//    space for an explicit parent index, making Parent and
-	//    Stack optimal.
-
 	events := c.events()
-	for i := c.index; i >= 0; i-- {
-		if j := events[i].index; j > i {
-			// push
-			stack = append(stack, Cursor{c.in, i})
-		} else {
-			// pop: skip to before corresponding push
-			i = j
-		}
+	for i := c.index; i >= 0; i = events[i].parent {
+		stack = append(stack, Cursor{c.in, i})
 	}
 	slices.Reverse(stack)
 	return stack
@@ -200,25 +176,12 @@ func (c Cursor) Stack(stack []Cursor) []Cursor {
 // Parent returns the parent of the current node.
 //
 // Parent must not be called on the Root node (whose [Cursor.Node] returns nil).
-//
-// The cost of this operation is proportional to the number of the
-// node's preceding siblings.
 func (c Cursor) Parent() Cursor {
 	if c.index < 0 {
 		panic("Cursor.Parent called on Root node")
 	}
 
-	events := c.events()
-	i := c.index - 1
-	for i >= 0 {
-		if j := events[i].index; j > i { // push?
-			break
-		} else {
-			// pop: skip to before corresponding push
-			i = j - 1
-		}
-	}
-	return Cursor{c.in, i}
+	return Cursor{c.in, c.events()[c.index].parent}
 }
 
 // NextSibling returns the cursor for the next sibling node in the
@@ -234,7 +197,7 @@ func (c Cursor) NextSibling() (Cursor, bool) {
 
 	events := c.events()
 	i := events[c.index].index + 1 // after corresponding pop
-	if i < len(events) {
+	if i < int32(len(events)) {
 		if events[i].index > i { // push?
 			return Cursor{c.in, i}, true
 		}
@@ -267,8 +230,8 @@ func (c Cursor) PrevSibling() (Cursor, bool) {
 // or zero if it has no children.
 func (c Cursor) FirstChild() (Cursor, bool) {
 	events := c.events()
-	i := c.index + 1                            // i=0 if c is root
-	if i < len(events) && events[i].index > i { // push?
+	i := c.index + 1                                   // i=0 if c is root
+	if i < int32(len(events)) && events[i].index > i { // push?
 		return Cursor{c.in, i}, true
 	}
 	return Cursor{}, false
@@ -359,7 +322,7 @@ func (c Cursor) FindPos(start, end token.Pos) (Cursor, bool) {
 	// This algorithm could be implemented using c.Inspect,
 	// but it is about 2.5x slower.
 
-	best := -1 // push index of latest (=innermost) node containing range
+	best := int32(-1) // push index of latest (=innermost) node containing range
 	for i, limit := c.indices(); i < limit; i++ {
 		ev := events[i]
 		if ev.index > i { // push?

internal/astutil/cursor/hooks.go

@@ -18,9 +18,10 @@ import (
 // Copied from inspector.event; must remain in sync.
 // (Note that the linkname effects a type coercion too.)
 type event struct {
-	node  ast.Node
-	typ   uint64 // typeOf(node) on push event, or union of typ strictly between push and pop events on pop events
-	index int    // index of corresponding push or pop event (relative to this event's index, +ve=push, -ve=pop)
+	node   ast.Node
+	typ    uint64 // typeOf(node) on push event, or union of typ strictly between push and pop events on pop events
+	index  int32  // index of corresponding push or pop event (relative to this event's index, +ve=push, -ve=pop)
+	parent int32  // index of parent's push node (defined for push nodes only)
 }
 
 //go:linkname maskOf golang.org/x/tools/go/ast/inspector.maskOf