internal/godoc: preserve sharing of unresolved identifiers

Preserve the sharing between the list of unresolved identifiers in an
ast.File and the identifiers in the tree. Doc rendering (specifically,
links to packages) requires that sharing to be preserved.

Change-Id: I88c82b6dcdf1806c8f5d59b6a717f7a0b5a1c5a6
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/259100
Trust: Jonathan Amsterdam <jba@google.com>
Run-TryBot: Jonathan Amsterdam <jba@google.com>
Reviewed-by: Julie Qiu <julie@golang.org>
diff --git a/internal/godoc/encode.go b/internal/godoc/encode.go
index 36fd2d0..366563e 100644
--- a/internal/godoc/encode.go
+++ b/internal/godoc/encode.go
@@ -97,7 +97,7 @@
 	}
 
 	for _, f := range p.Files {
-		removeCycles(f.AST)
+		removeCycles(f)
 	}
 
 	var buf bytes.Buffer
@@ -111,7 +111,7 @@
 		return nil, fmt.Errorf("enc.Encode: %v", err)
 	}
 	for _, f := range p.Files {
-		fixupObjects(f.AST)
+		fixupObjects(f)
 	}
 	return buf.Bytes(), nil
 }
@@ -133,13 +133,13 @@
 		return nil, err
 	}
 	for _, f := range p.Files {
-		fixupObjects(f.AST)
+		fixupObjects(f)
 	}
 	return p, nil
 }
 
 // removeCycles removes cycles from f. There are two sources of cycles
-// in an ast.File: Scopes and Objects.
+// in an ast.File: Scopes and Objects. Also, some Idents are shared.
 //
 // removeCycles removes all Scopes, since doc generation doesn't use them. Doc
 // generation does use Objects, and it needs object identity to be preserved
@@ -159,10 +159,14 @@
 // numbers into Ident.Objs. We take advantage of the fact that the Data and Decl
 // fields are of type interface{}, storing the object number into Data and the
 // Decl number into Decl.
