golang源码分析：抽象语法树

golangLeetcode

发布于 2023-03-14 20:50:02

3700

发布于 2023-03-14 20:50:02

文章被收录于专栏：golang算法架构leetcode技术php

golang提供了非常强大的工具集合，通过这些工具我们可以非常方便地进行源码的分析加工，在代码中插入我们想要的代码，或者提取源码中我们关心的信息。如何使用呢其实非常简单：

1，解析源码文件得到抽象语法树

2，定义我们自己需要的访问者

3，通过walk方法遍历语法树，提取我们需要的信息。

fset := token.NewFileSet()
f, err := parser.ParseFile(fset, path, nil, parser.ParseComments)
ast.Walk(&Visitor{
  fset: fset,
}, f)

其中需要我们自定义的就是定义一个访问者。它的定义如下，我们只需要实现一个Visit接口，它的入参是Node也就是抽象语法树上的一个节点，我们可以根据节点的不同类型实现我们需要的不同功能。返回的是一个Visitor，关于返回值的使用是这么约定的：如果我们想继续解析当前节点的子节点，就返回一个不是nil的Visitor，这样就可以继续递归解析，否则结束当前节点的遍历

type Visitor interface {
  Visit(node Node) (w Visitor)
}

其中Node也是一个接口，通过它我们可以方便得到节点在源码中的起止位置

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
}

其实上面Visitor这个约定基本上归纳了Walk方法的功能

func Walk(v Visitor, node Node) {
  if v = v.Visit(node); v == nil {
    return
  }

  // walk children
  // (the order of the cases matches the order
  // of the corresponding node types in ast.go)
  switch n := node.(type) {
  // Comments and fields
  case *Comment:
    // nothing to do

  case *CommentGroup:
    for _, c := range n.List {
      Walk(v, c)
    }

可以看到，它先调用Visit方法，访问当前节点，如果返回值是nil，直接结束返回，否则，判断节点的类型，针对不同节点类型，拆分出节点的孩子节点，然后继续遍历节点的孩子节点。

基于上述知识，我们就可以非常轻松地自定义我们需要的访问者，比如我想获取当前文件引用了哪些包：

type PkgVisitor struct {
  fset *token.FileSet
  Pkgs []string
}

func (v *PkgVisitor) Visit(node ast.Node) ast.Visitor {
  switch node.(type) {
  case *ast.Package:
    pkgs := node.(*ast.Package)
    fmt.Println("pkg name:", pkgs.Name)
    if pkgs != nil {
      v.Pkgs = append(v.Pkgs, pkgs.Name)
    }
   }
   return nil
  }

或者，我们在中golang源码分析:goc集成测试覆盖率实现原理（2）分析的，实现源码单测覆盖率统计

func (f *File) Visit(node ast.Node) ast.Visitor {
  switch n := node.(type) {
  case *ast.BlockStmt:
    // If it's a switch or select, the body is a list of case clauses; don't tag the block itself.
    if len(n.List) > 0 {
      switch n.List[0].(type) {
      case *ast.CaseClause: // switch
        for _, n := range n.List {
          clause := n.(*ast.CaseClause)
          f.addCounters(clause.Colon+1, clause.Colon+1, clause.End(), clause.Body, false)
        }
        return f
      case *ast.CommClause: // select
        for _, n := range n.List {
          clause := n.(*ast.CommClause)
          f.addCounters(clause.Colon+1, clause.Colon+1, clause.End(), clause.Body, false)
        }
        return f
      }
    }
    f.addCounters(n.Lbrace, n.Lbrace+1, n.Rbrace+1, n.List, true) // +1 to step past closing brace.
  case *ast.IfStmt:
    if n.Init != nil {
      ast.Walk(f, n.Init)
    }
    ast.Walk(f, n.Cond)
    ast.Walk(f, n.Body)
    if n.Else == nil {
      return nil
    }
    // The elses are special, because if we have
    //  if x {
    //  } else if y {
    //  }
    // we want to cover the "if y". To do this, we need a place to drop the counter,
    // so we add a hidden block:
    //  if x {
    //  } else {
    //    if y {
    //    }
    //  }
    elseOffset := f.findText(n.Body.End(), "else")
    if elseOffset < 0 {
      panic("lost else")
    }
    f.edit.Insert(elseOffset+4, "{")
    f.edit.Insert(f.offset(n.Else.End()), "}")

    // We just created a block, now walk it.
    // Adjust the position of the new block to start after
    // the "else". That will cause it to follow the "{"
    // we inserted above.
    pos := f.fset.File(n.Body.End()).Pos(elseOffset + 4)
    switch stmt := n.Else.(type) {
    case *ast.IfStmt:
      block := &ast.BlockStmt{
        Lbrace: pos,
        List:   []ast.Stmt{stmt},
        Rbrace: stmt.End(),
      }
      n.Else = block
    case *ast.BlockStmt:
      stmt.Lbrace = pos
    default:
      panic("unexpected node type in if")
    }
    ast.Walk(f, n.Else)
    return nil
  case *ast.SelectStmt:
    // Don't annotate an empty select - creates a syntax error.
    if n.Body == nil || len(n.Body.List) == 0 {
      return nil
    }
  case *ast.SwitchStmt:
    // Don't annotate an empty switch - creates a syntax error.
    if n.Body == nil || len(n.Body.List) == 0 {
      if n.Init != nil {
        ast.Walk(f, n.Init)
      }
      if n.Tag != nil {
        ast.Walk(f, n.Tag)
      }
      return nil
    }
  case *ast.TypeSwitchStmt:
    // Don't annotate an empty type switch - creates a syntax error.
    if n.Body == nil || len(n.Body.List) == 0 {
      if n.Init != nil {
        ast.Walk(f, n.Init)
      }
      ast.Walk(f, n.Assign)
      return nil
    }
  }
  return f
}

当然，我们也可以做得复杂点，比如实现函数内逻辑分支的可视化：https://github.com/xiazemin/jstree-go

本文参与腾讯云自媒体同步曝光计划，分享自微信公众号。

原始发表：2023-03-04，如有侵权请联系 cloudcommunity@tencent.com 删除

遍历