From 18c0563fcd9ffbca8f6b071e6a29078a4e8573e9 Mon Sep 17 00:00:00 2001
From: Justin Buchanan <justinbuchanan@google.com>
Date: Wed, 19 Feb 2020 10:41:36 -0800
Subject: [PATCH] feat(go extractor): relative paths against
 KYTHE_ROOT_DIRECTORY (#4380)

By default, the go extractor makes source files relative to GOPATH or
the module root if using modules. This change extends the extractor to
use the KYTHE_ROOT_DIRECTORY environment variable instead if it is
provided.
---
 kythe/go/extractors/cmd/gotool/gotool.go |  1 +
 kythe/go/extractors/golang/golang.go     | 11 ++++++++++-
 kythe/go/extractors/govname/govname.go   | 14 ++++++++++++--
 3 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/kythe/go/extractors/cmd/gotool/gotool.go b/kythe/go/extractors/cmd/gotool/gotool.go
index d2b28fe127..ba178a242a 100644
--- a/kythe/go/extractors/cmd/gotool/gotool.go
+++ b/kythe/go/extractors/cmd/gotool/gotool.go
@@ -120,6 +120,7 @@ func main() {
 			DefaultCorpus:             *corpus,
 			Rules:                     rules,
 			CanonicalizePackageCorpus: *canonicalizePackageCorpus,
+			RootDirectory:             os.Getenv("KYTHE_ROOT_DIRECTORY"),
 		},
 	}
 	if *extraFiles != "" {
diff --git a/kythe/go/extractors/golang/golang.go b/kythe/go/extractors/golang/golang.go
index 563f63d13a..fc8119e3ab 100644
--- a/kythe/go/extractors/golang/golang.go
+++ b/kythe/go/extractors/golang/golang.go
@@ -377,12 +377,21 @@ func (p *Package) EachUnit(ctx context.Context, f func(cu *apb.CompilationUnit,
 // The digest will be the complete path as written -- this will be replaced
 // with the content digest in the fetcher.
 func (p *Package) addFiles(cu *apb.CompilationUnit, root, base string, names []string) {
+	// If a root directory is specified, use it instead of the the root from the
+	// go package loader.
+	if p.ext.RootDirectory != "" {
+		root = p.ext.RootDirectory
+	}
+	if !strings.HasSuffix(root, "/") {
+		root += "/"
+	}
+
 	for _, name := range names {
 		path := name
 		if base != "" {
 			path = filepath.Join(base, name)
 		}
-		trimmed := strings.TrimPrefix(path, root+"/")
+		trimmed := strings.TrimPrefix(path, root)
 		vn := &spb.VName{
 			Corpus: p.ext.DefaultCorpus,
 			Path:   trimmed,
diff --git a/kythe/go/extractors/govname/govname.go b/kythe/go/extractors/govname/govname.go
index d6c7d1a656..1bf21bf3ab 100644
--- a/kythe/go/extractors/govname/govname.go
+++ b/kythe/go/extractors/govname/govname.go
@@ -54,6 +54,11 @@ type PackageVNameOptions struct {
 	// Rules optionally provides a list of rules to apply to go package and file
 	// paths to customize output vnames. See the vnameutil package for details.
 	Rules vnameutil.Rules
+
+	// If set, file and package paths are made relative to this directory before
+	// applying vname rules (if any). If unset, the module root (if using
+	// modules) or the gopath directory is used instead.
+	RootDirectory string
 }
 
 // ForPackage returns a VName for a Go package.
@@ -97,9 +102,14 @@ type PackageVNameOptions struct {
 //   }
 func ForPackage(pkg *build.Package, opts *PackageVNameOptions) *spb.VName {
 	if !pkg.Goroot && opts != nil && opts.Rules != nil {
-		relpath, err := filepath.Rel(pkg.Root, pkg.Dir)
+		root := pkg.Root
+		if opts.RootDirectory != "" {
+			root = opts.RootDirectory
+		}
+
+		relpath, err := filepath.Rel(root, pkg.Dir)
 		if err != nil {
-			log.Fatalf("relativizing path %q against dir %q: %v", pkg.Dir, pkg.Root, err)
+			log.Fatalf("relativizing path %q against dir %q: %v", pkg.Dir, root, err)
 		}
 		if relpath == "." {
 			relpath = ""