GO/ast编程(三)-ast生成代码

作者: 温岭夹糕 | 来源:发表于2024-01-13 10:56 被阅读0次

Golang标准库——go(1)
vue源码目录设计
AST与babel
编译原理术语解释
使用typescript来修改js文件
代码混淆
babel 与 AST
Go ast
AST
Swift day-1

1.常用节点学习

1.1Decl节点

demo1

func TestAstbuild(t *testing.T) {
    str :=
        `//@auther zjb
    package ast
    
    //@hi
    func main(){
    
    }
`
    src := []byte(str)
    fset := token.NewFileSet()
    f, err := parser.ParseFile(fset, "src.go", src, parser.ParseComments)
    if err != nil {
        t.Fatal(err)
    }
    var v AstVistor
    v.Com = make(map[string]string)

    ast.Inspect(f, func(n ast.Node) bool {
    //断点打在这里，不断查看n的类型和值
        return true
    })
}

ast.Inspect是ast.Walk的封装，内部都是使用深度优先遍历DFS，因此我们可以通过打断点的方式绘制出该ast结构

image.png 其中FuncDecl还可以继续展开成更复杂的形式

image.png
还记得前文提到的Decl是什么节点吗？是声明节点，那么这个FuncDecl是函数声明的节点，实际上，Decl声明节点共有三种：

FuncDecl
BadDecl 语法错误声明的占位符节点
GenDecl 通用声明节点(包括import导入、const常量、type类型、variable变量)，那么这里type可以是声明一个普通类型

type a int

也可以是声明一个接口或结构体类型

type a interface{}

为了验证我们的猜想，我们将上述demo变量修改

    str :=
        `
package ast

type a interface{}
`

经过打断点发现确实如此(实际上也是如此)

image.png

在GenDecl源码的注释中描述了以下3种声明类型

    //  token.IMPORT  *ImportSpec
    //  token.CONST   *ValueSpec
    //  token.TYPE    *TypeSpec
    //  token.VAR     *ValueSpec
    GenDecl struct {
        Doc    *CommentGroup // associated documentation; or nil
        TokPos token.Pos     // position of Tok
        Tok    token.Token   // IMPORT, CONST, TYPE, or VAR
        Lparen token.Pos     // position of '(', if any
        Specs  []Spec
        Rparen token.Pos // position of ')', if any
    }

那么如果有“通过接口定义来生成代码”这样的需求，我们就可以使用GenDecl来定位声明接口的位置

1.2 Spec节点类型

Spec在前文ast.File字段的学习中我们了解过(ImportSpec是导入声明)，我们在上文1.1学习到Spec接口是和GenDecl通用声明节点是强关联的（是以切片形式被保存在GenDecl.Spec中），其他两个很容易理解，这里主要看TypeSpec

    TypeSpec struct {
        Doc        *CommentGroup // associated documentation; or nil
        Name       *Ident        // type name
        TypeParams *FieldList    // type parameters; or nil
        Assign     token.Pos     // position of '=', if any
        Type       Expr          // *Ident, *ParenExpr, *SelectorExpr, *StarExpr, or any of the *XxxTypes
        Comment    *CommentGroup // line comments; or nil
    }

我们看到TypeSpec里面的Type是一个表达式节点，它可以是：

Ident 标识符表达式节点
ParenExpr 括号表达式节点
SelectorExpr 选择器表达式节点
StarExpr 索引表达式节点
XxxTypes 包括ArrayType,structType,FuncType,InterfacType,MapType,ChanType ，即可以和type关键字一起使用的那些

2.实战简单代码生成

2.1定位接口声明

myterface.go

package myast

type RouteAble interface {}

demo2

func TestMakeAst(t *testing.T) {
    fset := token.NewFileSet()
    f, err := parser.ParseFile(fset, "./myinterface.go", nil, 0)
    if err != nil {
        t.Fatal(err)
    }
    ast.Inspect(f, func(n ast.Node) bool {
        if genNode, ok := n.(*ast.GenDecl); ok {
            for _, spec := range genNode.Specs {
                if tp_spec, ok := spec.(*ast.TypeSpec); ok {
                    if _, ok := tp_spec.Type.(*ast.InterfaceType); ok {
                        t.Log(tp_spec.Name.Name)//TODO 代码生成
                    }
                }
            }
        }
        return true
    })
}

