cursor.go

     1  // Copyright 2025 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
     6  
     7  import (
     8  	"fmt"
     9  	"go/ast"
    10  	"go/token"
    11  	"iter"
    12  	"reflect"
    13  
    14  	"golang.org/x/tools/go/ast/edge"
    15  )
    16  
    17  // A Cursor represents an [ast.Node]. It is immutable.
    18  //
    19  // Two Cursors compare equal if they represent the same node.
    20  //
    21  // Call [Inspector.Root] to obtain a valid cursor for the virtual root
    22  // node of the traversal.
    23  //
    24  // Use the following methods to navigate efficiently around the tree:
    25  //   - for ancestors, use [Cursor.Parent] and [Cursor.Enclosing];
    26  //   - for children, use [Cursor.Child], [Cursor.Children],
    27  //     [Cursor.FirstChild], and [Cursor.LastChild];
    28  //   - for siblings, use [Cursor.PrevSibling] and [Cursor.NextSibling];
    29  //   - for descendants, use [Cursor.FindByPos], [Cursor.FindNode],
    30  //     [Cursor.Inspect], and [Cursor.Preorder].
    31  //
    32  // Use the [Cursor.ChildAt] and [Cursor.ParentEdge] methods for
    33  // information about the edges in a tree: which field (and slice
    34  // element) of the parent node holds the child.
    35  type Cursor struct {
    36  	in    *Inspector
    37  	index int32 // index of push node; -1 for virtual root node
    38  }
    39  
    40  // Root returns a cursor for the virtual root node,
    41  // whose children are the files provided to [New].
    42  //
    43  // Its [Cursor.Node] and [Cursor.Stack] methods return nil.
    44  func (in *Inspector) Root() Cursor {
    45  	return Cursor{in, -1}
    46  }
    47  
    48  // At returns the cursor at the specified index in the traversal,
    49  // which must have been obtained from [Cursor.Index] on a Cursor
    50  // belonging to the same Inspector (see [Cursor.Inspector]).
    51  func (in *Inspector) At(index int32) Cursor {
    52  	if index < 0 {
    53  		panic("negative index")
    54  	}
    55  	if int(index) >= len(in.events) {
    56  		panic("index out of range for this inspector")
    57  	}
    58  	if in.events[index].index < index {
    59  		panic("invalid index") // (a push, not a pop)
    60  	}
    61  	return Cursor{in, index}
    62  }
    63  
    64  // Inspector returns the cursor's Inspector.
    65  func (c Cursor) Inspector() *Inspector { return c.in }
    66  
    67  // Index returns the index of this cursor position within the package.
    68  //
    69  // Clients should not assume anything about the numeric Index value
    70  // except that it increases monotonically throughout the traversal.
    71  // It is provided for use with [At].
    72  //
    73  // Index must not be called on the Root node.
    74  func (c Cursor) Index() int32 {
    75  	if c.index < 0 {
    76  		panic("Index called on Root node")
    77  	}
    78  	return c.index
    79  }
    80  
    81  // Node returns the node at the current cursor position,
    82  // or nil for the cursor returned by [Inspector.Root].
    83  func (c Cursor) Node() ast.Node {
    84  	if c.index < 0 {
    85  		return nil
    86  	}
    87  	return c.in.events[c.index].node
    88  }
    89  
    90  // String returns information about the cursor's node, if any.
    91  func (c Cursor) String() string {
    92  	if c.in == nil {
    93  		return "(invalid)"
    94  	}
    95  	if c.index < 0 {
    96  		return "(root)"
    97  	}
    98  	return reflect.TypeOf(c.Node()).String()
    99  }
   100  
   101  // indices return the [start, end) half-open interval of event indices.
   102  func (c Cursor) indices() (int32, int32) {
   103  	if c.index < 0 {
   104  		return 0, int32(len(c.in.events)) // root: all events
   105  	} else {
   106  		return c.index, c.in.events[c.index].index + 1 // just one subtree
   107  	}
   108  }
   109  
   110  // Preorder returns an iterator over the nodes of the subtree
   111  // represented by c in depth-first order. Each node in the sequence is
   112  // represented by a Cursor that allows access to the Node, but may
   113  // also be used to start a new traversal, or to obtain the stack of
   114  // nodes enclosing the cursor.
   115  //
   116  // The traversal sequence is determined by [ast.Inspect]. The types
   117  // argument, if non-empty, enables type-based filtering of events. The
   118  // function f if is called only for nodes whose type matches an
   119  // element of the types slice.
   120  //
   121  // If you need control over descent into subtrees,
   122  // or need both pre- and post-order notifications, use [Cursor.Inspect]
   123  func (c Cursor) Preorder(types ...ast.Node) iter.Seq[Cursor] {
   124  	mask := maskOf(types)
   125  
   126  	return func(yield func(Cursor) bool) {
   127  		events := c.in.events
   128  
   129  		for i, limit := c.indices(); i < limit; {
   130  			ev := events[i]
   131  			if ev.index > i { // push?
   132  				if ev.typ&mask != 0 && !yield(Cursor{c.in, i}) {
   133  					break
   134  				}
   135  				pop := ev.index
   136  				if events[pop].typ&mask == 0 {
   137  					// Subtree does not contain types: skip.
   138  					i = pop + 1
   139  					continue
   140  				}
   141  			}
   142  			i++
   143  		}
   144  	}
   145  }
   146  
   147  // Inspect visits the nodes of the subtree represented by c in
   148  // depth-first order. It calls f(n) for each node n before it
   149  // visits n's children. If f returns true, Inspect invokes f
   150  // recursively for each of the non-nil children of the node.
   151  //
   152  // Each node is represented by a Cursor that allows access to the
   153  // Node, but may also be used to start a new traversal, or to obtain
   154  // the stack of nodes enclosing the cursor.
   155  //
   156  // The complete traversal sequence is determined by [ast.Inspect].
   157  // The types argument, if non-empty, enables type-based filtering of
   158  // events. The function f if is called only for nodes whose type
   159  // matches an element of the types slice.
   160  func (c Cursor) Inspect(types []ast.Node, f func(c Cursor) (descend bool)) {
   161  	mask := maskOf(types)
   162  	events := c.in.events
   163  	for i, limit := c.indices(); i < limit; {
   164  		ev := events[i]
   165  		if ev.index > i {
   166  			// push
   167  			pop := ev.index
   168  			if ev.typ&mask != 0 && !f(Cursor{c.in, i}) ||
   169  				events[pop].typ&mask == 0 {
   170  				// The user opted not to descend, or the
   171  				// subtree does not contain types:
   172  				// skip past the pop.
   173  				i = pop + 1
   174  				continue
   175  			}
   176  		}
   177  		i++
   178  	}
   179  }
   180  
   181  // Enclosing returns an iterator over the nodes enclosing the current
   182  // current node, starting with the Cursor itself.
   183  //
   184  // Enclosing must not be called on the Root node (whose [Cursor.Node] returns nil).
   185  //
   186  // The types argument, if non-empty, enables type-based filtering of
   187  // events: the sequence includes only enclosing nodes whose type
   188  // matches an element of the types slice.
   189  func (c Cursor) Enclosing(types ...ast.Node) iter.Seq[Cursor] {
   190  	if c.index < 0 {
   191  		panic("Cursor.Enclosing called on Root node")
   192  	}
   193  
   194  	mask := maskOf(types)
   195  
   196  	return func(yield func(Cursor) bool) {
   197  		events := c.in.events
   198  		for i := c.index; i >= 0; i = events[i].parent {
   199  			if events[i].typ&mask != 0 && !yield(Cursor{c.in, i}) {
   200  				break
   201  			}
   202  		}
   203  	}
   204  }
   205  
   206  // Parent returns the parent of the current node.
   207  //
   208  // Parent must not be called on the Root node (whose [Cursor.Node] returns nil).
   209  func (c Cursor) Parent() Cursor {
   210  	if c.index < 0 {
   211  		panic("Cursor.Parent called on Root node")
   212  	}
   213  
   214  	return Cursor{c.in, c.in.events[c.index].parent}
   215  }
   216  
   217  // ParentEdge returns the identity of the field in the parent node
   218  // that holds this cursor's node, and if it is a list, the index within it.
   219  //
   220  // For example, f(x, y) is a CallExpr whose three children are Idents.
   221  // f has edge kind [edge.CallExpr_Fun] and index -1.
   222  // x and y have kind [edge.CallExpr_Args] and indices 0 and 1, respectively.
   223  //
   224  // If called on a child of the Root node, it returns ([edge.Invalid], -1).
   225  //
   226  // ParentEdge must not be called on the Root node (whose [Cursor.Node] returns nil).
   227  func (c Cursor) ParentEdge() (edge.Kind, int) {
   228  	if c.index < 0 {
   229  		panic("Cursor.ParentEdge called on Root node")
   230  	}
   231  	events := c.in.events
   232  	pop := events[c.index].index
   233  	return unpackEdgeKindAndIndex(events[pop].parent)
   234  }
   235  
   236  // ChildAt returns the cursor for the child of the
   237  // current node identified by its edge and index.
   238  // The index must be -1 if the edge.Kind is not a slice.
   239  // The indicated child node must exist.
   240  //
   241  // ChildAt must not be called on the Root node (whose [Cursor.Node] returns nil).
   242  //
   243  // Invariant: c.Parent().ChildAt(c.ParentEdge()) == c.
   244  func (c Cursor) ChildAt(k edge.Kind, idx int) Cursor {
   245  	target := packEdgeKindAndIndex(k, idx)
   246  
   247  	// Unfortunately there's no shortcut to looping.
   248  	events := c.in.events
   249  	i := c.index + 1
   250  	for {
   251  		pop := events[i].index
   252  		if pop < i {
   253  			break
   254  		}
   255  		if events[pop].parent == target {
   256  			return Cursor{c.in, i}
   257  		}
   258  		i = pop + 1
   259  	}
   260  	panic(fmt.Sprintf("ChildAt(%v, %d): no such child of %v", k, idx, c))
   261  }
   262  
   263  // Child returns the cursor for n, which must be a direct child of c's Node.
   264  //
   265  // Child must not be called on the Root node (whose [Cursor.Node] returns nil).
   266  func (c Cursor) Child(n ast.Node) Cursor {
   267  	if c.index < 0 {
   268  		panic("Cursor.Child called on Root node")
   269  	}
   270  
   271  	if false {
   272  		// reference implementation
   273  		for child := range c.Children() {
   274  			if child.Node() == n {
   275  				return child
   276  			}
   277  		}
   278  
   279  	} else {
   280  		// optimized implementation
   281  		events := c.in.events
   282  		for i := c.index + 1; events[i].index > i; i = events[i].index + 1 {
   283  			if events[i].node == n {
   284  				return Cursor{c.in, i}
   285  			}
   286  		}
   287  	}
   288  	panic(fmt.Sprintf("Child(%T): not a child of %v", n, c))
   289  }
   290  
   291  // NextSibling returns the cursor for the next sibling node in the same list
   292  // (for example, of files, decls, specs, statements, fields, or expressions) as
   293  // the current node. It returns (zero, false) if the node is the last node in
   294  // the list, or is not part of a list.
   295  //
   296  // NextSibling must not be called on the Root node.
   297  //
   298  // See note at [Cursor.Children].
   299  func (c Cursor) NextSibling() (Cursor, bool) {
   300  	if c.index < 0 {
   301  		panic("Cursor.NextSibling called on Root node")
   302  	}
   303  
   304  	events := c.in.events
   305  	i := events[c.index].index + 1 // after corresponding pop
   306  	if i < int32(len(events)) {
   307  		if events[i].index > i { // push?
   308  			return Cursor{c.in, i}, true
   309  		}
   310  	}
   311  	return Cursor{}, false
   312  }
   313  
   314  // PrevSibling returns the cursor for the previous sibling node in the
   315  // same list (for example, of files, decls, specs, statements, fields,
   316  // or expressions) as the current node. It returns zero if the node is
   317  // the first node in the list, or is not part of a list.
   318  //
   319  // It must not be called on the Root node.
   320  //
   321  // See note at [Cursor.Children].
   322  func (c Cursor) PrevSibling() (Cursor, bool) {
   323  	if c.index < 0 {
   324  		panic("Cursor.PrevSibling called on Root node")
   325  	}
   326  
   327  	events := c.in.events
   328  	i := c.index - 1
   329  	if i >= 0 {
   330  		if j := events[i].index; j < i { // pop?
   331  			return Cursor{c.in, j}, true
   332  		}
   333  	}
   334  	return Cursor{}, false
   335  }
   336  
   337  // FirstChild returns the first direct child of the current node,
   338  // or zero if it has no children.
   339  func (c Cursor) FirstChild() (Cursor, bool) {
   340  	events := c.in.events
   341  	i := c.index + 1                                   // i=0 if c is root
   342  	if i < int32(len(events)) && events[i].index > i { // push?
   343  		return Cursor{c.in, i}, true
   344  	}
   345  	return Cursor{}, false
   346  }
   347  
   348  // LastChild returns the last direct child of the current node,
   349  // or zero if it has no children.
   350  func (c Cursor) LastChild() (Cursor, bool) {
   351  	events := c.in.events
   352  	if c.index < 0 { // root?
   353  		if len(events) > 0 {
   354  			// return push of final event (a pop)
   355  			return Cursor{c.in, events[len(events)-1].index}, true
   356  		}
   357  	} else {
   358  		j := events[c.index].index - 1 // before corresponding pop
   359  		// Inv: j == c.index if c has no children
   360  		//  or  j is last child's pop.
   361  		if j > c.index { // c has children
   362  			return Cursor{c.in, events[j].index}, true
   363  		}
   364  	}
   365  	return Cursor{}, false
   366  }
   367  
   368  // Children returns an iterator over the direct children of the
   369  // current node, if any.
   370  //
   371  // When using Children, NextChild, and PrevChild, bear in mind that a
   372  // Node's children may come from different fields, some of which may
   373  // be lists of nodes without a distinguished intervening container
   374  // such as [ast.BlockStmt].
   375  //
   376  // For example, [ast.CaseClause] has a field List of expressions and a
   377  // field Body of statements, so the children of a CaseClause are a mix
   378  // of expressions and statements. Other nodes that have "uncontained"
   379  // list fields include:
   380  //
   381  //   - [ast.ValueSpec] (Names, Values)
   382  //   - [ast.CompositeLit] (Type, Elts)
   383  //   - [ast.IndexListExpr] (X, Indices)
   384  //   - [ast.CallExpr] (Fun, Args)
   385  //   - [ast.AssignStmt] (Lhs, Rhs)
   386  //
   387  // So, do not assume that the previous sibling of an ast.Stmt is also
   388  // an ast.Stmt, or if it is, that they are executed sequentially,
   389  // unless you have established that, say, its parent is a BlockStmt
   390  // or its [Cursor.ParentEdge] is [edge.BlockStmt_List].
   391  // For example, given "for S1; ; S2 {}", the predecessor of S2 is S1,
   392  // even though they are not executed in sequence.
   393  func (c Cursor) Children() iter.Seq[Cursor] {
   394  	return func(yield func(Cursor) bool) {
   395  		c, ok := c.FirstChild()
   396  		for ok && yield(c) {
   397  			c, ok = c.NextSibling()
   398  		}
   399  	}
   400  }
   401  
   402  // Contains reports whether c contains or is equal to c2.
   403  //
   404  // Both Cursors must belong to the same [Inspector];
   405  // neither may be its Root node.
   406  func (c Cursor) Contains(c2 Cursor) bool {
   407  	if c.in != c2.in {
   408  		panic("different inspectors")
   409  	}
   410  	events := c.in.events
   411  	return c.index <= c2.index && events[c2.index].index <= events[c.index].index
   412  }
   413  
   414  // FindNode returns the cursor for node n if it belongs to the subtree
   415  // rooted at c. It returns zero if n is not found.
   416  func (c Cursor) FindNode(n ast.Node) (Cursor, bool) {
   417  
   418  	// FindNode is equivalent to this code,
   419  	// but more convenient and 15-20% faster:
   420  	if false {
   421  		for candidate := range c.Preorder(n) {
   422  			if candidate.Node() == n {
   423  				return candidate, true
   424  			}
   425  		}
   426  		return Cursor{}, false
   427  	}
   428  
   429  	// TODO(adonovan): opt: should we assume Node.Pos is accurate
   430  	// and combine type-based filtering with position filtering
   431  	// like FindByPos?
   432  
   433  	mask := maskOf([]ast.Node{n})
   434  	events := c.in.events
   435  
   436  	for i, limit := c.indices(); i < limit; i++ {
   437  		ev := events[i]
   438  		if ev.index > i { // push?
   439  			if ev.typ&mask != 0 && ev.node == n {
   440  				return Cursor{c.in, i}, true
   441  			}
   442  			pop := ev.index
   443  			if events[pop].typ&mask == 0 {
   444  				// Subtree does not contain type of n: skip.
   445  				i = pop
   446  			}
   447  		}
   448  	}
   449  	return Cursor{}, false
   450  }
   451  
   452  // FindByPos returns the cursor for the innermost node n in the tree
   453  // rooted at c such that n.Pos() <= start && end <= n.End().
   454  // (For an *ast.File, it uses the bounds n.FileStart-n.FileEnd.)
   455  //
   456  // It returns zero if none is found.
   457  // Precondition: start <= end.
   458  //
   459  // See also [astutil.PathEnclosingInterval], which
   460  // tolerates adjoining whitespace.
   461  func (c Cursor) FindByPos(start, end token.Pos) (Cursor, bool) {
   462  	if end < start {
   463  		panic("end < start")
   464  	}
   465  	events := c.in.events
   466  
   467  	// This algorithm could be implemented using c.Inspect,
   468  	// but it is about 2.5x slower.
   469  
   470  	best := int32(-1) // push index of latest (=innermost) node containing range
   471  	for i, limit := c.indices(); i < limit; i++ {
   472  		ev := events[i]
   473  		if ev.index > i { // push?
   474  			n := ev.node
   475  			var nodeEnd token.Pos
   476  			if file, ok := n.(*ast.File); ok {
   477  				nodeEnd = file.FileEnd
   478  				// Note: files may be out of Pos order.
   479  				if file.FileStart > start {
   480  					i = ev.index // disjoint, after; skip to next file
   481  					continue
   482  				}
   483  			} else {
   484  				nodeEnd = n.End()
   485  				if n.Pos() > start {
   486  					break // disjoint, after; stop
   487  				}
   488  			}
   489  			// Inv: node.{Pos,FileStart} <= start
   490  			if end <= nodeEnd {
   491  				// node fully contains target range
   492  				best = i
   493  			} else if nodeEnd < start {
   494  				i = ev.index // disjoint, before; skip forward
   495  			}
   496  		}
   497  	}
   498  	if best >= 0 {
   499  		return Cursor{c.in, best}, true
   500  	}
   501  	return Cursor{}, false
   502  }
   503
View as plain text