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  // [Tried this, wasn't faster. --adonovan]
    89  
    90  // Preorder visits all the nodes of the files supplied to New in
    91  // depth-first order. It calls f(n) for each node n before it visits
    92  // n's children.
    93  //
    94  // The complete traversal sequence is determined by [ast.Inspect].
    95  // The types argument, if non-empty, enables type-based filtering of
    96  // events. The function f is called only for nodes whose type
    97  // matches an element of the types slice.
    98  //
    99  // The [Cursor.Preorder] method provides a richer alternative interface.
   100  // Example:
   101  //
   102  //	for c := range in.Root().Preorder(types) { ... }
   103  func (in *Inspector) Preorder(types []ast.Node, f func(ast.Node)) {
   104  	// Because it avoids postorder calls to f, and the pruning
   105  	// check, Preorder is almost twice as fast as Nodes. The two
   106  	// features seem to contribute similar slowdowns (~1.4x each).
   107  
   108  	// This function is equivalent to the PreorderSeq call below,
   109  	// but to avoid the additional dynamic call (which adds 13-35%
   110  	// to the benchmarks), we expand it out.
   111  	//
   112  	// in.PreorderSeq(types...)(func(n ast.Node) bool {
   113  	// 	f(n)
   114  	// 	return true
   115  	// })
   116  
   117  	mask := maskOf(types)
   118  	for i := int32(0); i < int32(len(in.events)); {
   119  		ev := in.events[i]
   120  		if ev.index > i {
   121  			// push
   122  			if ev.typ&mask != 0 {
   123  				f(ev.node)
   124  			}
   125  			pop := ev.index
   126  			if in.events[pop].typ&mask == 0 {
   127  				// Subtrees do not contain types: skip them and pop.
   128  				i = pop + 1
   129  				continue
   130  			}
   131  		}
   132  		i++
   133  	}
   134  }
   135  
   136  // Nodes visits the nodes of the files supplied to New in depth-first
   137  // order. It calls f(n, true) for each node n before it visits n's
   138  // children. If f returns true, Nodes invokes f recursively for each
   139  // of the non-nil children of the node, followed by a call of
   140  // f(n, false).
   141  //
   142  // The complete traversal sequence is determined by [ast.Inspect].
   143  // The types argument, if non-empty, enables type-based filtering of
   144  // events. The function f if is called only for nodes whose type
   145  // matches an element of the types slice.
   146  //
   147  // The [Cursor.Inspect] method provides a richer alternative interface.
   148  // Example:
   149  //
   150  //	in.Root().Inspect(types, func(c Cursor) bool {
   151  //		...
   152  //		return true
   153  //	}
   154  func (in *Inspector) Nodes(types []ast.Node, f func(n ast.Node, push bool) (proceed bool)) {
   155  	mask := maskOf(types)
   156  	for i := int32(0); i < int32(len(in.events)); {
   157  		ev := in.events[i]
   158  		if ev.index > i {
   159  			// push
   160  			pop := ev.index
   161  			if ev.typ&mask != 0 {
   162  				if !f(ev.node, true) {
   163  					i = pop + 1 // jump to corresponding pop + 1
   164  					continue
   165  				}
   166  			}
   167  			if in.events[pop].typ&mask == 0 {
   168  				// Subtrees do not contain types: skip them.
   169  				i = pop
   170  				continue
   171  			}
   172  		} else {
   173  			// pop
   174  			push := ev.index
   175  			if in.events[push].typ&mask != 0 {
   176  				f(ev.node, false)
   177  			}
   178  		}
   179  		i++
   180  	}
   181  }
   182  
   183  // WithStack visits nodes in a similar manner to Nodes, but it
   184  // supplies each call to f an additional argument, the current
   185  // traversal stack. The stack's first element is the outermost node,
   186  // an *ast.File; its last is the innermost, n.
   187  //
   188  // The [Cursor.Inspect] method provides a richer alternative interface.
   189  // Example:
   190  //
   191  //	in.Root().Inspect(types, func(c Cursor) bool {
   192  //		stack := slices.Collect(c.Enclosing())
   193  //		...
   194  //		return true
   195  //	})
   196  func (in *Inspector) WithStack(types []ast.Node, f func(n ast.Node, push bool, stack []ast.Node) (proceed bool)) {
   197  	mask := maskOf(types)
   198  	var stack []ast.Node
   199  	for i := int32(0); i < int32(len(in.events)); {
   200  		ev := in.events[i]
   201  		if ev.index > i {
   202  			// push
   203  			pop := ev.index
   204  			stack = append(stack, ev.node)
   205  			if ev.typ&mask != 0 {
   206  				if !f(ev.node, true, stack) {
   207  					i = pop + 1
   208  					stack = stack[:len(stack)-1]
   209  					continue
   210  				}
   211  			}
   212  			if in.events[pop].typ&mask == 0 {
   213  				// Subtrees does not contain types: skip them.
   214  				i = pop
   215  				continue
   216  			}
   217  		} else {
   218  			// pop
   219  			push := ev.index
   220  			if in.events[push].typ&mask != 0 {
   221  				f(ev.node, false, stack)
   222  			}
   223  			stack = stack[:len(stack)-1]
   224  		}
   225  		i++
   226  	}
   227  }
   228  
   229  // traverse builds the table of events representing a traversal.
   230  func traverse(files []*ast.File) []event {
   231  	// Preallocate approximate number of events
   232  	// based on source file extent of the declarations.
   233  	// (We use End-Pos not FileStart-FileEnd to neglect
   234  	// the effect of long doc comments.)
   235  	// This makes traverse faster by 4x (!).
   236  	var extent int
   237  	for _, f := range files {
   238  		extent += int(f.End() - f.Pos())
   239  	}
   240  	// This estimate is based on the net/http package.
   241  	capacity := min(extent*33/100, 1e6) // impose some reasonable maximum (1M)
   242  
   243  	v := &visitor{
   244  		events: make([]event, 0, capacity),
   245  		stack:  []item{{index: -1}}, // include an extra event so file nodes have a parent
   246  	}
   247  	for _, file := range files {
   248  		walk(v, edge.Invalid, -1, file)
   249  	}
   250  	return v.events
   251  }
   252  
   253  type visitor struct {
   254  	events []event
   255  	stack  []item
   256  }
   257  
   258  type item struct {
   259  	index            int32  // index of current node's push event
   260  	parentIndex      int32  // index of parent node's push event
   261  	typAccum         uint64 // accumulated type bits of current node's descendants
   262  	edgeKindAndIndex int32  // edge.Kind and index, bit packed
   263  }
   264  
   265  func (v *visitor) push(ek edge.Kind, eindex int, node ast.Node) {
   266  	var (
   267  		index       = int32(len(v.events))
   268  		parentIndex = v.stack[len(v.stack)-1].index
   269  	)
   270  	v.events = append(v.events, event{
   271  		node:   node,
   272  		parent: parentIndex,
   273  		typ:    typeOf(node),
   274  		index:  0, // (pop index is set later by visitor.pop)
   275  	})
   276  	v.stack = append(v.stack, item{
   277  		index:            index,
   278  		parentIndex:      parentIndex,
   279  		edgeKindAndIndex: packEdgeKindAndIndex(ek, eindex),
   280  	})
   281  
   282  	// 2B nodes ought to be enough for anyone!
   283  	if int32(len(v.events)) < 0 {
   284  		panic("event index exceeded int32")
   285  	}
   286  
   287  	// 32M elements in an []ast.Node ought to be enough for anyone!
   288  	if ek2, eindex2 := unpackEdgeKindAndIndex(packEdgeKindAndIndex(ek, eindex)); ek2 != ek || eindex2 != eindex {
   289  		panic("Node slice index exceeded uint25")
   290  	}
   291  }
   292  
   293  func (v *visitor) pop(node ast.Node) {
   294  	top := len(v.stack) - 1
   295  	current := v.stack[top]
   296  
   297  	push := &v.events[current.index]
   298  	parent := &v.stack[top-1]
   299  
   300  	push.index = int32(len(v.events))              // make push event refer to pop
   301  	parent.typAccum |= current.typAccum | push.typ // accumulate type bits into parent
   302  
   303  	v.stack = v.stack[:top]
   304  
   305  	v.events = append(v.events, event{
   306  		node:   node,
   307  		typ:    current.typAccum,
   308  		index:  current.index,
   309  		parent: current.edgeKindAndIndex, // see [unpackEdgeKindAndIndex]
   310  	})
   311  }
   312
View as plain text