Imported Upstream version 1.3upstream/1.3

1 files changed, 74 insertions, 0 deletions

diff --git a/src/pkg/go/ast/example_test.go b/src/pkg/go/ast/example_test.go
index 632bfcfd0..d2e734f2c 100644
--- a/src/pkg/go/ast/example_test.go
+++ b/src/pkg/go/ast/example_test.go
@@ -5,8 +5,10 @@
 package ast_test
 
 import (
+	"bytes"
 	"fmt"
 	"go/ast"
+	"go/format"
 	"go/parser"
 	"go/token"
 )
@@ -134,3 +136,75 @@ func main() {
 	//     57  .  }
 	//     58  }
 }
+
+// This example illustrates how to remove a variable declaration
+// in a Go program while maintaining correct comment association
+// using an ast.CommentMap.
+func ExampleCommentMap() {
+	// src is the input for which we create the AST that we
+	// are going to manipulate.
+	src := `
+// This is the package comment.
+package main
+
+// This comment is associated with the hello constant.
+const hello = "Hello, World!" // line comment 1
+
+// This comment is associated with the foo variable.
+var foo = hello // line comment 2 
+
+// This comment is associated with the main function.
+func main() {
+	fmt.Println(hello) // line comment 3
+}
+`
+
+	// Create the AST by parsing src.
+	fset := token.NewFileSet() // positions are relative to fset
+	f, err := parser.ParseFile(fset, "src.go", src, parser.ParseComments)
+	if err != nil {
+		panic(err)
+	}
+
+	// Create an ast.CommentMap from the ast.File's comments.
+	// This helps keeping the association between comments
+	// and AST nodes.
+	cmap := ast.NewCommentMap(fset, f, f.Comments)
+
+	// Remove the first variable declaration from the list of declarations.
+	f.Decls = removeFirstVarDecl(f.Decls)
+
+	// Use the comment map to filter comments that don't belong anymore
+	// (the comments associated with the variable declaration), and create
+	// the new comments list.
+	f.Comments = cmap.Filter(f).Comments()
+
+	// Print the modified AST.
+	var buf bytes.Buffer
+	if err := format.Node(&buf, fset, f); err != nil {
+		panic(err)
+	}
+	fmt.Printf("%s", buf.Bytes())
+
+	// output:
+	// // This is the package comment.
+	// package main
+	//
+	// // This comment is associated with the hello constant.
+	// const hello = "Hello, World!" // line comment 1
+	//
+	// // This comment is associated with the main function.
+	// func main() {
+	// 	fmt.Println(hello) // line comment 3
+	// }
+}
+
+func removeFirstVarDecl(list []ast.Decl) []ast.Decl {
+	for i, decl := range list {
+		if gen, ok := decl.(*ast.GenDecl); ok && gen.Tok == token.VAR {
+			copy(list[i:], list[i+1:])
+			return list[:len(list)-1]
+		}
+	}
+	panic("variable declaration not found")
+}