inspector.go

     1  // Copyright 2018 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package inspector provides helper functions for traversal over the
     6  // syntax trees of a package, including node filtering by type, and
     7  // materialization of the traversal stack.
     8  //
     9  // During construction, the inspector does a complete traversal and
    10  // builds a list of push/pop events and their node type. Subsequent
    11  // method calls that request a traversal scan this list, rather than walk
    12  // the AST, and perform type filtering using efficient bit sets.
    13  // This representation is sometimes called a "balanced parenthesis tree."
    14  //
    15  // Experiments suggest the inspector's traversals are about 2.5x faster
    16  // than [ast.Inspect], but it may take around 5 traversals for this
    17  // benefit to amortize the inspector's construction cost.
    18  // If efficiency is the primary concern, do not use Inspector for
    19  // one-off traversals.
    20  //
    21  // The [Cursor] type provides a more flexible API for efficient
    22  // navigation of syntax trees in all four "cardinal directions". For
    23  // example, traversals may be nested, so you can find each node of
    24  // type A and then search within it for nodes of type B. Or you can
    25  // traverse from a node to its immediate neighbors: its parent, its
    26  // previous and next sibling, or its first and last child. We
    27  // recommend using methods of Cursor in preference to Inspector where
    28  // possible.
    29  package inspector
    30  
    31  // There are four orthogonal features in a traversal:
    32  //  1 type filtering
    33  //  2 pruning
    34  //  3 postorder calls to f
    35  //  4 stack
    36  // Rather than offer all of them in the API,
    37  // only a few combinations are exposed:
    38  // - Preorder is the fastest and has fewest features,
    39  //   but is the most commonly needed traversal.
    40  // - Nodes and WithStack both provide pruning and postorder calls,
    41  //   even though few clients need it, because supporting two versions
    42  //   is not justified.
    43  // More combinations could be supported by expressing them as
    44  // wrappers around a more generic traversal, but this was measured
    45  // and found to degrade performance significantly (30%).
    46  
    47  import (
    48  	"go/ast"
    49  
    50  	"golang.org/x/tools/go/ast/edge"
    51  )
    52  
    53  // An Inspector provides methods for inspecting
    54  // (traversing) the syntax trees of a package.
    55  type Inspector struct {
    56  	events []event
    57  }
    58  
    59  func packEdgeKindAndIndex(ek edge.Kind, index int) int32 {
    60  	return int32(uint32(index+1)<<7 | uint32(ek))
    61  }
    62  
    63  // unpackEdgeKindAndIndex unpacks the edge kind and edge index (within
    64  // an []ast.Node slice) from the parent field of a pop event.
    65  func unpackEdgeKindAndIndex(x int32) (edge.Kind, int) {
    66  	// The "parent" field of a pop node holds the
    67  	// edge Kind in the lower 7 bits and the index+1
    68  	// in the upper 25.
    69  	return edge.Kind(x & 0x7f), int(x>>7) - 1
    70  }
    71  
    72  // New returns an Inspector for the specified syntax trees.
    73  func New(files []*ast.File) *Inspector {
    74  	return &Inspector{traverse(files)}
    75  }
    76  
    77  // An event represents a push or a pop
    78  // of an ast.Node during a traversal.
    79  type event struct {
    80  	node   ast.Node
    81  	typ    uint64 // typeOf(node) on push event, or union of typ strictly between push and pop events on pop events
    82  	index  int32  // index of corresponding push or pop event
    83  	parent int32  // index of parent's push node (push nodes only), or packed edge kind/index (pop nodes only)
    84  }
    85  
    86  // TODO: Experiment with storing only the second word of event.node (unsafe.Pointer).
    87  // Type can be recovered from the sole bit in typ.
    88  
    89  // Preorder visits all the nodes of the files supplied to New in
    90  // depth-first order. It calls f(n) for each node n before it visits
    91  // n's children.
    92  //
    93  // The complete traversal sequence is determined by [ast.Inspect].
    94  // The types argument, if non-empty, enables type-based filtering of
    95  // events. The function f is called only for nodes whose type
    96  // matches an element of the types slice.
    97  //
    98  // The [Cursor.Preorder] method provides a richer alternative interface.
    99  // Example:
   100  //
   101  //	for c := range in.Root().Preorder(types) { ... }
   102  func (in *Inspector) Preorder(types []ast.Node, f func(ast.Node)) {
   103  	// Because it avoids postorder calls to f, and the pruning
   104  	// check, Preorder is almost twice as fast as Nodes. The two
   105  	// features seem to contribute similar slowdowns (~1.4x each).
   106  
   107  	// This function is equivalent to the PreorderSeq call below,
   108  	// but to avoid the additional dynamic call (which adds 13-35%
   109  	// to the benchmarks), we expand it out.
   110  	//
   111  	// in.PreorderSeq(types...)(func(n ast.Node) bool {
   112  	// 	f(n)
   113  	// 	return true
   114  	// })
   115  
   116  	mask := maskOf(types)
   117  	for i := int32(0); i < int32(len(in.events)); {
   118  		ev := in.events[i]
   119  		if ev.index > i {
   120  			// push
   121  			if ev.typ&mask != 0 {
   122  				f(ev.node)
   123  			}
   124  			pop := ev.index
   125  			if in.events[pop].typ&mask == 0 {
   126  				// Subtrees do not contain types: skip them and pop.
   127  				i = pop + 1
   128  				continue
   129  			}
   130  		}
   131  		i++
   132  	}
   133  }
   134  
   135  // Nodes visits the nodes of the files supplied to New in depth-first
   136  // order. It calls f(n, true) for each node n before it visits n's
   137  // children. If f returns true, Nodes invokes f recursively for each
   138  // of the non-nil children of the node, followed by a call of
   139  // f(n, false).
   140  //
   141  // The complete traversal sequence is determined by [ast.Inspect].
   142  // The types argument, if non-empty, enables type-based filtering of
   143  // events. The function f if is called only for nodes whose type
   144  // matches an element of the types slice.
   145  //
   146  // The [Cursor.Inspect] method provides a richer alternative interface.
   147  // Example:
   148  //
   149  //	in.Root().Inspect(types, func(c Cursor) bool {
   150  //		...
   151  //		return true
   152  //	}
   153  func (in *Inspector) Nodes(types []ast.Node, f func(n ast.Node, push bool) (proceed bool)) {
   154  	mask := maskOf(types)
   155  	for i := int32(0); i < int32(len(in.events)); {
   156  		ev := in.events[i]
   157  		if ev.index > i {
   158  			// push
   159  			pop := ev.index
   160  			if ev.typ&mask != 0 {
   161  				if !f(ev.node, true) {
   162  					i = pop + 1 // jump to corresponding pop + 1
   163  					continue
   164  				}
   165  			}
   166  			if in.events[pop].typ&mask == 0 {
   167  				// Subtrees do not contain types: skip them.
   168  				i = pop
   169  				continue
   170  			}
   171  		} else {
   172  			// pop
   173  			push := ev.index
   174  			if in.events[push].typ&mask != 0 {
   175  				f(ev.node, false)
   176  			}
   177  		}
   178  		i++
   179  	}
   180  }
   181  
   182  // WithStack visits nodes in a similar manner to Nodes, but it
   183  // supplies each call to f an additional argument, the current
   184  // traversal stack. The stack's first element is the outermost node,
   185  // an *ast.File; its last is the innermost, n.
   186  //
   187  // The [Cursor.Inspect] method provides a richer alternative interface.
   188  // Example:
   189  //
   190  //	in.Root().Inspect(types, func(c Cursor) bool {
   191  //		stack := slices.Collect(c.Enclosing())
   192  //		...
   193  //		return true
   194  //	})
   195  func (in *Inspector) WithStack(types []ast.Node, f func(n ast.Node, push bool, stack []ast.Node) (proceed bool)) {
   196  	mask := maskOf(types)
   197  	var stack []ast.Node
   198  	for i := int32(0); i < int32(len(in.events)); {
   199  		ev := in.events[i]
   200  		if ev.index > i {
   201  			// push
   202  			pop := ev.index
   203  			stack = append(stack, ev.node)
   204  			if ev.typ&mask != 0 {
   205  				if !f(ev.node, true, stack) {
   206  					i = pop + 1
   207  					stack = stack[:len(stack)-1]
   208  					continue
   209  				}
   210  			}
   211  			if in.events[pop].typ&mask == 0 {
   212  				// Subtrees does not contain types: skip them.
   213  				i = pop
   214  				continue
   215  			}
   216  		} else {
   217  			// pop
   218  			push := ev.index
   219  			if in.events[push].typ&mask != 0 {
   220  				f(ev.node, false, stack)
   221  			}
   222  			stack = stack[:len(stack)-1]
   223  		}
   224  		i++
   225  	}
   226  }
   227  
   228  // traverse builds the table of events representing a traversal.
   229  func traverse(files []*ast.File) []event {
   230  	// Preallocate approximate number of events
   231  	// based on source file extent of the declarations.
   232  	// (We use End-Pos not FileStart-FileEnd to neglect
   233  	// the effect of long doc comments.)
   234  	// This makes traverse faster by 4x (!).
   235  	var extent int
   236  	for _, f := range files {
   237  		extent += int(f.End() - f.Pos())
   238  	}
   239  	// This estimate is based on the net/http package.
   240  	capacity := min(extent*33/100, 1e6) // impose some reasonable maximum (1M)
   241  
   242  	v := &visitor{
   243  		events: make([]event, 0, capacity),
   244  		stack:  []item{{index: -1}}, // include an extra event so file nodes have a parent
   245  	}
   246  	for _, file := range files {
   247  		walk(v, edge.Invalid, -1, file)
   248  	}
   249  	return v.events
   250  }
   251  
   252  type visitor struct {
   253  	events []event
   254  	stack  []item
   255  }
   256  
   257  type item struct {
   258  	index            int32  // index of current node's push event
   259  	parentIndex      int32  // index of parent node's push event
   260  	typAccum         uint64 // accumulated type bits of current node's descendants
   261  	edgeKindAndIndex int32  // edge.Kind and index, bit packed
   262  }
   263  
   264  func (v *visitor) push(ek edge.Kind, eindex int, node ast.Node) {
   265  	var (
   266  		index       = int32(len(v.events))
   267  		parentIndex = v.stack[len(v.stack)-1].index
   268  	)
   269  	v.events = append(v.events, event{
   270  		node:   node,
   271  		parent: parentIndex,
   272  		typ:    typeOf(node),
   273  		index:  0, // (pop index is set later by visitor.pop)
   274  	})
   275  	v.stack = append(v.stack, item{
   276  		index:            index,
   277  		parentIndex:      parentIndex,
   278  		edgeKindAndIndex: packEdgeKindAndIndex(ek, eindex),
   279  	})
   280  
   281  	// 2B nodes ought to be enough for anyone!
   282  	if int32(len(v.events)) < 0 {
   283  		panic("event index exceeded int32")
   284  	}
   285  
   286  	// 32M elements in an []ast.Node ought to be enough for anyone!
   287  	if ek2, eindex2 := unpackEdgeKindAndIndex(packEdgeKindAndIndex(ek, eindex)); ek2 != ek || eindex2 != eindex {
   288  		panic("Node slice index exceeded uint25")
   289  	}
   290  }
   291  
   292  func (v *visitor) pop(node ast.Node) {
   293  	top := len(v.stack) - 1
   294  	current := v.stack[top]
   295  
   296  	push := &v.events[current.index]
   297  	parent := &v.stack[top-1]
   298  
   299  	push.index = int32(len(v.events))              // make push event refer to pop
   300  	parent.typAccum |= current.typAccum | push.typ // accumulate type bits into parent
   301  
   302  	v.stack = v.stack[:top]
   303  
   304  	v.events = append(v.events, event{
   305  		node:   node,
   306  		typ:    current.typAccum,
   307  		index:  current.index,
   308  		parent: current.edgeKindAndIndex, // see [unpackEdgeKindAndIndex]
   309  	})
   310  }
   311
View as plain text