github.com/golangci/go-tools@v0.0.0-20190318060251-af6baa5dc196/ssa/ssautil/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 ssautil // import "github.com/golangci/go-tools/ssa/ssautil"
     6  
     7  import "github.com/golangci/go-tools/ssa"
     8  
     9  // This file defines utilities for visiting the SSA representation 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  //
    22  func AllFunctions(prog *ssa.Program) map[*ssa.Function]bool {
    23  	visit := visitor{
    24  		prog: prog,
    25  		seen: make(map[*ssa.Function]bool),
    26  	}
    27  	visit.program()
    28  	return visit.seen
    29  }
    30  
    31  type visitor struct {
    32  	prog *ssa.Program
    33  	seen map[*ssa.Function]bool
    34  }
    35  
    36  func (visit *visitor) program() {
    37  	for _, pkg := range visit.prog.AllPackages() {
    38  		for _, mem := range pkg.Members {
    39  			if fn, ok := mem.(*ssa.Function); ok {
    40  				visit.function(fn)
    41  			}
    42  		}
    43  	}
    44  	for _, T := range visit.prog.RuntimeTypes() {
    45  		mset := visit.prog.MethodSets.MethodSet(T)
    46  		for i, n := 0, mset.Len(); i < n; i++ {
    47  			visit.function(visit.prog.MethodValue(mset.At(i)))
    48  		}
    49  	}
    50  }
    51  
    52  func (visit *visitor) function(fn *ssa.Function) {
    53  	if !visit.seen[fn] {
    54  		visit.seen[fn] = true
    55  		var buf [10]*ssa.Value // avoid alloc in common case
    56  		for _, b := range fn.Blocks {
    57  			for _, instr := range b.Instrs {
    58  				for _, op := range instr.Operands(buf[:0]) {
    59  					if fn, ok := (*op).(*ssa.Function); ok {
    60  						visit.function(fn)
    61  					}
    62  				}
    63  			}
    64  		}
    65  	}
    66  }
    67  
    68  // MainPackages returns the subset of the specified packages
    69  // named "main" that define a main function.
    70  // The result may include synthetic "testmain" packages.
    71  func MainPackages(pkgs []*ssa.Package) []*ssa.Package {
    72  	var mains []*ssa.Package
    73  	for _, pkg := range pkgs {
    74  		if pkg.Pkg.Name() == "main" && pkg.Func("main") != nil {
    75  			mains = append(mains, pkg)
    76  		}
    77  	}
    78  	return mains
    79  }