feat: connect method call with method definition using typeinfo (cloudwego#13)

* feat: connect function-call with method defination

* deduplicate

Author: AsterDY

src/compress/golang/go.mod

@@ -6,3 +6,5 @@ require (
 	github.com/davecgh/go-spew v1.1.1
 	golang.org/x/tools v0.16.1
 )
+
+require golang.org/x/mod v0.14.0 // indirect

src/compress/golang/go.sum

@@ -1,4 +1,7 @@
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+golang.org/x/mod v0.14.0 h1:dGoOF9QVLYng8IHTm7BAyWqCqSheQ5pYWGhzW00YJr0=
+golang.org/x/mod v0.14.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
+golang.org/x/sync v0.5.0 h1:60k92dhOjHxJkrqnwsfl8KuaHbn/5dl0lUPUklKo3qE=
 golang.org/x/tools v0.16.1 h1:TLyB3WofjdOEepBHAU20JdNC1Zbg87elYofWYAY5oZA=
 golang.org/x/tools v0.16.1/go.mod h1:kYVVN6I1mBNoB1OX+noeBjbRk4IUEPa7JJ+TJMEooJ0=

src/compress/golang/plugin/file.go

@@ -19,6 +19,7 @@ package main
 import (
 	"go/ast"
 	"go/token"
+	"go/types"
 	"path/filepath"
 	"strings"
 )
@@ -66,6 +67,31 @@ func isGoBuiltinFunc(name string) bool {
 	}
 }
 