-func removeCycles(f *ast.File) {
-	f.Scope.Objects = nil // doc doesn't use scopes
+//
+// The AST includes a list of unresolved Idents, which are shared with Idents
+// in the tree itself. We assign these numbers as well, and store the numbers
+// in a separate field of File.
+func removeCycles(f *File) {
+	f.AST.Scope.Objects = nil // doc doesn't use scopes
 
-	// First pass: assign every Decl and Spec a number.
+	// First pass: assign every Decl, Spec and Ident a number.
 	// Since these aren't shared and Inspect is deterministic,
 	// this walk will produce the same sequence of Decls after encoding/decoding.
 	// Also assign a unique number to each Object we find in an Ident.
@@ -170,15 +174,18 @@
 	// produce the same sequence. So we store their numbers separately.
 	declNums := map[interface{}]int{}
 	objNums := map[*ast.Object]int{}
-	ast.Inspect(f, func(n ast.Node) bool {
+	ast.Inspect(f.AST, func(n ast.Node) bool {
 		if isRelevantDecl(n) {
 			if _, ok := declNums[n]; ok {
 				panic(fmt.Sprintf("duplicate decl %+v", n))
 			}
 			declNums[n] = len(declNums)
-		} else if id, ok := n.(*ast.Ident); ok && id.Obj != nil {
-			if _, ok := objNums[id.Obj]; !ok {
-				objNums[id.Obj] = len(objNums)
+		} else if id, ok := n.(*ast.Ident); ok {
+			declNums[id] = len(declNums) // remember Idents for Unresolved list.
+			if id.Obj != nil {
+				if _, ok := objNums[id.Obj]; !ok {
+					objNums[id.Obj] = len(objNums)
+				}
 			}
 		}
 		return true
@@ -189,7 +196,7 @@
 	// if it's not a relevant Decl.
 	// The Data field gets a number from the objNums map. (This destroys
 	// whatever might be in the Data field, but doc generation doesn't care.)
-	ast.Inspect(f, func(n ast.Node) bool {
+	ast.Inspect(f.AST, func(n ast.Node) bool {
 		id, ok := n.(*ast.Ident)
 		if !ok || id.Obj == nil {
 			return true
@@ -209,20 +216,31 @@
 		id.Obj.Decl = d
 		return true
 	})
+
+	// Replace the unresolved identifiers with their numbers.
+	f.UnresolvedNums = nil
+	for _, id := range f.AST.Unresolved {
+		// If we can't find an identifier, assume it was in a part of the AST
+		// deleted by removeUnusedASTNodes, and ignore it.
+		if num, ok := declNums[id]; ok {
+			f.UnresolvedNums = append(f.UnresolvedNums, num)
+		}
+	}
+	f.AST.Unresolved = nil
 }
 
 // fixupObjects re-establishes the original Object and Decl relationships of the
-// ast.File f.
+// File.
 //
-// f is the result of EncodeASTFiles, which uses removeCycles (see above) to
-// modify ast.Objects so that they are uniquely identified by their Data field,
-// and refer to their Decl via a number in the Decl field. fixupObjects uses
-// those values to reconstruct the same set of relationships.
-func fixupObjects(f *ast.File) {
-	// First pass: reconstruct the numbers of every Decl.
+// f is the result of Encode, which uses removeCycles (see above) to modify
+// ast.Objects so that they are uniquely identified by their Data field, and
+// refer to their Decl via a number in the Decl field. fixupObjects uses those
+// values to reconstruct the same set of relationships.
+func fixupObjects(f *File) {
+	// First pass: reconstruct the numbers of every Decl and Ident.
 	var decls []ast.Node
-	ast.Inspect(f, func(n ast.Node) bool {
-		if isRelevantDecl(n) {
+	ast.Inspect(f.AST, func(n ast.Node) bool {
+		if _, ok := n.(*ast.Ident); ok || isRelevantDecl(n) {
 			decls = append(decls, n)
 		}
 		return true
@@ -230,7 +248,7 @@
 
 	// Second pass: replace the numbers in Ident.Objs with the right Nodes.
 	var objs []*ast.Object
-	ast.Inspect(f, func(n ast.Node) bool {
+	ast.Inspect(f.AST, func(n ast.Node) bool {
 		id, ok := n.(*ast.Ident)
 		if !ok || id.Obj == nil {
 			return true
@@ -260,6 +278,13 @@
 		}
 		return true
 	})
+
+	// Fix up unresolved identifiers.
+	f.AST.Unresolved = make([]*ast.Ident, len(f.UnresolvedNums))
+	for i, num := range f.UnresolvedNums {
+		f.AST.Unresolved[i] = decls[num].(*ast.Ident)
+	}
+	f.UnresolvedNums = nil
 }
 
 // isRelevantDecl reports whether n is a Node for a declaration relevant to
diff --git a/internal/godoc/godoc.go b/internal/godoc/godoc.go
index 9faea12..f1af28a 100644
--- a/internal/godoc/godoc.go
+++ b/internal/godoc/godoc.go
@@ -32,8 +32,9 @@
 
 // A File contains everything needed about a source file to render documentation.
 type File struct {
-	Name string // full file pathname relative to zip content directory
-	AST  *ast.File
+	Name           string // full file pathname relative to zip content directory
+	AST            *ast.File
+	UnresolvedNums []int // used to handle sharing of unresolved identifiers
 }
 
 // NewPackage returns a new Package with the given fset and set of module package paths.
diff --git a/internal/godoc/render_test.go b/internal/godoc/render_test.go
index 6f30eb1..6287504 100644
--- a/internal/godoc/render_test.go
+++ b/internal/godoc/render_test.go
@@ -30,6 +30,7 @@
 	if err != nil {
 		t.Fatal(err)
 	}
+
 	wantSyn, wantImports, wantDoc, err := p.Render(ctx, "p", si, mi, "", "")
 	if err != nil {
 		t.Fatal(err)
@@ -58,7 +59,7 @@
 	}
 
 	// Verify that removing AST nodes doesn't change the doc.
-	p, err = packageForDir(filepath.Join("testdata", "p"), false)
+	p, err = packageForDir(filepath.Join("testdata", "p"), true)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -66,7 +67,7 @@
 
 	// Verify that encoding then decoding generates the same doc.
 	// We can't re-use p to encode because it's been rendered.
-	p, err = packageForDir(filepath.Join("testdata", "p"), false)
+	p, err = packageForDir(filepath.Join("testdata", "p"), true)
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/internal/godoc/testdata/p/p.go b/internal/godoc/testdata/p/p.go
index 42fcc93..8c873aa 100644
--- a/internal/godoc/testdata/p/p.go
+++ b/internal/godoc/testdata/p/p.go
@@ -6,7 +6,10 @@
 // of other things to say, but that's the gist of it.
 package p
 
-import "fmt"
+import (
+	"fmt"
+	"time"
+)
 
 // const
 const C = 1
@@ -15,8 +18,8 @@
 var V = 2
 
 // exported func
-func F() {
-	fmt.Println("ok")
+func F(t time.Time) {
+	fmt.Println(t)
 }
 
 // unexported func