/
scparser.go
363 lines (301 loc) · 9.56 KB
/
scparser.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
package scparser
import (
"fmt"
"go/ast"
"go/token"
"go/types"
"os"
"os/exec"
"strings"
"golang.org/x/tools/go/packages"
)
// Parse retrieves the source code of the specified function and its underlying functions
// within the Go module packages. It takes the package path and function name as input
// arguments and returns a formatted string containing the combined source code.
// The function will panic if the provided function is not found in the package path.
func Parse(funcPkgPath, funcName string, excludeRoot, codeOnly bool) string {
// Change the working directory to the given package directory
changeBack := changeDir(funcPkgPath)
defer changeBack()
funcSig, funcToFileAndPkg := initialize(funcName)
p := newParser(funcToFileAndPkg)
// Process the function and its underlying functions up to a depth of 5 (or 6 if root is excluded)
if excludeRoot {
p.processFunction(funcSig, 6)
} else {
p.processFunction(funcSig, 5)
}
return p.toString(excludeRoot, codeOnly)
}
type parser struct {
// rootPkg is the root package of the Go module
rootPkg *packages.Package
// funcToFileAndPkg is a map that stores the file and package for each function signature
funcToFileAndPkg map[*types.Signature]fileAndPkg
// functions is a map of packages to their function source code
functions map[*packages.Package]string
// pkgOrder is an ordered list of processed packages to maintain the order of processing
pkgOrder []*packages.Package
// seen is a map to keep track of already processed functions
seen map[*types.Signature]bool
}
// fileAndPkg is a struct that contains a pointer to an ast.File and a pointer to a packages.Package
type fileAndPkg struct {
file *ast.File
pkg *packages.Package
}
func newParser(funcToFileAndPkg map[*types.Signature]fileAndPkg) *parser {
return &parser{
funcToFileAndPkg: funcToFileAndPkg,
functions: make(map[*packages.Package]string),
seen: make(map[*types.Signature]bool),
}
}
// Convert functions into one string
func (p *parser) toString(excludeRoot, codeOnly bool) string {
var result string
for k, pkg := range p.pkgOrder {
if k == 0 && excludeRoot {
continue
}
if k > 1 || (k == 1 && !excludeRoot) {
result += formatPkg(pkg.Name, codeOnly) + "\n"
}
result += formatFunctions(p.functions[pkg], codeOnly)
if k < len(p.pkgOrder)-1 {
result += "\n\n"
}
}
return result
}
func formatPkg(pkgName string, codeOnly bool) string {
if codeOnly {
return `// ` + pkgName
}
return pkgName
}
func formatFunctions(functions string, codeOnly bool) string {
if codeOnly {
return functions
}
return "```" + functions + "```"
}
// processFunction processes a function with the provided package path and signature, and its underlying functions up to the specified depth
func (p *parser) processFunction(funcSig *types.Signature, depth int) {
// Check if the function signature has already been processed
// If so, return early to avoid processing it again
if p.seen[funcSig] {
return
}
// Try to get the file and package information associated with the function signature
// Return if the function signature is not found in the map (i.e. not in a go mod package)
f, ok := p.funcToFileAndPkg[funcSig]
if !ok {
return
}
// Inspect the AST (Abstract Syntax Tree) of the file
ast.Inspect(f.file, func(n ast.Node) bool {
// Check if the node is a function declaration
fn, ok := n.(*ast.FuncDecl)
if !ok || fn.Name == nil {
return true
}
// Check if the function signature matches the target function signature
sig, ok := f.pkg.TypesInfo.ObjectOf(fn.Name).Type().(*types.Signature)
if !ok || sig != funcSig {
return true
}
// Extract the source code of the function
funcSrc, err := extractSourceCode(f.pkg.Fset, f.file, fn)
if err != nil {
panic(err)
}
// If the package is not yet in the functions map, add it to the pkgOrder list
if _, ok := p.functions[f.pkg]; !ok {
p.pkgOrder = append(p.pkgOrder, f.pkg)
}
// Append the extracted function source code to the existing source code for the package, separated by a newline
p.functions[f.pkg] += "\n" + funcSrc
// Add the function to the map of processed functions
p.seen[funcSig] = true
// Process the underlying functions
p.processUnderlyingFunctions(f.pkg, fn, depth-1)
// Return false to stop AST traversal once the target function is found and processed
return false
})
}
// processUnderlyingFunctions processes the underlying functions called within the given function up to a specified depth
func (p *parser) processUnderlyingFunctions(pkg *packages.Package, fn *ast.FuncDecl, depth int) {
if depth <= 0 {
return
}
// Check if function decleration has body
if fn.Body == nil {
return
}
// Inspect the AST of the function body
ast.Inspect(fn.Body, func(n ast.Node) bool {
// Check if the node is a call expression (function call)
ce, ok := n.(*ast.CallExpr)
if !ok {
return true
}
var funcNode *ast.Ident
// Get the function node from the call expression
switch fun := ce.Fun.(type) {
case *ast.Ident:
funcNode = fun
case *ast.SelectorExpr:
funcNode = fun.Sel
default:
return true
}
if funcNode == nil {
return true
}
obj := pkg.TypesInfo.ObjectOf(funcNode)
if obj == nil {
return true
}
funcPkg := obj.Pkg()
if funcPkg == nil {
return true
}
// Get the function signature from the function node
funcSig, ok := obj.Type().(*types.Signature)
if !ok {
return true
}
// Process the underlying functions recursively
p.processFunction(funcSig, depth)
return true
})
}
// extractSourceCode extracts the source code of a function, including comments, from the provided file and function declaration
func extractSourceCode(fset *token.FileSet, file *ast.File, fn *ast.FuncDecl) (string, error) {
var sb strings.Builder
// Read the content of the file containing the function
fileContent, err := os.ReadFile(fset.Position(fn.Pos()).Filename)
if err != nil {
return "", err
}
// Split the file content into lines
lines := strings.Split(string(fileContent), "\n")
start := fset.Position(fn.Pos()).Line - 1
// Include comments above the function
if fn.Doc != nil {
for _, comment := range fn.Doc.List {
if comment == nil {
continue
}
commentStart := fset.Position(comment.Pos()).Line - 1
commentEnd := fset.Position(comment.End()).Line - 1
for i := commentStart; i <= commentEnd; i++ {
sb.WriteString(lines[i])
sb.WriteString("\n")
}
}
}
// Extract the function source code from the start to end line
end := fset.Position(fn.End()).Line - 1
for i := start; i <= end; i++ {
sb.WriteString(lines[i])
sb.WriteString("\n")
}
return sb.String(), nil
}
// parseGoModFile parses the go.mod file and returns a slice of package paths.
func parseGoModFile() []string {
content, err := os.ReadFile("go.mod")
if err != nil {
panic(err)
}
var goModPaths []string
lines := strings.Split(string(content), "\n")
for _, line := range lines {
fields := strings.Fields(line)
if len(fields) > 0 {
pkg := fields[0]
if pkg == `module` || pkg == `require` {
pkg = fields[1]
}
if pkg != `` && pkg != `)` && pkg != `(` && pkg != `require` && pkg != `go` {
goModPaths = append(goModPaths, pkg)
}
}
}
return goModPaths
}
// loadPackages loads and returns the (sub)packages in the current working directory.
func loadPackages() []*packages.Package {
err := exec.Command(`go`, `mod`, `vendor`).Run()
if err != nil {
fmt.Println("Warning: go mod vendor failed:", err)
}
pkgs, err := packages.Load(&packages.Config{
Mode: packages.NeedName | packages.NeedFiles | packages.NeedSyntax | packages.NeedTypes | packages.NeedTypesInfo,
}, "...")
if err != nil {
panic(err)
}
if len(pkgs) == 0 {
panic(`no packages found`)
}
return pkgs
}
// initialize searches for the target function with the provided name in the root package loads the go.mod packages
func initialize(funcName string) (*types.Signature, map[*types.Signature]fileAndPkg) {
goModPaths := parseGoModFile()
pkgs := loadPackages()
// Return the signature of the intial target function
var funcSig *types.Signature
// Collect all function signatures and their respective files
funcToFileAndPkg := make(map[*types.Signature]fileAndPkg)
for _, pkg := range pkgs {
// Skip packages not listed in go.mod
if !isGoModPkg(goModPaths, pkg.PkgPath) {
continue
}
for _, file := range pkg.Syntax {
ast.Inspect(file, func(n ast.Node) bool {
// Check if the node is a function declaration
fn, ok := n.(*ast.FuncDecl)
if !ok || fn.Name == nil {
return true
}
// Get the function signature from the TypesInfo of the package
obj := pkg.TypesInfo.ObjectOf(fn.Name)
if obj == nil {
return true
}
sig, ok := obj.Type().(*types.Signature)
if !ok || sig == nil {
return true
}
funcToFileAndPkg[sig] = fileAndPkg{
file: file,
pkg: pkg,
}
// If the function is the initial target function, store the signature
if pkg.PkgPath == goModPaths[0] && fn.Name.Name == funcName {
funcSig = sig
}
return true
})
}
if pkg.PkgPath == goModPaths[0] && funcSig == nil {
panic(fmt.Sprintf("Function %s not found in package path", funcName))
}
}
return funcSig, funcToFileAndPkg
}
// isGoModPkg checks if the provided package path is listed in the go.mod file
func isGoModPkg(goModPaths []string, pkgPath string) bool {
// Iterate through the goModPaths to check if the given package path matches or is a subpackage of any listed package
for _, path := range goModPaths {
if pkgPath == path || strings.HasPrefix(pkgPath, path+`/`) {
return true
}
}
return false
}