...
1
2
3
4
5 package ssautil
6
7 import "golang.org/x/tools/go/ssa"
8
9
10
11
12
13
14
15
16
17
18
19
20
21 func AllFunctions(prog *ssa.Program) map[*ssa.Function]bool {
22 visit := visitor{
23 prog: prog,
24 seen: make(map[*ssa.Function]bool),
25 }
26 visit.program()
27 return visit.seen
28 }
29
30 type visitor struct {
31 prog *ssa.Program
32 seen map[*ssa.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.(*ssa.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 *ssa.Function) {
52 if !visit.seen[fn] {
53 visit.seen[fn] = true
54 var buf [10]*ssa.Value
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).(*ssa.Function); ok {
59 visit.function(fn)
60 }
61 }
62 }
63 }
64 }
65 }
66
67
68
69
70 func MainPackages(pkgs []*ssa.Package) []*ssa.Package {
71 var mains []*ssa.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 }
79
View as plain text