cmd/gofmt: simplify map key literals

Simplify map key literals in "gofmt -s"

Fixes #16461.

Change-Id: Ia61739b34a30ac27f6696f94a98809109a8a7b61
Reviewed-on: https://go-review.googlesource.com/25530
Run-TryBot: Russ Cox <rsc@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Russ Cox <rsc@golang.org>

3 files changed, 64 insertions, 24 deletions

diff --git a/src/cmd/gofmt/simplify.go b/src/cmd/gofmt/simplify.go
index 2ebf4cde0b..1a0e8174af 100644
--- a/src/cmd/gofmt/simplify.go
+++ b/src/cmd/gofmt/simplify.go
@@ -17,47 +17,33 @@ func (s simplifier) Visit(node ast.Node) ast.Visitor {
 	case *ast.CompositeLit:
 		// array, slice, and map composite literals may be simplified
 		outer := n
-		var eltType ast.Expr
+		var keyType, eltType ast.Expr
 		switch typ := outer.Type.(type) {
 		case *ast.ArrayType:
 			eltType = typ.Elt
 		case *ast.MapType:
+			keyType = typ.Key
 			eltType = typ.Value
 		}
 
 		if eltType != nil {
+			var ktyp reflect.Value
+			if keyType != nil {
+				ktyp = reflect.ValueOf(keyType)
+			}
 			typ := reflect.ValueOf(eltType)
 			for i, x := range outer.Elts {
 				px := &outer.Elts[i]
 				// look at value of indexed/named elements
 				if t, ok := x.(*ast.KeyValueExpr); ok {
+					if keyType != nil {
+						s.simplifyLiteral(ktyp, keyType, t.Key, &t.Key)
+					}
 					x = t.Value
 					px = &t.Value
 				}
-				ast.Walk(s, x) // simplify x
-				// if the element is a composite literal and its literal type
-				// matches the outer literal's element type exactly, the inner
-				// literal type may be omitted
-				if inner, ok := x.(*ast.CompositeLit); ok {
-					if match(nil, typ, reflect.ValueOf(inner.Type)) {
-						inner.Type = nil
-					}
-				}
-				// if the outer literal's element type is a pointer type *T
-				// and the element is & of a composite literal of type T,
-				// the inner &T may be omitted.
-				if ptr, ok := eltType.(*ast.StarExpr); ok {
-					if addr, ok := x.(*ast.UnaryExpr); ok && addr.Op == token.AND {
-						if inner, ok := addr.X.(*ast.CompositeLit); ok {
-							if match(nil, reflect.ValueOf(ptr.X), reflect.ValueOf(inner.Type)) {
-								inner.Type = nil // drop T
-								*px = inner      // drop &
-							}
-						}
-					}
-				}
+				s.simplifyLiteral(typ, eltType, x, px)
 			}
-
 			// node was simplified - stop walk (there are no subnodes to simplify)
 			return nil
 		}
@@ -113,6 +99,32 @@ func (s simplifier) Visit(node ast.Node) ast.Visitor {
 	return s
 }
 
+func (s simplifier) simplifyLiteral(typ reflect.Value, astType, x ast.Expr, px *ast.Expr) {
+	ast.Walk(s, x) // simplify x
+
+	// if the element is a composite literal and its literal type
+	// matches the outer literal's element type exactly, the inner
+	// literal type may be omitted
+	if inner, ok := x.(*ast.CompositeLit); ok {
+		if match(nil, typ, reflect.ValueOf(inner.Type)) {
+			inner.Type = nil
+		}
+	}
+	// if the outer literal's element type is a pointer type *T
+	// and the element is & of a composite literal of type T,
+	// the inner &T may be omitted.
+	if ptr, ok := astType.(*ast.StarExpr); ok {
+		if addr, ok := x.(*ast.UnaryExpr); ok && addr.Op == token.AND {
+			if inner, ok := addr.X.(*ast.CompositeLit); ok {
+				if match(nil, reflect.ValueOf(ptr.X), reflect.ValueOf(inner.Type)) {
+					inner.Type = nil // drop T
+					*px = inner      // drop &
+				}
+			}
+		}
+	}
+}
+
 func isBlank(x ast.Expr) bool {
 	ident, ok := x.(*ast.Ident)
 	return ok && ident.Name == "_"
diff --git a/src/cmd/gofmt/testdata/composites.golden b/src/cmd/gofmt/testdata/composites.golden
index fc9c98e625..a06a69d096 100644
--- a/src/cmd/gofmt/testdata/composites.golden
+++ b/src/cmd/gofmt/testdata/composites.golden
@@ -6,6 +6,10 @@ type T struct {
 	x, y int
 }
 
+type T2 struct {
+	w, z int
+}
+
 var _ = [42]T{
 	{},
 	{1, 2},
@@ -202,3 +206,13 @@ var pieces4 = []*Piece{
 	{2, 0, Point{4, 1}, []Point{{0, 0}, {1, 0}, {1, 0}, {1, 0}}, nil, nil},
 	{3, 0, Point{1, 4}, []Point{{0, 0}, {0, 1}, {0, 1}, {0, 1}}, nil, nil},
 }
+
+var _ = map[T]T2{
+	{1, 2}: {3, 4},
+	{5, 6}: {7, 8},
+}
+
+var _ = map[*T]*T2{
+	{1, 2}: {3, 4},
+	{5, 6}: {7, 8},
+}
diff --git a/src/cmd/gofmt/testdata/composites.input b/src/cmd/gofmt/testdata/composites.input
index fc7598af99..9d28ac7ed3 100644
--- a/src/cmd/gofmt/testdata/composites.input
+++ b/src/cmd/gofmt/testdata/composites.input
@@ -6,6 +6,10 @@ type T struct {
 	x, y int
 }
 
+type T2 struct {
+	w, z int
+}
+
 var _ = [42]T{
 	T{},
 	T{1, 2},
@@ -202,3 +206,13 @@ var pieces4 = []*Piece{
 	&Piece{2, 0, Point{4, 1}, []Point{Point{0, 0}, Point{1, 0}, Point{1, 0}, Point{1, 0}}, nil, nil},
 	&Piece{3, 0, Point{1, 4}, []Point{Point{0, 0}, Point{0, 1}, Point{0, 1}, Point{0, 1}}, nil, nil},
 }
+
+var _ = map[T]T2{
+	T{1, 2}: T2{3, 4},
+	T{5, 6}: T2{7, 8},
+}
+
+var _ = map[*T]*T2{
+	&T{1, 2}: &T2{3, 4},
+	&T{5, 6}: &T2{7, 8},
+}