From 116ee62a3b2ec3da71560c3dc6cae5fbaa004f9c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vytautas=20=C5=A0altenis?= <vytas@rtfb.lt>
Date: Mon, 14 Aug 2017 21:43:52 +0300
Subject: [PATCH] Implement import comment extraction

Look for import comments following the package declaration. The
implementation of findImportComment() very closely follows go/build's
findImportComment().

Tests and validation against importRoot to come in a later commit.
---
 internal/gps/pkgtree/pkgtree.go | 77 +++++++++++++++++++++++++++++++--
 1 file changed, 74 insertions(+), 3 deletions(-)

diff --git a/internal/gps/pkgtree/pkgtree.go b/internal/gps/pkgtree/pkgtree.go
index 4fc4e710a2..3a6c53297b 100644
--- a/internal/gps/pkgtree/pkgtree.go
+++ b/internal/gps/pkgtree/pkgtree.go
@@ -5,7 +5,9 @@
 package pkgtree
 
 import (
+	"bytes"
 	"fmt"
+	"go/ast"
 	"go/build"
 	"go/parser"
 	gscan "go/scanner"
@@ -125,7 +127,8 @@ func ListPackages(fileRoot, importRoot string) (PackageTree, error) {
 		// Find all the imports, across all os/arch combos
 		//p, err := fullPackageInDir(wp)
 		p := &build.Package{
-			Dir: wp,
+			Dir:        wp,
+			ImportPath: ip,
 		}
 		err = fillPackage(p)
 
@@ -140,7 +143,7 @@ func ListPackages(fileRoot, importRoot string) (PackageTree, error) {
 			}
 		} else {
 			switch err.(type) {
-			case gscan.ErrorList, *gscan.Error, *build.NoGoError:
+			case gscan.ErrorList, *gscan.Error, *build.NoGoError, *ConflictingImportComments:
 				// This happens if we encounter malformed or nonexistent Go
 				// source code
 				ptree.Packages[ip] = PackageOrErr{
@@ -208,6 +211,7 @@ func fillPackage(p *build.Package) error {
 
 	var testImports []string
 	var imports []string
+	var importComments []string
 	for _, file := range gofiles {
 		// Skip underscore-led or dot-led files, in keeping with the rest of the toolchain.
 		bPrefix := filepath.Base(file)[0]
@@ -232,6 +236,10 @@ func fillPackage(p *build.Package) error {
 
 		var ignored bool
 		for _, c := range pf.Comments {
+			ic := findImportComment(pf.Name, c)
+			if ic != "" {
+				importComments = append(importComments, ic)
+			}
 			if c.Pos() > pf.Package { // +build comment must come before package
 				continue
 			}
@@ -280,7 +288,16 @@ func fillPackage(p *build.Package) error {
 			}
 		}
 	}
-
+	importComments = uniq(importComments)
+	if len(importComments) > 1 {
+		return &ConflictingImportComments{
+			ImportPath:     p.ImportPath,
+			ImportComments: importComments,
+		}
+	}
+	if len(importComments) > 0 {
+		p.ImportComment = importComments[0]
+	}
 	imports = uniq(imports)
 	testImports = uniq(testImports)
 	p.Imports = imports
@@ -288,6 +305,60 @@ func fillPackage(p *build.Package) error {
 	return nil
 }
 
+var (
+	slashSlash = []byte("//")
+	slashStar  = []byte("/*")
+	starSlash  = []byte("*/")
+	newline    = []byte("\n")
+	importKwd  = []byte("import ")
+)
+
+func findImportComment(pkgName *ast.Ident, c *ast.CommentGroup) string {
+	afterPkg := pkgName.NamePos + token.Pos(len(pkgName.Name)) + 1
+	commentSlash := c.List[0].Slash
+	if afterPkg != commentSlash {
+		return ""
+	}
+	text := []byte(c.List[0].Text)
+	switch {
+	case bytes.HasPrefix(text, slashSlash):
+		eol := bytes.IndexByte(text, '\n')
+		if eol < 0 {
+			eol = len(text)
+		}
+		text = text[2:eol]
+	case bytes.HasPrefix(text, slashStar):
+		text = text[2:]
+		end := bytes.Index(text, starSlash)
+		if end < 0 {
+			// malformed comment
+			return ""
+		}
+		text = text[:end]
+		if bytes.IndexByte(text, '\n') > 0 {
+			// multiline comment, can't be an import comment
+			return ""
+		}
+	}
+	text = bytes.TrimSpace(text)
+	if !bytes.HasPrefix(text, importKwd) {
+		return ""
+	}
+	quotedPath := bytes.TrimSpace(text[len(importKwd):])
+	return string(bytes.Trim(quotedPath, `"`))
+}
+
+// ConflictingImportComments conflict
+type ConflictingImportComments struct {
+	ImportPath     string
+	ImportComments []string
+}
+
+func (e *ConflictingImportComments) Error() string {
+	return fmt.Sprintf("import path %s had conflicting import comments: %q",
+		e.ImportPath, strings.Join(e.ImportComments, "\", \""))
+}
+
 // LocalImportsError indicates that a package contains at least one relative
 // import that will prevent it from compiling.
 //