Repository: kythe/kythe
Pull request #5314 (2022-06-22 17:19:47)

feat(go_indexer): add --override_stdlib_corpus flag
Problem: we index the go std library with the "google3/golib" corpus and the rest of go code with "google3". We want the cross-references to work between these, so std library elements in both need to have the same corpus.

For many projects we index, we want stdlib nodes to be assigned the compilation unit's corpus (enabled with --use_compilation_corpus_as_default), however in some cases we want to give stdlib items an explicit corpus so we can get cross-references to a shared instance of the std library.

shahms: Is it feasible to add tests for the overwriting behavior?
justbuchanan: I added a new verifier test showing the main corpus being used for functions defined in the file and the override corpus used for imported stdlib items.

Patch:
diff --git a/kythe/go/extractors/govname/govname.go b/kythe/go/extractors/govname/govname.go
index fd91c9867..b9a6a421c 100644
--- a/kythe/go/extractors/govname/govname.go
+++ b/kythe/go/extractors/govname/govname.go
@@ -34,7 +34,7 @@ import (
 const (
 	pathTail     = `(?:/(?P<path>.+))?$`
 	packageSig   = "package"
-	golangCorpus = "golang.org"
+	GolangCorpus = "golang.org"
 )
 
 // Language is the language string to use for Go VNames.
@@ -82,33 +82,34 @@ type PackageVNameOptions struct {
 // applicable. Note that vname rules are ignored for go stdlib packages.
 //
 // Examples:
-//   ForPackage(<kythe.io/kythe/go/util/schema>, &{CanonicalizePackageCorpus: false}) => {
-//     Corpus: "kythe.io",
-//     Path: "kythe/go/util/schema",
-//		 Language: "go",
-//		 Signature: "package",
-//   }
 //
-//   ForPackage( <kythe.io/kythe/go/util/schema>, &{CanonicalizePackageCorpus: true}) => {
-//     Corpus: "github.com/kythe/kythe",
-//     Path: "kythe/go/util/schema",
-//		 Language: "go",
-//		 Signature: "package",
-//   }
+//	  ForPackage(<kythe.io/kythe/go/util/schema>, &{CanonicalizePackageCorpus: false}) => {
+//	    Corpus: "kythe.io",
+//	    Path: "kythe/go/util/schema",
+//			 Language: "go",
+//			 Signature: "package",
+//	  }
 //
-//   ForPackage(<github.com/kythe/kythe/kythe/go/util/schema>, &{CanonicalizePackageCorpus: false}) => {
-//     Corpus: "github.com/kythe/kythe",
-//     Path: "kythe/go/util/schema",
-//		 Language: "go",
-//		 Signature: "package",
-//   }
+//	  ForPackage( <kythe.io/kythe/go/util/schema>, &{CanonicalizePackageCorpus: true}) => {
+//	    Corpus: "github.com/kythe/kythe",
+//	    Path: "kythe/go/util/schema",
+//			 Language: "go",
+//			 Signature: "package",
+//	  }
 //
-//   ForPackage(<github.com/kythe/kythe/kythe/go/util/schema>, &{CanonicalizePackageCorpus: true}) => {
-//     Corpus: "github.com/kythe/kythe",
-//     Path: "kythe/go/util/schema",
-//		 Language: "go",
-//		 Signature: "package",
-//   }
+//	  ForPackage(<github.com/kythe/kythe/kythe/go/util/schema>, &{CanonicalizePackageCorpus: false}) => {
+//	    Corpus: "github.com/kythe/kythe",
+//	    Path: "kythe/go/util/schema",
+//			 Language: "go",
+//			 Signature: "package",
+//	  }
+//
+//	  ForPackage(<github.com/kythe/kythe/kythe/go/util/schema>, &{CanonicalizePackageCorpus: true}) => {
+//	    Corpus: "github.com/kythe/kythe",
+//	    Path: "kythe/go/util/schema",
+//			 Language: "go",
+//			 Signature: "package",
+//	  }
 func ForPackage(pkg *build.Package, opts *PackageVNameOptions) *spb.VName {
 	if !pkg.Goroot && opts != nil && opts.Rules != nil {
 		root := pkg.Root
@@ -166,7 +167,7 @@ func ForPackage(pkg *build.Package, opts *PackageVNameOptions) *spb.VName {
 		if opts.UseDefaultCorpusForStdLib {
 			v.Corpus = opts.DefaultCorpus
 		} else {
-			v.Corpus = golangCorpus
+			v.Corpus = GolangCorpus
 		}
 	} else if strings.HasPrefix(ip, ".") {
 		// Local import; no corpus
@@ -185,7 +186,7 @@ func ForPackage(pkg *build.Package, opts *PackageVNameOptions) *spb.VName {
 // ForBuiltin returns a VName for a Go built-in with the given signature.
 func ForBuiltin(signature string) *spb.VName {
 	return &spb.VName{
-		Corpus:    golangCorpus,
+		Corpus:    GolangCorpus,
 		Language:  Language,
 		Root:      "ref/spec",
 		Signature: signature,
@@ -196,7 +197,7 @@ func ForBuiltin(signature string) *spb.VName {
 // given import path.
 func ForStandardLibrary(importPath string) *spb.VName {
 	return &spb.VName{
-		Corpus:    golangCorpus,
+		Corpus:    GolangCorpus,
 		Language:  Language,
 		Path:      importPath,
 		Signature: "package",
@@ -207,7 +208,7 @@ func ForStandardLibrary(importPath string) *spb.VName {
 // This includes the "golang.org" corpus but excludes the "golang.org/x/..."
 // extension repositories.  If v == nil, the answer is false.
 func IsStandardLibrary(v *spb.VName) bool {
-	return v != nil && (v.Language == "go" || v.Language == "") && v.Corpus == golangCorpus
+	return v != nil && (v.Language == "go" || v.Language == "") && v.Corpus == GolangCorpus
 }
 
 // ImportPath returns the putative Go import path corresponding to v.  The
@@ -230,7 +231,7 @@ func ImportPath(v *spb.VName, goRoot string) string {
 
 // rootRelative reports whether path has the form
 //
-//     root[/pkg/os_arch/]tail
+//	root[/pkg/os_arch/]tail
 //
 // and if so, returns the tail. It returns path, false if path does not have
 // this form.
diff --git a/kythe/go/extractors/govname/types.go b/kythe/go/extractors/govname/types.go
index 66b002882..3c1af7fb2 100644
--- a/kythe/go/extractors/govname/types.go
+++ b/kythe/go/extractors/govname/types.go
@@ -25,38 +25,38 @@ import (
 
 // BasicType returns the VName for a basic builtin Go type.
 func BasicType(b *types.Basic) *spb.VName {
-	return &spb.VName{Corpus: golangCorpus, Language: Language, Signature: b.Name() + "#builtin"}
+	return &spb.VName{Corpus: GolangCorpus, Language: Language, Signature: b.Name() + "#builtin"}
 }
 
 // FunctionConstructorType returns the VName for the builtin Go function type constructor.
 func FunctionConstructorType() *spb.VName {
-	return &spb.VName{Corpus: golangCorpus, Language: Language, Signature: "fn#builtin"}
+	return &spb.VName{Corpus: GolangCorpus, Language: Language, Signature: "fn#builtin"}
 }
 
 // TupleConstructorType returns the VName for the builtin Go tuple type constructor.
 func TupleConstructorType() *spb.VName {
-	return &spb.VName{Corpus: golangCorpus, Language: Language, Signature: "tuple#builtin"}
+	return &spb.VName{Corpus: GolangCorpus, Language: Language, Signature: "tuple#builtin"}
 }
 
 // MapConstructorType returns the VName for the builtin Go map type constructor.
 func MapConstructorType() *spb.VName {
-	return &spb.VName{Corpus: golangCorpus, Language: Language, Signature: "map#builtin"}
+	return &spb.VName{Corpus: GolangCorpus, Language: Language, Signature: "map#builtin"}
 }
 
 // ArrayConstructorType returns the VName for the builtin Go array type
 // constructor of a given length.
 func ArrayConstructorType(length int64) *spb.VName {
-	return &spb.VName{Corpus: golangCorpus, Language: Language, Signature: fmt.Sprintf("array%d#builtin", length)}
+	return &spb.VName{Corpus: GolangCorpus, Language: Language, Signature: fmt.Sprintf("array%d#builtin", length)}
 }
 
 // SliceConstructorType returns the VName for the builtin Go slice type constructor.
 func SliceConstructorType() *spb.VName {
-	return &spb.VName{Corpus: golangCorpus, Language: Language, Signature: "slice#builtin"}
+	return &spb.VName{Corpus: GolangCorpus, Language: Language, Signature: "slice#builtin"}
 }
 
 // PointerConstructorType returns the VName for the builtin Go pointer type constructor.
 func PointerConstructorType() *spb.VName {
-	return &spb.VName{Corpus: golangCorpus, Language: Language, Signature: "pointer#builtin"}
+	return &spb.VName{Corpus: GolangCorpus, Language: Language, Signature: "pointer#builtin"}
 }
 
 // ChanConstructorType returns the VName for the builtin Go chan type
@@ -71,10 +71,10 @@ func ChanConstructorType(dir types.ChanDir) *spb.VName {
 	default:
 		chanType = "chan"
 	}
-	return &spb.VName{Corpus: golangCorpus, Language: Language, Signature: chanType + "#builtin"}
+	return &spb.VName{Corpus: GolangCorpus, Language: Language, Signature: chanType + "#builtin"}
 }
 
 // VariadicConstructorType returns the VName for the builtin Go variadic type constructor.
 func VariadicConstructorType() *spb.VName {
-	return &spb.VName{Corpus: golangCorpus, Language: Language, Signature: "variadic#builtin"}
+	return &spb.VName{Corpus: GolangCorpus, Language: Language, Signature: "variadic#builtin"}
 }
diff --git a/kythe/go/indexer/BUILD b/kythe/go/indexer/BUILD
index 695f3e844..bacbd0af3 100644
--- a/kythe/go/indexer/BUILD
+++ b/kythe/go/indexer/BUILD
@@ -296,6 +296,16 @@ empty_corpus_test(
     entries = ":stdlibimport_test_entries.entries.gz",
 )
 
+go_indexer_test(
+    name = "stdliboverride_test",
+    srcs = ["testdata/basic/stdliboverride.go"],
+    extra_extractor_args = [
+        "--corpus=kythe",
+    ],
+    override_stdlib_corpus = "STDLIB_OVERRIDE",
+    use_compilation_corpus_as_default = True,
+)
+
 # load(":testdata/go_indexer_test.bzl", "go_integration_test")
 # TODO(#2375): (closed?) requires MarkedSource resolution in pipeline
 # go_integration_test(
diff --git a/kythe/go/indexer/cmd/go_indexer/go_indexer.go b/kythe/go/indexer/cmd/go_indexer/go_indexer.go
index 907fe0384..c7391c781 100644
--- a/kythe/go/indexer/cmd/go_indexer/go_indexer.go
+++ b/kythe/go/indexer/cmd/go_indexer/go_indexer.go
@@ -52,6 +52,7 @@ var (
 	verbose                        = flag.Bool("verbose", false, "Emit verbose log information")
 	contOnErr                      = flag.Bool("continue", false, "Log errors encountered during analysis but do not exit unsuccessfully")
 	useCompilationCorpusAsDefault  = flag.Bool("use_compilation_corpus_as_default", false, "Nodes that otherwise wouldn't have a corpus (such as tapps) are given the corpus of the compilation unit being indexed.")
+	overrideStdlibCorpus           = flag.String("override_stdlib_corpus", "", "If set, all stdlib nodes are assigned this corpus")
 
 	writeEntry func(context.Context, *spb.Entry) error
 	docURL     *url.URL
@@ -160,6 +161,7 @@ func indexGo(ctx context.Context, unit *apb.CompilationUnit, f indexer.Fetcher)
 		DocBase:                        docURL,
 		OnlyEmitDocURIsForStandardLibs: *onlyEmitDocURIsForStandardLibs,
 		UseCompilationCorpusAsDefault:  *useCompilationCorpusAsDefault,
+		OverrideStdlibCorpus:           *overrideStdlibCorpus,
 	})
 }
 
diff --git a/kythe/go/indexer/emit.go b/kythe/go/indexer/emit.go
index afcfc0015..671e6151e 100644
--- a/kythe/go/indexer/emit.go
+++ b/kythe/go/indexer/emit.go
@@ -68,6 +68,9 @@ type EmitOptions struct {
 	// Nodes that otherwise wouldn't have a corpus (such as tapps) are given the
 	// corpus of the compilation unit being indexed.
 	UseCompilationCorpusAsDefault bool
+
+	// If set, all stdlib nodes are assigned this corpus.
+	OverrideStdlibCorpus string
 }
 
 func (e *EmitOptions) emitMarkedSource() bool {
@@ -300,6 +303,18 @@ func (e *emitter) visitFuncDecl(decl *ast.FuncDecl, stack stackFunc) {
 	e.emitParameters(decl.Type, sig, info)
 }
 
+// rewrittenCorpusForVName returns the new corpus that should be assigned to the
+// given vname based on the OverrideStdlibCorpus and UseCompilationCorpusAsDefault options
+func (e *emitter) rewrittenCorpusForVName(v *spb.VName) string {
+	if e.opts.OverrideStdlibCorpus != "" && v.GetCorpus() == govname.GolangCorpus {
+		return e.opts.OverrideStdlibCorpus
+	}
+	if e.opts.UseCompilationCorpusAsDefault {
+		return e.pi.VName.GetCorpus()
+	}
+	return v.GetCorpus()
+}
+
 // emitTApp emits a tapp node and returns its VName.  The new tapp is emitted
 // with given constructor and parameters.  The constructor's kind is also
 // emitted if this is the first time seeing it.
@@ -315,8 +330,8 @@ func (e *emitter) emitTApp(ms *cpb.MarkedSource, ctorKind string, ctor *spb.VNam
 		components = append(components, p)
 	}
 	v := &spb.VName{Language: govname.Language, Signature: hashSignature(components)}
-	if e.opts.UseCompilationCorpusAsDefault {
-		v.Corpus = e.pi.VName.GetCorpus()
+	if e.opts.UseCompilationCorpusAsDefault || e.opts.OverrideStdlibCorpus != "" {
+		v.Corpus = e.rewrittenCorpusForVName(v)
 	}
 	if e.pi.typeEmitted.Add(v.Signature) {
 		e.writeFact(v, facts.NodeKind, nodes.TApp)
@@ -985,27 +1000,27 @@ func (e *emitter) writeSatisfies(src, tgt types.Object) {
 }
 
 func (e *emitter) writeFact(src *spb.VName, name, value string) {
-	if e.opts.UseCompilationCorpusAsDefault {
+	if e.opts.UseCompilationCorpusAsDefault || e.opts.OverrideStdlibCorpus != "" {
 		src = proto.Clone(src).(*spb.VName)
-		src.Corpus = e.pi.VName.GetCorpus()
+		src.Corpus = e.rewrittenCorpusForVName(src)
 	}
 	e.check(e.sink.writeFact(e.ctx, src, name, value))
 }
 
 func (e *emitter) writeEdge(src, tgt *spb.VName, kind string) {
-	if e.opts.UseCompilationCorpusAsDefault {
+	if e.opts.UseCompilationCorpusAsDefault || e.opts.OverrideStdlibCorpus != "" {
 		src = proto.Clone(src).(*spb.VName)
-		src.Corpus = e.pi.VName.GetCorpus()
+		src.Corpus = e.rewrittenCorpusForVName(src)
 		tgt = proto.Clone(tgt).(*spb.VName)
-		tgt.Corpus = e.pi.VName.GetCorpus()
+		tgt.Corpus = e.rewrittenCorpusForVName(tgt)
 	}
 	e.check(e.sink.writeEdge(e.ctx, src, tgt, kind))
 }
 
 func (e *emitter) writeAnchor(node ast.Node, src *spb.VName, start, end int) {
-	if e.opts.UseCompilationCorpusAsDefault {
+	if e.opts.UseCompilationCorpusAsDefault || e.opts.OverrideStdlibCorpus != "" {
 		src = proto.Clone(src).(*spb.VName)
-		src.Corpus = e.pi.VName.GetCorpus()
+		src.Corpus = e.rewrittenCorpusForVName(src)
 	}
 	if _, ok := e.anchored[node]; ok {
 		return // this node already has an anchor
@@ -1015,9 +1030,9 @@ func (e *emitter) writeAnchor(node ast.Node, src *spb.VName, start, end int) {
 }
 
 func (e *emitter) writeDiagnostic(src *spb.VName, d diagnostic) {
-	if e.opts.UseCompilationCorpusAsDefault {
+	if e.opts.UseCompilationCorpusAsDefault || e.opts.OverrideStdlibCorpus != "" {
 		src = proto.Clone(src).(*spb.VName)
-		src.Corpus = e.pi.VName.GetCorpus()
+		src.Corpus = e.rewrittenCorpusForVName(src)
 	}
 	e.check(e.sink.writeDiagnostic(e.ctx, src, d))
 }


Test patch:
diff --git a/kythe/go/extractors/govname/govname_test.go b/kythe/go/extractors/govname/govname_test.go
index 94d145657..bf64b6c1b 100644
--- a/kythe/go/extractors/govname/govname_test.go
+++ b/kythe/go/extractors/govname/govname_test.go
@@ -142,7 +142,7 @@ func TestNotStandardLib(t *testing.T) {
 func TestForBuiltin(t *testing.T) {
 	const signature = "blah"
 	want := &spb.VName{
-		Corpus:    golangCorpus,
+		Corpus:    GolangCorpus,
 		Language:  Language,
 		Root:      "ref/spec",
 		Signature: signature,
diff --git a/kythe/go/indexer/testdata/basic/stdliboverride.go b/kythe/go/indexer/testdata/basic/stdliboverride.go
new file mode 100644
index 000000000..d62ee6e4b
--- /dev/null
+++ b/kythe/go/indexer/testdata/basic/stdliboverride.go
@@ -0,0 +1,16 @@
+// Package stdliboverride tests that corpus labels are assigned correctly when
+// the --override_stdlib_corpus and --use_compilation_corpus_as_default flags
+// are enabled.
+package stdliboverride
+
+import (
+    //- @"\"bytes\"" ref/imports BYTES=vname("package","STDLIB_OVERRIDE",_,"bytes","go")
+	"bytes"
+)
+
+//- @myFunc=vname(_,"kythe",_,_,"go") defines/binding _MY_FUNC
+func myFunc() {
+    //- @bytes ref BYTES
+    //- @NewBuffer ref vname(_,"STDLIB_OVERRIDE",_,_,"go")
+	bytes.NewBuffer(nil)
+}
diff --git a/kythe/go/indexer/testdata/go_indexer_test.bzl b/kythe/go/indexer/testdata/go_indexer_test.bzl
index 6735ce223..0baedc414 100644
--- a/kythe/go/indexer/testdata/go_indexer_test.bzl
+++ b/kythe/go/indexer/testdata/go_indexer_test.bzl
@@ -172,6 +172,9 @@ def _go_entries(ctx):
     if ctx.attr.use_compilation_corpus_as_default:
         iargs.append("-use_compilation_corpus_as_default")
 
+    if ctx.attr.override_stdlib_corpus:
+        iargs.append("-override_stdlib_corpus=%s" % ctx.attr.override_stdlib_corpus)
+
     # If the test wants linkage metadata, enable support for it in the indexer.
     if ctx.attr.metadata_suffix:
         iargs += ["-meta", ctx.attr.metadata_suffix]
@@ -207,6 +210,7 @@ go_entries = rule(
         # The suffix used to recognize linkage metadata files, if non-empty.
         "metadata_suffix": attr.string(default = ""),
         "use_compilation_corpus_as_default": attr.bool(default = False),
+        "override_stdlib_corpus": attr.string(default = ""),
 
         # The location of the Go indexer binary.
         "_indexer": attr.label(
@@ -255,6 +259,7 @@ def _go_indexer(
         emit_anchor_scopes = False,
         allow_duplicates = False,
         use_compilation_corpus_as_default = False,
+        override_stdlib_corpus= "",
         metadata_suffix = "",
         extra_extractor_args = []):
     if importpath == None:
@@ -279,6 +284,7 @@ def _go_indexer(
         has_marked_source = has_marked_source,
         emit_anchor_scopes = emit_anchor_scopes,
         use_compilation_corpus_as_default = use_compilation_corpus_as_default,
+        override_stdlib_corpus = override_stdlib_corpus,
         kzip = ":" + kzip,
         metadata_suffix = metadata_suffix,
     )
@@ -299,6 +305,7 @@ def go_indexer_test(
         emit_anchor_scopes = False,
         allow_duplicates = False,
         use_compilation_corpus_as_default = False,
+        override_stdlib_corpus= "",
         metadata_suffix = "",
         extra_extractor_args = []):
     entries = _go_indexer(
@@ -308,6 +315,7 @@ def go_indexer_test(
         has_marked_source = has_marked_source,
         emit_anchor_scopes = emit_anchor_scopes,
         use_compilation_corpus_as_default = use_compilation_corpus_as_default,
+        override_stdlib_corpus=override_stdlib_corpus,
         importpath = import_path,
         metadata_suffix = metadata_suffix,
         deps = deps,