在进入这一阶段之前可以先阅读下这篇 Understanding Go programs with go/parser ，对go的AST有一定的感性认识。
比如这个我们一直在默默测试的 main.go
1: package main
2:
3: import (
4: "context"
5: "fmt"
6: )
7:
8: func hello(ctx context.Context) {
9: fmt.Println("hello world")
10: }
11:
12: func main() {
13: hello(context.Background())
14: }
15:
经过 go/parser.ParseFile
1: package main
2:
3: import (
4: "go/parser"
5: "go/token"
6: "io/ioutil"
7: "os"
8: "path/filepath"
9:
10: "github.com/davecgh/go-spew/spew"
11: )
12:
13: func main() {
14: home, err := os.UserHomeDir()
15: if err != nil {
16: panic(err)
17: }
18:
19: file := filepath.Join(home, "src", "github.com", "shanexu", "go-playground", "helloworld", "main.go")
20: src, err := ioutil.ReadFile(file)
21: if err != nil {
22: panic(err)
23: }
24:
25: fset := token.NewFileSet()
26: f, err := parser.ParseFile(fset, "main.go", src, parser.AllErrors)
27: if err != nil {
28: panic(err)
29: }
30:
31: spew.Dump(f)
32: }
得到AST：
而 cmd/compile/internal/syntax.Parse 的结果则都是 cmd/compile/internal/syntax 包下的类型。基本都能和 go/ast 下的类型一一对应。比如：ast.File和syntax.File，ast.ExprStmt和syntax.ExprStmt，ast.Ident和syntax.Name，ast.CallExpr和syntax.CallExpr，ast.SelectorExpr和syntax.SelectorExpr。
我们现在可以对AST出手。
例如我们想针对所有main包下以hello开头的函数在进入方法时打印“start ${方法名}...”，在离开方法是打印“stop ${方法名}...”， 则可以在文件 cmd/compile/internal/gc/noder.go Parse结束后，修改p.file的值。代码如下：
52: p.file, _ = syntax.Parse(base, f, p.error, p.pragma, syntax.CheckBranches) // errors are tracked via p.error
53: if p.file.PkgName.Value == "main" {
54: for _, d := range p.file.DeclList {
55: d, _ := d.(*syntax.FuncDecl)
56: if d == nil {
57: continue
58: }
59: if !strings.HasPrefix(d.Name.Value, "hello") {
60: continue
61: }
62: d.Body.List = append([]syntax.Stmt{
63: &syntax.ExprStmt{
64: X: &syntax.CallExpr{
65: Fun: &syntax.SelectorExpr{
66: X: &syntax.Name{Value: "fmt"},
67: Sel: &syntax.Name{Value: "Println"},
68: },
69: ArgList: []syntax.Expr{
70: &syntax.BasicLit{
71: Value: strconv.Quote("start " + d.Name.Value + "..."),
72: Kind: syntax.StringLit,
73: },
74: },
75: },
76: },
77: &syntax.CallStmt{
78: Tok: syntax.Defer,
79: Call: &syntax.CallExpr{
80: Fun: &syntax.SelectorExpr{
81: X: &syntax.Name{Value: "fmt"},
82: Sel: &syntax.Name{Value: "Println"},
83: },
84: ArgList: []syntax.Expr{
85: &syntax.BasicLit{
86: Value: strconv.Quote("stop " + d.Name.Value + "..."),
87: Kind: syntax.StringLit,
88: },
89: },
90: },
91: },
92: }, d.Body.List...)
93: }
94: }
修改完后需要重新编译golang。
cd $HOME/src/github.com/golang/go/src
./make.bash
cd $HOME/src/github.com/shanexu/go-playground
go clean -cache
go run helloworld/main.go
看到如下结果：
start hello...
hello world
stop hello...
在修改了AST之后，实际上hello方法的源码应该长这样：
func hello(ctx context.Context) {
fmt.Println("start hello...")
defer fmt.Println("stop hello...")
fmt.Println("hello world")
}
这里有个问题，插入的代码新引入了fmt包，如果原始代码里面没有引入fmt包会怎样？
$ go clean -cache; go run helloworld/main.go
# command-line-arguments
helloworld/main.go:7:16: undefined: fmt in fmt.Println
果然编译失败了。
在语法树中按需插入import呢？
52: if p.file.PkgName.Value == "main" {
53: for _, d := range p.file.DeclList {
54: d, _ := d.(*syntax.FuncDecl)
55: if d == nil {
56: continue
57: }
58: if !strings.HasPrefix(d.Name.Value, "hello") {
59: continue
60: }
61: hasHello = true
62: d.Body.List = append([]syntax.Stmt{
63: &syntax.ExprStmt{
64: X: &syntax.CallExpr{
65: Fun: &syntax.SelectorExpr{
66: X: &syntax.Name{Value: "fmt"},
67: Sel: &syntax.Name{Value: "Println"},
68: },
69: ArgList: []syntax.Expr{
70: &syntax.BasicLit{
71: Value: strconv.Quote("start " + d.Name.Value + "..."),
72: Kind: syntax.StringLit,
73: },
74: },
75: },
76: },
77: &syntax.CallStmt{
78: Tok: syntax.Defer,
79: Call: &syntax.CallExpr{
80: Fun: &syntax.SelectorExpr{
81: X: &syntax.Name{Value: "fmt"},
82: Sel: &syntax.Name{Value: "Println"},
83: },
84: ArgList: []syntax.Expr{
85: &syntax.BasicLit{
86: Value: strconv.Quote("stop " + d.Name.Value + "..."),
87: Kind: syntax.StringLit,
88: },
89: },
90: },
91: },
92: }, d.Body.List...)
93: }
94: }
95: if hasHello {
96: hasFmtImport := false
97: for _, d := range p.file.DeclList {
98: d, _ := d.(*syntax.ImportDecl)
99: if d == nil {
100: continue
101: }
102: if d.Path.Value != "fmt" {
103: continue
104: }
105: hasFmtImport = true
106: break
107: }
108: if !hasFmtImport {
109: p.file.DeclList = append([]syntax.Decl{
110: &syntax.ImportDecl{
111: Path: &syntax.BasicLit{
112: Value: `"fmt"`, Kind: syntax.StringLit,
113: },
114: },
115: }, p.file.DeclList...)
116: }
117: }
重新编译golang，执行go run
$ go clean -cache; go run helloworld/main.go
# command-line-arguments
helloworld/main.go:1:9: can't find import: "fmt"
根据错误信息找到，抛错位置，发现parseFile过程中，有findpkg的过程，起调用栈如下。
包变量 packageFile 这个map中获取相关包的信息。
988: if packageFile != nil {
989: file, ok = packageFile[name]
990: return file, ok
991: }
包变量 packageFile 在 readImportCfg 方法中初始化。
所以 readImportCfg 方法的入参就是，调用compile命令时，选项 -importcfg 的值。
805: func readImportCfg(file string) {
806: packageFile = map[string]string{}
807: data, err := ioutil.ReadFile(file)
808: if err != nil {
809: log.Fatalf("-importcfg: %v", err)
810: }
要解决这个问题，就需要在 importcfg 文件中加入新加的包。
问题又回到的了 importcfg 文件的生成过程了。
文件 cmd/go/internal/work/exec.go 第634行，有关importcfg的内容的逻辑。
634: // Prepare Go import config.
635: // We start it off with a comment so it can't be empty, so icfg.Bytes() below is never nil.
636: // It should never be empty anyway, but there have been bugs in the past that resulted
637: // in empty configs, which then unfortunately turn into "no config passed to compiler",
638: // and the compiler falls back to looking in pkg itself, which mostly works,
639: // except when it doesn't.
640: var icfg bytes.Buffer
641: fmt.Fprintf(&icfg, "# import config\n")
642: for i, raw := range a.Package.Internal.RawImports {
643: final := a.Package.Imports[i]
644: if final != raw {
645: fmt.Fprintf(&icfg, "importmap %s=%s\n", raw, final)
646: }
647: }
现在归结于 Package.Imports 的值的设置了。
文件 go/build/build.go=，先用 =go/parser.ParseFile 解析源文件，然后获取其中的imports。
847: pf, err := parser.ParseFile(fset, filename, data, parser.ImportsOnly|parser.ParseComments)
848: if err != nil {
849: badFile(err)
850: continue
851: }
修改代码如下：
847: pf, err := parser.ParseFile(fset, filename, data, parser.ImportsOnly|parser.ParseComments)
848: if err != nil {
849: badFile(err)
850: continue
851: }
852:
853: if strings.HasSuffix(filename, "helloworld/main.go") {
854: hasFmtImport := false
855: for _, i := range pf.Imports {
856: if i.Path.Value == `"fmt"` {
857: hasFmtImport = true
858: break
859: }
860: }
861: if !hasFmtImport {
862: pf.Imports = append(pf.Imports, &ast.ImportSpec{
863: Path: &ast.BasicLit{
864: Value: `"fmt"`,
865: Kind: token.STRING,
866: },
867: })
868: if len(pf.Decls) > 0 {
869: d, ok := pf.Decls[0].(*ast.GenDecl)
870: if ok {
871: d.Specs = append(d.Specs, &ast.ImportSpec{
872: Path: &ast.BasicLit{
873: Kind: token.STRING,
874: Value: `"fmt"`,
875: },
876: })
877: }
878: }
879: }
880: }
重新编译golang，执行go run
$ go clean -cache; go run helloworld/main.go
start hello...
stop hello...
成功加上了import。
