func()

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
	})
}