From 6902c24fa18bfc9328421ca85a2774afbaff7dbb Mon Sep 17 00:00:00 2001
From: win5do <fffeng566@gmail.com>
Date: Fri, 10 Jan 2025 15:14:55 +0800
Subject: [PATCH] cmd/goimports: add regroup flag support regroup imports

Updates golang/go#64271
---
 cmd/goimports/goimports.go   |  1 +
 internal/imports/fix_test.go | 67 ++++++++++++++++++++++++++++++++++++
 internal/imports/imports.go  | 56 +++++++++++++++++++++++++++++-
 3 files changed, 123 insertions(+), 1 deletion(-)

diff --git a/cmd/goimports/goimports.go b/cmd/goimports/goimports.go
index dcb5023a2e7..103eacb8991 100644
--- a/cmd/goimports/goimports.go
+++ b/cmd/goimports/goimports.go
@@ -53,6 +53,7 @@ func init() {
 	flag.BoolVar(&options.AllErrors, "e", false, "report all errors (not just the first 10 on different lines)")
 	flag.StringVar(&options.LocalPrefix, "local", "", "put imports beginning with this string after 3rd-party packages; comma-separated list")
 	flag.BoolVar(&options.FormatOnly, "format-only", false, "if true, don't fix imports and only format. In this mode, goimports is effectively gofmt, with the addition that imports are grouped into sections.")
+	flag.BoolVar(&options.RegroupImports, "regroup", false, "remove blank line and regroup imports")
 }
 
 func report(err error) {
diff --git a/internal/imports/fix_test.go b/internal/imports/fix_test.go
index 02ddd480dfd..a2c2a94acaf 100644
--- a/internal/imports/fix_test.go
+++ b/internal/imports/fix_test.go
@@ -3085,3 +3085,70 @@ func BenchmarkMatchesPath(b *testing.B) {
 		})
 	}
 }
+
+func TestRegroupImports(t *testing.T) {
+	const input = `package foo
+
+import (
+	"fmt"
+
+	"github.com/foo/a"
+	"github.com/foo/b"
+
+	"context"
+	"go.pkg.com/bar/x"
+	"go.pkg.com/bar/y"
+
+	"github.com/foo/c"
+	"go.pkg.com/bar/z"
+)
+
+var (
+	ctx  context.Context
+	fmt1 fmt.Formatter
+	a1   a.A
+	b1   b.A
+	c1   c.A
+	x1   x.A
+	y1   y.A
+	z1   z.A
+)
+`
+
+	const want = `package foo
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/foo/a"
+	"github.com/foo/b"
+	"github.com/foo/c"
+	"go.pkg.com/bar/x"
+	"go.pkg.com/bar/y"
+	"go.pkg.com/bar/z"
+)
+
+var (
+	ctx  context.Context
+	fmt1 fmt.Formatter
+	a1   a.A
+	b1   b.A
+	c1   c.A
+	x1   x.A
+	y1   y.A
+	z1   z.A
+)
+`
+
+	testConfig{
+		module: packagestest.Module{
+			Name: "foo.com",
+			Files: fm{
+				"p/test.go": input,
+			},
+		},
+	}.processTest(t, "foo.com", "p/test.go", nil, &Options{
+		RegroupImports: true,
+	}, want)
+}
diff --git a/internal/imports/imports.go b/internal/imports/imports.go
index 2215a12880a..d7147ed483c 100644
--- a/internal/imports/imports.go
+++ b/internal/imports/imports.go
@@ -41,11 +41,19 @@ type Options struct {
 	TabIndent bool // Use tabs for indent (true if nil *Options provided)
 	TabWidth  int  // Tab width (8 if nil *Options provided)
 
-	FormatOnly bool // Disable the insertion and deletion of imports
+	FormatOnly     bool // Disable the insertion and deletion of imports
+	RegroupImports bool // Remove blank line in imports.
 }
 
 // Process implements golang.org/x/tools/imports.Process with explicit context in opt.Env.
 func Process(filename string, src []byte, opt *Options) (formatted []byte, err error) {
+	if opt.RegroupImports {
+		src, err = removeBlankLineInImport(filename, src)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	fileSet := token.NewFileSet()
 	var parserMode parser.Mode
 	if opt.Comments {
@@ -357,3 +365,49 @@ func addImportSpaces(r io.Reader, breaks []string) ([]byte, error) {
 	}
 	return out.Bytes(), nil
 }
+
+// removeBlankLineInImport remove blank line in import block.
+func removeBlankLineInImport(filename string, src []byte) (out []byte, rerr error) {
+	tokenSet := token.NewFileSet()
+	astFile, err := parser.ParseFile(tokenSet, filename, src, parser.ParseComments)
+	if err != nil {
+		return src, err
+	}
+
+	if len(astFile.Decls) <= 1 {
+		return src, nil
+	}
+	tokenFile := tokenSet.File(astFile.Pos())
+
+	for i := 0; i < len(astFile.Decls); i++ {
+		decl := astFile.Decls[i]
+		gen, ok := decl.(*ast.GenDecl)
+		if !ok || gen.Tok != token.IMPORT || declImports(gen, "C") {
+			continue
+		}
+
+		if !gen.Lparen.IsValid() {
+			// Not a block: sorted by default.
+			continue
+		}
+
+		lpLine := tokenFile.Line(gen.Lparen)
+		rpLine := tokenFile.Line(gen.Rparen)
+		src = removeEmptyLine(src, lpLine, rpLine)
+	}
+
+	return src, nil
+}
+
+func removeEmptyLine(src []byte, l, r int) []byte {
+	lines := bytes.Split(src, []byte("\n"))
+	for i := l; i < r; i++ {
+		if i < len(lines) && len(bytes.TrimSpace(lines[i])) == 0 {
+			lines = append(lines[:i], lines[i+1:]...)
+			i--
+			r--
+		}
+	}
+
+	return bytes.Join(lines, []byte("\n"))
+}