From 96a609c2d7fcb8e1dc370e0cda9e4eb3bfb77412 Mon Sep 17 00:00:00 2001
From: Robert Griesemer <gri@golang.org>
Date: Mon, 25 Jun 2012 13:58:28 -0700
Subject: [PATCH] gofmt: handle comments correctly in rewrites

R=rsc
CC=golang-dev
https://golang.org/cl/6294076
---
 src/cmd/gofmt/gofmt_test.go            |  2 +-
 src/cmd/gofmt/rewrite.go               |  5 ++++-
 src/cmd/gofmt/testdata/rewrite5.golden | 15 +++++++++++++++
 src/cmd/gofmt/testdata/rewrite5.input  | 15 +++++++++++++++
 4 files changed, 35 insertions(+), 2 deletions(-)
 create mode 100644 src/cmd/gofmt/testdata/rewrite5.golden
 create mode 100644 src/cmd/gofmt/testdata/rewrite5.input

diff --git a/src/cmd/gofmt/gofmt_test.go b/src/cmd/gofmt/gofmt_test.go
index 4b280500976..903ba2177d0 100644
--- a/src/cmd/gofmt/gofmt_test.go
+++ b/src/cmd/gofmt/gofmt_test.go
@@ -66,7 +66,6 @@ func runTest(t *testing.T, in, out, flags string) {
 	}
 }
 
-// TODO(gri) Add more test cases!
 var tests = []struct {
 	in, flags string
 }{
@@ -78,6 +77,7 @@ var tests = []struct {
 	{"testdata/rewrite2.input", "-r=int->bool"},
 	{"testdata/rewrite3.input", "-r=x->x"},
 	{"testdata/rewrite4.input", "-r=(x)->x"},
+	{"testdata/rewrite5.input", "-r=x+x->2*x"},
 	{"testdata/stdin*.input", "-stdin"},
 	{"testdata/comments.input", ""},
 	{"testdata/import.input", ""},
diff --git a/src/cmd/gofmt/rewrite.go b/src/cmd/gofmt/rewrite.go
index 3c7861f0d11..dfabb61983c 100644
--- a/src/cmd/gofmt/rewrite.go
+++ b/src/cmd/gofmt/rewrite.go
@@ -55,6 +55,7 @@ func dump(msg string, val reflect.Value) {
 
 // rewriteFile applies the rewrite rule 'pattern -> replace' to an entire file.
 func rewriteFile(pattern, replace ast.Expr, p *ast.File) *ast.File {
+	cmap := ast.NewCommentMap(fileSet, p, p.Comments)
 	m := make(map[string]reflect.Value)
 	pat := reflect.ValueOf(pattern)
 	repl := reflect.ValueOf(replace)
@@ -73,7 +74,9 @@ func rewriteFile(pattern, replace ast.Expr, p *ast.File) *ast.File {
 		}
 		return val
 	}
-	return apply(f, reflect.ValueOf(p)).Interface().(*ast.File)
+	r := apply(f, reflect.ValueOf(p)).Interface().(*ast.File)
+	r.Comments = cmap.Filter(r).Comments() // recreate comments list
+	return r
 }
 
 // setValue is a wrapper for x.SetValue(y); it protects
diff --git a/src/cmd/gofmt/testdata/rewrite5.golden b/src/cmd/gofmt/testdata/rewrite5.golden
new file mode 100644
index 00000000000..5a448a63d37
--- /dev/null
+++ b/src/cmd/gofmt/testdata/rewrite5.golden
@@ -0,0 +1,15 @@
+// Copyright 2011 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Rewriting of expressions containing nodes with associated comments to
+// expressions without those nodes must also eliminate the associated
+// comments.
+
+package p
+
+func f(x int) int {
+	_ = 2 * x // this comment remains in the rewrite
+	_ = 2 * x
+	return 2 * x
+}
diff --git a/src/cmd/gofmt/testdata/rewrite5.input b/src/cmd/gofmt/testdata/rewrite5.input
new file mode 100644
index 00000000000..0d759e69b6d
--- /dev/null
+++ b/src/cmd/gofmt/testdata/rewrite5.input
@@ -0,0 +1,15 @@
+// Copyright 2011 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Rewriting of expressions containing nodes with associated comments to
+// expressions without those nodes must also eliminate the associated
+// comments.
+
+package p
+
+func f(x int) int {
+	_ = x + x // this comment remains in the rewrite
+	_ = x /* this comment must not be in the rewrite */ + x
+	return x /* this comment must not be in the rewrite */ + x
+}