go.undefinedlabs.com/scopeagent@v0.4.2/ast/sources.go (about)

     1  package ast
     2  
     3  import (
     4  	"fmt"
     5  	"go/ast"
     6  	"go/parser"
     7  	"go/token"
     8  	"path/filepath"
     9  	"runtime"
    10  	"strings"
    11  	"sync"
    12  )
    13  
    14  var (
    15  	methodCodes map[string]map[string]*MethodCodeBoundaries
    16  	mutex       sync.Mutex
    17  )
    18  
    19  type MethodCodeBoundaries struct {
    20  	Package string
    21  	Name    string
    22  	File    string
    23  	Start   CodePos
    24  	End     CodePos
    25  }
    26  type CodePos struct {
    27  	Line   int
    28  	Column int
    29  }
    30  
    31  // Gets the function source code boundaries from the caller method
    32  func GetFuncSourceFromCaller(skip int) (*MethodCodeBoundaries, error) {
    33  	pc, _, _, _ := runtime.Caller(skip + 1)
    34  	return GetFuncSource(pc)
    35  }
    36  
    37  // Gets the function source code boundaries from a method
    38  func GetFuncSourceForName(pc uintptr, name string) (*MethodCodeBoundaries, error) {
    39  	mFunc := runtime.FuncForPC(pc)
    40  	mFile, _ := mFunc.FileLine(pc)
    41  	mFile = filepath.Clean(mFile)
    42  	fileCode, err := getCodesForFile(mFile)
    43  	if err != nil {
    44  		return nil, err
    45  	}
    46  	return fileCode[name], nil
    47  }
    48  
    49  // Gets the function source code boundaries from a method
    50  func GetFuncSource(pc uintptr) (*MethodCodeBoundaries, error) {
    51  	mFunc := runtime.FuncForPC(pc)
    52  	mFile, _ := mFunc.FileLine(pc)
    53  	mFile = filepath.Clean(mFile)
    54  	fileCode, err := getCodesForFile(mFile)
    55  	if err != nil {
    56  		return nil, err
    57  	}
    58  
    59  	parts := strings.Split(mFunc.Name(), ".")
    60  	funcName := parts[len(parts)-1]
    61  	return fileCode[funcName], nil
    62  }
    63  
    64  func getCodesForFile(file string) (map[string]*MethodCodeBoundaries, error) {
    65  	mutex.Lock()
    66  	defer mutex.Unlock()
    67  	if methodCodes == nil {
    68  		methodCodes = map[string]map[string]*MethodCodeBoundaries{}
    69  	}
    70  	if methodCodes[file] == nil {
    71  		methodCodes[file] = map[string]*MethodCodeBoundaries{}
    72  
    73  		fSet := token.NewFileSet()
    74  		f, err := parser.ParseFile(fSet, file, nil, 0)
    75  		if err != nil {
    76  			return nil, err
    77  		}
    78  
    79  		packageName := f.Name.String()
    80  		for _, decl := range f.Decls {
    81  			if fDecl, ok := decl.(*ast.FuncDecl); ok {
    82  				bPos := fDecl.Pos()
    83  				if fDecl.Body != nil {
    84  					bEnd := fDecl.Body.End()
    85  					if bPos.IsValid() && bEnd.IsValid() {
    86  						pos := fSet.PositionFor(bPos, true)
    87  						end := fSet.PositionFor(bEnd, true)
    88  						var instTypes []string
    89  						if fDecl.Recv != nil && fDecl.Recv.List != nil {
    90  							for _, recv := range fDecl.Recv.List {
    91  								if recVStarExpr, ok := recv.Type.(*ast.StarExpr); ok {
    92  									if recVStarExprIdent, ok := recVStarExpr.X.(*ast.Ident); ok {
    93  										instTypes = append(instTypes, fmt.Sprintf("*%s", recVStarExprIdent.Name))
    94  									}
    95  								} else if recVIdent, ok := recv.Type.(*ast.Ident); ok {
    96  									instTypes = append(instTypes, recVIdent.Name)
    97  								}
    98  							}
    99  						}
   100  						name := ""
   101  						if instTypes != nil {
   102  							name = fmt.Sprintf("(%s).", strings.Join(instTypes, ", "))
   103  						}
   104  						name = name + fDecl.Name.String()
   105  						methodCode := MethodCodeBoundaries{
   106  							Package: packageName,
   107  							Name:    name,
   108  							File:    file,
   109  							Start: CodePos{
   110  								Line:   pos.Line,
   111  								Column: pos.Column,
   112  							},
   113  							End: CodePos{
   114  								Line:   end.Line,
   115  								Column: end.Column,
   116  							},
   117  						}
   118  						methodCodes[file][methodCode.Name] = &methodCode
   119  					}
   120  				}
   121  			}
   122  		}
   123  	}
   124  	return methodCodes[file], nil
   125  }