Golang ast.Walk函数代码示例

func (v *visitor) Visit(node ast.Node) ast.Visitor {
	switch node := node.(type) {
	case *ast.Ident:
		v.use(v.pkg.Info.Uses[node])

	case *ast.ValueSpec:
		if !v.insideFunc {
			for _, ident := range node.Names {
				if !isReserved(ident.Name) {
					v.decl(v.pkg.Info.Defs[ident])
				}
			}
		}
		for _, val := range node.Values {
			ast.Walk(v, val)
		}
		return nil

	case *ast.FuncDecl:
		if node.Body != nil {
			v.insideFunc = true
			ast.Walk(v, node.Body)
			v.insideFunc = false
		}

		return nil
	}

	return v
}

func (this *astEquals) Check(params []interface{}, names []string) (bool, string) {
	var obtainedNodes, expectedNodes chan *NodeWithBreadcrumbs
	var obtainedNode, expectedNode *NodeWithBreadcrumbs

	obtainedNodes = make(chan *NodeWithBreadcrumbs)
	expectedNodes = make(chan *NodeWithBreadcrumbs)

	go ast.Walk(&AstChannelWalker{Out: obtainedNodes}, params[0].(ast.Node))
	go ast.Walk(&AstChannelWalker{Out: expectedNodes}, params[1].(ast.Node))

	for {
		obtainedNode = <-obtainedNodes
		expectedNode = <-expectedNodes

		if obtainedNode == nil && expectedNode == nil {
			break
		}

		if obtainedNode == nil || expectedNode == nil {
			return false, fmt.Sprintf("\n%+v\ndid not match\n%+v", obtainedNode, expectedNode)
		}

		if !this.nodeEquals(obtainedNode.Node, expectedNode.Node) {
			return false, fmt.Sprintf("\n%+v\ndid not match\n%+v", obtainedNode, expectedNode)
		}
	}

	return true, ""
}

func (x *Indexer) visitSpec(kind SpotKind, spec ast.Spec) {
	switch n := spec.(type) {
	case *ast.ImportSpec:
		x.visitIdent(ImportDecl, n.Name)
		if n.Path != nil {
			if imp, err := strconv.Unquote(n.Path.Value); err == nil {
				x.importCount[x.intern(imp)]++
			}
		}

	case *ast.AliasSpec:
		x.visitIdent(kind, n.Name)
		ast.Walk(x, n.Orig)

	case *ast.ValueSpec:
		for _, n := range n.Names {
			x.visitIdent(kind, n)
		}
		ast.Walk(x, n.Type)
		for _, v := range n.Values {
			ast.Walk(x, v)
		}

	case *ast.TypeSpec:
		x.visitIdent(TypeDecl, n.Name)
		ast.Walk(x, n.Type)
	}
}

// Visits files for used nodes.
func (p *Package) Visit(node ast.Node) ast.Visitor {
	u := usedWalker(*p) // hopefully p fields are references.
	switch n := node.(type) {
	// don't walk whole file, but only:
	case *ast.ValueSpec:
		// - variable initializers
		for _, value := range n.Values {
			ast.Walk(&u, value)
		}
		// variable types.
		if n.Type != nil {
			ast.Walk(&u, n.Type)
		}
	case *ast.BlockStmt:
		// - function bodies
		for _, stmt := range n.List {
			ast.Walk(&u, stmt)
		}
	case *ast.FuncDecl:
		// - function signatures
		ast.Walk(&u, n.Type)
	case *ast.TypeSpec:
		// - type declarations
		ast.Walk(&u, n.Type)
	}
	return p
}

// rewriteOptionalMethods makes specific mutations to marshaller methods that belong to types identified
// as being "optional" (they may be nil on the wire). This allows protobuf to serialize a map or slice and
// properly discriminate between empty and nil (which is not possible in protobuf).
// TODO: move into upstream gogo-protobuf once https://github.com/gogo/protobuf/issues/181
//   has agreement
func rewriteOptionalMethods(decl ast.Decl, isOptional OptionalFunc) {
	switch t := decl.(type) {
	case *ast.FuncDecl:
		ident, ptr, ok := receiver(t)
		if !ok {
			return
		}

		// correct initialization of the form `m.Field = &OptionalType{}` to
		// `m.Field = OptionalType{}`
		if t.Name.Name == "Unmarshal" {
			ast.Walk(optionalAssignmentVisitor{fn: isOptional}, t.Body)
		}

		if !isOptional(ident.Name) {
			return
		}

		switch t.Name.Name {
		case "Unmarshal":
			ast.Walk(&optionalItemsVisitor{}, t.Body)
		case "MarshalTo", "Size", "String":
			ast.Walk(&optionalItemsVisitor{}, t.Body)
			fallthrough
		case "Marshal":
			// if the method has a pointer receiver, set it back to a normal receiver
			if ptr {
				t.Recv.List[0].Type = ident
			}
		}
	}
}

