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>

diff --git a/src/cmd/gofmt/simplify.go b/src/cmd/gofmt/simplify.go
index 2ebf4cde0beb1d0b2b8bb718eed12e96ebb2eb3b..1a0e8174afa384fad243c46fa56c8cd845d761d3 100644 (file)
--- 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 fc9c98e625b4b869250a464fc5fd6b08fc42b613..a06a69d0965ed26233fe0a8ff243581e84456703 100644 (file)
--- 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 fc7598af99edf1fc729f0e51f459ceda11ca73c0..9d28ac7ed31256463839d13e357e9557b29d3dea 100644 (file)
--- 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},
+}