cmd/guru: compute referrers as packages are type checked

This CL makes little observable difference to the behavior but paves the
way for streaming 'referrers' and (later) 'implements' queries which
scan the entire workspace, but print each result as soon as it is found.

The go/loader package now exposes a hook, AfterTypeCheck, that lets
clients inspect each package as soon as it is type-checked, and also
modify it, for instance to release unneeded data structures.

A 'referrers' query applied to an exported object must scan the entire
workspace.  It uses this hook so to gather uses of the query object in
streaming fashion.  However, for now, it still accumulates the results
and prints them all at the end, though I propose to change that in a
follow-up.

Code details:
- The referrers logic had a 2-iteration loop to load first the query
  package and then if necessary the enlarged program.  The second
  iteration has now been unrolled and split into globalReferrers.
- Queries for package names (whether in a package declaration or
  a qualified identifier) have been split off into packageReferrers.
  It now loads all direct importers of the query package,
  which catches some references that were missing before.
  (It used to inspect only the forward dependencies of the query
  package.)

Also:
- Referrers.Pos (the position of query identifier) was removed from the
  JSON output.  It's a nuisance to compute now, and it's already
  absent from the plain output.
  (In a follow-up, I plan to simplify the information content of the
  JSON output exactly what is currently printed in the plain output.)

Change-Id: Ia5677636dc7b0fe4461a5d393107665757fb9a97
Reviewed-on: https://go-review.googlesource.com/19794
Reviewed-by: Daniel Morsing <daniel.morsing@gmail.com>

cmd/guru/referrers.go

