gofmt: handle comments correctly in rewrites

R=rsc
CC=golang-dev
https://golang.org/cl/6294076

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", ""},

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

src/cmd/gofmt/testdata/rewrite5.golden

+// 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
+}

src/cmd/gofmt/testdata/rewrite5.input

+// 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
+}