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>

3 files changed, 62 insertions, 47 deletions

diff --git a/src/cmd/gofmt/gofmt.go b/src/cmd/gofmt/gofmt.go
index 4bba44489d..d7a77a9682 100644
--- a/src/cmd/gofmt/gofmt.go
+++ b/src/cmd/gofmt/gofmt.go
@@ -335,66 +335,49 @@ func backupFile(filename string, data []byte, perm os.FileMode) (string, error)
 // 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 abefcb6c58..0184aaa6ce 100644
--- a/src/cmd/gofmt/testdata/go2numbers.golden
+++ b/src/cmd/gofmt/testdata/go2numbers.golden
@@ -163,8 +163,24 @@ const (
 	_ = 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 51a9f8eaf6..f3e7828d94 100644
--- a/src/cmd/gofmt/testdata/go2numbers.input
+++ b/src/cmd/gofmt/testdata/go2numbers.input
@@ -163,8 +163,24 @@ const (
 	_ = 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
 )