internal/godoc: add render test, fix node removal

Add a test that verifies that removeUnusedASTNodes doesn't
change the doc.

Fix some issues with removeUnusedASTNodes: it can't remove function
bodies that may be part of examples, and it can't remove the top-level
Comments field.

Also, two changes that prevent clobbering the AST too early:

- Remove cycles from a Package when encoding, not when creating; otherwise
  we can't use the Package to render doc before we encode it.

- Tell go/doc to preserve the AST, so we can render, then encode.

Change-Id: I2725d278572cd3de3f2019157faf224ab4b592a1
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/258424
Reviewed-by: Julie Qiu <julie@golang.org>
Trust: Jonathan Amsterdam <jba@google.com>
diff --git a/internal/godoc/encode.go b/internal/godoc/encode.go
index 93f4445..46070df 100644
--- a/internal/godoc/encode.go
+++ b/internal/godoc/encode.go
@@ -54,6 +54,7 @@
 		&ast.ParenExpr{},
 		&ast.RangeStmt{},
 		&ast.ReturnStmt{},
+		&ast.Scope{},
 		&ast.SelectorExpr{},
 		&ast.SliceExpr{},
 		&ast.StarExpr{},
@@ -74,6 +75,10 @@
 func (p *Package) Encode() (_ []byte, err error) {
 	defer derrors.Wrap(&err, "godoc.Package.Encode()")
 
+	for _, f := range p.Files {
+		removeCycles(f.AST)
+	}
+
 	var buf bytes.Buffer
 	io.WriteString(&buf, encodingType)
 	enc := gob.NewEncoder(&buf)
diff --git a/internal/godoc/godoc.go b/internal/godoc/godoc.go
index cf6579a..a262b5a 100644
--- a/internal/godoc/godoc.go
+++ b/internal/godoc/godoc.go
@@ -8,6 +8,7 @@
 import (
 	"go/ast"
 	"go/token"
+	"strings"
 
 	"golang.org/x/pkgsite/internal/godoc/dochtml"
 )
@@ -49,7 +50,6 @@
 	if removeNodes {
 		removeUnusedASTNodes(f)
 	}
-	removeCycles(f)
 	p.Files = append(p.Files, &File{
 		Name: p.Fset.Position(f.Package).Filename,
 		AST:  f,
@@ -59,6 +59,11 @@
 // 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) {
+	// Don't trim anything from a file in a XXX_test package; it
+	// may be part of a playable example.
+	if strings.HasSuffix(pf.Name.Name, "_test") {
+		return
+	}
 	var decls []ast.Decl
 	for _, d := range pf.Decls {
 		if f, ok := d.(*ast.FuncDecl); ok {
@@ -66,10 +71,14 @@
 			if f.Name == nil || !ast.IsExported(f.Name.Name) {
 				continue
 			}
-			f.Body = nil
+			// Remove the function body, unless it's an example.
+			// The doc contains example bodies.
+			if !strings.HasPrefix(f.Name.Name, "Example") {
+				f.Body = nil
+			}
 		}
 		decls = append(decls, d)
 	}
-	pf.Comments = nil
+	// Don't remove pf.Comments; they may contain Notes.
 	pf.Decls = decls
 }
diff --git a/internal/godoc/godoc_test.go b/internal/godoc/godoc_test.go
index 876f96d..1627361 100644
--- a/internal/godoc/godoc_test.go
+++ b/internal/godoc/godoc_test.go
@@ -40,13 +40,13 @@
 // Exp is exported.
 func Exp() {}
 
-// unexp is not exported.
+// unexp is not exported, but the comment is preserved for notes.
 func unexp() {}
 
 // M is exported.
 func (t T) M() int {}
 
-// m isn't.
+// m isn't, but the comment is preserved for notes.
 func (T) m() {}
 
 // U is an exported method of an unexported type.
@@ -80,9 +80,13 @@
 // Exp is exported.
 func Exp()
 
+// unexp is not exported, but the comment is preserved for notes.
+
 // M is exported.
 func (t T) M() int
 
+// m isn't, but the comment is preserved for notes.
+
 // 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
diff --git a/internal/godoc/render.go b/internal/godoc/render.go
index dc48c56..496ec63 100644
--- a/internal/godoc/render.go
+++ b/internal/godoc/render.go
@@ -44,6 +44,10 @@
 	if modInfo.ModulePath == stdlib.ModulePath {
 		importPath = innerPath
 	}
+	if modInfo.ModulePackages == nil {
+		modInfo.ModulePackages = p.ModulePackagePaths
+	}
+
 	// The "builtin" package in the standard library is a special case.
 	// We want to show documentation for all globals (not just exported ones),
 	// and avoid association of consts, vars, and factory functions with types
@@ -55,7 +59,7 @@
 	}
 
 	// Compute package documentation.
-	var m doc.Mode
+	m := doc.PreserveAST // so we can render doc, then save the AST afterwards.
 	if noFiltering {
 		m |= doc.AllDecls
 	}
@@ -67,6 +71,7 @@
 	if err != nil {
 		return "", nil, safehtml.HTML{}, fmt.Errorf("doc.NewFromFiles: %v", err)
 	}
+
 	if d.ImportPath != importPath {
 		panic(fmt.Errorf("internal error: *doc.Package has an unexpected import path (%q != %q)", d.ImportPath, importPath))
 	}
diff --git a/internal/godoc/render_test.go b/internal/godoc/render_test.go
new file mode 100644
index 0000000..370b254
--- /dev/null
+++ b/internal/godoc/render_test.go
@@ -0,0 +1,62 @@
+// Copyright 2020 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package godoc
+
+import (
+	"context"
+	"path/filepath"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"golang.org/x/pkgsite/internal/source"
+	"golang.org/x/pkgsite/internal/testing/sample"
+)
+
+func TestRender(t *testing.T) {
+	ctx := context.Background()
+
+	si := source.NewGitHubInfo(sample.ModulePath, "", "abcde")
+	mi := &ModuleInfo{
+		ModulePath:      sample.ModulePath,
+		ResolvedVersion: sample.VersionString,
+		ModulePackages:  nil,
+	}
+
+	// Use a Package created locally and without nodes removed as the desired doc.
+	p, err := packageForDir(filepath.Join("testdata", "p"), false)
+	if err != nil {
+		t.Fatal(err)
+	}
+	wantSyn, wantImports, wantDoc, err := p.Render(ctx, "p", si, mi, "GOOS", "GOARCH")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	check := func(p *Package) {
+		t.Helper()
+		gotSyn, gotImports, gotDoc, err := p.Render(ctx, "p", si, mi, "GOOS", "GOARCH")
+		if err != nil {
+			t.Fatal(err)
+		}
+		if gotSyn != wantSyn {
+			t.Errorf("synopsis: got %q, want %q", gotSyn, wantSyn)
+		}
+		if !cmp.Equal(gotImports, wantImports) {
+			t.Errorf("imports: got %v, want %v", gotImports, wantImports)
+		}
+		if diff := cmp.Diff(wantDoc.String(), gotDoc.String()); diff != "" {
+			t.Errorf("doc mismatch (-want, +got):\n%s", diff)
+			t.Logf("---- want ----\n%s", wantDoc)
+			t.Logf("---- got ----\n%s", gotDoc)
+		}
+	}
+
+	// Verify that removing AST nodes doesn't change the doc.
+	p, err = packageForDir(filepath.Join("testdata", "p"), false)
+	if err != nil {
+		t.Fatal(err)
+	}
+	check(p)
+}
diff --git a/internal/godoc/testdata/p/example_test.go b/internal/godoc/testdata/p/example_test.go
new file mode 100644
index 0000000..78d00af
--- /dev/null
+++ b/internal/godoc/testdata/p/example_test.go
@@ -0,0 +1,23 @@
+// Copyright 2020 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package p_test
+
+import "golang.org/x/pkgsite/internal/godoc/testdata/p"
+
+// non-executable example
+func ExampleTF() {
+	// example comment
+	app := App{}
+	app.Name = "greet"
+	_ = app.Run([]string{"greet"})
+}
+
+// executable example
+func ExampleF() {
+	// example comment
+	p.F()
+	// Output:
+	// ok
+}
diff --git a/internal/godoc/testdata/p/p.go b/internal/godoc/testdata/p/p.go
new file mode 100644
index 0000000..42fcc93
--- /dev/null
+++ b/internal/godoc/testdata/p/p.go
@@ -0,0 +1,64 @@
+// Copyright 202 0The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package p is for testing godoc.Render. There are a lot
+// of other things to say, but that's the gist of it.
+package p
+
+import "fmt"
+
+// const
+const C = 1
+
+// var
+var V = 2
+
+// exported func
+func F() {
+	fmt.Println("ok")
+}
+
+// unexported func
+func unexp() {
+}
+
+// type
+type T int
+
+// typeConstant
+const CT T = 3
+
+// typeVariable
+var VT T
+
+// typeFunc
+func TF() T {
+	return T(0)
+}
+
+// method
+// BUG(uid): this verifies that notes are rendered
+func (T) M() {}
+
+// unexported method
+func (T) m() {}
+
+type S1 struct {
+	F int // field
+}
+
+type us struct {
+	G bool
+	u int
+}
+type S2 struct {
+	S1 // embedded struct
+	us // embedded, unexported struct
+	H  int
+}
+
+// I is an interface.
+type I interface {
+	M1()
+}