Golang ast.NewCommentMap函数代码示例

func RewriteGeneratedGogoProtobufFile(name string, extractFn ExtractFunc, optionalFn OptionalFunc, header []byte) error {
	return rewriteFile(name, header, func(fset *token.FileSet, file *ast.File) error {
		cmap := ast.NewCommentMap(fset, file, file.Comments)

		// transform methods that point to optional maps or slices
		for _, d := range file.Decls {
			rewriteOptionalMethods(d, optionalFn)
		}

		// remove types that are already declared
		decls := []ast.Decl{}
		for _, d := range file.Decls {
			if dropExistingTypeDeclarations(d, extractFn) {
				continue
			}
			if dropEmptyImportDeclarations(d) {
				continue
			}
			decls = append(decls, d)
		}
		file.Decls = decls

		// remove unmapped comments
		file.Comments = cmap.Filter(file).Comments()
		return nil
	})
}

// ParseFile will create a Build from the file path that
// was passed. FileSet of the Build will only contain a
// single file.
func ParseFile(path string) (*Build, error) {
	var fileSet token.FileSet

	astTree, err := parser.ParseFile(&fileSet, path, nil, parser.AllErrors|parser.ParseComments)
	if err != nil {
		return nil, err
	}

	fileName := filepath.Base(path)

	// create a comment map from file
	commentMap := ast.NewCommentMap(&fileSet, astTree, astTree.Comments)

	// create new build for the file
	build := NewBuild()
	fileAST, err := ParseFileAST(fileName, astTree, commentMap)
	if err != nil {
		return nil, err
	}

	// add parsed file to the build file set
	build.AddFile(fileName, fileAST)

	return build, nil
}

// rewriteFile applies the rewrite rule 'pattern -> replace' to an entire file.
func rewriteFile(pattern, replace ast.Expr, p *ast.File) *ast.File {
	cmap := ast.NewCommentMap(fileSet, p, p.Comments)
	m := make(map[string]reflect.Value)
	pat := reflect.ValueOf(pattern)
	repl := reflect.ValueOf(replace)

	var rewriteVal func(val reflect.Value) reflect.Value
	rewriteVal = func(val reflect.Value) reflect.Value {
		// don't bother if val is invalid to start with
		if !val.IsValid() {
			return reflect.Value{}
		}
		for k := range m {
			delete(m, k)
		}
		val = apply(rewriteVal, val)
		if match(m, pat, val) {
			val = subst(m, repl, reflect.ValueOf(val.Interface().(ast.Node).Pos()))
		}
		return val
	}

	r := apply(rewriteVal, reflect.ValueOf(p)).Interface().(*ast.File)
	r.Comments = cmap.Filter(r).Comments() // recreate comments list
	return r
}

func processFile(filename string) []Function {
	var res []Function
	fset := token.NewFileSet() // positions are relative to fset
	f, err := parser.ParseFile(fset, filename, nil, parser.ParseComments)
	if err != nil {
		panic(err)
	}

	cmap := ast.NewCommentMap(fset, f, f.Comments)
	ast.Inspect(f, func(n ast.Node) bool {
		switch x := n.(type) {
		case *ast.FuncDecl:
			fun := Function{Begin: fset.Position(x.Pos()).Line,
				Package:  f.Name.String(),
				Name:     x.Name.String(),
				End:      fset.Position(x.End()).Line,
				Filepath: fset.Position(x.Pos()).Filename}

			fun.getTSpec(cmap[n])
			res = append(res, fun)
		}
		return true
	})
	return res
}

func (gas *Analyzer) process(filename string, source interface{}) error {
	mode := parser.ParseComments
	root, err := parser.ParseFile(gas.context.FileSet, filename, source, mode)
	if err == nil {
		gas.context.Comments = ast.NewCommentMap(gas.context.FileSet, root, root.Comments)
		gas.context.Root = root

		// here we get type info
		gas.context.Info = &types.Info{
			Types:      make(map[ast.Expr]types.TypeAndValue),
			Defs:       make(map[*ast.Ident]types.Object),
			Uses:       make(map[*ast.Ident]types.Object),
			Selections: make(map[*ast.SelectorExpr]*types.Selection),
			Scopes:     make(map[ast.Node]*types.Scope),
			Implicits:  make(map[ast.Node]types.Object),
		}

		conf := types.Config{Importer: importer.Default()}
		gas.context.Pkg, _ = conf.Check("pkg", gas.context.FileSet, []*ast.File{root}, gas.context.Info)
		if err != nil {
			gas.logger.Println("failed to check imports")
			return err
		}

		ast.Walk(gas, root)
		gas.Stats.NumFiles++
	}
	return err
}

