基于AST分析的源码定制化修复扩展方案

基于AST分析的源码定制化修复扩展方案

背景：

• 编程语言: golang
• IDE: jetbrains系列，goland
• 业务需求:

某种代码（例：日志打印）写作规范需要修改。修改规则略复杂，简单替换无法完成，现存代码量极大。
进阶需求：由于DFX等原因，造成代码引入额外代码，比较丑陋。一般开发人员不关注额外的格式，可自由编写修改前规范的代码，在提交/保存文件时自动格式化为目标规范代码。

例如：

// 修改前
log.Logger().Info("format", parameters1, parameters2...)
log.Logger().Error("format", parameters1, parameters2...)
log.Logger().EventXX("format", parameters1, parameters2...)

// 修改后
_ = log.I() && log.Logger().Info("format", parameters1, parameters2...)
_ = log.E() && log.Logger().Error("format", parameters1, parameters2...)
_ = log.I() && log.Logger().EventXX("format", parameters1, parameters2...)

替换说明：根据日志级别，在每个日志前面加入短路控制，达到性能优化效果。

预备知识

AST基本知识

goland的FileWatcher

A File Watcher is a GoLand system that tracks changes to your files and runs a third-party standalone application. GoLand provides predefined File Watcher templates for a number of such standard popular third-party tools (compilers, compressors, prettifiers, and others).

解决方案

简单一句话：分析抽象语法树，并根据要求修改节点，重新生成文件替换原文件。
进阶需求——自动修改：搭配goland的FileWatcher功能，完成保存时自动格式化（可扩展定制更多功能）。

AST分析思路

1. 待修改代码筛选filter

模式识别：
本例中，待修改的代码都是在log包中。
整句是ast.ExprStmt类型。
代码位置一般存在于ast.BlockStmt和ast.CaseClause中。

为了便于替换，需要有“替换槽”，所以从ast.BlockStmt和ast.CaseClause作为根节点进行分析。

// A BlockStmt node represents a braced statement list.
type BlockStmt struct {
	Lbrace token.Pos // position of "{"
	List   []Stmt
	Rbrace token.Pos // position of "}"
}

// A CaseClause represents a case of an expression or type switch statement.
type CaseClause struct {
	Case  token.Pos // position of "case" or "default" keyword
	List  []Expr    // list of expressions or types; nil means default case
	Colon token.Pos // position of ":"
	Body  []Stmt    // statement list; or nil
}

我们取出他们中的[]ast.Stmt切片。
每一个日志都是ast.Stmt类型。
打出来一个样例，这里删掉了一些不重要的信息，可以识别出我们要替换的模式：

  1: *ast.ExprStmt {
 .  X: *ast.CallExpr {
 .  .  Fun: *ast.SelectorExpr {
 .  .  .  X: *ast.CallExpr {
 .  .  .  .  Fun: *ast.SelectorExpr {
 .  .  .  .  .  X: *ast.Ident {
 .  .  .  .  .  .  Name: "log" // 包名在这
 .  .  .  .  .  }
 .  .  .  .  .  Sel: *ast.Ident {
 .  .  .  .  .  .  Name: "Logger"
 .  .  .  .  .  }
 .  .  .  .  }
 .  .  .  .  Lparen: xxx.go:110:16
 .  .  .  .  Ellipsis: -
 .  .  .  .  Rparen: xxx.go:110:17
 .  .  .  }
 .  .  .  Sel: *ast.Ident {
 .  .  .  .  Name: "Info" // 体现日志级别的函数在这
 .  .  .  }
 .  .  }
 .  .  Lparen: xxx.go:110:23
 .  .  Args: []ast.Expr (len = 1) {
 .  .  .  0: *ast.BasicLit {
 .  .  .  .  Kind: STRING
 .  .  .  .  Value: "format"
 .  .  .  }
 .  .  }
 .  .  Ellipsis: -
 .  .  Rparen: xxx.go:110:56
 .  }
 }

2. 原代码与目标代码之间的差异分析

_ = log.I() && log.Logger().Info("format", parameters1, parameters2...)
目标代码列于上。
它整体是一个ast.AssignStmt，可替代原ast.ExprStmt的位置。

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

赋值语句，由左操作，赋值符号，右操作三个部分组成。
我们的Lhs是_，Tok是= :token.ASSIGN，Rhs是一个二元表达式ast.BinaryExpr即Level() && 原代码那一行。

根据分析，我们可以写出一个构造函数，

// Assignment语句构造函数
func buildAssignStmt(lvl string, caller *ast.CallExpr, logidentPos token.Pos) *ast.AssignStmt {
	assign := &ast.AssignStmt{
		Lhs: []ast.Expr{
			&ast.Ident{
				NamePos: logidentPos, // 位置就用原来那一行的位置
				Name:    "_", // 忽略返回值
			},
		},
		TokPos: logidentPos,
		Tok:    token.ASSIGN,
		Rhs: []ast.Expr{
			&ast.BinaryExpr{ // X Op Y 二元表达式
				X: &ast.CallExpr{ // X是 log.Level()
					Fun: &ast.SelectorExpr{
						X: &ast.Ident{
							Name:    "log",
							NamePos: logidentPos,
						},
						Sel: &ast.Ident{Name: lvl},
					},
				},
				OpPos: logidentPos,
				Op:    token.LAND, // 操作符是 &&
				Y:     caller, // Y是原语句
			},
		},
	}
	return assign
}

3. AST Node替换

只需要把原block Stmt切片中的内容换成新构造的Assignment即可。

4. 重新生成文件

使用format.Node()，向dst写入格式化好的ast节点。

// Node formats node in canonical gofmt style and writes the result to dst.
//
// The node type must be *ast.File, *printer.CommentedNode, []ast.Decl,
// []ast.Stmt, or assignment-compatible to ast.Expr, ast.Decl, ast.Spec,
// or ast.Stmt. Node does not modify node. Imports are not sorted for
// nodes representing partial source files (for instance, if the node is
// not an *ast.File or a *printer.CommentedNode not wrapping an *ast.File).
//
// The function may return early (before the entire result is written)
// and return a formatting error, for instance due to an incorrect AST.
//
func Node(dst io.Writer, fset *token.FileSet, node interface{}) error {...}

goland自动保存

需要把文件名作为main函数的参数，并在goland里增加如下配置。即可在保存文件时触发自动保存。

效果演示：