internal/godoc: restore AST after Encode

After encoding, put the AST back as it was (with a few irrelevant things missing).

That lets us Encode, then Render. The other way won't work. go/doc
removes things from the AST. It has a mode (doc.PreserveAST) that
doesn't change the AST, but has the effect of rendering function
bodies into the doc.

Change-Id: I02a2464bf514a876d9ab6bc0007deb7e246e999a
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/258740
Trust: Jonathan Amsterdam <jba@google.com>
Run-TryBot: Jonathan Amsterdam <jba@google.com>
TryBot-Result: kokoro <noreply+kokoro@google.com>
Reviewed-by: Julie Qiu <julie@golang.org>
diff --git a/internal/fetch/load.go b/internal/fetch/load.go
index 9c4f836..2818e61 100644
--- a/internal/fetch/load.go
+++ b/internal/fetch/load.go
@@ -121,10 +121,8 @@
 		docPkg.AddFile(pf, removeNodes)
 	}
 
-	synopsis, imports, docHTML, err := docPkg.Render(ctx, innerPath, sourceInfo, modInfo, goos, goarch)
-	if err != nil && !errors.Is(err, godoc.ErrTooLarge) {
-		return nil, err
-	}
+	// Encode before rendering: both operations mess with the AST, but Encode restores
+	// it enough to make Render work.
 	var src []byte
 	if experiment.IsActive(ctx, internal.ExperimentInsertPackageSource) {
 		src, err = docPkg.Encode()
@@ -132,6 +130,11 @@
 			return nil, err
 		}
 	}
+
+	synopsis, imports, docHTML, err := docPkg.Render(ctx, innerPath, sourceInfo, modInfo, goos, goarch)
+	if err != nil && !errors.Is(err, godoc.ErrTooLarge) {
+		return nil, err
+	}
 	importPath := path.Join(modulePath, innerPath)
 	if modulePath == stdlib.ModulePath {
 		importPath = innerPath
diff --git a/internal/godoc/encode.go b/internal/godoc/encode.go
index 46070df..004221d 100644
--- a/internal/godoc/encode.go
+++ b/internal/godoc/encode.go
@@ -72,6 +72,9 @@
 }
 
 // Encode encodes a Package into a byte slice.
+// During its operation, Encode modifies the AST,
+// but it restores it to a state suitable for
+// rendering before it returns.
 func (p *Package) Encode() (_ []byte, err error) {
 	defer derrors.Wrap(&err, "godoc.Package.Encode()")
 
@@ -89,6 +92,9 @@
 	if err := enc.Encode(p.gobPackage); err != nil {
 		return nil, err
 	}
+	for _, f := range p.Files {
+		fixupObjects(f.AST)
+	}
 	return buf.Bytes(), nil
 }
 
diff --git a/internal/godoc/render.go b/internal/godoc/render.go
index de2b831..1199efe 100644
--- a/internal/godoc/render.go
+++ b/internal/godoc/render.go
@@ -36,8 +36,9 @@
 var MaxDocumentationHTML = 20 * megabyte
 
 // Render renders the documentation for the package.
-// This is mostly copied from internal/fetch/fetch.go.
+// Rendering destroys p's AST; do not call any methods of p after it returns.
 func (p *Package) Render(ctx context.Context, innerPath string, sourceInfo *source.Info, modInfo *ModuleInfo, goos, goarch string) (synopsis string, imports []string, html safehtml.HTML, err error) {
+	// This is mostly copied from internal/fetch/fetch.go.
 	defer derrors.Wrap(&err, "godoc.Package.Render(%q, %q, %q, %q, %q)", modInfo.ModulePath, modInfo.ResolvedVersion, innerPath, goos, goarch)
 
 	importPath := path.Join(modInfo.ModulePath, innerPath)