cmd/gofmt: fix normalization of imaginary number literals

The old code only normalized decimal integer imaginary number
literals. But with the generalized imaginary number syntax,
the number value may be decimal, binary, octal, or hexadecimal,
integer or floating-point.

The new code only looks at the number pattern. Only for decimal
integer imaginary literals do we need to strip leading zeroes.
The remaining normalization code simply ignore the 'i' suffix.
As a result, the new code is both simpler and shorter.

Fixes #32718.

Change-Id: If43fc962a48ed62002e65d5c81fddbb9bd283984
Reviewed-on: https://go-review.googlesource.com/c/go/+/183378
Reviewed-by: Ian Lance Taylor <iant@golang.org>
diff --git a/src/cmd/gofmt/gofmt.go b/src/cmd/gofmt/gofmt.go
index 4bba444..d7a77a9 100644
--- a/src/cmd/gofmt/gofmt.go
+++ b/src/cmd/gofmt/gofmt.go
@@ -335,66 +335,49 @@
 // alone.
 func normalizeNumbers(n ast.Node) bool {
 	lit, _ := n.(*ast.BasicLit)
-	if lit == nil {
+	if lit == nil || (lit.Kind != token.INT && lit.Kind != token.FLOAT && lit.Kind != token.IMAG) {
 		return true
 	}
 	if len(lit.Value) < 2 {
-		return false // only one digit - nothing to do
+		return false // only one digit (common case) - nothing to do
 	}
 	// len(lit.Value) >= 2
 
+	// We ignore lit.Kind because for lit.Kind == token.IMAG the literal may be an integer
+	// or floating-point value, decimal or not. Instead, just consider the literal pattern.
 	x := lit.Value
-	switch lit.Kind {
-	case token.INT:
-		switch x[:2] {
-		case "0X":
-			lit.Value = "0x" + x[2:]
-		case "0O":
-			lit.Value = "0o" + x[2:]
-		case "0B":
-			lit.Value = "0b" + x[2:]
+	switch x[:2] {
+	default:
+		// 0-prefix octal, decimal int, or float (possibly with 'i' suffix)
+		if i := strings.LastIndexByte(x, 'E'); i >= 0 {
+			x = x[:i] + "e" + x[i+1:]
+			break
 		}
-
-	case token.FLOAT:
-		switch lit.Value[:2] {
-		default:
-			if i := strings.LastIndexByte(x, 'E'); i >= 0 {
-				lit.Value = x[:i] + "e" + x[i+1:]
-			}
-		case "0x":
-			if i := strings.LastIndexByte(x, 'P'); i >= 0 {
-				lit.Value = x[:i] + "p" + x[i+1:]
-			}
-		case "0X":
-			if i := strings.LastIndexByte(x, 'P'); i >= 0 {
-				lit.Value = "0x" + x[2:i] + "p" + x[i+1:]
-			} else {
-				lit.Value = "0x" + x[2:]
-			}
-		}
-
-	case token.IMAG:
-		// Note that integer imaginary literals may contain
-		// any decimal digit even if they start with zero.
-		// Imaginary literals should always end in 'i' but be
-		// conservative and check anyway before proceeding.
-		if x[0] == '0' && x[len(x)-1] == 'i' && isDecimals(x[1:len(x)-1]) {
+		// remove leading 0's from integer (but not floating-point) imaginary literals
+		if x[len(x)-1] == 'i' && strings.IndexByte(x, '.') < 0 && strings.IndexByte(x, 'e') < 0 {
 			x = strings.TrimLeft(x, "0_")
 			if x == "i" {
 				x = "0i"
 			}
-			lit.Value = x
 		}
+	case "0X":
+		x = "0x" + x[2:]
+		fallthrough
+	case "0x":
+		// possibly a hexadecimal float
+		if i := strings.LastIndexByte(x, 'P'); i >= 0 {
+			x = x[:i] + "p" + x[i+1:]
+		}
+	case "0O":
+		x = "0o" + x[2:]
+	case "0o":
+		// nothing to do
+	case "0B":
+		x = "0b" + x[2:]
+	case "0b":
+		// nothing to do
 	}
 
+	lit.Value = x
 	return false
 }
-
-// isDecimals reports whether x consists entirely of decimal digits and underscores.
-func isDecimals(x string) bool {
-	i := 0
-	for i < len(x) && ('0' <= x[i] && x[i] <= '9' || x[i] == '_') {
-		i++
-	}
-	return i == len(x)
-}
diff --git a/src/cmd/gofmt/testdata/go2numbers.golden b/src/cmd/gofmt/testdata/go2numbers.golden
index abefcb6..0184aaa 100644
--- a/src/cmd/gofmt/testdata/go2numbers.golden
+++ b/src/cmd/gofmt/testdata/go2numbers.golden
@@ -163,8 +163,24 @@
 	_ = 0.i
 	_ = 123.i
 	_ = 0123.i
+	_ = 000123.i
+
+	_ = 0e0i
+	_ = 123e0i
+	_ = 0123e0i
+	_ = 000123e0i
 
 	_ = 0.e+1i
-	_ = 123.E-1_0i
+	_ = 123.e-1_0i
 	_ = 01_23.e123i
+	_ = 00_01_23.e123i
+
+	_ = 0b1010i
+	_ = 0b1010i
+	_ = 0o660i
+	_ = 0o660i
+	_ = 0xabcDEFi
+	_ = 0xabcDEFi
+	_ = 0xabcDEFp0i
+	_ = 0xabcDEFp0i
 )
diff --git a/src/cmd/gofmt/testdata/go2numbers.input b/src/cmd/gofmt/testdata/go2numbers.input
index 51a9f8e..f3e7828 100644
--- a/src/cmd/gofmt/testdata/go2numbers.input
+++ b/src/cmd/gofmt/testdata/go2numbers.input
@@ -163,8 +163,24 @@
 	_ = 0.i
 	_ = 123.i
 	_ = 0123.i
+	_ = 000123.i
+
+	_ = 0e0i
+	_ = 123e0i
+	_ = 0123E0i
+	_ = 000123E0i
 
 	_ = 0.e+1i
 	_ = 123.E-1_0i
 	_ = 01_23.e123i
+	_ = 00_01_23.e123i
+
+	_ = 0b1010i
+	_ = 0B1010i
+	_ = 0o660i
+	_ = 0O660i
+	_ = 0xabcDEFi
+	_ = 0XabcDEFi
+	_ = 0xabcDEFP0i
+	_ = 0XabcDEFp0i
 )
diff --git a/src/math/cmplx/cmath_test.go b/src/math/cmplx/cmath_test.go
index fbb49fd..57ba76a 100644
--- a/src/math/cmplx/cmath_test.go
+++ b/src/math/cmplx/cmath_test.go
@@ -222,7 +222,7 @@
 	(-2.499956739197529585028819e+00 + 1.759751724335650228957144e+00i),
 	(7.357094338218116311191939e+04 - 5.089973412479151648145882e+04i),
 	(1.320777296067768517259592e+01 - 3.165621914333901498921986e+01i),
-	(-3.123287828297300934072149e-07 - 1.9849567521490553032502223E-7i),
+	(-3.123287828297300934072149e-07 - 1.9849567521490553032502223e-7i),
 	(8.0622651468477229614813e+04 - 7.80028727944573092944363e+04i),
 	(-1.0268824572103165858577141e+00 - 4.716844738244989776610672e-01i),
 	(-4.35953819012244175753187e+01 + 2.2036445974645306917648585e+02i),