fix: adjust line number via //line directive

Author: y1yang0

tool/instrument/inst_func.go

@@ -260,15 +260,17 @@ func (rp *RuleProcessor) insertRaw(r *resource.InstFuncRule, decl *dst.FuncDecl)
 	util.Assert(r.OnEnter != "" || r.OnExit != "", "sanity check")
 	if r.OnEnter != "" {
 		// Prepend raw code snippet to function body for onEnter
-		onEnterSnippet, err := util.ParseAstFromSnippet(r.OnEnter)
+		p := util.NewAstParser()
+		onEnterSnippet, err := p.ParseSnippet(r.OnEnter)
 		if err != nil {
 			return err
 		}
 		decl.Body.List = append(onEnterSnippet, decl.Body.List...)
 	}
 	if r.OnExit != "" {
 		// Use defer func(){ raw_code_snippet }() for onExit
-		onExitSnippet, err := util.ParseAstFromSnippet(
+		p := util.NewAstParser()
+		onExitSnippet, err := p.ParseSnippet(
 			fmt.Sprintf("defer func(){ %s }()", r.OnExit),
 		)
 		if err != nil {
@@ -301,8 +303,8 @@ func sortFuncRules(fnRules []*resource.InstFuncRule) []*resource.InstFuncRule {
 
 func (rp *RuleProcessor) writeTrampoline(pkgName string) error {
 	// Prepare trampoline code header
-	code := "package " + pkgName
-	trampoline, err := util.ParseAstFromSource(code)
+	p := util.NewAstParser()
+	trampoline, err := p.ParseSource("package " + pkgName)
 	if err != nil {
 		return err
 	}

tool/instrument/optimize.go

@@ -150,7 +150,8 @@ func (rp *RuleProcessor) newCallContextImpl(tjump *TJump) (dst.Expr, error) {
 	// One line please, otherwise debugging line number will be a nightmare
 	tmpl := fmt.Sprintf("&CallContextImpl%s{Params:[]interface{}{},ReturnVals:[]interface{}{}}",
 		rp.rule2Suffix[tjump.rule])
-	astRoot, err := util.ParseAstFromSnippet(tmpl)
+	p := util.NewAstParser()
+	astRoot, err := p.ParseSnippet(tmpl)
 	if err != nil {
 		return nil, err
 	}

tool/instrument/trampoline.go

@@ -71,7 +71,8 @@ var trampolineTemplate string
 func (rp *RuleProcessor) materializeTemplate() error {
 	// Read trampoline template and materialize onEnter and onExit function
 	// declarations based on that
-	astRoot, err := util.ParseAstFromSource(trampolineTemplate)
+	p := util.NewAstParser()
+	astRoot, err := p.ParseSource(trampolineTemplate)
 	if err != nil {
 		return err
 	}

tool/preprocess/preprocess.go

@@ -18,6 +18,8 @@ import (
 	"bufio"
 	"embed"
 	"fmt"
+	"go/parser"
+	"go/token"
 	"os"
 	"os/exec"
 	"os/signal"
@@ -31,6 +33,7 @@ import (
 	"github.com/alibaba/opentelemetry-go-auto-instrumentation/tool/errc"
 	"github.com/alibaba/opentelemetry-go-auto-instrumentation/tool/resource"
 	"github.com/alibaba/opentelemetry-go-auto-instrumentation/tool/util"
+	"github.com/dave/dst"
 	"golang.org/x/mod/modfile"
 	"golang.org/x/tools/go/packages"
 )
@@ -536,7 +539,8 @@ func (dp *DepProcessor) addExplicitImport(importPaths ...string) (err error) {
 		if !util.IsGoFile(file) {
 			continue
 		}
-		astRoot, err := util.ParseAstFromFile(file)
+		p := util.NewAstParser()
+		astRoot, err := p.ParseFile(file, parser.ParseComments)
 		if err != nil {
 			return err
 		}
@@ -549,6 +553,18 @@ func (dp *DepProcessor) addExplicitImport(importPaths ...string) (err error) {
 			}
 		}
 
+		for _, decl := range astRoot.Decls {
+			decl, ok := decl.(*dst.GenDecl)
+			if !ok {
+				continue
+			}
+			if decl.Tok != token.IMPORT {
+				continue
+			}
+			decl.Decs.Before = dst.EmptyLine
+			decl.Decs.Start.Append("//wow")
+		}
+
 		// Prepend import path to the file
 		for _, importPath := range importPaths {
 			util.AddImportForcely(astRoot, importPath)

tool/preprocess/setup.go

@@ -171,7 +171,8 @@ func (dp *DepProcessor) copyRule(path, target string,
 		return err
 	}
 	text = util.RemoveGoBuildComment(text)
-	astRoot, err := util.ParseAstFromSource(text)
+	p := util.NewAstParser()
+	astRoot, err := p.ParseSource(text)
 	if err != nil {
 		return err
 	}

tool/util/ast.go

@@ -394,31 +394,40 @@ func MatchStructDecl(decl dst.Decl, structType string) bool {
 // To solve this issue, we disable DST's automatic Import management
 // and use plain AST manipulation to add imports.
 
-// ParseAstFromSnippet parses the AST from incomplete source code snippet.
-func ParseAstFromSnippet(codeSnippnet string) ([]dst.Stmt, error) {
-	fset := token.NewFileSet()
+type AstParser struct {
+	fset *token.FileSet
+}
+
+func NewAstParser() *AstParser {
+	return &AstParser{
+		fset: token.NewFileSet(),
+	}
+}
+
+// ParseSnippet parses the AST from incomplete source code snippet.
+func (ap *AstParser) ParseSnippet(codeSnippnet string) ([]dst.Stmt, error) {
+	Assert(codeSnippnet != "", "empty code snippet")
 	snippet := "package main; func _() {" + codeSnippnet + "}"
-	file, err := decorator.ParseFile(fset, "", snippet, 0)
+	file, err := decorator.ParseFile(ap.fset, "", snippet, 0)
 	if err != nil {
 		return nil, errc.New(errc.ErrParseCode, err.Error())
 	}
 	return file.Decls[0].(*dst.FuncDecl).Body.List, nil
 }
 
-// ParseAstFromSource parses the AST from complete source code.
-func ParseAstFromSource(source string) (*dst.File, error) {
+// ParseSource parses the AST from complete source code.
+func (ap *AstParser) ParseSource(source string) (*dst.File, error) {
 	Assert(source != "", "empty source")
-	dec := decorator.NewDecorator(token.NewFileSet())
+	dec := decorator.NewDecorator(ap.fset)
 	dstRoot, err := dec.Parse(source)
 	if err != nil {
 		return nil, errc.New(errc.ErrParseCode, err.Error())
 	}
 	return dstRoot, nil
 }
 
-func parseAstMode(filePath string, mode parser.Mode) (*dst.File, error) {
+func (ap *AstParser) ParseFile(filePath string, mode parser.Mode) (*dst.File, error) {
 	name := filepath.Base(filePath)
-	fset := token.NewFileSet()
 	file, err := os.Open(filePath)
 	if err != nil {
 		return nil, errc.New(errc.ErrOpenFile, err.Error())
@@ -429,11 +438,11 @@ func parseAstMode(filePath string, mode parser.Mode) (*dst.File, error) {
 			LogFatal("failed to close file %s: %v", file.Name(), err)
 		}
 	}(file)
-	astFile, err := parser.ParseFile(fset, name, file, mode)
+	astFile, err := parser.ParseFile(ap.fset, name, file, mode)
 	if err != nil {
 		return nil, errc.New(errc.ErrParseCode, err.Error())
 	}
-	dec := decorator.NewDecorator(fset)
+	dec := decorator.NewDecorator(ap.fset)
 	dstFile, err := dec.DecorateFile(astFile)
 	if err != nil {
 		return nil, errc.New(errc.ErrParseCode, err.Error())
@@ -442,16 +451,16 @@ func parseAstMode(filePath string, mode parser.Mode) (*dst.File, error) {
 }
 
 func ParseAstFromFileOnlyPackage(filePath string) (*dst.File, error) {
-	return parseAstMode(filePath, parser.PackageClauseOnly)
+	return NewAstParser().ParseFile(filePath, parser.PackageClauseOnly)
 }
 
 func ParseAstFromFileFast(filePath string) (*dst.File, error) {
-	return parseAstMode(filePath, parser.SkipObjectResolution)
+	return NewAstParser().ParseFile(filePath, parser.SkipObjectResolution)
 }
 
 // ParseAstFromFile parses the AST from complete source file.
 func ParseAstFromFile(filePath string) (*dst.File, error) {
-	return parseAstMode(filePath, parser.ParseComments)
+	return NewAstParser().ParseFile(filePath, parser.ParseComments)
 }
 
 // WriteAstToFile writes the AST to source file.