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