// packageExports is a local implementation of ast.PackageExports
// which correctly updates each package file's comment list.
// (The ast.PackageExports signature is frozen, hence the local
// implementation).
//
func packageExports(fset *token.FileSet, pkg *ast.Package) {
	for _, src := range pkg.Files {
		cmap := ast.NewCommentMap(fset, src, src.Comments)
		ast.FileExports(src)
		src.Comments = cmap.Filter(src).Comments()
	}
}

// filterInfo updates info to include only the nodes that match the given
// filter args.
func filterInfo(args []string, info *PageInfo) {
	rx, err := makeRx(args)
	if err != nil {
		log.Fatalf("illegal regular expression from %v: %v", args, err)
	}

	filter := func(s string) bool { return rx.MatchString(s) }
	switch {
	case info.PAst != nil:
		newPAst := map[string]*ast.File{}
		for name, a := range info.PAst {
			cmap := ast.NewCommentMap(info.FSet, a, a.Comments)
			a.Comments = []*ast.CommentGroup{} // remove all comments.
			ast.FilterFile(a, filter)
			if len(a.Decls) > 0 {
				newPAst[name] = a
			}
			for _, d := range a.Decls {
				// add back the comments associated with d only
				comments := cmap.Filter(d).Comments()
				a.Comments = append(a.Comments, comments...)
			}
		}
		info.PAst = newPAst // add only matching files.
	case info.PDoc != nil:
		info.PDoc.Filter(filter)
	}
}

// ParseDir will create a Build from the directory that
// was passed into the function.
func ParseDir(path string) (*Build, error) {
	var fileSet token.FileSet

	packages, err := parser.ParseDir(&fileSet, path, nil, parser.AllErrors)
	if err != nil {
		return nil, err
	}

	// create new build for the file set
	build := NewBuild()

	// iterate over all packages in the directory
	for _, pkg := range packages {
		// iterate over all files within the package
		for name, astTree := range pkg.Files {
			baseName := filepath.Base(name)

			// create a comment map from file
			commentMap := ast.NewCommentMap(&fileSet, astTree, astTree.Comments)

			fileAST, err := ParseFileAST(baseName, astTree, commentMap)
			if err != nil {
				return nil, err
			}
			build.AddFile(baseName, fileAST)
		}
	}

	return build, nil
}

func getCommentMap(fileName string) (cm map[string]string) {
	fileSet := token.NewFileSet() // positions are relative to fset

	// Parse the file containing this very example
	// but stop after processing the imports.
	f, err := parser.ParseFile(fileSet, fileName, nil, parser.ParseComments)
	if err != nil {
		fmt.Println(err)
		return
	}

	commentMap := ast.NewCommentMap(fileSet, f, f.Comments)
	cm = make(map[string]string)
	for n, cgs := range commentMap {
		fmt.Printf("%#v,%#v --- %#v\n", n.(ast.Node).Pos(), n.(ast.Node).End(), toText(cgs))
		comment := toText(cgs)
		if len(strings.TrimSpace(comment)) == 0 {
			continue
		}
		split := strings.SplitN(comment, " ", 2)
		godoc := split[1]
		key := split[0]
		cm[key] = godoc
	}
	return cm
}

// 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 functions(f *ast.File, info types.Info, fset *token.FileSet) ([]Function, error) {
	fns := exportedFuncs(f, fset)
	fns = errorOrVoid(fns, info)

	cmtMap := ast.NewCommentMap(fset, f, f.Comments)

	functions := make([]Function, len(fns))

	for i, fn := range fns {
		fun := Function{Name: fn.Name.Name}
		fun.Comment = combine(cmtMap[fn])

		// we only support null returns or error returns, so if there's a
		// return, it's an error.
		if len(fn.Type.Results.List) > 0 {
			fun.IsError = true
		}
		params := fn.Type.Params.List
		fun.Params = make([]Param, 0, len(params))
		for _, field := range params {
			t := info.TypeOf(field.Type)
			pointer := false
			if p, ok := t.(*types.Pointer); ok {
				t = p.Elem()
				pointer = true
			}
			if b, ok := t.(*types.Basic); ok {
				if b.Kind() == types.UnsafePointer {
					log.Printf(
						"Can't create command for function %q because its parameter %q is an unsafe.Pointer.",
						fn.Name.Name,
						field.Names[0])
					break
				}

				fieldCmt := combine(cmtMap[field])
				// handle a, b, c int
				for _, name := range field.Names {
					nameCmt := combine(cmtMap[name])
					if nameCmt == "" {
						nameCmt = fieldCmt
					}
					param := Param{
						Name:      name.Name,
						Type:      b.Kind(),
						IsPointer: pointer,
						Comment:   nameCmt,
					}
					fun.Params = append(fun.Params, param)
				}
				continue
			}
		}
		functions[i] = fun
	}
	return functions, nil
}