func (x *Indexer) visitSpec(spec ast.Spec, isVarDecl bool) {
	switch n := spec.(type) {
	case *ast.ImportSpec:
		x.visitComment(n.Doc)
		x.visitIdent(ImportDecl, n.Name)
		ast.Walk(x, n.Path)
		x.visitComment(n.Comment)

	case *ast.ValueSpec:
		x.visitComment(n.Doc)
		kind := ConstDecl
		if isVarDecl {
			kind = VarDecl
		}
		for _, n := range n.Names {
			x.visitIdent(kind, n)
		}
		ast.Walk(x, n.Type)
		for _, v := range n.Values {
			ast.Walk(x, v)
		}
		x.visitComment(n.Comment)

	case *ast.TypeSpec:
		x.visitComment(n.Doc)
		x.visitIdent(TypeDecl, n.Name)
		ast.Walk(x, n.Type)
		x.visitComment(n.Comment)
	}
}

// Visits files for used nodes.
func (p *Package) Visit(node ast.Node) ast.Visitor {
	u := usedWalker(*p) // hopefully p fields are references.
	switch n := node.(type) {
	// don't walk whole file, but only:
	case *ast.ValueSpec:
		// - variable initializers
		for _, value := range n.Values {
			ast.Walk(&u, value)
		}
		// variable types.
		if n.Type != nil {
			ast.Walk(&u, n.Type)
		}
	case *ast.BlockStmt:
		// - function bodies
		for _, stmt := range n.List {
			ast.Walk(&u, stmt)
		}
	case *ast.FuncDecl:
		// - function signatures
		if *includeTests {
			// Test* functions are always used
			if strings.HasPrefix(n.Name.String(), "Test") {
				u.used[n.Name.String()] = true
			}
		}
		ast.Walk(&u, n.Type)
	case *ast.TypeSpec:
		// - type declarations
		ast.Walk(&u, n.Type)
	}
	return p
}

func instrument(pkg, shortName, fullName string, fset *token.FileSet, parsedFile *ast.File, info *types.Info, out io.Writer, lits map[Literal]struct{}, blocks *[]CoverBlock, sonar *[]CoverBlock) {
	file := &File{
		fset:      fset,
		pkg:       pkg,
		shortName: shortName,
		fullName:  fullName,
		astFile:   parsedFile,
		blocks:    blocks,
		info:      info,
	}
	file.addImport("github.com/dvyukov/go-fuzz/go-fuzz-dep", fuzzdepPkg, "Main")

	if lits != nil {
		ast.Walk(&LiteralCollector{lits}, file.astFile)
	}

	ast.Walk(file, file.astFile)

	if sonar != nil {
		s := &Sonar{
			fset:      fset,
			shortName: shortName,
			fullName:  fullName,
			pkg:       pkg,
			blocks:    sonar,
			info:      info,
		}
		ast.Walk(s, file.astFile)
	}

	file.print(out)
}

func restoreIdentMapping(oldFile, newFile *ast.File, identMap st.IdentifierMap) {
	comm := make(chan *ast.Ident)
	source := &restoreIMSourceVisitor{comm}
	dest := &restoreIMDestVisitor{identMap, comm}
	go ast.Walk(source, oldFile)
	ast.Walk(dest, newFile)
}

func main() {
	goopt.Parse(func() []string { return nil })
	if len(goopt.Args) > 0 {
		x, err := parser.ParseFiles(myfiles, goopt.Args, 0)
		die(err)
		fmt.Fprintln(os.Stderr, "Parsed: ", *x["main"])
		die(typechecker.CheckPackage(myfiles, x["main"], nil))
		fmt.Fprintln(os.Stderr, "Checked: ", *x["main"])
		//for _,a := range x["main"].Files {
		//	die(printer.Fprint(os.Stdout, a))
		//}

		aaa := x86.StartData
		var bbb *Stack
		var cv = CompileVisitor{&aaa, make(map[string]string), bbb.New("global")}
		ast.Walk(StringVisitor(cv), x["main"])

		cv.Append(x86.StartText...)
		ast.Walk(&cv, x["main"])

		// Here we just add a crude debug library
		cv.Append(x86.Debugging...)
		ass := x86.Assembly(*cv.assembly)
		//fmt.Println(ass)
		die(elf.AssembleAndLink(goopt.Args[0][:len(goopt.Args[0])-3], []byte(ass)))
	}
}

