go/ast/astutil: match prefix segments when adding imports

AddImport and AddNamedImport attempt to place new
imports in roughly the correct place--and thus the
correct group--by matching prefixes. Matching prefixes
byte-by-byte led to "regexp" being grouped with "rsc.io/p".
Instead, match prefixes by segments.

Fixes golang/go#9961.

Change-Id: I52b7c58a9a2fbe85c2b5297e50c87d409364bda3
Reviewed-on: https://go-review.googlesource.com/8090
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>

go/ast/astutil/imports.go

@@ -308,13 +308,15 @@ func declImports(gen *ast.GenDecl, path string) bool {
 	return false
 }
 
-// matchLen returns the length of the longest prefix shared by x and y.
+// matchLen returns the length of the longest path segment prefix shared by x and y.
 func matchLen(x, y string) int {
-	i := 0
-	for i < len(x) && i < len(y) && x[i] == y[i] {
-		i++
+	n := 0
+	for i := 0; i < len(x) && i < len(y) && x[i] == y[i]; i++ {
+		if x[i] == '/' {
+			n++
 		}
-	return i
+	}
+	return n
 }
 
 // isTopName returns true if n is a top-level unresolved identifier with the given name.

go/ast/astutil/imports_test.go

@@ -313,6 +313,31 @@ package main // comment on
 import "time"
 
 type T time.Time
+`,
+	},
+
+	// Issue 9961: Match prefixes using path segments rather than bytes
+	{
+		name: "issue 9961",
+		pkg:  "regexp",
+		in: `package main
+
+import (
+	"flag"
+	"testing"
+
+	"rsc.io/p"
+)
+`,
+		out: `package main
+
+import (
+	"flag"
+	"regexp"
+	"testing"
+
+	"rsc.io/p"
+)
 `,
 	},
 }

imports/fix_test.go

@@ -669,21 +669,61 @@ import (
 )
 
 func main() { fmt.Println("pi:", math.Pi) }
+`,
+	},
+
+	// Too aggressive prefix matching
+	// golang.org/issue/9961
+	{
+		name: "issue 9961",
+		in: `package p
+
+import (
+	"zip"
+
+	"rsc.io/p"
+)
+
+var (
+	_ = fmt.Print
+	_ = zip.Store
+	_ p.P
+	_ = regexp.Compile
+)
+`,
+		out: `package p
+
+import (
+	"fmt"
+	"regexp"
+	"zip"
+
+	"rsc.io/p"
+)
+
+var (
+	_ = fmt.Print
+	_ = zip.Store
+	_ p.P
+	_ = regexp.Compile
+)
 `,
 	},
 }
 
 func TestFixImports(t *testing.T) {
 	simplePkgs := map[string]string{
-		"fmt":       "fmt",
-		"os":        "os",
-		"math":      "math",
 		"appengine": "appengine",
-		"user":      "appengine/user",
-		"zip":       "archive/zip",
 		"bytes":     "bytes",
+		"fmt":       "fmt",
+		"math":      "math",
+		"os":        "os",
+		"p":         "rsc.io/p",
+		"regexp":    "regexp",
 		"snappy":    "code.google.com/p/snappy-go/snappy",
 		"str":       "strings",
+		"user":      "appengine/user",
+		"zip":       "archive/zip",
 	}
 	findImport = func(pkgName string, symbols map[string]bool) (string, bool, error) {
 		return simplePkgs[pkgName], pkgName == "str", nil