iter.go (2971B)
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 }