From 1b71478939bb0cdc49959d9447c53d0488942e49 Mon Sep 17 00:00:00 2001 From: Dimitar Petrov Date: Sun, 6 Sep 2020 21:13:28 +0300 Subject: [PATCH] Add import aliases --- cmd/apply.go | 2 +- cmd/revert.go | 2 +- tracing/deinstrument.go | 4 +- tracing/deinstrument_test.go | 4 +- tracing/instrument.go | 2 +- tracing/instrument_test.go | 92 +++++++++++++++++----------------- tracing/interfaces.go | 6 +-- tracing/unused_imports.go | 10 ++-- tracing/unused_imports_test.go | 4 +- tracing/util.go | 8 ++- 10 files changed, 70 insertions(+), 64 deletions(-) diff --git a/cmd/apply.go b/cmd/apply.go index 674a4b9..5487280 100644 --- a/cmd/apply.go +++ b/cmd/apply.go @@ -41,6 +41,6 @@ func (ac *ApplyCmd) Run() error { if err != nil { return err } - return ac.importsGroomer.RemoveUnusedImportFromDirectory(path, []string{"fmt", "runtime", "rand"}) + return ac.importsGroomer.RemoveUnusedImportFromDirectory(path, map[string]string{"fmt":"", "runtime":"rt", "rand":""}) }) } diff --git a/cmd/revert.go b/cmd/revert.go index 07add8d..5cdb142 100644 --- a/cmd/revert.go +++ b/cmd/revert.go @@ -41,6 +41,6 @@ func (rc *RevertCmd) Run() error { if err != nil { return err } - return rc.importsGroomer.RemoveUnusedImportFromDirectory(path, []string{"fmt", "runtime", "rand"}) + return rc.importsGroomer.RemoveUnusedImportFromDirectory(path, map[string]string{"fmt":"", "runtime":"rt", "rand":""}) }) } diff --git a/tracing/deinstrument.go b/tracing/deinstrument.go index 058b85e..5a6454b 100644 --- a/tracing/deinstrument.go +++ b/tracing/deinstrument.go @@ -68,7 +68,9 @@ func (cd *codeDeinstrumenter) DeinstrumentFile(fset *token.FileSet, file *ast.Fi if checkInstrumentationStatementsIntegrity(t) { t.Body.List = t.Body.List[instrumentationStmtsCount:] - t.Body.List[0].Decorations().Before = dst.None + if len(t.Body.List) > 0 { + t.Body.List[0].Decorations().Before = dst.None + } } } } diff --git a/tracing/deinstrument_test.go b/tracing/deinstrument_test.go index decf1a2..2746c0f 100644 --- a/tracing/deinstrument_test.go +++ b/tracing/deinstrument_test.go @@ -43,7 +43,7 @@ func TestDeinstrumentFile(t *testing.T) { } var buff2 bytes.Buffer - if err := NewImportsGroomer().RemoveUnusedImportFromFile(fset, file, &buff2, []string{"fmt", "runtime", "rand"}); err != nil { + if err := NewImportsGroomer().RemoveUnusedImportFromFile(fset, file, &buff2, map[string]string{"fmt":"", "runtime":"rt", "rand":""}); err != nil { t.Fatal(err) } @@ -87,7 +87,7 @@ func TestDeinstrumentDirectory(t *testing.T) { t.Fatal(err) } - if err := NewImportsGroomer().RemoveUnusedImportFromDirectory("test", []string{"fmt", "runtime", "rand"}); err != nil { + if err := NewImportsGroomer().RemoveUnusedImportFromDirectory("test", map[string]string{"fmt":"", "runtime":"rt", "rand":""}); err != nil { t.Fatal(err) } diff --git a/tracing/instrument.go b/tracing/instrument.go index effe985..e2161d6 100644 --- a/tracing/instrument.go +++ b/tracing/instrument.go @@ -83,7 +83,7 @@ func (ci *codeInstrumenter) InstrumentPackage(fset *token.FileSet, pkg *ast.Pack func (ci *codeInstrumenter) InstrumentFile(fset *token.FileSet, file *ast.File, out io.Writer) error { astutil.AddImport(fset, file, "fmt") - astutil.AddImport(fset, file, "runtime") + astutil.AddNamedImport(fset, file, "rt", "runtime") astutil.AddImport(fset, file, "rand") // Needed because ast does not support floating comments and deletes them. diff --git a/tracing/instrument_test.go b/tracing/instrument_test.go index dcf4c06..78c44e2 100644 --- a/tracing/instrument_test.go +++ b/tracing/instrument_test.go @@ -29,7 +29,7 @@ const resultCodeWithoutImports = `package a import ( "fmt" "rand" - "runtime" + rt "runtime" ) func test(i int, b bool) int { @@ -37,11 +37,11 @@ func test(i int, b bool) int { /* prinTracer */ funcName := "test" caller := "unknown" - if funcPC, _, _, ok := runtime.Caller(0); ok { - funcName = runtime.FuncForPC(funcPC).Name() + if funcPC, _, _, ok := rt.Caller(0); ok { + funcName = rt.FuncForPC(funcPC).Name() } - if callerPC, _, _, ok := runtime.Caller(1); ok { - caller = runtime.FuncForPC(callerPC).Name() + if callerPC, _, _, ok := rt.Caller(1); ok { + caller = rt.FuncForPC(callerPC).Name() } idBytes := make([]byte, 16) _, _ = rand.Read(idBytes) @@ -60,11 +60,11 @@ func main() { /* prinTracer */ funcName := "main" caller := "unknown" - if funcPC, _, _, ok := runtime.Caller(0); ok { - funcName = runtime.FuncForPC(funcPC).Name() + if funcPC, _, _, ok := rt.Caller(0); ok { + funcName = rt.FuncForPC(funcPC).Name() } - if callerPC, _, _, ok := runtime.Caller(1); ok { - caller = runtime.FuncForPC(callerPC).Name() + if callerPC, _, _, ok := rt.Caller(1); ok { + caller = rt.FuncForPC(callerPC).Name() } idBytes := make([]byte, 16) _, _ = rand.Read(idBytes) @@ -81,7 +81,7 @@ const editedResultCodeWithoutImports = `package a import ( "fmt" "rand" - "runtime" + rt "runtime" ) func test(i int, b bool) int { @@ -89,11 +89,11 @@ func test(i int, b bool) int { /* prinTracer */ funcName := "test2" caller := "unknown2" - if funcPC, _, _, ok := runtime.Caller(0); ok { - funcName = runtime.FuncForPC(funcPC).Name() + if funcPC, _, _, ok := rt.Caller(0); ok { + funcName = rt.FuncForPC(funcPC).Name() } - if callerPC, _, _, ok := runtime.Caller(1); ok { - caller = runtime.FuncForPC(callerPC).Name() + if callerPC, _, _, ok := rt.Caller(1); ok { + caller = rt.FuncForPC(callerPC).Name() } fmt.Println("test") idBytes := make([]byte, 16) @@ -112,11 +112,11 @@ func main() { funcName := "main" caller := "unknown" - if funcPC, _, _, ok := runtime.Caller(0); ok { - funcName = runtime.FuncForPC(funcPC).Name() + if funcPC, _, _, ok := rt.Caller(0); ok { + funcName = rt.FuncForPC(funcPC).Name() } - if callerPC, _, _, ok := runtime.Caller(1); ok { - caller = runtime.FuncForPC(callerPC).Name() + if callerPC, _, _, ok := rt.Caller(1); ok { + caller = rt.FuncForPC(callerPC).Name() } idBytes := make([]byte, 16) _, _ = rand.Read(idBytes) @@ -151,7 +151,7 @@ const resultCodeWithFmtImport = `package a import ( "fmt" "rand" - "runtime" + rt "runtime" ) func test(i int, b bool) int { @@ -159,11 +159,11 @@ func test(i int, b bool) int { /* prinTracer */ funcName := "test" caller := "unknown" - if funcPC, _, _, ok := runtime.Caller(0); ok { - funcName = runtime.FuncForPC(funcPC).Name() + if funcPC, _, _, ok := rt.Caller(0); ok { + funcName = rt.FuncForPC(funcPC).Name() } - if callerPC, _, _, ok := runtime.Caller(1); ok { - caller = runtime.FuncForPC(callerPC).Name() + if callerPC, _, _, ok := rt.Caller(1); ok { + caller = rt.FuncForPC(callerPC).Name() } idBytes := make([]byte, 16) _, _ = rand.Read(idBytes) @@ -182,11 +182,11 @@ func main() { /* prinTracer */ funcName := "main" caller := "unknown" - if funcPC, _, _, ok := runtime.Caller(0); ok { - funcName = runtime.FuncForPC(funcPC).Name() + if funcPC, _, _, ok := rt.Caller(0); ok { + funcName = rt.FuncForPC(funcPC).Name() } - if callerPC, _, _, ok := runtime.Caller(1); ok { - caller = runtime.FuncForPC(callerPC).Name() + if callerPC, _, _, ok := rt.Caller(1); ok { + caller = rt.FuncForPC(callerPC).Name() } idBytes := make([]byte, 16) _, _ = rand.Read(idBytes) @@ -224,7 +224,7 @@ const resultCodeWithMultipleImports = `package a import ( "fmt" "rand" - "runtime" + rt "runtime" "strconv" ) @@ -233,11 +233,11 @@ func test(i int, b bool) int { /* prinTracer */ funcName := "test" caller := "unknown" - if funcPC, _, _, ok := runtime.Caller(0); ok { - funcName = runtime.FuncForPC(funcPC).Name() + if funcPC, _, _, ok := rt.Caller(0); ok { + funcName = rt.FuncForPC(funcPC).Name() } - if callerPC, _, _, ok := runtime.Caller(1); ok { - caller = runtime.FuncForPC(callerPC).Name() + if callerPC, _, _, ok := rt.Caller(1); ok { + caller = rt.FuncForPC(callerPC).Name() } idBytes := make([]byte, 16) _, _ = rand.Read(idBytes) @@ -256,11 +256,11 @@ func main() { /* prinTracer */ funcName := "main" caller := "unknown" - if funcPC, _, _, ok := runtime.Caller(0); ok { - funcName = runtime.FuncForPC(funcPC).Name() + if funcPC, _, _, ok := rt.Caller(0); ok { + funcName = rt.FuncForPC(funcPC).Name() } - if callerPC, _, _, ok := runtime.Caller(1); ok { - caller = runtime.FuncForPC(callerPC).Name() + if callerPC, _, _, ok := rt.Caller(1); ok { + caller = rt.FuncForPC(callerPC).Name() } idBytes := make([]byte, 16) _, _ = rand.Read(idBytes) @@ -297,7 +297,7 @@ const resultCodeWithImportsWithoutFmt = `package a import ( "fmt" "rand" - "runtime" + rt "runtime" "strconv" ) @@ -306,11 +306,11 @@ func test(i int, b bool) int { /* prinTracer */ funcName := "test" caller := "unknown" - if funcPC, _, _, ok := runtime.Caller(0); ok { - funcName = runtime.FuncForPC(funcPC).Name() + if funcPC, _, _, ok := rt.Caller(0); ok { + funcName = rt.FuncForPC(funcPC).Name() } - if callerPC, _, _, ok := runtime.Caller(1); ok { - caller = runtime.FuncForPC(callerPC).Name() + if callerPC, _, _, ok := rt.Caller(1); ok { + caller = rt.FuncForPC(callerPC).Name() } idBytes := make([]byte, 16) _, _ = rand.Read(idBytes) @@ -329,11 +329,11 @@ func main() { /* prinTracer */ funcName := "main" caller := "unknown" - if funcPC, _, _, ok := runtime.Caller(0); ok { - funcName = runtime.FuncForPC(funcPC).Name() + if funcPC, _, _, ok := rt.Caller(0); ok { + funcName = rt.FuncForPC(funcPC).Name() } - if callerPC, _, _, ok := runtime.Caller(1); ok { - caller = runtime.FuncForPC(callerPC).Name() + if callerPC, _, _, ok := rt.Caller(1); ok { + caller = rt.FuncForPC(callerPC).Name() } idBytes := make([]byte, 16) _, _ = rand.Read(idBytes) @@ -358,7 +358,7 @@ const resultCodeWithoutFunction = `package a import ( "fmt" "rand" - "runtime" + rt "runtime" ) type test struct { diff --git a/tracing/interfaces.go b/tracing/interfaces.go index 57d18d2..8e2bb24 100644 --- a/tracing/interfaces.go +++ b/tracing/interfaces.go @@ -22,7 +22,7 @@ type CodeDeinstrumenter interface { //go:generate counterfeiter . ImportsGroomer type ImportsGroomer interface { - RemoveUnusedImportFromFile(fset *token.FileSet, file *ast.File, out io.Writer, importsToRemove []string) error - RemoveUnusedImportFromPackage(fset *token.FileSet, pkg *ast.Package, importsToRemove []string) error - RemoveUnusedImportFromDirectory(path string, importsToRemove []string) error + RemoveUnusedImportFromFile(fset *token.FileSet, file *ast.File, out io.Writer, importsToRemove map[string]string) error + RemoveUnusedImportFromPackage(fset *token.FileSet, pkg *ast.Package, importsToRemove map[string]string) error + RemoveUnusedImportFromDirectory(path string, importsToRemove map[string]string) error } diff --git a/tracing/unused_imports.go b/tracing/unused_imports.go index b250854..9a5ef31 100644 --- a/tracing/unused_imports.go +++ b/tracing/unused_imports.go @@ -18,7 +18,7 @@ func NewImportsGroomer() ImportsGroomer { return &importsGroomer{} } -func (ig *importsGroomer) RemoveUnusedImportFromDirectory(path string, importsToRemove []string) error { +func (ig *importsGroomer) RemoveUnusedImportFromDirectory(path string, importsToRemove map[string]string) error { fset := token.NewFileSet() filter := func(info os.FileInfo) bool { return testsFilter(info) && generatedFilter(path, info) @@ -36,7 +36,7 @@ func (ig *importsGroomer) RemoveUnusedImportFromDirectory(path string, importsTo return nil } -func (ig *importsGroomer) RemoveUnusedImportFromPackage(fset *token.FileSet, pkg *ast.Package, importsToRemove []string) error { +func (ig *importsGroomer) RemoveUnusedImportFromPackage(fset *token.FileSet, pkg *ast.Package, importsToRemove map[string]string) error { for fileName, file := range pkg.Files { sourceFile, err := os.OpenFile(fileName, os.O_TRUNC|os.O_WRONLY, 0664) if err != nil { @@ -49,10 +49,10 @@ func (ig *importsGroomer) RemoveUnusedImportFromPackage(fset *token.FileSet, pkg return nil } -func (ig *importsGroomer) RemoveUnusedImportFromFile(fset *token.FileSet, file *ast.File, out io.Writer, importsToRemove []string) error { - for _, importToRemove := range importsToRemove { +func (ig *importsGroomer) RemoveUnusedImportFromFile(fset *token.FileSet, file *ast.File, out io.Writer, importsToRemove map[string]string) error { + for importToRemove, alias := range importsToRemove { if !astutil.UsesImport(file, importToRemove) { - astutil.DeleteImport(fset, file, importToRemove) + astutil.DeleteNamedImport(fset, file, alias, importToRemove) } } // Needed because ast does not support floating comments and deletes them. diff --git a/tracing/unused_imports_test.go b/tracing/unused_imports_test.go index c0ee1a1..2522ddf 100644 --- a/tracing/unused_imports_test.go +++ b/tracing/unused_imports_test.go @@ -27,7 +27,7 @@ func TestRemoveUnusedImportsFromFile(t *testing.T) { t.Fatal(err) } var buff bytes.Buffer - if err := NewImportsGroomer().RemoveUnusedImportFromFile(fset, file, &buff, []string{"fmt", "runtime", "rand"}); err != nil { + if err := NewImportsGroomer().RemoveUnusedImportFromFile(fset, file, &buff, map[string]string{"fmt":"", "runtime":"rt", "rand":""}); err != nil { t.Fatal(err) } @@ -63,7 +63,7 @@ func TestRemoveUnusedImportsFromDirectory(t *testing.T) { i++ } - if err := NewImportsGroomer().RemoveUnusedImportFromDirectory("test", []string{"fmt", "runtime", "rand"}); err != nil { + if err := NewImportsGroomer().RemoveUnusedImportFromDirectory("test", map[string]string{"fmt":"", "runtime":"rt", "rand":""}); err != nil { t.Fatal(err) } diff --git a/tracing/util.go b/tracing/util.go index 308bace..dca6232 100644 --- a/tracing/util.go +++ b/tracing/util.go @@ -165,7 +165,7 @@ func newGetFuncNameIfStatement(funcIndex, funcPcVarName, funcNameVarName string) &dst.CallExpr{ Fun: &dst.SelectorExpr{ X: &dst.Ident{ - Name: "runtime", + Name: "rt", }, Sel: &dst.Ident{ Name: "Caller", @@ -198,7 +198,7 @@ func newGetFuncNameIfStatement(funcIndex, funcPcVarName, funcNameVarName string) X: &dst.CallExpr{ Fun: &dst.SelectorExpr{ X: &dst.Ident{ - Name: "runtime", + Name: "rt", }, Sel: &dst.Ident{ Name: "FuncForPC", @@ -253,6 +253,8 @@ func newMakeByteSliceStmt() *dst.AssignStmt { } } +// Returns dst statement like: +// _, _ = rand.Read(idBytes) func newRandReadStmt() *dst.AssignStmt { return &dst.AssignStmt{ Lhs: []dst.Expr{ @@ -284,6 +286,8 @@ func newRandReadStmt() *dst.AssignStmt { } } +// Returns dst statement like: +// callID := fmt.Sprintf("%x-%x-%x-%x-%x", idBytes[0:4], idBytes[4:6], idBytes[6:8], idBytes[8:10], idBytes[10:]) func newParseUUIDFromByteSliceStmt(callIDVarName string) *dst.AssignStmt { return &dst.AssignStmt{ Lhs: []dst.Expr{