Skip to content

Commit aeaf596

Browse files
welkeyeverAsterDY
authored andcommitted
refactor: getMain also return function node (cloudwego#12)
1 parent 59d460b commit aeaf596

1 file changed

Lines changed: 5 additions & 3 deletions

File tree

src/compress/golang/plugin/go_ast.go

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -394,17 +394,19 @@ type SingleFunction struct {
394394
// }
395395
// }
396396

397-
func (p *goParser) getMain() *MainStream {
397+
func (p *goParser) getMain() (*MainStream, *Function) {
398398
m := &MainStream{
399399
RelatedFunctions: make([]SingleFunction, 0),
400400
}
401401

402402
var functionCalledInMain = map[string]*Function{}
403+
var mainFunc *Function
403404

404405
Out:
405406
for _, v := range p.processedPkgFunctions {
406407
for _, vv := range v {
407408
if vv.Name == "main" {
409+
mainFunc = vv
408410
m.MainFunc = vv.Content
409411
for k, v := range vv.InternalFunctionCalls {
410412
functionCalledInMain[k] = v
@@ -418,7 +420,7 @@ Out:
418420
}
419421
}
420422
p.fillFunctionContent(functionCalledInMain, &m.RelatedFunctions)
421-
return m
423+
return m, mainFunc
422424
}
423425

424426
func (p *goParser) fillFunctionContent(f map[string]*Function, fl *[]SingleFunction) {
@@ -490,7 +492,7 @@ func main() {
490492
}
491493

492494
// p.generateStruct()
493-
m := p.getMain()
495+
m, _ := p.getMain()
494496

495497
out := bytes.NewBuffer(nil)
496498
encoder := json.NewEncoder(out)

0 commit comments

Comments
 (0)