iter.go

     1  // Copyright 2024 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  //go:build go1.23
     6  
     7  package inspector
     8  
     9  import (
    10  	"go/ast"
    11  	"iter"
    12  )
    13  
    14  // PreorderSeq returns an iterator that visits all the
    15  // nodes of the files supplied to [New] in depth-first order.
    16  // It visits each node n before n's children.
    17  // The complete traversal sequence is determined by ast.Inspect.
    18  //
    19  // The types argument, if non-empty, enables type-based filtering:
    20  // only nodes whose type matches an element of the types slice are
    21  // included in the sequence.
    22  //
    23  // Example:
    24  //
    25  //	for call := range in.PreorderSeq((*ast.CallExpr)(nil)) { ... }
    26  //
    27  // The [All] function is more convenient if there is exactly one node type:
    28  //
    29  //	for call := range All[*ast.CallExpr](in) { ... }
    30  //
    31  // See also the newer and more flexible [Cursor] API, which lets you
    32  // start the traversal at an arbitrary node, and reports each matching
    33  // node by its Cursor, enabling easier navigation.
    34  // The above example would be written thus:
    35  //
    36  //	for curCall := range in.Root().Preorder((*ast.CallExpr)(nil)) {
    37  //		call := curCall.Node().(*ast.CallExpr)
    38  //		...
    39  //	}
    40  func (in *Inspector) PreorderSeq(types ...ast.Node) iter.Seq[ast.Node] {
    41  
    42  	// This implementation is identical to Preorder,
    43  	// except that it supports breaking out of the loop.
    44  
    45  	return func(yield func(ast.Node) bool) {
    46  		mask := maskOf(types)
    47  		for i := int32(0); i < int32(len(in.events)); {
    48  			ev := in.events[i]
    49  			if ev.index > i {
    50  				// push
    51  				if ev.typ&mask != 0 {
    52  					if !yield(ev.node) {
    53  						break
    54  					}
    55  				}
    56  				pop := ev.index
    57  				if in.events[pop].typ&mask == 0 {
    58  					// Subtrees do not contain types: skip them and pop.
    59  					i = pop + 1
    60  					continue
    61  				}
    62  			}
    63  			i++
    64  		}
    65  	}
    66  }
    67  
    68  // All[N] returns an iterator over all the nodes of type N.
    69  // N must be a pointer-to-struct type that implements ast.Node.
    70  //
    71  // Example:
    72  //
    73  //	for call := range All[*ast.CallExpr](in) { ... }
    74  //
    75  // See also the newer and more flexible [Cursor] API, which lets you
    76  // start the traversal at an arbitrary node, and reports each matching
    77  // node by its Cursor, enabling easier navigation.
    78  // The above example would be written thus:
    79  //
    80  //	for curCall := range in.Root().Preorder((*ast.CallExpr)(nil)) {
    81  //		call := curCall.Node().(*ast.CallExpr)
    82  //		...
    83  //	}
    84  func All[N interface {
    85  	*S
    86  	ast.Node
    87  }, S any](in *Inspector) iter.Seq[N] {
    88  
    89  	// To avoid additional dynamic call overheads,
    90  	// we duplicate rather than call the logic of PreorderSeq.
    91  
    92  	mask := typeOf((N)(nil))
    93  	return func(yield func(N) bool) {
    94  		for i := int32(0); i < int32(len(in.events)); {
    95  			ev := in.events[i]
    96  			if ev.index > i {
    97  				// push
    98  				if ev.typ&mask != 0 {
    99  					if !yield(ev.node.(N)) {
   100  						break
   101  					}
   102  				}
   103  				pop := ev.index
   104  				if in.events[pop].typ&mask == 0 {
   105  					// Subtrees do not contain types: skip them and pop.
   106  					i = pop + 1
   107  					continue
   108  				}
   109  			}
   110  			i++
   111  		}
   112  	}
   113  }
   114
View as plain text