in tools/go-agent/instrument/plugins/rewrite/func.go [96:217]
func (c *Context) enhanceFuncStmt(stmt dst.Stmt) {
// for the variables created in the sub statement, ex: if, func(), the temporary variable count should be recorded
subCallTypes := []reflect.Type{
reflect.TypeOf(&dst.IfStmt{}),
reflect.TypeOf(&dst.BlockStmt{}),
reflect.TypeOf(&dst.TypeSwitchStmt{}),
}
dstutil.Apply(stmt, func(cursor *dstutil.Cursor) bool {
for _, t := range subCallTypes {
if reflect.TypeOf(cursor.Node()) == t {
c.rewriteMapping.pushBlockStack()
}
}
switch n := cursor.Node().(type) {
case *dst.BlockStmt:
for _, tmp := range n.List {
c.enhanceFuncStmt(tmp)
}
case *dst.AssignStmt:
for _, l := range n.Lhs {
if k, v := c.enhanceVarNameWhenRewrite(l); k != "" {
c.rewriteMapping.addVarMapping(k, v)
}
}
for i, r := range n.Rhs {
if k, v := c.enhanceTypeNameWhenRewrite(r, nil, i); k != "" {
c.rewriteMapping.addTypeMapping(k, v)
}
}
case *dst.BinaryExpr:
c.rewriteVarIfExistingMapping(n.X, n)
c.rewriteVarIfExistingMapping(n.Y, n)
case *dst.CallExpr:
c.enhanceTypeNameWhenRewrite(n.Fun, n, -1)
for inx, arg := range n.Args {
c.enhanceTypeNameWhenRewrite(arg, n, inx)
}
case *dst.ReturnStmt:
for inx, arg := range n.Results {
c.enhanceTypeNameWhenRewrite(arg, n, inx)
}
case *dst.FuncType:
c.enhanceFuncParameter(n.Params)
c.enhanceFuncParameter(n.Results)
case *dst.ExprStmt:
c.enhanceTypeNameWhenRewrite(n.X, n, -1)
case *dst.TypeAssertExpr:
c.enhanceTypeNameWhenRewrite(n.X, n, -1)
c.enhanceTypeNameWhenRewrite(n.Type, n, -1)
case *dst.IfStmt:
c.enhanceFuncStmt(n.Init)
c.enhanceTypeNameWhenRewrite(n.Cond, n, -1)
if n.Body != nil {
for _, stmt := range n.Body.List {
c.enhanceFuncStmt(stmt)
}
}
if n.Else != nil {
c.enhanceFuncStmt(n.Else)
}
case *dst.RangeStmt:
c.enhanceTypeNameWhenRewrite(n.X, n, -1)
if k, v := c.enhanceVarNameWhenRewrite(n.Key); k != "" {
c.rewriteMapping.addVarMapping(k, v)
}
if k, v := c.enhanceVarNameWhenRewrite(n.Value); k != "" {
c.rewriteMapping.addVarMapping(k, v)
}
if n.Body != nil {
for _, stmt := range n.Body.List {
c.enhanceFuncStmt(stmt)
}
}
case *dst.ValueSpec:
for _, n := range n.Names {
if k, v := c.enhanceVarNameWhenRewrite(n); k != "" {
c.rewriteMapping.addVarMapping(k, v)
}
}
c.enhanceTypeNameWhenRewrite(n.Type, n, -1)
for _, subVal := range n.Values {
c.enhanceTypeNameWhenRewrite(subVal, n, -1)
}
case *dst.TypeSwitchStmt:
c.enhanceFuncStmt(n.Init)
c.enhanceFuncStmt(n.Assign)
if n.Body != nil {
for _, stmt := range n.Body.List {
c.enhanceFuncStmt(stmt)
}
}
case *dst.SwitchStmt:
c.enhanceFuncStmt(n.Init)
c.enhanceTypeNameWhenRewrite(n.Tag, n, -1)
if n.Body != nil {
for _, stmt := range n.Body.List {
c.enhanceFuncStmt(stmt)
}
}
case *dst.CaseClause:
for i, stmt := range n.List {
c.enhanceTypeNameWhenRewrite(stmt, n, i)
}
for _, stmt := range n.Body {
c.enhanceFuncStmt(stmt)
}
default:
return true
}
return false
}, func(cursor *dstutil.Cursor) bool {
// all templates variables should be removed
for _, t := range subCallTypes {
if reflect.TypeOf(cursor.Node()) == t {
c.rewriteMapping.popBlockStack()
break
}
}
return true
})
}