optimize detect import blocks (#120)

Signed-off-by: Loong <loong.dai@intel.com>

pkg/gci/gci.go

@@ -139,20 +139,6 @@ func LoadFormatGoFile(file io.FileObj, cfg config.Config) (src, dist []byte, err
 		return nil, nil, err
 	}
 
-	var head []byte
-	// use `import` as key to locate
-	importEnd := bytes.LastIndex(src[:headEnd], []byte{'t'})
-	head = make([]byte, importEnd+1)
-	copy(head, src[:importEnd+1])
-	// append `(\n\t`
-	head = append(head, []byte{utils.LeftParenthesis, utils.Linebreak, utils.Indent}...)
-	log.L().Debug(fmt.Sprintf("head: %s", head))
-	tail := src[tailStart:]
-	// for test
-	if len(tail) == 0 {
-		tail = []byte{utils.RightParenthesis, utils.Linebreak}
-	}
-
 	firstWithIndex := true
 
 	var body []byte
@@ -170,10 +156,22 @@ func LoadFormatGoFile(file io.FileObj, cfg config.Config) (src, dist []byte, err
 		}
 	}
 
-	if tail[0] != utils.Linebreak {
-		body = append(body, utils.Linebreak)
+	head := src[:headEnd]
+	tail := src[tailStart:]
+
+	head = append(head, utils.Linebreak)
+	// add beginning of import block
+	head = append(head, `import (`...)
+	// add end of import block
+	body = append(body, []byte{utils.RightParenthesis, utils.Linebreak}...)
+
+	log.L().Debug(fmt.Sprintf("head:\n%s", head))
+	log.L().Debug(fmt.Sprintf("body:\n%s", body))
+	if len(tail) > 20 {
+		log.L().Debug(fmt.Sprintf("tail:\n%s", tail[:20]))
+	} else {
+		log.L().Debug(fmt.Sprintf("tail:\n%s", tail))
 	}
-	log.L().Debug(fmt.Sprintf("body: %s", body))
 
 	var totalLen int
 	slices := [][]byte{head, body, tail}
@@ -185,6 +183,7 @@ func LoadFormatGoFile(file io.FileObj, cfg config.Config) (src, dist []byte, err
 	for _, s := range slices {
 		i += copy(dist[i:], s)
 	}
+	log.L().Debug(fmt.Sprintf("raw:\n%s", dist))
 	dist, err = goFormat.Source(dist)
 	if err != nil {
 		return nil, nil, err

pkg/gci/gci_test.go

@@ -36,7 +36,7 @@ func TestRun(t *testing.T) {
 	}
 	for _, testFile := range testFiles {
 		fileBaseName := strings.TrimSuffix(testFile, ".in.go")
-		t.Run(fileBaseName, func(t *testing.T) {
+		t.Run("pkg/gci/"+testFile, func(t *testing.T) {
 			t.Parallel()
 
 			gciCfg, err := config.InitializeGciConfigFromYAML(fileBaseName + ".cfg.yaml")

pkg/gci/internal/testdata/multiple-imports.cfg.yaml

@@ -0,0 +1 @@
+common.cfg.yaml

pkg/gci/internal/testdata/multiple-imports.in.go

@@ -9,3 +9,9 @@ import (
 
 	"github.com/daixiang0/test"
 )
+
+import "math"
+
+
+func main() {
+}

pkg/gci/internal/testdata/multiple-imports.out.go

@@ -3,7 +3,11 @@ package main
 import (
 	"context"
 	"fmt"
+	"math"
 	"os"
 
 	"github.com/daixiang0/test"
 )
+
+func main() {
+}

pkg/parse/parse.go

@@ -81,23 +81,53 @@ func ParseFile(src []byte, filename string) (ImportList, int, int, error) {
 		return nil, 0, 0, NoImportError{}
 	}
 
-	var headEnd, tailStart int
+	var (
+		// headEnd means the start of import block
+		headEnd int
+		// tailStart means the end + 1 of import block
+		tailStart int
+		// lastImportStart means the start of last import block
+		lastImportStart int
+		data            ImportList
+	)
 
-	var data ImportList
-	for i, imp := range f.Imports {
+	for i, d := range f.Decls {
+		switch d.(type) {
+		case *ast.GenDecl:
+			dd := d.(*ast.GenDecl)
+			if dd.Tok == token.IMPORT {
+				// there are two cases, both end with linebreak:
+				// 1.
+				// import (
+				//	 "xxxx"
+				// )
+				// 2.
+				// import "xxx"
+				if headEnd == 0 {
+					headEnd = int(d.Pos()) - 1
+				}
+				tailStart = int(d.End())
+				lastImportStart = i
+			}
+		}
+	}
+
+	if len(f.Decls) > lastImportStart+1 {
+		tailStart = int(f.Decls[lastImportStart+1].Pos() - 1)
+	}
+
+	for _, imp := range f.Imports {
 		if imp.Path.Value == C {
+			if imp.Comment != nil {
+				headEnd = int(imp.Comment.End())
+			} else {
+				headEnd = int(imp.Path.End())
+			}
 			continue
 		}
 
 		start, end, name := getImports(imp)
 
-		if headEnd == 0 {
-			headEnd = start
-		}
-		if i == len(f.Imports)-1 {
-			tailStart = end
-		}
-
 		data = append(data, &GciImports{
 			Start: start,
 			End:   end,