+func (p *goParser) inspectFile(ctx *fileContext, f *ast.File) (map[string]*Function, map[string]*Struct, error) {
+	fileStructs := map[string]*Struct{}
+	fileFuncs := map[string]*Function{}
+	cont := true
+	ast.Inspect(f, func(node ast.Node) bool {
+		if funcDecl, ok := node.(*ast.FuncDecl); ok {
+			// parse funcs
+			f, ct := p.parseFunc(ctx, funcDecl)
+			fileFuncs[f.Name] = f
+			cont = ct
+		} else if typDecl, ok := node.(*ast.TypeSpec); ok {
+			// parse structs
+			struName := typDecl.Name.Name
+			struDecl, ok := typDecl.Type.(*ast.StructType)
+			if ok {
+				st, ct := p.parseStruct(ctx, struName, struDecl)
+				fileStructs[struName] = st
+				cont = ct
+			}
+		}
+		return cont
+	})
+	return fileFuncs, fileStructs, nil
+}
+
 // getOrSetFunc get a function in the map, or alloc and set a new one if not exists
 func (p *goParser) getOrSetFunc(pkg, name string) *Function {
 	pkgFuncs := p.processedPkgFunctions[pkg]
@@ -127,7 +153,7 @@ func (p *goParser) seprateImports(impts []*ast.ImportSpec) (map[string]string, m
 
 // parseFunc parses all function declaration in one file
 func (p *goParser) parseFunc(ctx *fileContext, funcDecl *ast.FuncDecl) (*Function, bool) {
-
+	// funcObj := ctx.pkgTypeInfo.Defs[funcDecl.Name]
 	var associatedStruct *Struct
 	isMethod := funcDecl.Recv != nil
 	if isMethod {
@@ -150,6 +176,7 @@ func (p *goParser) parseFunc(ctx *fileContext, funcDecl *ast.FuncDecl) (*Functio
 	var functionCalls, methodCalls = map[string]*Function{}, map[string]*Function{}
 
 	ast.Inspect(funcDecl.Body, func(node ast.Node) bool {
+		// scope := ctx.pkgTypeInfo.Scopes[node]
 		call, ok := node.(*ast.CallExpr)
 		if ok {
 			var funcName string
@@ -181,20 +208,40 @@ func (p *goParser) parseFunc(ctx *fileContext, funcDecl *ast.FuncDecl) (*Functio
 					thirdPartyFunctionCalls[funcName] = &ThirdPartyIdentity{PkgPath: impt, Identity: expr.Sel.Name}
 					return true
 				}
-				// WHY: skip sys imports?
-				if _, ok := ctx.sysImports[x.(*ast.Ident).Name]; ok {
+				// NOTICE: skip sys imports?
+				if ctx.IsSysImport(x.(*ast.Ident).Name) {
 					// internalFunctionCalls[funcName] = p.getOrSetFunc(impt, expr.Sel.Name)
 					return true
 				}
+				// check if it's method calls
+				sel, ok := ctx.pkgTypeInfo.Selections[expr]
+				if ok && (sel.Kind() == types.MethodExpr || sel.Kind() == types.MethodVal) {
+					// builtin or std libs, just ignore
+					m := sel.Obj()
+					if m.Pkg() == nil || ctx.IsSysImport(m.Pkg().Name()) {
+						return true
+					}
+					// try assert as named type
+					obj := getTypeNamed(sel.Recv())
+					if obj == nil {
+						return true
+					}
 
-				// Fallback must be method calls
-				// FIXME: get type info of object
-				f := p.getOrSetFunc(ctx.pkgPath, funcName)
-				f.IsMethod = true
-				f.AssociatedStruct = &Struct{Name: x.(*ast.Ident).Name}
-				methodCalls[funcName] = f
-				// TODO: seperate internal method and third-party method
-
+					mpkg := m.Pkg().Path()
+					//NOTICE: use {structName.methodName} as method key
+					mname := obj.Name() + "." + m.Name()
+					f := p.getOrSetFunc(mpkg, mname)
+					f.AssociatedStruct = p.getOrSetStruct(mpkg, obj.Name())
+					f.IsMethod = true
+
+					if strings.HasPrefix(mpkg, p.modName) {
+						// internal pkg
+						methodCalls[funcName] = f
+					} else {
+						// external pkg
+						thirdPartyMethodCalls[funcName] = &ThirdPartyIdentity{mpkg, mname}
+					}
+				}
 				return true
 			case *ast.Ident:
 				funcName = expr.Name
@@ -227,6 +274,22 @@ func (p *goParser) parseFunc(ctx *fileContext, funcDecl *ast.FuncDecl) (*Functio
 	return f, true
 }
 
+func getTypeNamed(typ types.Type) types.Object {
+	if pt, ok := typ.(*types.Pointer); ok {
+		typ = pt.Elem()
+	}
+	name, ok := typ.(*types.Named)
+	if ok {
+		return name.Obj()
+	}
+	return nil
+}
+
+func (ctx *fileContext) IsSysImport(alias string) bool {
+	_, ok := ctx.sysImports[alias]
+	return ok
+}
+
 // Struct holds the information about a struct
 type Struct struct {
 	Name     string // Name of the struct
@@ -257,6 +320,7 @@ type fileContext struct {
 	sysImports        map[string]string
 	projectImports    map[string]string
 	thirdPartyImports map[string]string
+	pkgTypeInfo       *types.Info
 }
 
 // parse a ast.StructType node and renturn allocated *Struct
@@ -287,7 +351,7 @@ func (p *goParser) parseStruct(ctx *fileContext, struName string, struDecl *ast.
 		}
 
 		types := []ThirdPartyIdentity{}
-		isFunc := getTypeName(ctx.bs, fieldDecl.Type, &types)
+		isFunc := getTypeName(ctx.fset, ctx.bs, fieldDecl.Type, &types)
 
 		for _, ty := range types {
 			if isFunc {
@@ -328,35 +392,35 @@ func (p *goParser) parseStruct(ctx *fileContext, struName string, struDecl *ast.
 
 // handle typ expr and return not-builtin type identity and return if the type if a func signature.
 // ret is used to store results.
-func getTypeName(file []byte, typ ast.Expr, ret *[]ThirdPartyIdentity) bool {
+func getTypeName(fset *token.FileSet, file []byte, typ ast.Expr, ret *[]ThirdPartyIdentity) bool {
 	switch ty := typ.(type) {
 	case *ast.Ident:
 		if !isGoBuiltinFunc(ty.Name) {
 			*ret = append(*ret, ThirdPartyIdentity{Identity: ty.Name})
 		}
 		return false
 	case *ast.StarExpr:
-		return getTypeName(file, ty.X, ret)
+		return getTypeName(fset, file, ty.X, ret)
 	case *ast.ArrayType:
-		return getTypeName(file, ty.Elt, ret)
+		return getTypeName(fset, file, ty.Elt, ret)
 	case *ast.MapType:
-		a := getTypeName(file, ty.Key, ret)
-		b := getTypeName(file, ty.Value, ret)
+		a := getTypeName(fset, file, ty.Key, ret)
+		b := getTypeName(fset, file, ty.Value, ret)
 		return a || b
 	case *ast.ChanType:
-		return getTypeName(file, ty.Value, ret)
+		return getTypeName(fset, file, ty.Value, ret)
 	case *ast.SelectorExpr:
 		pkg, ok := ty.X.(*ast.Ident)
 		if ok {
 			*ret = append(*ret, ThirdPartyIdentity{Identity: ty.Sel.Name, PkgPath: pkg.Name})
 		}
 		return false
 	case *ast.FuncType:
-		name := string(file[ty.Func:typ.End()])
+		name := string(file[fset.Position(ty.Func).Offset:fset.Position(typ.End()).Offset])
 		*ret = append(*ret, ThirdPartyIdentity{Identity: name})
 		return true
 	case *ast.InterfaceType:
-		name := string(file[ty.Interface:typ.End()])
+		name := string(file[fset.Position(ty.Interface).Offset:fset.Position(typ.End()).Offset])
 		*ret = append(*ret, ThirdPartyIdentity{Identity: name})
 		return false
 	}

src/compress/golang/plugin/go_ast.go

@@ -5,9 +5,6 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
-	"go/ast"
-	"go/parser"
-	"go/token"
 	"io/ioutil"
 	"os"
 	"path/filepath"
@@ -91,83 +88,6 @@ func (p *goParser) associateStructWithMethods() {
 	}
 }
 
-// parseFile parse single go file and return all functions in it
-// warning: this function has no cache, do not call it with repeated file
-func (p *goParser) parseFile(filePath string) (map[string]*Function, error) {
-	fset := token.NewFileSet()
-
-	bs, err := ioutil.ReadFile(filePath)
-	if err != nil {
-		return nil, err
-	}
-
-	f, err := parser.ParseFile(fset, filePath, bs, 0)
-	if err != nil {
-		return nil, err
-	}
-
-	sysImports, projectImports, thirdPartyImports := p.seprateImports(f.Imports)
-	pkgPath := p.pkgPathFromABS(filepath.Dir(filePath))
-	ctx := &fileContext{
-		filePath:          filePath,
-		pkgPath:           pkgPath,
-		bs:                bs,
-		fset:              fset,
-		sysImports:        sysImports,
-		projectImports:    projectImports,
-		thirdPartyImports: thirdPartyImports,
-	}
-
-	fileFuncs, _, err := p.inspectFile(ctx, f)
-	return fileFuncs, err
-}
-
-func (p *goParser) inspectFile(ctx *fileContext, f *ast.File) (map[string]*Function, map[string]*Struct, error) {
-	fileStructs := map[string]*Struct{}
-	fileFuncs := map[string]*Function{}
-	cont := true
-	ast.Inspect(f, func(node ast.Node) bool {
-		if funcDecl, ok := node.(*ast.FuncDecl); ok {
-			// parse funcs
-			f, ct := p.parseFunc(ctx, funcDecl)
-			fileFuncs[f.Name] = f
-			cont = ct
-		} else if typDecl, ok := node.(*ast.TypeSpec); ok {
-			// parse structs
-			struName := typDecl.Name.Name
-			struDecl, ok := typDecl.Type.(*ast.StructType)
-			if ok {
-				st, ct := p.parseStruct(ctx, struName, struDecl)
-				fileStructs[struName] = st
-				cont = ct
-			}
-		}
-		return cont
-	})
-	return fileFuncs, fileStructs, nil
-}
-
-// ParseDir parse a single go package in the dir
-// If the dir has been visted, return cache
-func (p *goParser) ParseDir(dir string) (map[string]*Function, error) {
-	// unify dir ./xxx/xxx -> xxx/xxx
-	if !strings.HasPrefix(dir, "/") {
-		dir = filepath.Join(p.homePageDir, dir)
-	}
-	pkgPath := p.pkgPathFromABS(dir)
-	if p.visited[pkgPath] {
-		return p.processedPkgFunctions[pkgPath], nil
-	}
-	for _, f := range getGoFilesInDir(dir) {
-		_, err := p.parseFile(f)
-		if err != nil {
-			return nil, err
-		}
-	}
-	p.visited[pkgPath] = true
-	return p.processedPkgFunctions[pkgPath], nil
-}
-
 // TODO: Parallel transformation
 // ParseTilTheEnd parse the all go files from the starDir,
 // and their related go files in the project recursively
@@ -179,18 +99,19 @@ func (p *goParser) ParseTilTheEnd(startDir string) error {
 			return err
 		}
 	}
-	functionList, err := p.ParseDir(startDir)
-	if err != nil {
+	if err := p.ParseDir(startDir); err != nil {
 		return err
 	}
-	for _, f := range functionList {
-		// Notice: local funcs has been parsed in ParseDir
-		for _, fc := range f.InternalFunctionCalls {
-			if fc.FilePath != "" || fc.IsMethod {
-				continue
-			}
-			if err := p.ParseTilTheEnd(p.pkgPathToABS(fc.PkgPath)); err != nil {
-				return err
+	for _, pkg := range p.processedPkgFunctions {
+		for _, f := range pkg {
+			// Notice: local funcs has been parsed in ParseDir
+			for _, fc := range f.InternalFunctionCalls {
+				if fc.FilePath != "" {
+					continue
+				}
+				if err := p.ParseTilTheEnd(p.pkgPathToABS(fc.PkgPath)); err != nil {
+					return err
+				}
 			}
 		}
 	}
@@ -207,6 +128,32 @@ type MainStream struct {
 	RelatedStruct []SingleStruct
 }
 
+func (m *MainStream) Dedup() {
+	fs := map[string]string{}
+	for _, f := range m.RelatedFunctions {
+		fs[f.CallName] = f.Content
+	}
+	m.RelatedFunctions = m.RelatedFunctions[:len(fs)]
+	i := 0
+	for k, v := range fs {
+		m.RelatedFunctions[i].CallName = k
+		m.RelatedFunctions[i].Content = v
+		i++
+	}
+
+	fs = map[string]string{}
+	for _, f := range m.RelatedStruct {
+		fs[f.Name] = f.Content
+	}
+	m.RelatedStruct = m.RelatedStruct[:len(fs)]
+	i = 0
+	for k, v := range fs {
+		m.RelatedStruct[i].Name = k
+		m.RelatedStruct[i].Content = v
+		i++
+	}
+}
+
 type SingleFunction struct {
 	CallName string
 	Content  string
@@ -257,7 +204,7 @@ func (p *goParser) fillRelatedContent(f *Function, fl *[]SingleFunction, sl *[]S
 		// for method which has been associated with struct, push the struct
 		if ff.AssociatedStruct != nil && ff.AssociatedStruct.Content != "" {
 			ss := SingleStruct{
-				Name:    ff.AssociatedStruct.Name,
+				Name:    ff.PkgPath + "." + ff.AssociatedStruct.Name,
 				Content: ff.AssociatedStruct.Content,
 			}
 			*sl = append(*sl, ss)
@@ -309,6 +256,7 @@ func main() {
 
 	// p.generateStruct()
 	m, _ := p.getMain()
+	m.Dedup()
 
 	out := bytes.NewBuffer(nil)
 	encoder := json.NewEncoder(out)

src/compress/golang/plugin/go_ast_test.go

@@ -42,6 +42,7 @@ func Test_goParser_ParseTilTheEnd(t *testing.T) {
 			if fun.Name != "main" {
 				t.Fail()
 			}
+			out.Dedup()
 			if out, err := json.MarshalIndent(out, "", "  "); err != nil {
 				t.Fatalf("json.Marshal() error = %v", err)
 			} else {