Golang AST語法樹使用教程及示例

阿新 • • 發佈：2019-12-31

許多自動化程式碼生成工具都離不開語法樹分析，例如goimport，gomock，wire等專案都離不開語法樹分析。基於語法樹分析，可以實現許多有趣實用的工具。本篇將結合示例，展示如何基於ast標準包操作語法樹。

本篇中的程式碼的完整示例可以在這裡找到：ast-example

Quick Start

首先我們看下語法樹長什麼樣子，以下程式碼將列印./demo.go檔案的語法樹：

package main

import (
	"go/ast"
	"go/parser"
	"go/token"
	"log"
	"path/filepath"
)

func main() {
	fset := token.NewFileSet()
	// 這裡取絕對路徑，方便打印出來的語法樹可以轉跳到編輯器
	path,_ := filepath.Abs("./demo.go" 
)
	f,err := parser.ParseFile(fset,path,nil,parser.AllErrors)
	if err != nil {
		log.Println(err)
		return
	}
	// 列印語法樹
	ast.Print(fset,f)
}

複製程式碼

package main

import (
	"context"
)

// Foo 結構體
type Foo struct {
	i int
}

// Bar 介面
type Bar interface {
	Do(ctx context.Context) error
}

// main方法
func main 
() {
    a := 1
}
複製程式碼

demo.go檔案已儘量簡化，但其語法樹的輸出內容依舊十分龐大。我們擷取部分來做一些簡要的說明。