@@ -12,7 +12,10 @@ import (
 	"go/token"
 	"go/types"
 	"io"
+	"log"
 	"sort"
+	"strings"
+	"sync"
 
 	"golang.org/x/tools/cmd/guru/serial"
 	"golang.org/x/tools/go/buildutil"
@@ -21,104 +24,311 @@ import (
 )
 
 // Referrers reports all identifiers that resolve to the same object
-// as the queried identifier, within any package in the analysis scope.
+// as the queried identifier, within any package in the workspace.
 func referrers(q *Query) error {
-	lconf := loader.Config{Build: q.Build}
+	fset := token.NewFileSet()
+	lconf := loader.Config{Fset: fset, Build: q.Build}
 	allowErrors(&lconf)
 
 	if _, err := importQueryPackage(q.Pos, &lconf); err != nil {
 		return err
 	}
 
-	var id *ast.Ident
-	var obj types.Object
-	var lprog *loader.Program
-	var pass2 bool
-	var qpos *queryPos
-	for {
-		// Load/parse/type-check the program.
-		var err error
-		lprog, err = lconf.Load()
-		if err != nil {
-			return err
-		}
-		q.Fset = lprog.Fset
-
-		qpos, err = parseQueryPos(lprog, q.Pos, false)
-		if err != nil {
-			return err
-		}
-
-		id, _ = qpos.path[0].(*ast.Ident)
-		if id == nil {
-			return fmt.Errorf("no identifier here")
-		}
-
-		obj = qpos.info.ObjectOf(id)
-		if obj == nil {
-			// Happens for y in "switch y := x.(type)",
-			// the package declaration,
-			// and unresolved identifiers.
-			if _, ok := qpos.path[1].(*ast.File); ok { // package decl?
-				pkg := qpos.info.Pkg
-				obj = types.NewPkgName(id.Pos(), pkg, pkg.Name(), pkg)
-			} else {
-				return fmt.Errorf("no object for identifier: %T", qpos.path[1])
-			}
-		}
-
-		if pass2 {
-			break
-		}
-
-		// If the identifier is exported, we must load all packages that
-		// depend transitively upon the package that defines it.
-		// Treat PkgNames as exported, even though they're lowercase.
-		if _, isPkg := obj.(*types.PkgName); !(isPkg || obj.Exported()) {
-			break // not exported
-		}
-
-		// Scan the workspace and build the import graph.
-		// Ignore broken packages.
-		_, rev, _ := importgraph.Build(q.Build)
-
-		// Re-load the larger program.
-		// Create a new file set so that ...
-		// External test packages are never imported,
-		// so they will never appear in the graph.
-		// (We must reset the Config here, not just reset the Fset field.)
-		lconf = loader.Config{
-			Fset:  token.NewFileSet(),
-			Build: q.Build,
-		}
-		allowErrors(&lconf)
-		for path := range rev.Search(obj.Pkg().Path()) {
-			lconf.ImportWithTests(path)
-		}
-		pass2 = true
+	// Load/parse/type-check the query package.
+	lprog, err := lconf.Load()
+	if err != nil {
+		return err
 	}
 
-	// Iterate over all go/types' Uses facts for the entire program.
-	var refs []*ast.Ident
-	for _, info := range lprog.AllPackages {
-		for id2, obj2 := range info.Uses {
-			if sameObj(obj, obj2) {
-				refs = append(refs, id2)
-			}
-		}
+	qpos, err := parseQueryPos(lprog, q.Pos, false)
+	if err != nil {
+		return err
 	}
-	sort.Sort(byNamePos{q.Fset, refs})
 
+	id, _ := qpos.path[0].(*ast.Ident)
+	if id == nil {
+		return fmt.Errorf("no identifier here")
+	}
+
+	obj := qpos.info.ObjectOf(id)
+	if obj == nil {
+		// Happens for y in "switch y := x.(type)",
+		// the package declaration,
+		// and unresolved identifiers.
+		if _, ok := qpos.path[1].(*ast.File); ok { // package decl?
+			return packageReferrers(q, qpos.info.Pkg.Path())
+		}
+		return fmt.Errorf("no object for identifier: %T", qpos.path[1])
+	}
+
+	// Imported package name?
+	if pkgname, ok := obj.(*types.PkgName); ok {
+		return packageReferrers(q, pkgname.Imported().Path())
+	}
+
+	if obj.Pkg() == nil {
+		return fmt.Errorf("references to predeclared %q are everywhere!", obj.Name())
+	}
+
+	// For a globally accessible object defined in package P, we
+	// must load packages that depend on P.  Specifically, for a
+	// package-level object, we need load only direct importers
+	// of P, but for a field or interface method, we must load
+	// any package that transitively imports P.
+	if global, pkglevel := classify(obj); global {
+		// We'll use the the object's position to identify it in the larger program.
+		objposn := fset.Position(obj.Pos())
+		defpkg := obj.Pkg().Path() // defining package
+		return globalReferrers(q, qpos.info.Pkg.Path(), defpkg, objposn, pkglevel)
+	}
+
+	// Find uses of obj within the query package.
+	refs := usesOf(obj, qpos.info)
+	sort.Sort(byNamePos{fset, refs})
+	q.Fset = fset
 	q.result = &referrersResult{
 		build: q.Build,
-		qpos:  qpos,
-		query: id,
+		fset:  fset,
+		qinfo: qpos.info,
+		obj:   obj,
+		refs:  refs,
+	}
+	return nil // success
+}
+
+// classify classifies objects by how far
+// we have to look to find references to them.
+func classify(obj types.Object) (global, pkglevel bool) {
+	if obj.Exported() {
+		if obj.Parent() == nil {
+			// selectable object (field or method)
+			return true, false
+		}
+		if obj.Parent() == obj.Pkg().Scope() {
+			// lexical object (package-level var/const/func/type)
+			return true, true
+		}
+	}
+	// object with unexported named or defined in local scope
+	return false, false
+}
+
+// packageReferrers finds all references to the specified package
+// throughout the workspace and populates q.result.
+func packageReferrers(q *Query, path string) error {
+	// Scan the workspace and build the import graph.
+	// Ignore broken packages.
+	_, rev, _ := importgraph.Build(q.Build)
+
+	// Find the set of packages that directly import the query package.
+	// Only those packages need typechecking of function bodies.
+	users := rev[path]
+
+	// Load the larger program.
+	fset := token.NewFileSet()
+	lconf := loader.Config{
+		Fset:  fset,
+		Build: q.Build,
+		TypeCheckFuncBodies: func(p string) bool {
+			return users[strings.TrimSuffix(p, "_test")]
+		},
+	}
+	allowErrors(&lconf)
+	for path := range users {
+		lconf.ImportWithTests(path)
+	}
+	lprog, err := lconf.Load()
+	if err != nil {
+		return err
+	}
+
+	// Find uses of [a fake PkgName that imports] the package.
+	//
+	// TODO(adonovan): perhaps more useful would be to show imports
+	// of the package instead of qualified identifiers.
+	qinfo := lprog.Package(path)
+	obj := types.NewPkgName(token.NoPos, qinfo.Pkg, qinfo.Pkg.Name(), qinfo.Pkg)
+	refs := usesOf(obj, lprog.InitialPackages()...)
+	sort.Sort(byNamePos{fset, refs})
+	q.Fset = fset
+	q.result = &referrersResult{
+		build: q.Build,
+		fset:  fset,
+		qinfo: qinfo,
 		obj:   obj,
 		refs:  refs,
 	}
 	return nil
 }
 
+// globalReferrers finds references throughout the entire workspace to the
+// object at the specified source position.  Its defining package is defpkg,
+// and the query package is qpkg.  isPkgLevel indicates whether the object
+// is defined at package-level.
+func globalReferrers(q *Query, qpkg, defpkg string, objposn token.Position, isPkgLevel bool) error {
+	// Scan the workspace and build the import graph.
+	// Ignore broken packages.
+	_, rev, _ := importgraph.Build(q.Build)
+
+	// Find the set of packages that depend on defpkg.
+	// Only function bodies in those packages need type-checking.
+	var users map[string]bool
+	if isPkgLevel {
+		users = rev[defpkg]  // direct importers
+		users[defpkg] = true // plus the defining package itself
+	} else {
+		users = rev.Search(defpkg) // transitive importers
+	}
+
+	// Prepare to load the larger program.
+	fset := token.NewFileSet()
+	lconf := loader.Config{
+		Fset:  fset,
+		Build: q.Build,
+		TypeCheckFuncBodies: func(p string) bool {
+			return users[strings.TrimSuffix(p, "_test")]
+		},
+	}
+	allowErrors(&lconf)
+
+	// The importgraph doesn't treat external test packages
+	// as separate nodes, so we must use ImportWithTests.
+	for path := range users {
+		lconf.ImportWithTests(path)
+	}
+
+	// The remainder of this function is somewhat tricky because it
+	// operates on the concurrent stream of packages observed by the
+	// loader's AfterTypeCheck hook.  Most of guru's helper
+	// functions assume the entire program has already been loaded,
+	// so we can't use them here.
+	// TODO(adonovan): smooth things out once the other changes have landed.
+
+	var (
+		mu    sync.Mutex
+		qobj  types.Object
+		qinfo *loader.PackageInfo // info for qpkg
+	)
+
+	// For efficiency, we scan each package for references
+	// just after it has been type-checked.  The loader calls
+	// AfterTypeCheck (concurrently), providing us with a stream of
+	// packages.
+	ch := make(chan []*ast.Ident)
+	lconf.AfterTypeCheck = func(info *loader.PackageInfo, files []*ast.File) {
+		// Only inspect packages that depend on the declaring package
+		// (and thus were type-checked).
+		if lconf.TypeCheckFuncBodies(info.Pkg.Path()) {
+			// Record the query object and its package when we see it.
+			mu.Lock()
+			if qobj == nil && info.Pkg.Path() == defpkg {
+				// Find the object by its position (slightly ugly).
+				qobj = findObject(fset, &info.Info, objposn)
+				if qobj == nil {
+					// It really ought to be there;
+					// we found it once already.
+					log.Fatalf("object at %s not found in package %s",
+						objposn, defpkg)
+				}
+				qinfo = info
+			}
+			obj := qobj
+			mu.Unlock()
+
+			// Look for references to the query object.
+			if obj != nil {
+				ch <- usesOf(obj, info)
+			}
+		}
+
+		// TODO(adonovan): opt: save memory by eliminating unneeded scopes/objects.
+		// (Requires go/types change for Go 1.7.)
+		//   info.Pkg.Scope().ClearChildren()
+
+		// Discard the file ASTs and their accumulated type
+		// information to save memory.
+		info.Files = nil
+		info.Defs = make(map[*ast.Ident]types.Object)
+		info.Uses = make(map[*ast.Ident]types.Object)
+		info.Implicits = make(map[ast.Node]types.Object)
+
+		// Also, disable future collection of wholly unneeded
+		// type information for the package in case there is
+		// more type-checking to do (augmentation).
+		info.Types = nil
+		info.Scopes = nil
+		info.Selections = nil
+	}
+
+	go func() {
+		lconf.Load() // ignore error
+		close(ch)
+	}()
+
+	var refs []*ast.Ident
+	for ids := range ch {
+		refs = append(refs, ids...)
+	}
+	sort.Sort(byNamePos{fset, refs})
+
+	if qobj == nil {
+		log.Fatal("query object not found during reloading")
+	}
+
+	// TODO(adonovan): in a follow-up, do away with the
+	// analyze/display split so we can print a stream of output
+	// directly from the AfterTypeCheck hook.
+	// (We should not assume that users let the program run long
+	// enough for Load to return.)
+
+	q.Fset = fset
+	q.result = &referrersResult{
+		build: q.Build,
+		fset:  fset,
+		qinfo: qinfo,
+		obj:   qobj,
+		refs:  refs,
+	}
+
+	return nil // success
+}
+
+// findObject returns the object defined at the specified position.
+func findObject(fset *token.FileSet, info *types.Info, objposn token.Position) types.Object {
+	good := func(obj types.Object) bool {
+		if obj == nil {
+			return false
+		}
+		posn := fset.Position(obj.Pos())
+		return posn.Filename == objposn.Filename && posn.Offset == objposn.Offset
+	}
+	for _, obj := range info.Defs {
+		if good(obj) {
+			return obj
+		}
+	}
+	for _, obj := range info.Implicits {
+		if good(obj) {
+			return obj
+		}
+	}
+	return nil
+}
+
+// usesOf returns all identifiers in the packages denoted by infos
+// that refer to queryObj.
+func usesOf(queryObj types.Object, infos ...*loader.PackageInfo) []*ast.Ident {
+	var refs []*ast.Ident
+	for _, info := range infos {
+		for id, obj := range info.Uses {
+			if sameObj(queryObj, obj) {
+				refs = append(refs, id)
+			}
+		}
+	}
+	return refs
+}
+
 // same reports whether x and y are identical, or both are PkgNames
 // that import the same Package.
 //
@@ -160,14 +370,16 @@ func (p byNamePos) Less(i, j int) bool {
 
 type referrersResult struct {
 	build *build.Context
+	fset  *token.FileSet
+	qinfo *loader.PackageInfo
 	qpos  *queryPos
-	query *ast.Ident   // identifier of query
 	obj   types.Object // object it denotes
 	refs  []*ast.Ident // set of all other references to it
 }
 
 func (r *referrersResult) display(printf printfFunc) {
-	printf(r.obj, "%d references to %s", len(r.refs), r.qpos.objectString(r.obj))
+	printf(r.obj, "%d references to %s",
+		len(r.refs), types.ObjectString(r.obj, types.RelativeTo(r.qinfo.Pkg)))
 
 	// Show referring lines, like grep.
 	type fileinfo struct {
@@ -181,7 +393,7 @@ func (r *referrersResult) display(printf printfFunc) {
 	// First pass: start the file reads concurrently.
 	sema := make(chan struct{}, 20) // counting semaphore to limit I/O concurrency
 	for _, ref := range r.refs {
-		posn := r.qpos.fset.Position(ref.Pos())
+		posn := r.fset.Position(ref.Pos())
 		fi := fileinfosByName[posn.Filename]
 		if fi == nil {
 			fi = &fileinfo{data: make(chan interface{})}
@@ -243,11 +455,8 @@ func readFile(ctxt *build.Context, filename string) ([]byte, error) {
 	return buf.Bytes(), nil
 }
 
-// TODO(adonovan): encode extent, not just Pos info, in Serial form.
-
 func (r *referrersResult) toSerial(res *serial.Result, fset *token.FileSet) {
 	referrers := &serial.Referrers{
-		Pos:  fset.Position(r.query.Pos()).String(),
 		Desc: r.obj.String(),
 	}
 	if pos := r.obj.Pos(); pos != token.NoPos { // Package objects have no Pos()

cmd/guru/serial/serial.go

@@ -27,7 +27,6 @@ type Peers struct {
 
 // A Referrers is the result of a 'referrers' query.
 type Referrers struct {
-	Pos    string   `json:"pos"`              // location of the query reference
 	ObjPos string   `json:"objpos,omitempty"` // location of the definition
 	Desc   string   `json:"desc"`             // description of the denoted object
 	Refs   []string `json:"refs,omitempty"`   // locations of all references

cmd/guru/testdata/src/referrers-json/main.golden

@@ -2,12 +2,20 @@
 {
 	"mode": "referrers",
 	"referrers": {
-		"pos": "testdata/src/referrers-json/main.go:14:8",
-		"objpos": "testdata/src/referrers-json/main.go:7:8",
 		"desc": "package lib",
 		"refs": [
+			"testdata/src/describe/main.go:91:8",
+			"testdata/src/imports/main.go:18:12",
+			"testdata/src/imports/main.go:19:2",
+			"testdata/src/imports/main.go:20:2",
+			"testdata/src/imports/main.go:21:8",
+			"testdata/src/imports/main.go:26:8",
 			"testdata/src/referrers-json/main.go:14:8",
-			"testdata/src/referrers-json/main.go:14:19"
+			"testdata/src/referrers-json/main.go:14:19",
+			"testdata/src/referrers/ext_test.go:10:7",
+			"testdata/src/referrers/int_test.go:7:7",
+			"testdata/src/referrers/main.go:16:8",
+			"testdata/src/referrers/main.go:16:19"
 		]
 	}
 }
@@ -15,7 +23,6 @@
 {
 	"mode": "referrers",
 	"referrers": {
-		"pos": "testdata/src/referrers-json/main.go:15:8",
 		"objpos": "testdata/src/lib/lib.go:5:13",
 		"desc": "func (lib.Type).Method(x *int) *int",
 		"refs": [
@@ -33,7 +40,6 @@
 {
 	"mode": "referrers",
 	"referrers": {
-		"pos": "testdata/src/referrers-json/main.go:17:2",
 		"objpos": "testdata/src/referrers-json/main.go:14:6",
 		"desc": "var v lib.Type",
 		"refs": [
@@ -48,7 +54,6 @@
 {
 	"mode": "referrers",
 	"referrers": {
-		"pos": "testdata/src/referrers-json/main.go:20:10",
 		"objpos": "testdata/src/referrers-json/main.go:10:2",
 		"desc": "field f int",
 		"refs": [

cmd/guru/testdata/src/referrers/main.golden

@@ -8,14 +8,22 @@
 	var s2 s
 
 -------- @referrers ref-package --------
-4 references to package lib
+12 references to package lib
+	var _ lib.Outer // @describe lib-outer "Outer"
+	const c = lib.Const // @describe ref-const "Const"
+	lib.Func()          // @describe ref-func "Func"
+	lib.Var++           // @describe ref-var "Var"
+	var t lib.Type      // @describe ref-type "Type"
+	var _ lib.Type // @describe ref-pkg "lib"
+	var v lib.Type = lib.Const // @referrers ref-package "lib"
+	var v lib.Type = lib.Const // @referrers ref-package "lib"
 	_ = (lib.Type).Method // ref from external test package
 	_ = (lib.Type).Method // ref from internal test package
 	var v lib.Type = lib.Const // @referrers ref-package "lib"
 	var v lib.Type = lib.Const // @referrers ref-package "lib"
 
 -------- @referrers ref-method --------
-7 references to func (lib.Type).Method(x *int) *int
+7 references to func (Type).Method(x *int) *int
 	p := t.Method(&a)   // @describe ref-method "Method"
 	_ = v.Method               // @referrers ref-method "Method"
 	_ = v.Method

go/loader/loader.go

@@ -107,6 +107,25 @@ type Config struct {
 	//
 	// It must be safe to call concurrently from multiple goroutines.
 	FindPackage func(ctxt *build.Context, fromDir, importPath string, mode build.ImportMode) (*build.Package, error)
+
+	// AfterTypeCheck is called immediately after a list of files
+	// has been type-checked and appended to info.Files.
+	//
+	// This optional hook function is the earliest opportunity for
+	// the client to observe the output of the type checker,
+	// which may be useful to reduce analysis latency when loading
+	// a large program.
+	//
+	// The function is permitted to modify info.Info, for instance
+	// to clear data structures that are no longer needed, which can
+	// dramatically reduce peak memory consumption.
+	//
+	// The function may be called twice for the same PackageInfo:
+	// once for the files of the package and again for the
+	// in-package test files.
+	//
+	// It must be safe to call concurrently from multiple goroutines.
+	AfterTypeCheck func(info *PackageInfo, files []*ast.File)
 }
 
 // A PkgSpec specifies a non-importable package to be created by Load.
@@ -971,8 +990,6 @@ func (imp *importer) load(bp *build.Package) *PackageInfo {
 // dependency edges that should be checked for potential cycles.
 //
 func (imp *importer) addFiles(info *PackageInfo, files []*ast.File, cycleCheck bool) {
-	info.Files = append(info.Files, files...)
-
 	// Ensure the dependencies are loaded, in parallel.
 	var fromPath string
 	if cycleCheck {
@@ -990,6 +1007,11 @@ func (imp *importer) addFiles(info *PackageInfo, files []*ast.File, cycleCheck b
 	// Ignore the returned (first) error since we
 	// already collect them all in the PackageInfo.
 	info.checker.Files(files)
+	info.Files = append(info.Files, files...)
+
+	if imp.conf.AfterTypeCheck != nil {
+		imp.conf.AfterTypeCheck(info, files)
+	}
 
 	if trace {
 		fmt.Fprintf(os.Stderr, "%s: stop %q\n",