//output:
// RouteAble

实际上就是遍历ast并进行类型断言，有点前端js回调地狱的那味道了。

2.2手动构建ast

那既然定位好了，接口的信息也都能获取了，接下来就是要生成对应实现了，我们回想一下第一篇文章提到的gofmt核心是把源代码变成ast，再由ast生成代码，意味着我们也要构建对应的ast，那肯定不能在当前树操作

    out,err := parser.ParseFile(fset,"out.go","package myast",0)
    if err != nil {
        t.Fatal(err)
    }

接口生成对应的实现也应该是一个声明，它应该是一个*ast.TypeSpec，且字段Type是structType

                        structName := tp_spec.Name.Name+"Default"
                        new_typeSpec := &ast.TypeSpec{
                            Name: ast.NewIdent(structName),
                            Type: &ast.StructType{
                                Fields: &ast.FieldList{},
                            },
                        }

注意这里StructType.Fileds表示结构体成员，我们默认没有

    StructType struct {
        Struct     token.Pos  // position of "struct" keyword
        Fields     *FieldList // list of field declarations
        Incomplete bool       // true if (source) fields are missing in the Fields list
    }

type FieldList struct {
    Opening token.Pos //  左括号位置
    List    []*Field  // field list; or nil
    Closing token.Pos // 右括号位置
}

然后我们创建的TypeSpec需要一个GenDecl载体并挂载上去，同时GenDecl也要挂载到根节点ast.File上

                        new_decl := &ast.GenDecl{
                            Tok: token.TYPE,
                            Specs: []ast.Spec{new_typeSpec},
                        }
                        out.Decls = append(out.Decls, new_decl)

此时代码为

func TestMakeAst(t *testing.T) {
    fset := token.NewFileSet()
    f, err := parser.ParseFile(fset, "./myinterface.go", nil, 0)
    if err != nil {
        t.Fatal(err)
    }
    out,err := parser.ParseFile(fset,"out.go","package myast",0)
    if err != nil {
        t.Fatal(err)
    }
    ast.Inspect(f, func(n ast.Node) bool {
        if genNode, ok := n.(*ast.GenDecl); ok {
            for _, spec := range genNode.Specs {
                if tp_spec, ok := spec.(*ast.TypeSpec); ok {
                    if _, ok := tp_spec.Type.(*ast.InterfaceType); ok {
                        //TODO
                        structName := tp_spec.Name.Name+"Default"
                        new_typeSpec := &ast.TypeSpec{
                            Name: ast.NewIdent(structName),
                            Type: &ast.StructType{
                                Fields: &ast.FieldList{},
                            },
                        }
                        new_decl := &ast.GenDecl{
                            Tok: token.TYPE,
                            Specs: []ast.Spec{new_typeSpec},
                        }
                        out.Decls = append(out.Decls, new_decl)
                    }
                }
            }
        }
        return true
    })
}

2.3生成代码

使用go/format来翻译ast树成字节流，这里使用format.Node

    file, err := os.Create("out.go")
    if err != nil {
        t.Fatal(err)
    }
    defer file.Close()

    var buf bytes.Buffer
    err = format.Node(&buf, fset, out)
    if err != nil {
        t.Fatal(err)
    }
    _, _ = file.WriteString(buf.String())

效果如下

image.png
那么问题来了，函数怎么办，也就是最复杂的FuncDecl，而且节点类型还有stmt语句节点还没解析，那是不是两者有一些关系？我们回头看开头demo画的ast

image.png ast.FuncDecl下有一个ast.BlockStmt，两者之间确实有联系

3.函数

