github.com/amarpal/go-tools@v0.0.0-20240422043104-40142f59f616/go/ir/irutil/visit.go (about) 1 // Copyright 2013 The Go Authors. All rights reserved. 2 // Use of this source code is governed by a BSD-style 3 // license that can be found in the LICENSE file. 4 5 package irutil 6 7 import "github.com/amarpal/go-tools/go/ir" 8 9 // This file defines utilities for visiting the IR of 10 // a Program. 11 // 12 // TODO(adonovan): test coverage. 13 14 // AllFunctions finds and returns the set of functions potentially 15 // needed by program prog, as determined by a simple linker-style 16 // reachability algorithm starting from the members and method-sets of 17 // each package. The result may include anonymous functions and 18 // synthetic wrappers. 19 // 20 // Precondition: all packages are built. 21 func AllFunctions(prog *ir.Program) map[*ir.Function]bool { 22 visit := visitor{ 23 prog: prog, 24 seen: make(map[*ir.Function]bool), 25 } 26 visit.program() 27 return visit.seen 28 } 29 30 type visitor struct { 31 prog *ir.Program 32 seen map[*ir.Function]bool 33 } 34 35 func (visit *visitor) program() { 36 for _, pkg := range visit.prog.AllPackages() { 37 for _, mem := range pkg.Members { 38 if fn, ok := mem.(*ir.Function); ok { 39 visit.function(fn) 40 } 41 } 42 } 43 for _, T := range visit.prog.RuntimeTypes() { 44 mset := visit.prog.MethodSets.MethodSet(T) 45 for i, n := 0, mset.Len(); i < n; i++ { 46 visit.function(visit.prog.MethodValue(mset.At(i))) 47 } 48 } 49 } 50 51 func (visit *visitor) function(fn *ir.Function) { 52 if !visit.seen[fn] { 53 visit.seen[fn] = true 54 var buf [10]*ir.Value // avoid alloc in common case 55 for _, b := range fn.Blocks { 56 for _, instr := range b.Instrs { 57 for _, op := range instr.Operands(buf[:0]) { 58 if fn, ok := (*op).(*ir.Function); ok { 59 visit.function(fn) 60 } 61 } 62 } 63 } 64 } 65 } 66 67 // MainPackages returns the subset of the specified packages 68 // named "main" that define a main function. 69 // The result may include synthetic "testmain" packages. 70 func MainPackages(pkgs []*ir.Package) []*ir.Package { 71 var mains []*ir.Package 72 for _, pkg := range pkgs { 73 if pkg.Pkg.Name() == "main" && pkg.Func("main") != nil { 74 mains = append(mains, pkg) 75 } 76 } 77 return mains 78 }