gofmt: handle comments correctly in rewrites

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

diff --git a/src/cmd/gofmt/gofmt_test.go b/src/cmd/gofmt/gofmt_test.go
index 4b280500976a29c9848fb7df3e65428ca9cfba02..903ba2177d0e8d0fae80db5e3754562962d64d41 100644 (file)
--- 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 3c7861f0d1152e7ad8e1753854e8fe22eafd9b23..dfabb61983ce1379e761836cfd11fe48c2151ebc 100644 (file)
--- 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 (file)
index 0000000..5a448a6
--- /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 (file)
index 0000000..0d759e6
--- /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
+}