3.1InterfaceType

    // An InterfaceType node represents an interface type.
    InterfaceType struct {
        Interface  token.Pos  // position of "interface" keyword
        Methods    *FieldList // list of embedded interfaces, methods, or types
        Incomplete bool       // true if (source) methods or types are missing in the Methods list
    }

InterfaceType.Methods 是一个列表或集合，可以是内嵌的接口、类型或方法，这很容易理解，GO不支持继承，使用都是嵌入，FildList一些Expr表达式节点的装饰集合

type FieldList struct {
    Opening token.Pos //  左括号位置
    List    []*Field  // field list; or nil
    Closing token.Pos // 右括号位置
}
type Field struct {
    Doc     *CommentGroup // associated documentation; or nil
    Names   []*Ident      // field/method/(type) parameter names; or nil
    Type    Expr          // field/method/parameter type; or nil
    Tag     *BasicLit     // field tag; or nil
    Comment *CommentGroup // line comments; or nil
}

需要注意Field.Names是一个切片，如果她是一个方法节点，即Field.Type== *ast.FuncType，那么切片中只有函数名一个值
InterfaceType.Incomplete 如果没有方法或类型，即空接口eface，那么则为true

那么我们就可以利用methods字段提取出接口的方法了，我们给我们的接口添加方法

//myinterface.go
type RouteAble interface {
    Route(method, path string) string
}

获取接口方法demo2.1