func trans(pks map[string]*ast.Package) {
	ts := Trans{}
	for _, pk := range pks {
		ast.Walk(normalizer{}, pk)
		ast.Walk(ts, pk)
	}
	fmt.Println("main();")
}

// 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)
	}
}

func (vis *pointerCandidatesVisitor) Visit(node ast.Node) ast.Visitor {
	switch t := node.(type) {
	case *ast.AssignStmt:
		for _, expr := range t.Rhs {
			ast.Walk(vis, expr)
		}
		if t.Tok == token.DEFINE {
			return nil
		}
		for _, ee := range t.Lhs {
			depth := 1
			stop := false
			for !stop {
				fmt.Printf("%T\n", ee)
				switch e := ee.(type) {
				case *ast.Ident:
					s := vis.identMap.GetSymbol(e)
					if vis.params.Contains(s) && depth > 0 {
						if i, ok := vis.result[s]; !ok || i < depth {
							vis.result[s] = depth
						}
					}
					stop = true
				case *ast.StarExpr:
					depth--
					ee = e.X
				case *ast.UnaryExpr:
					if e.Op == token.AND {
						depth++
					}
					ee = e.X
				case *ast.ParenExpr:
					ee = e.X
				case *ast.IndexExpr:
					ast.Walk(vis, e.Index)
					stop = true
				case *ast.SelectorExpr:
					ast.Walk(vis, e.X)
					stop = true
				default:
					stop = true
				}
			}
		}
		return nil
	case *ast.UnaryExpr:
		if t.Op != token.AND {
			return vis
		}
		vis.checkAddrOperators(t)
		return nil
	case *ast.StarExpr:
		vis.checkAddrOperators(t)
		return nil
	}
	return vis
}

func compareTwoTrees(src string) {
	v1 := &channelPusher{}
	v1.fileSet = token.NewFileSet()
	v1.queue = make(chan *ast.Node)

	v2 := &channelPusher{}
	v2.fileSet = token.NewFileSet()
	v2.queue = make(chan *ast.Node)

	tree1, err := parser.ParseExpr(src)
	if err != nil {
		panic(err)
	}

	src2 := "x + 2*y"
	tree2, err := parser.ParseExpr(src2)
	if err != nil {
		panic(err)
	}
	done := make(chan struct{})
	defer close(done)

	go func() {
		ast.Walk(v1, tree1)
		close(v1.queue)
		done <- struct{}{}
	}()
	go func() {
		ast.Walk(v2, tree2)
		close(v2.queue)
		done <- struct{}{}
	}()

	var n1, n2 *ast.Node
	quit := false
	for !quit {
		select {
		case n1 = <-v1.queue:
		case n2 = <-v2.queue:
		case <-done:
			quit = true
		}
		if n1 != nil && n2 != nil {
			if !equalNodes(n1, n2) {
				println("!equalNodes")
				break
			}
			println("equalNodes")
			n1 = nil
			n2 = nil
		}
	}
}

func (v *idVisitor) Visit(n ast.Node) (w ast.Visitor) {
	switch node := n.(type) {
	case *ast.Ident, *ast.SelectorExpr:
		id, err := str(node.(ast.Expr))
		if err == nil {
			v.rIDs[id] = n
		}
	case *ast.StarExpr:
		ast.Walk(v, node.X)
		return nil
	case *ast.CallExpr:
		stmt, ok := node.Fun.(*ast.SelectorExpr)
		if ok {
			ast.Walk(v, stmt.X)
		}

		if node.Args == nil {
			return nil
		}

		for _, arg := range node.Args {
			ast.Walk(v, arg)
		}
		return nil
	case *ast.AssignStmt:
		for _, e := range node.Lhs {
			id, err := str(e)
			if err != nil {
				return
			}
			v.wIDs[id] = n
		}

		for _, e := range node.Rhs {
			ast.Walk(v, e)
		}
		return nil
	case *ast.GenDecl:
		for _, s := range node.Specs {
			val, ok := s.(*ast.ValueSpec)
			if !ok {
				continue
			}

			for _, n := range val.Names {
				id, _ := str(n)
				v.wIDs[id] = node
			}
		}
	}
	w = v
	return
}