首先是檔案所屬的包名，和其宣告在檔案中的位置：

 0  *ast.File {
     1  .  Package: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:1:1
     2  .  Name: *ast.Ident {
     3  .  .  NamePos: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:1:9
     4  .  .  Name: "main" 

     5  .  }
     ...
複製程式碼

緊接著是Decls，也就是Declarations，其包含了宣告的一些變數，方法，介面等：

...
     6  .  Decls: []ast.Decl (len = 4) {
     7  .  .  0: *ast.GenDecl {
     8  .  .  .  TokPos: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:3:1
     9  .  .  .  Tok: import
    10  .  .  .  Lparen: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:3:8
    11  .  .  .  Specs: []ast.Spec (len = 1) {
    12  .  .  .  .  0: *ast.ImportSpec {
    13  .  .  .  .  .  Path: *ast.BasicLit {
    14  .  .  .  .  .  .  ValuePos: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:4:2
    15  .  .  .  .  .  .  Kind: STRING
    16  .  .  .  .  .  .  Value: "\"context\""
    17  .  .  .  .  .  }
    18  .  .  .  .  .  EndPos: -
    19  .  .  .  .  }
    20  .  .  .  }
    21  .  .  .  Rparen: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:5:1
    22  .  .  }
 ....
複製程式碼

可以看到該語法樹包含了4條Decl記錄，我們取第一條記錄為例，該記錄為*ast.GenDecl型別。不難看出這條記錄對應的是我們的import程式碼段。始位置(TokPos)，左右括號的位置(Lparen,Rparen)，和import的包（Specs）等資訊都能從語法樹中得到。

語法樹的列印信來自ast.File結構體：

// 該結構體位於標準包 go/ast/ast.go 中，有興趣可以轉跳到原始碼閱讀更詳盡的註釋
type File struct {
	Doc        *CommentGroup   // associated documentation; or nil
	Package    token.Pos       // position of "package" keyword
	Name       *Ident          // package name
	Decls      []Decl          // top-level declarations; or nil
	Scope      *Scope          // package scope (this file only)
	Imports    []*ImportSpec   // imports in this file
	Unresolved []*Ident        // unresolved identifiers in this file
	Comments   []*CommentGroup // list of all comments in the source file
}
複製程式碼

結合註釋和欄位名我們大概知道每個欄位的含義，接下來我們詳細梳理一下語法樹的組成結構。

Node節點

整個語法樹由不同的node組成，從原始碼註釋中可以得知主要有如下三種node：

但實際在程式碼，出現了第四種node：Spec Node，每種node都有專門的介面定義：

...
// All node types implement the Node interface.
type Node interface {
	Pos() token.Pos // position of first character belonging to the node
	End() token.Pos // position of first character immediately after the node
}

// All expression nodes implement the Expr interface.
type Expr interface {
	Node
	exprNode()
}

// All statement nodes implement the Stmt interface.
type Stmt interface {
	Node
	stmtNode()
}

// All declaration nodes implement the Decl interface.
type Decl interface {
	Node
	declNode()
}
...

// A Spec node represents a single (non-parenthesized) import,// constant,type,or variable declaration.
//
type (
	// The Spec type stands for any of *ImportSpec,*ValueSpec,and *TypeSpec.
	Spec interface {
		Node
		specNode()
	}
....
)
複製程式碼

可以看到所有的node都繼承Node介面，記錄了node的開始和結束位置。還記得Quick Start示例中的Decls嗎？它正是declaration nodes。除去上述四種使用介面進行分類的node，還有些node沒有再額外定義介面細分類別，僅實現了Node介面，為了方便描述，在本篇中我把這些節點稱為common node。 $GOROOT/src/go/ast/ast.go列舉了所有所有節點的實現，我們從中挑選幾個作為例子，感受一下它們的區別。

Expression and Type

先來看expression node。

...
	// An Ident node represents an identifier.
	Ident struct {
		NamePos token.Pos // identifier position
		Name    string    // identifier name
		Obj     *Object   // denoted object; or nil
	}
...
複製程式碼

Indent（identifier）表示一個識別符號，比如Quick Start示例中表示包名的Name欄位就是一個expression node：

 0  *ast.File {
     1  .  Package: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:1:1
     2  .  Name: *ast.Ident { <----
     3  .  .  NamePos: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:1:9
     4  .  .  Name: "main"
     5  .  }
     ...
複製程式碼

接下來是type node。

...
	// A StructType node represents a struct type.
	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
	}

	// Pointer types are represented via StarExpr nodes.

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

	// An InterfaceType node represents an interface type.
	InterfaceType struct {
		Interface  token.Pos  // position of "interface" keyword
		Methods    *FieldList // list of methods
		Incomplete bool       // true if (source) methods are missing in the Methods list
	}
...
複製程式碼

type node很好理解，它包含一些複合型別，例如在Quick Start中出現的StructType,FuncType和InterfaceType。

Statement

賦值語句，控制語句（if，else,for，select...）等均屬於statement node。

...
	// An AssignStmt node represents an assignment or
	// a short variable declaration.
	//
	AssignStmt struct {
		Lhs    []Expr
		TokPos token.Pos   // position of Tok
		Tok    token.Token // assignment token,DEFINE
		Rhs    []Expr
	}
...

	// An IfStmt node represents an if statement.
	IfStmt struct {
		If   token.Pos // position of "if" keyword
		Init Stmt      // initialization statement; or nil
		Cond Expr      // condition
		Body *BlockStmt
		Else Stmt // else branch; or nil
	}
...
複製程式碼

例如Quick Start中，我們在main函式中對變數a賦值的程式片段就屬於AssignStmt:

...
 174  .  .  .  Body: *ast.BlockStmt {
   175  .  .  .  .  Lbrace: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:18:13
   176  .  .  .  .  List: []ast.Stmt (len = 1) {
   177  .  .  .  .  .  0: *ast.AssignStmt { <--- 這裡
   178  .  .  .  .  .  .  Lhs: []ast.Expr (len = 1) {
   179  .  .  .  .  .  .  .  0: *ast.Ident {
   180  .  .  .  .  .  .  .  .  NamePos: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:19:2
   181  .  .  .  .  .  .  .  .  Name: "a"
...
複製程式碼

Spec Node

Spec node只有3種，分別是ImportSpec，ValueSpec和TypeSpec：

	// An ImportSpec node represents a single package import.
	ImportSpec struct {
		Doc     *CommentGroup // associated documentation; or nil
		Name    *Ident        // local package name (including "."); or nil
		Path    *BasicLit     // import path
		Comment *CommentGroup // line comments; or nil
		EndPos  token.Pos     // end of spec (overrides Path.Pos if nonzero)
	}

	// A ValueSpec node represents a constant or variable declaration
	// (ConstSpec or VarSpec production).
	//
	ValueSpec struct {
		Doc     *CommentGroup // associated documentation; or nil
		Names   []*Ident      // value names (len(Names) > 0)
		Type    Expr          // value type; or nil
		Values  []Expr        // initial values; or nil
		Comment *CommentGroup // line comments; or nil
	}

	// A TypeSpec node represents a type declaration (TypeSpec production).
	TypeSpec struct {
		Doc     *CommentGroup // associated documentation; or nil
		Name    *Ident        // type name
		Assign  token.Pos     // position of '=',if any
		Type    Expr          // *Ident,*ParenExpr,*SelectorExpr,*StarExpr,or any of the *XxxTypes
		Comment *CommentGroup // line comments; or nil
	}
複製程式碼

ImportSpec表示一個單獨的import，ValueSpec表示一個常量或變數的宣告，TypeSpec則表示一個type宣告。例如在Quick Start示例中，出現了ImportSpec和TypeSpec

import (
	"context" // <--- 這裡是一個ImportSpec node
)

// Foo 結構體
type Foo struct { // <--- 這裡是一個TypeSpec node
	i int
}
複製程式碼

在語法樹的列印結果中可以看到對應的輸出，小夥伴們可自行查詢。

Declaration Node

Declaration node也只有三種：

...
type (
	// A BadDecl node is a placeholder for declarations containing
	// syntax errors for which no correct declaration nodes can be
	// created.
	//
	BadDecl struct {
		From,To token.Pos // position range of bad declaration
	}

	// A GenDecl node (generic declaration node) represents an import,type or variable declaration. A valid Lparen position
	// (Lparen.IsValid()) indicates a parenthesized declaration.
	//
	// Relationship between Tok value and Specs element type:
	//
	//	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,VAR
		Lparen token.Pos     // position of '(',if any
		Specs  []Spec
		Rparen token.Pos // position of ')',if any
	}

	// A FuncDecl node represents a function declaration.
	FuncDecl struct {
		Doc  *CommentGroup // associated documentation; or nil
		Recv *FieldList    // receiver (methods); or nil (functions)
		Name *Ident        // function/method name
		Type *FuncType     // function signature: parameters,results,and position of "func" keyword
		Body *BlockStmt    // function body; or nil for external (non-Go) function
	}
)
...
複製程式碼

BadDecl表示一個有語法錯誤的節點； GenDecl用於表示import,const，type或變數宣告；FunDecl用於表示函式宣告。 GenDecl和FunDecl在Quick Start例子中均有出現，小夥伴們可自行查詢。

Common Node

除去上述四種類別劃分的node,還有一些node不屬於上面四種類別：

// Comment 註釋節點，代表單行的 //-格式 或 /*-格式的註釋.
type Comment struct {
    ...
}
...
// CommentGroup 註釋塊節點，包含多個連續的Comment
type CommentGroup struct {
    ...
}

// Field 欄位節點,可以代表結構體定義中的欄位，介面定義中的方法列表，函式前面中的入參和返回值欄位
type Field struct {
    ...
}
...
// FieldList 包含多個Field
type FieldList struct {
    ...
}

// File 表示一個檔案節點
type File struct {
	...
}

// Package 表示一個包節點
type Package struct {
    ...
}
複製程式碼

Quick Start示例包含了上面列舉的所有node，小夥伴們可以自行查詢。更為詳細的註釋和具體的結構體欄位請查閱原始碼。

所有的節點型別大致列舉完畢，其中還有許多具體的節點型別未能一一列舉，但基本上都是大同小異，原始碼註釋也比較清晰，等用到的時候再細看也不遲。現在我們對整個語法樹的構造有了基本的瞭解，接下來通過幾個示例來演示具體用法。

示例

為檔案中所有介面方法新增context引數

實現這個功能我們需要四步：

遍歷整個語法樹
判斷是否已經importcontext包，如果沒有則import
遍歷所有的介面方法，判斷方法列表中是否有context.Context型別的入參，如果沒有我們將其新增到方法的第一個引數
將修改過後的語法樹轉換成Go程式碼並輸出

遍歷語法樹

語法樹層級較深，巢狀關係複雜，如果不能完全掌握node之間的關係和巢狀規則，我們很難自己寫出正確的遍歷方法。不過好在ast包已經為我們提供了遍歷方法：

func Walk(v Visitor,node Node) 
複製程式碼

type Visitor interface {
	Visit(node Node) (w Visitor)
}
複製程式碼

Walk方法會按照深度優先搜尋方法（depth-first order）遍歷整個語法樹，我們只需按照我們的業務需要，實現Visitor介面即可。 Walk每遍歷一個節點就會呼叫Visitor.Visit方法，傳入當前節點。如果Visit返回nil，則停止遍歷當前節點的子節點。本示例的Visitor實現如下：

// Visitor
type Visitor struct {
}
func (v *Visitor) Visit(node ast.Node) ast.Visitor {
	switch node.(type) {
	case *ast.GenDecl:
		genDecl := node.(*ast.GenDecl)
		// 查詢有沒有import context包
		// Notice：沒有考慮沒有import任何包的情況
		if genDecl.Tok == token.IMPORT {
			v.addImport(genDecl)
			// 不需要再遍歷子樹
			return nil
		}
	case *ast.InterfaceType:
		// 遍歷所有的介面型別
		iface := node.(*ast.InterfaceType)
		addContext(iface)
		// 不需要再遍歷子樹
		return nil
	}
	return v
}
複製程式碼

新增import

// addImport 引入context包
func (v *Visitor) addImport(genDecl *ast.GenDecl) {
	// 是否已經import
	hasImported := false
	for _,v := range genDecl.Specs {
		imptSpec := v.(*ast.ImportSpec)
		// 如果已經包含"context"
		if imptSpec.Path.Value == strconv.Quote("context") {
			hasImported = true
		}
	}
	// 如果沒有import context，則import
	if !hasImported {
		genDecl.Specs = append(genDecl.Specs,&ast.ImportSpec{
			Path: &ast.BasicLit{
				Kind:  token.STRING,Value: strconv.Quote("context"),},})
	}
}

複製程式碼

為介面方法新增引數

// addContext 新增context引數
func addContext(iface *ast.InterfaceType) {
	// 介面方法不為空時，遍歷介面方法
	if iface.Methods != nil || iface.Methods.List != nil {
		for _,v := range iface.Methods.List {
			ft := v.Type.(*ast.FuncType)
			hasContext := false
			// 判斷引數中是否包含context.Context型別
			for _,v := range ft.Params.List {
				if expr,ok := v.Type.(*ast.SelectorExpr); ok {
					if ident,ok := expr.X.(*ast.Ident); ok {
						if ident.Name == "context" {
							hasContext = true
						}
					}
				}
			}
			// 為沒有context引數的方法新增context引數
			if !hasContext {
				ctxField := &ast.Field{
					Names: []*ast.Ident{
						ast.NewIdent("ctx"),// Notice: 沒有考慮import別名的情況
					Type: &ast.SelectorExpr{
						X:   ast.NewIdent("context"),Sel: ast.NewIdent("Context"),}
				list := []*ast.Field{
					ctxField,}
				ft.Params.List = append(list,ft.Params.List...)
			}
		}
	}
}
複製程式碼

將語法樹轉換成Go程式碼

format包為我們提供了轉換函式，format.Node會將語法樹按照gofmt的格式輸出：

...
	var output []byte
	buffer := bytes.NewBuffer(output)
	err = format.Node(buffer,fset,f)
	if err != nil {
		log.Fatal(err)
	}
	// 輸出Go程式碼
	fmt.Println(buffer.String())
...
複製程式碼

輸出結果如下：

package main

import (
        "context"
)

type Foo interface {
        FooA(ctx context.Context,i int)
        FooB(ctx context.Context,j int)
        FooC(ctx context.Context)
}

type Bar interface {
        BarA(ctx context.Context,i int)
        BarB(ctx context.Context)
        BarC(ctx context.Context)
}
複製程式碼

可以看到我們所有的介面方的第一個引數都變成了context.Context。建議將示例中的語法樹先打印出來，再對照著程式碼看，方便理解。

一些坑與不足

至此我們已經完成了語法樹的解析，遍歷，修改以及輸出。但細心的小夥伴可能已經發現：示例中的檔案並沒有出現一行註釋。這的確是有意為之，如果我們加上註釋，會發現最終生成檔案的註釋就像迷途的羔羊，完全找不到自己的位置。比如這樣：

//修改前
type Foo interface {
	FooA(i int)
	// FooB
	FooB(j int)
	FooC(ctx context.Context)
}

// 修改後
type Foo interface {
    FooA(ctx context.
            // FooB
            Context,i int)

    FooB(ctx context.Context,j int)
    FooC(ctx context.Context)
}
複製程式碼

導致這種現象的原因在於：ast包生成的語法樹中的註釋是"free-floating"的。還記得每個node都有Pos()和End()方法來標識其位置嗎？對於非註釋節點，語法樹能夠正確的調整他們的位置，但卻不能自動調整註釋節點的位置。如果我們想要讓註釋出現在正確的位置上，我們必須手動設定節點Pos和End。原始碼註釋中提到了這個問題：

issue中有具體的討論，官方承認這是一個設計缺陷，但還是遲遲未能改進。其中有位迫不及待的小哥提供了自己的方案：

如果實在是要對有註釋的語法樹進行修改，可以嘗試一下。雖然語法樹的確存在修改困難問題，但其還是能滿足大部分基於語法樹分析的程式碼生成工作了(gomock,wire等等)。

參考

Golang AST語法樹使用教程及示例

Quick Start

Node節點

Expression and Type

Statement

Spec Node

Declaration Node

Common Node

示例

為檔案中所有介面方法新增context引數

遍歷語法樹

新增import

為介面方法新增引數

將語法樹轉換成Go程式碼

一些坑與不足

參考

Golang AST語法樹使用教程及示例

Java HashSet使用方法教程及示例程式碼

教你如何用AST語法樹對程式碼“動手腳”

從零寫一個編譯器（九）：語義分析之構造抽象語法樹(AST)

JDT.AST:抽象語法樹-原始碼管理

MongoDB使用者許可權操作語法及示例

python AST 抽象語法樹

PHP程式執行流程：語法分析（Parse）與抽象語法樹（AST），PHP-Parser

抖音資料採集Frida教程，Frida Java Hook 詳解：程式碼及示例（下）

LaTex公式語法教程及手冊（附emlogpro公式顯示外掛katex說明）

Typora使用教程及MarkDown語法

常用語法及示例

Razor常用語法介紹及示例

js逆向--基於babel學習AST(Abstract Syntax Tree, 抽象語法樹)

AST抽象語法樹的基本思想

[WebKit] JavaScriptCore解析--基礎篇（一）位元組碼的生成及抽象語法樹的構建詳情分析

【轉載】Linux下裝置樹內容（詳細）總結及示例解析

CMake學習筆記（一）基本概念介紹、入門教程及CLion安裝配置

Ognl 語法基礎教程

golang定時任務踩坑及終極解決方案

Golang AST語法樹使用教程及示例

Quick Start

Node節點

Expression and Type

Statement

Spec Node

Declaration Node

Common Node

示例

為檔案中所有介面方法新增context引數

遍歷語法樹

新增import

為介面方法新增引數

將語法樹轉換成Go程式碼

一些坑與不足

參考

相關推薦