//实战代码片段 inf 是通过断言*ast.interfaceType 获得的变量
                        if !inf.Incomplete {
                            for _, method := range inf.Methods.List {
                                //如果Interface.Methods.List中获取的Filed不是方法
                                //则暂时不处理
                                _, ok := method.Type.(*ast.FuncType)
                                if !ok {
                                    continue
                                }

                                var method_name string
                                for _, name := range method.Names {
                                    method_name = name.Name
                                }
                                if method_name == "" {
                                    panic("no method name found")
                                }
                                t.Log(method_name)

                            }                       
//output:Route

3.2FuncType

表达式节点FuncType 代表了一整个函数

func TestXxx(t string) string

源码

    FuncType struct {
        Func       token.Pos  // position of "func" keyword (token.NoPos if there is no "func")
        TypeParams *FieldList // type parameters; or nil
        Params     *FieldList // (incoming) parameters; non-nil
        Results    *FieldList // (outgoing) results; or nil
    }

params和Results顾名思义表示参数和返回值

3.3FuncDecl

我们想要生成的ast是什么样的，是不是类似下面代码的ast?
demo3

package myast

type H struct {
}

func (h *H) Hello(a string) string {
    return "hello"
}

我们通过ast.Inspect和打断点画出ast.FuncDecl分支

image.png

观察到它有一个方法接收者、函数名、参数、返回类型、函数体和返回值这些要素，我们构建FuncDecl也需要这些，查看FuncDecl源码

    FuncDecl struct {
        Doc  *CommentGroup // associated documentation; or nil
        Recv *FieldList    // receiver (methods); or nil (functions)
        Name *Ident        // function/method name
        Type *FuncType     // function signature: type and value parameters, results, and position of "func" keyword
        Body *BlockStmt    // function body; or nil for external (non-Go) function
    }

Recv是方法的接收者，如果不是结构体的函数，则为nil；如果是则它的Field.Names包含接受指针名，Field.Type==ast.StarExpr，即下面这个分支 image.png
Name函数名属性
Type FuncType见上一节3.2中分析
Body 是BlockStmt类型,stmt是一个接口，函数体里面有哪些句子就往里面塞各自的语句节点，如for循环就塞forstmt，if就塞ifstmt

    BlockStmt struct {
        Lbrace token.Pos // position of "{"
        List   []Stmt
        Rbrace token.Pos // position of "}", if any (may be absent due to syntax error)
    }

demo3只有返回语句

image.png

    ReturnStmt struct {
        Return  token.Pos // position of "return" keyword
        Results []Expr    // result expressions; or nil
    }

我们看到ReturnStmt.Results是切片类型，这也是GO能实现多返回值的条件之一（主要还是函数栈）

3.4构建我们自己的FuncDecl

又要返回demo2.1继续扩展，注意FiledList不要为nil，需要初始化
demo2.2

                        out.Decls = append(out.Decls, new_decl)
                        //先不支持内嵌 获取接口方法
                        if !inf.Incomplete {
                            for _, method := range inf.Methods.List {
                                //如果Interface.Methods.List中获取的Filed不是方法
                                //则暂时不处理
                                method_type, ok := method.Type.(*ast.FuncType)
                                if !ok {
                                    continue
                                }

                                var method_name string
                                for _, name := range method.Names {
                                    method_name = name.Name
                                }
                                if method_name == "" {
                                    panic("no method name found")
                                }
                                //方法接收者是存储在FuncDecl的FieldList类型中
                                f_recv := &ast.FieldList{
                                    List: []*ast.Field{
                                        {
                                            // 这里我们取接口的首字母小写
                                            Names: []*ast.Ident{ast.NewIdent(strings.ToLower(string(structName[0])))},
                                            Type:  &ast.StarExpr{X: ast.NewIdent(structName)},
                                        },
                                    },
                                }

                                //方法参数存储在FuncDecl的FuncType的FieldList中
                                f_type := &ast.FuncType{
                                    Params: &ast.FieldList{
                                        List: method_type.Params.List,
                                    },
                                    Results: &ast.FieldList{
                                        List: method_type.Results.List,
                                    },
                                }
                                //给每个返回类型创建默认返回值
                                var res []ast.Expr
                                for _, field := range method_type.Results.List {
                                    switch ident := reflect.ValueOf(field.Type).Elem().Interface().(type) {
                                    case ast.Ident:
                                        res = append(res, CreateZeroExpr(ident.Name))
                                    default:
                                        res = append(res, &ast.Ident{Name: "nil"})
                                    }
                                }
                                //方法的函数体
                                f_body := &ast.BlockStmt{
                                    List: []ast.Stmt{
                                        &ast.ReturnStmt{Results: res},
                                    },
                                }

                                //注解
                                f_comment := &ast.CommentGroup{
                                    List: []*ast.Comment{
                                        {
                                            Text: "//@auther xxxx",
                                        },
                                        {
                                            Text: "//TODO",
                                        },
                                    },
                                }

                                //组装
                                new_funcDecl := &ast.FuncDecl{
                                    Name: ast.NewIdent(method_name),
                                    Doc:  f_comment,
                                    Recv: f_recv,
                                    Type: f_type,
                                    Body: f_body,
                                }
                                //添加到根节点下
                                out.Decls = append(out.Decls, new_funcDecl)
                            }

func CreateZeroExpr(name string) ast.Expr {
    switch name {
    case "bool":
        return &ast.Ident{Name: "false"}
    case "int","uint":
        return &ast.BasicLit{Kind: token.INT, Value: "0"}
    case "float32", "float64":
        return &ast.BasicLit{Kind: token.FLOAT, Value: "0.0"}
    case "string":
        return &ast.BasicLit{Kind: token.STRING, Value: `""`}
    default:
        return &ast.Ident{Name: "nil"}
    }
}

效果如下

image.png

需要注意如果接口有包导入，我们这里也要添加，但demo并未写出，需要自行添加，至于怎么加，自己分析ast结构照猫画虎去吧(需要注意import也是ast.GenDecl类型)

至此我们用3节内容已经基本掌握了ast编程，最后一节把编译内容前端的语法分析和后端部分补全

补充

1.包导入

// 解析import
func parseImports(file, out *ast.File) {
    import_decl := &ast.GenDecl{
        Tok:   token.IMPORT,
        Specs: []ast.Spec{},
    }
    for _, imp := range file.Imports {
        new_imp := &ast.ImportSpec{
            Path: &ast.BasicLit{
                Kind:  token.STRING,
                Value: imp.Path.Value,
            },
        }
        import_decl.Specs = append(import_decl.Specs, new_imp)
    }
    out.Decls = append(out.Decls, import_decl)
}