internal/fetch: strip unused data from AST

If the experiment "remove-unused-ast" is active,
remove parts of each file's AST that are not
needed to generate documentation.

Change-Id: Ida48ea1e94d8ccde670bd2318412b8a3fa524ce7
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/257659
Trust: Jonathan Amsterdam <jba@google.com>
Run-TryBot: Jonathan Amsterdam <jba@google.com>
Reviewed-by: Julie Qiu <julie@golang.org>
diff --git a/internal/experiment.go b/internal/experiment.go
index 64aadc8..64fca5a 100644
--- a/internal/experiment.go
+++ b/internal/experiment.go
@@ -6,21 +6,23 @@
 package internal
 
 const (
-	ExperimentAltRequeue   = "alt-requeue"
-	ExperimentAutocomplete = "autocomplete"
-	ExperimentSidenav      = "sidenav"
-	ExperimentUnitPage     = "unit-page"
-	ExperimentUseUnits     = "use-units"
+	ExperimentAltRequeue      = "alt-requeue"
+	ExperimentAutocomplete    = "autocomplete"
+	ExperimentRemoveUnusedAST = "remove-unused-ast"
+	ExperimentSidenav         = "sidenav"
+	ExperimentUnitPage        = "unit-page"
+	ExperimentUseUnits        = "use-units"
 )
 
 // Experiments represents all of the active experiments in the codebase and
 // a description of each experiment.
 var Experiments = map[string]string{
-	ExperimentAltRequeue:   "Requeue modules for reprocessing in a different order.",
-	ExperimentAutocomplete: "Enable autocomplete with search.",
-	ExperimentSidenav:      "Display documentation index on the left sidenav.",
-	ExperimentUnitPage:     "Enable the redesigned details page.",
-	ExperimentUseUnits:     "Read from paths, documentation, readmes, and package_imports tables.",
+	ExperimentAltRequeue:      "Requeue modules for reprocessing in a different order.",
+	ExperimentAutocomplete:    "Enable autocomplete with search.",
+	ExperimentRemoveUnusedAST: "Prune AST prior to rendering documentation HTML.",
+	ExperimentSidenav:         "Display documentation index on the left sidenav.",
+	ExperimentUnitPage:        "Enable the redesigned details page.",
+	ExperimentUseUnits:        "Read from paths, documentation, readmes, and package_imports tables.",
 }
 
 // Experiment holds data associated with an experimental feature for frontend
diff --git a/internal/fetch/load.go b/internal/fetch/load.go
index 48f918e..25f31c8 100644
--- a/internal/fetch/load.go
+++ b/internal/fetch/load.go
@@ -31,6 +31,7 @@
 	"golang.org/x/pkgsite/internal"
 	"golang.org/x/pkgsite/internal/config"
 	"golang.org/x/pkgsite/internal/derrors"
+	"golang.org/x/pkgsite/internal/experiment"
 	"golang.org/x/pkgsite/internal/fetch/dochtml"
 	"golang.org/x/pkgsite/internal/fetch/internal/doc"
 	"golang.org/x/pkgsite/internal/log"
@@ -112,6 +113,13 @@
 	}
 	var allGoFiles []*ast.File
 	for _, pf := range goFiles {
+		if experiment.IsActive(ctx, internal.ExperimentRemoveUnusedAST) {
+			// Don't strip the seemingly unexported functions from the builtin package;
+			// they are actually Go builtins like make, new, etc.
+			if !(modulePath == stdlib.ModulePath && innerPath == "builtin") {
+				removeUnusedASTNodes(pf)
+			}
+		}
 		allGoFiles = append(allGoFiles, pf)
 	}
 	d, err := loadPackageWithFiles(modulePath, innerPath, packageName, allGoFiles, fset)
@@ -380,3 +388,21 @@
 func ZipLoadShedStats() LoadShedStats {
 	return zipLoadShedder.stats()
 }
+
+// removeUnusedASTNodes removes parts of the AST not needed for documentation.
+// It doesn't remove unexported consts, vars or types, although it probably could.
+func removeUnusedASTNodes(pf *ast.File) {
+	var decls []ast.Decl
+	for _, d := range pf.Decls {
+		if f, ok := d.(*ast.FuncDecl); ok {
+			// Remove all unexported functions and function bodies.
+			if f.Name == nil || !ast.IsExported(f.Name.Name) {
+				continue
+			}
+			f.Body = nil
+		}
+		decls = append(decls, d)
+	}
+	pf.Comments = nil
+	pf.Decls = decls
+}
diff --git a/internal/fetch/load_test.go b/internal/fetch/load_test.go
index 94f90d2..669f19f 100644
--- a/internal/fetch/load_test.go
+++ b/internal/fetch/load_test.go
@@ -7,6 +7,9 @@
 import (
 	"archive/zip"
 	"bytes"
+	"go/format"
+	"go/parser"
+	"go/token"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -95,3 +98,96 @@
 		})
 	}
 }
+
+func TestRemoveUnusedASTNodes(t *testing.T) {
+	const file = `
+// Package-level comment.
+package p
+
+// const C
+const C = 1
+
+// leave unexported consts
+const c = 1
+
+// var V
+var V int
+
+// leave unexported vars
+var v int
+
+// type T
+type T int
+
+// leave unexported types
+type t int
+
+// Exp is exported.
+func Exp() {}
+
+// unexp is not exported.
+func unexp() {}
+
+// M is exported.
+func (t T) M() int {}
+
+// m isn't.
+func (T) m() {}
+
+// U is an exported method of an unexported type.
+// Its doc is not shown, unless t is embedded
+// in an exported type. We don't remove it to
+// be safe.
+func (t) U() {}
+`
+	////////////////
+	const want = `// Package-level comment.
+package p
+
+// const C
+const C = 1
+
+// leave unexported consts
+const c = 1
+
+// var V
+var V int
+
+// leave unexported vars
+var v int
+
+// type T
+type T int
+
+// leave unexported types
+type t int
+
+// Exp is exported.
+func Exp()
+
+// M is exported.
+func (t T) M() int
+
+// U is an exported method of an unexported type.
+// Its doc is not shown, unless t is embedded
+// in an exported type. We don't remove it to
+// be safe.
+func (t) U()
+`
+	////////////////
+
+	fset := token.NewFileSet()
+	astFile, err := parser.ParseFile(fset, "tst.go", file, parser.ParseComments)
+	if err != nil {
+		t.Fatal(err)
+	}
+	removeUnusedASTNodes(astFile)
+	var buf bytes.Buffer
+	if err := format.Node(&buf, fset, astFile); err != nil {
+		t.Fatal(err)
+	}
+	got := buf.String()
+	if diff := cmp.Diff(want, got); diff != "" {
+		t.Errorf("mismatch (-want, +got):\n%s", diff)
+	}
+}