❤️ set pos

process.go

@@ -44,7 +44,7 @@ func Process(filename string, src []byte) ([]byte, error) {
 				var ds ast.Stmt
 				switch t {
 				case TypeContext:
-					ds = buildDeferStmt(pkg, vn, sn)
+					ds = buildDeferStmt(fd.Body.Lbrace, pkg, vn, sn)
 				case TypeHttpRequest:
 					ds = buildDeferStmtWithHttpRequest(pkg, vn, sn)
 				case TypeUnknown:
@@ -95,25 +95,30 @@ func addImport(fs *token.FileSet, f *ast.File) (string, error) {
 // buildDeferStmt builds the defer statement with args.
 // ex:
 //    defer newrelic.FromContext(ctx).StartSegment("slow").End()
-func buildDeferStmt(pkgName, ctxName, segName string) *ast.DeferStmt {
+func buildDeferStmt(pos token.Pos, pkgName, ctxName, segName string) *ast.DeferStmt {
 	return &ast.DeferStmt{
 		Call: &ast.CallExpr{
 			Fun: &ast.SelectorExpr{
 				X: &ast.CallExpr{
 					Fun: &ast.SelectorExpr{
 						X: &ast.CallExpr{
 							Fun: &ast.SelectorExpr{
-								X:   &ast.Ident{Name: pkgName},
-								Sel: &ast.Ident{Name: "FromContext"},
+								X:   &ast.Ident{NamePos: pos, Name: pkgName},
+								Sel: &ast.Ident{NamePos: pos, Name: "FromContext"},
 							},
-							Args: []ast.Expr{&ast.Ident{Name: ctxName}},
+							Args: []ast.Expr{&ast.Ident{NamePos: pos, Name: ctxName}},
 						},
-						Sel: &ast.Ident{Name: "StartSegment"},
+						Sel: &ast.Ident{NamePos: pos, Name: "StartSegment"},
 					},
-					Args: []ast.Expr{&ast.BasicLit{Kind: token.STRING, Value: strconv.Quote(segName)}},
+					Args: []ast.Expr{&ast.BasicLit{
+						ValuePos: pos,
+						Kind:     token.STRING,
+						Value:    strconv.Quote(segName),
+					}},
 				},
-				Sel: &ast.Ident{Name: "End"},
+				Sel: &ast.Ident{NamePos: pos, Name: "End"},
 			},
+			Rparen: pos,
 		},
 	}
 }