|
1 | | -// Package spawner handles //goroutinectx:spawner directives. |
| 1 | +// Package spawner handles //goroutinectx:spawner directives and -external-spawner flag. |
2 | 2 | package spawner |
3 | 3 |
|
4 | 4 | import ( |
5 | 5 | "go/ast" |
6 | 6 | "go/types" |
7 | 7 | "strings" |
| 8 | + "unicode" |
8 | 9 |
|
9 | 10 | "golang.org/x/tools/go/analysis" |
10 | 11 | ) |
11 | 12 |
|
| 13 | +// FuncSpec holds parsed components of a spawner function specification. |
| 14 | +type FuncSpec struct { |
| 15 | + PkgPath string |
| 16 | + TypeName string // empty for package-level functions |
| 17 | + FuncName string |
| 18 | +} |
| 19 | + |
12 | 20 | // Map tracks functions marked with //goroutinectx:spawner. |
13 | 21 | // These functions are expected to spawn goroutines with their func arguments. |
14 | | -type Map map[*types.Func]struct{} |
| 22 | +type Map struct { |
| 23 | + local map[*types.Func]struct{} // from directives |
| 24 | + external []FuncSpec // from -external-spawner flag |
| 25 | +} |
15 | 26 |
|
16 | 27 | // IsSpawner checks if a function is marked as a spawner. |
17 | | -func (m Map) IsSpawner(fn *types.Func) bool { |
18 | | - _, ok := m[fn] |
| 28 | +func (m *Map) IsSpawner(fn *types.Func) bool { |
| 29 | + if m == nil { |
| 30 | + return false |
| 31 | + } |
| 32 | + |
| 33 | + // Check local map first (directive-based) |
| 34 | + if _, ok := m.local[fn]; ok { |
| 35 | + return true |
| 36 | + } |
| 37 | + |
| 38 | + // Check external specs (flag-based) |
| 39 | + return m.matchesExternal(fn) |
| 40 | +} |
| 41 | + |
| 42 | +// Len returns the total number of spawners (local + external). |
| 43 | +func (m *Map) Len() int { |
| 44 | + if m == nil { |
| 45 | + return 0 |
| 46 | + } |
| 47 | + |
| 48 | + return len(m.local) + len(m.external) |
| 49 | +} |
| 50 | + |
| 51 | +// matchesExternal checks if fn matches any external spec. |
| 52 | +func (m *Map) matchesExternal(fn *types.Func) bool { |
| 53 | + for _, spec := range m.external { |
| 54 | + if matchesSpec(fn, spec) { |
| 55 | + return true |
| 56 | + } |
| 57 | + } |
19 | 58 |
|
20 | | - return ok |
| 59 | + return false |
21 | 60 | } |
22 | 61 |
|
23 | | -// Build scans files for functions marked with the directive. |
24 | | -func Build(pass *analysis.Pass) Map { |
25 | | - m := make(Map) |
| 62 | +// matchesSpec checks if a function matches a FuncSpec. |
| 63 | +func matchesSpec(fn *types.Func, spec FuncSpec) bool { |
| 64 | + if fn.Name() != spec.FuncName { |
| 65 | + return false |
| 66 | + } |
| 67 | + |
| 68 | + pkg := fn.Pkg() |
| 69 | + if pkg == nil || pkg.Path() != spec.PkgPath { |
| 70 | + return false |
| 71 | + } |
| 72 | + |
| 73 | + // Check if it's a method |
| 74 | + sig := fn.Type().(*types.Signature) |
| 75 | + recv := sig.Recv() |
| 76 | + |
| 77 | + if spec.TypeName == "" { |
| 78 | + // Package-level function: should have no receiver |
| 79 | + return recv == nil |
| 80 | + } |
| 81 | + |
| 82 | + // Method: should have receiver of correct type |
| 83 | + if recv == nil { |
| 84 | + return false |
| 85 | + } |
| 86 | + |
| 87 | + recvType := recv.Type() |
| 88 | + // Handle pointer receivers |
| 89 | + if ptr, ok := recvType.(*types.Pointer); ok { |
| 90 | + recvType = ptr.Elem() |
| 91 | + } |
| 92 | + |
| 93 | + named, ok := recvType.(*types.Named) |
| 94 | + if !ok { |
| 95 | + return false |
| 96 | + } |
| 97 | + |
| 98 | + return named.Obj().Name() == spec.TypeName |
| 99 | +} |
| 100 | + |
| 101 | +// Build scans files for functions marked with the directive |
| 102 | +// and parses the external spawner flag. |
| 103 | +func Build(pass *analysis.Pass, externalSpawners string) *Map { |
| 104 | + m := &Map{ |
| 105 | + local: make(map[*types.Func]struct{}), |
| 106 | + external: parseExternal(externalSpawners), |
| 107 | + } |
26 | 108 |
|
27 | 109 | for _, file := range pass.Files { |
28 | | - buildSpawnersForFile(pass, file, m) |
| 110 | + buildSpawnersForFile(pass, file, m.local) |
29 | 111 | } |
30 | 112 |
|
31 | 113 | return m |
32 | 114 | } |
33 | 115 |
|
| 116 | +// parseExternal parses the -external-spawner flag value. |
| 117 | +// Format: comma-separated list of "pkg/path.Func" or "pkg/path.Type.Method". |
| 118 | +func parseExternal(s string) []FuncSpec { |
| 119 | + if s == "" { |
| 120 | + return nil |
| 121 | + } |
| 122 | + |
| 123 | + var specs []FuncSpec |
| 124 | + |
| 125 | + for part := range strings.SplitSeq(s, ",") { |
| 126 | + part = strings.TrimSpace(part) |
| 127 | + if part == "" { |
| 128 | + continue |
| 129 | + } |
| 130 | + |
| 131 | + spec := parseFunc(part) |
| 132 | + specs = append(specs, spec) |
| 133 | + } |
| 134 | + |
| 135 | + return specs |
| 136 | +} |
| 137 | + |
| 138 | +// parseFunc parses a single spawner function string into components. |
| 139 | +// Format: "pkg/path.Func" or "pkg/path.Type.Method". |
| 140 | +func parseFunc(s string) FuncSpec { |
| 141 | + spec := FuncSpec{} |
| 142 | + |
| 143 | + lastDot := strings.LastIndex(s, ".") |
| 144 | + if lastDot == -1 { |
| 145 | + spec.FuncName = s |
| 146 | + |
| 147 | + return spec |
| 148 | + } |
| 149 | + |
| 150 | + spec.FuncName = s[lastDot+1:] |
| 151 | + prefix := s[:lastDot] |
| 152 | + |
| 153 | + // Check if there's another dot (indicating Type.Method) |
| 154 | + // Type names start with uppercase in Go. |
| 155 | + secondLastDot := strings.LastIndex(prefix, ".") |
| 156 | + if secondLastDot != -1 { |
| 157 | + possibleType := prefix[secondLastDot+1:] |
| 158 | + if len(possibleType) > 0 && unicode.IsUpper(rune(possibleType[0])) { |
| 159 | + spec.TypeName = possibleType |
| 160 | + spec.PkgPath = prefix[:secondLastDot] |
| 161 | + |
| 162 | + return spec |
| 163 | + } |
| 164 | + } |
| 165 | + |
| 166 | + spec.PkgPath = prefix |
| 167 | + |
| 168 | + return spec |
| 169 | +} |
| 170 | + |
34 | 171 | // buildSpawnersForFile scans a single file for spawner directives. |
35 | | -func buildSpawnersForFile(pass *analysis.Pass, file *ast.File, m Map) { |
| 172 | +func buildSpawnersForFile(pass *analysis.Pass, file *ast.File, m map[*types.Func]struct{}) { |
36 | 173 | // Build a map of line -> comment for quick lookup |
37 | 174 | lineComments := make(map[int]string) |
38 | 175 |
|
|
0 commit comments