// Trim trims the AST rooted at node based on the coverage profile,
// removing irrelevant and unreached parts of the program.
// If the node is an *ast.File, comments are updated as well using
// an ast.CommentMap.
func (p *Profile) Trim(node ast.Node) {
	if f, ok := node.(*ast.File); ok {
		cmap := ast.NewCommentMap(p.Fset, f, f.Comments)
		ast.Walk(&trimVisitor{p}, f)
		f.Comments = cmap.Filter(f).Comments()
	} else {
		ast.Walk(&trimVisitor{p}, node)
	}
}

// Inline replaces each instance of identifier k with v.Ident in ast.File f,
// for k, v := range m.
// For all inlines that were triggeres it also adds imports from v.Imports to f.
// In addition, it removes top level type declarations of the form
// type k ...
// for all k in m.
//
// Every k in m should be a valid identifier.
// Every v.Ident should be a valid expression.
func Inline(fset *token.FileSet, f *ast.File, m map[string]Target) error {
	// Build the inline map.
	im := map[string]reflect.Value{}
	for k, v := range m {
		expr, err := parser.ParseExpr(k)
		if err != nil {
			return fmt.Errorf("failed to parse `%s`: %s", k, err)
		}
		if _, ok := expr.(*ast.Ident); !ok {
			return fmt.Errorf("expected identifier, got %s which is %T", k, expr)
		}
		expr, err = parser.ParseExpr(v.Ident)
		if err != nil {
			return fmt.Errorf("failed to parse `%s`: %s", v.Ident, err)
		}
		s := v.Ident
		if _, ok := expr.(*ast.StarExpr); ok {
			s = fmt.Sprintf("(%s)", s)
		}
		im[k] = reflect.ValueOf(ast.Ident{Name: s})
	}
	// Filter `type XXX ...` declarations out if we are inlining XXX.
	cmap := ast.NewCommentMap(fset, f, f.Comments)
	to := 0
	for _, d := range f.Decls {
		skip := false
		if t, ok := d.(*ast.GenDecl); ok {
			for _, s := range t.Specs {
				ts, ok := s.(*ast.TypeSpec)
				if !ok {
					continue
				}
				if _, ok = im[ts.Name.String()]; ok {
					skip = true
				}
			}
		}
		if !skip {
			f.Decls[to] = d
			to++
		}
	}
	if to != len(f.Decls) {
		f.Decls = f.Decls[:to]
		// Remove comments for the declarations that were filtered out.
		f.Comments = cmap.Filter(f).Comments()
	}
	// Add imports for the inlines that were triggered.
	for k := range inline(im, f) {
		for _, imp := range m[k].Imports {
			astutil.AddImport(fset, f, imp)
		}
	}
	return nil
}

func (file *File) initCMap() {
	goCMap := ast.NewCommentMap(file.FSet, file.AST, file.AST.Comments)

	myCMap := make(map[int]cmapel, len(file.AST.Comments))
	for _, comment := range file.AST.Comments {
		end := file.FSet.Position(comment.End())
		myCMap[end.Line] = cmapel{
			Comment: comment,
			End:     end,
		}
	}
	file.cmap = cmap{
		GoCMap: goCMap,
		MyCMap: myCMap,
	}
}

func AnalyzePkg(files []*ast.File, fileSet *token.FileSet, typesInfo *types.Info, typesPkg *types.Package, isBlocking func(*types.Func) bool) *Info {
	info := &Info{
		Info:          typesInfo,
		Pkg:           typesPkg,
		HasPointer:    make(map[*types.Var]bool),
		comments:      make(ast.CommentMap),
		IsBlocking:    isBlocking,
		FuncDeclInfos: make(map[*types.Func]*FuncInfo),
		FuncLitInfos:  make(map[*ast.FuncLit]*FuncInfo),
	}
	info.InitFuncInfo = info.newFuncInfo()

	for _, file := range files {
		for k, v := range ast.NewCommentMap(fileSet, file, file.Comments) {
			info.comments[k] = v
		}
		ast.Walk(info.InitFuncInfo, file)
	}

	for {
		done := true
		for _, funcInfo := range info.allInfos {
			for obj, calls := range funcInfo.LocalCalls {
				if len(info.FuncDeclInfos[obj].Blocking) != 0 {
					for _, call := range calls {
						funcInfo.markBlocking(call)
					}
					delete(funcInfo.LocalCalls, obj)
					done = false
				}
			}
		}
		if done {
			break
		}
	}

	for _, funcInfo := range info.allInfos {
		for _, continueStmt := range funcInfo.ContinueStmts {
			if funcInfo.Blocking[continueStmt.forStmt.Post] {
				funcInfo.markBlocking(continueStmt.analyzeStack)
			}
		}
	}

	return info
}