Forbid multiple instrumentations

README.md

@@ -106,6 +106,10 @@ You can also easily revert all the changes done by `printracer` by just executin
 printracer revert
 ```
 
+> NOTE: `printracer revert` reverts changes only if code block enclosed by /* prinTracer */ comment is not modified by hand. If you modify the instrumentation block then it should be manually reverted afterwards. 
+
+> NOTE: `printracer apply` will not apply any changes if find /* prinTracer */ comment directly above first statement if the function. This is needed to mitigate accidental multiple instrumentation which will then affect deinstrumentation and visualization negatively.
+You also can use it to signal that a particular function should not be instrumented.
 ### Visualization
 
 Let's say you have instrumented your code and captured the flow that is so hard to follow even the textual trace is confusing as hell.

tracing/deinstrument_test.go

@@ -73,6 +73,7 @@ func TestDeinstrumentDirectory(t *testing.T) {
 		{InputCode: resultCodeWithMultipleImports, OutputCode: codeWithMultipleImports},
 		{InputCode: resultCodeWithImportsWithoutFmt, OutputCode: codeWithImportsWithoutFmt},
 		{InputCode: resultCodeWithoutFunction, OutputCode: codeWithoutFunction},
+		{InputCode: editedResultCodeWithoutImports, OutputCode: editedResultCodeWithoutImports},
 	}
 
 	i := 0

tracing/instrument.go

@@ -96,15 +96,27 @@ func (ci *codeInstrumenter) InstrumentFile(fset *token.FileSet, file *ast.File,
 	dst.Inspect(f, func(n dst.Node) bool {
 		switch t := n.(type) {
 		case *dst.FuncDecl:
-			instrumentationStmts := buildInstrumentationStmts(t)
-			t.Body.List = append(instrumentationStmts[:], t.Body.List...)
-
-			t.Body.List[0].Decorations().Before = dst.EmptyLine
-			t.Body.List[0].Decorations().Start.Append(printracerCommentWatermark)
-			t.Body.List[instrumentationStmtsCount-1].Decorations().After = dst.EmptyLine
-			t.Body.List[instrumentationStmtsCount-1].Decorations().End.Append(printracerCommentWatermark)
+			if !ci.hasInstrumentationWatermark(t) {
+				instrumentationStmts := buildInstrumentationStmts(t)
+				t.Body.List = append(instrumentationStmts[:], t.Body.List...)
+
+				t.Body.List[0].Decorations().Before = dst.EmptyLine
+				t.Body.List[0].Decorations().Start.Append(printracerCommentWatermark)
+				t.Body.List[instrumentationStmtsCount-1].Decorations().After = dst.EmptyLine
+				t.Body.List[instrumentationStmtsCount-1].Decorations().End.Append(printracerCommentWatermark)
+			}
 		}
 		return true
 	})
 	return decorator.Fprint(out, f)
 }
+
+func (ci *codeInstrumenter) hasInstrumentationWatermark(f *dst.FuncDecl) bool {
+	if len(f.Body.List) > 0 {
+		firstStmntDecorations := f.Body.List[0].Decorations().Start.All()
+		if len(firstStmntDecorations) > 0 && firstStmntDecorations[0] == printracerCommentWatermark {
+			return true
+		}
+	}
+	return false
+}

tracing/instrument_test.go

@@ -346,6 +346,28 @@ func main() {
 }
 `
 
+const codeWithWatermarks = `package a
+
+import (
+	"crypto/rand"
+	"fmt"
+	rt "runtime"
+)
+
+func test(i int, b bool) int {
+	/* prinTracer */
+	if b {
+		return i
+	}
+	return 0
+}
+
+func main() {
+	/* prinTracer */
+	i := test(2, false)
+}
+`
+
 const codeWithoutFunction = `package a
 
 type test struct {
@@ -377,6 +399,8 @@ func TestInstrumentFile(t *testing.T) {
 		{Name: "InstrumentFileWithMultipleImports", InputCode: codeWithMultipleImports, OutputCode: resultCodeWithMultipleImports},
 		{Name: "InstrumentFileWithoutFmtImport", InputCode: codeWithImportsWithoutFmt, OutputCode: resultCodeWithImportsWithoutFmt},
 		{Name: "InstrumentFileWithoutFunctions", InputCode: codeWithoutFunction, OutputCode: resultCodeWithoutFunction},
+		{Name: "InstrumentFileDoesNotAffectAlreadyInstrumentedFiles", InputCode: resultCodeWithFmtImport, OutputCode: resultCodeWithFmtImport},
+		{Name: "FunctionsWithWatermarksShouldNotBeInstrumented", InputCode: codeWithWatermarks, OutputCode: codeWithWatermarks},
 	}
 
 	for _, test := range tests {
@@ -417,6 +441,7 @@ func TestInstrumentDirectory(t *testing.T) {
 		{InputCode: codeWithMultipleImports, OutputCode: resultCodeWithMultipleImports},
 		{InputCode: codeWithImportsWithoutFmt, OutputCode: resultCodeWithImportsWithoutFmt},
 		{InputCode: codeWithoutFunction, OutputCode: resultCodeWithoutFunction},
+		{InputCode: resultCodeWithFmtImport, OutputCode: resultCodeWithFmtImport},
 	}
 
 	i := 0