fix: associate methods (cloudwego#14)

Author: AsterDY

src/compress/golang/plugin/file.go

@@ -78,14 +78,19 @@ func (p *goParser) inspectFile(ctx *fileContext, f *ast.File) (map[string]*Funct
 			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
+			var st *Struct
+			var ct bool
+			// parse structs
+			if struDecl, ok := typDecl.Type.(*ast.StructType); ok {
+				st, ct = p.parseStruct(ctx, struName, struDecl)
+				// parse interface
+			} else if interDecl, ok := typDecl.Type.(*ast.InterfaceType); ok {
+				struName := typDecl.Name.Name
+				st, ct = p.parseInterface(ctx, struName, interDecl)
 			}
+			fileStructs[struName] = st
+			cont = ct
 		}
 		return cont
 	})
@@ -175,6 +180,9 @@ func (p *goParser) parseFunc(ctx *fileContext, funcDecl *ast.FuncDecl) (*Functio
 	var thirdPartyMethodCalls, thirdPartyFunctionCalls = map[string]*ThirdPartyIdentity{}, map[string]*ThirdPartyIdentity{}
 	var functionCalls, methodCalls = map[string]*Function{}, map[string]*Function{}
 
+	if funcDecl.Body == nil {
+		goto set_func
+	}
 	ast.Inspect(funcDecl.Body, func(node ast.Node) bool {
 		// scope := ctx.pkgTypeInfo.Scopes[node]
 		call, ok := node.(*ast.CallExpr)
@@ -230,12 +238,9 @@ func (p *goParser) parseFunc(ctx *fileContext, funcDecl *ast.FuncDecl) (*Functio
 					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
+						f := p.getOrSetFunc(mpkg, mname)
 						methodCalls[funcName] = f
 					} else {
 						// external pkg
@@ -253,6 +258,8 @@ func (p *goParser) parseFunc(ctx *fileContext, funcDecl *ast.FuncDecl) (*Functio
 		}
 		return true
 	})
+
+set_func:
 	name := funcDecl.Name.Name
 	if isMethod {
 		name = associatedStruct.Name + "." + name
@@ -292,10 +299,11 @@ func (ctx *fileContext) IsSysImport(alias string) bool {
 
 // Struct holds the information about a struct
 type Struct struct {
-	Name     string // Name of the struct
-	PkgPath         // Path to the package where the struct is defined
-	FilePath string // File where the struct is defined
-	Content  string // struct declaration content
+	IsInterface bool   //maybe a interface type decl
+	Name        string // Name of the struct
+	PkgPath            // Path to the package where the struct is defined
+	FilePath    string // File where the struct is defined
+	Content     string // struct declaration content
 
 	// related local structs in fields, key is {{pkgName.typName}} or {{typeName}}, val is declaration of the struct
 	InternalStructs map[string]*Struct
@@ -325,8 +333,7 @@ type fileContext struct {
 
 // parse a ast.StructType node and renturn allocated *Struct
 func (p *goParser) parseStruct(ctx *fileContext, struName string, struDecl *ast.StructType) (*Struct, bool) {
-	pkgPath := p.pkgPathFromABS(filepath.Dir(ctx.filePath))
-	st := p.getOrSetStruct(pkgPath, struName)
+	st := p.getOrSetStruct(ctx.pkgPath, struName)
 	st.FilePath = ctx.filePath
 
 	pos := ctx.fset.PositionFor(struDecl.Pos(), false).Offset
@@ -347,7 +354,7 @@ func (p *goParser) parseStruct(ctx *fileContext, struName string, struDecl *ast.
 			// TODO: combine all names
 			name = fieldDecl.Names[0].String()
 		} else {
-			name = string(ctx.bs[fieldDecl.Type.Pos():fieldDecl.Type.End()])
+			name = string(ctx.GetRawContent(fieldDecl))
 		}
 
 		types := []ThirdPartyIdentity{}
@@ -374,7 +381,7 @@ func (p *goParser) parseStruct(ctx *fileContext, struName string, struDecl *ast.
 				}
 			} else {
 				// local package
-				sub := p.getOrSetStruct(pkgPath, ty.Identity)
+				sub := p.getOrSetStruct(ctx.pkgPath, ty.Identity)
 				inStructs[name] = sub
 			}
 
@@ -390,6 +397,52 @@ func (p *goParser) parseStruct(ctx *fileContext, struName string, struDecl *ast.
 	return st, true
 }
 
+func (ctx *fileContext) GetRawContent(node ast.Node) []byte {
+	return ctx.bs[ctx.fset.Position(node.Pos()).Offset:ctx.fset.Position(node.End()).Offset]
+}
+
+func (p *goParser) parseInterface(ctx *fileContext, name string, decl *ast.InterfaceType) (*Struct, bool) {
+	if decl == nil || decl.Incomplete {
+		return nil, true
+	}
+
+	st := p.getOrSetStruct(ctx.pkgPath, name)
+	st.FilePath = ctx.filePath
+	st.IsInterface = true
+	st.Content = string(ctx.GetRawContent(decl))
+
+	methods := map[string]*Function{}
+	ast.Inspect(decl.Methods, func(n ast.Node) bool {
+		fieldDecl, ok := n.(*ast.Field)
+		if !ok {
+			return true
+		}
+		fname := ""
+		if len(fieldDecl.Names) > 0 {
+			// TODO: combine all names
+			fname = fieldDecl.Names[0].String()
+		} else {
+			fname = string(ctx.GetRawContent(fieldDecl.Type))
+		}
+
+		types := []ThirdPartyIdentity{}
+		isFunc := getTypeName(ctx.fset, ctx.bs, fieldDecl.Type, &types)
+		if !isFunc {
+			return true
+		}
+
+		f := p.getOrSetFunc(ctx.pkgPath, name+"."+fname)
+		f.IsMethod = true
+		f.AssociatedStruct = st
+		f.FilePath = ctx.filePath
+		methods[fname] = f
+		return true
+	})
+
+	st.Methods = methods
+	return st, true
+}
+
 // 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(fset *token.FileSet, file []byte, typ ast.Expr, ret *[]ThirdPartyIdentity) bool {

src/compress/golang/plugin/go_ast.go

@@ -28,13 +28,22 @@ func newGoParser(modName, homePageDir string) *goParser {
 	if err != nil {
 		panic(fmt.Sprintf("cannot get absolute path form homePageDir:%v", err))
 	}
-	return &goParser{
+
+	p := &goParser{
 		modName:               modName,
 		homePageDir:           abs,
 		processedPkgFunctions: map[PkgPath]map[string]*Function{},
 		processedPkgStruct:    map[PkgPath]map[string]*Struct{},
 		visited:               map[string]bool{},
 	}
+	if p.modName == "" {
+		var err error
+		p.modName, err = getModuleName(p.homePageDir + "/go.mod")
+		if err != nil {
+			panic(err.Error())
+		}
+	}
+	return p
 }
 
 // ToABS converts a local package path to absolute path
@@ -92,21 +101,26 @@ func (p *goParser) associateStructWithMethods() {
 // ParseTilTheEnd parse the all go files from the starDir,
 // and their related go files in the project recursively
 func (p *goParser) ParseTilTheEnd(startDir string) error {
-	if p.modName == "" {
-		var err error
-		p.modName, err = getModuleName(p.homePageDir + "/go.mod")
-		if err != nil {
-			return err
-		}
-	}
 	if err := p.ParseDir(startDir); err != nil {
 		return err
 	}
-	for _, pkg := range p.processedPkgFunctions {
+	for path, pkg := range p.processedPkgFunctions {
+		// ignore third-party packages
+		if !strings.Contains(path, p.modName) {
+			continue
+		}
 		for _, f := range pkg {
 			// Notice: local funcs has been parsed in ParseDir
 			for _, fc := range f.InternalFunctionCalls {
-				if fc.FilePath != "" {
+				if p.visited[fc.PkgPath] {
+					continue
+				}
+				if err := p.ParseTilTheEnd(p.pkgPathToABS(fc.PkgPath)); err != nil {
+					return err
+				}
+			}
+			for _, fc := range f.InternalMethodCalls {
+				if p.visited[fc.PkgPath] {
 					continue
 				}
 				if err := p.ParseTilTheEnd(p.pkgPathToABS(fc.PkgPath)); err != nil {
@@ -164,7 +178,7 @@ type SingleStruct struct {
 	Content string
 }
 
-func (p *goParser) getMain() (*MainStream, *Function) {
+func (p *goParser) getMain(depth int) (*MainStream, *Function) {
 	m := &MainStream{
 		RelatedFunctions: make([]SingleFunction, 0),
 	}
@@ -180,36 +194,61 @@ func (p *goParser) getMain() (*MainStream, *Function) {
 			}
 		}
 	}
-	p.fillRelatedContent(mainFunc, &m.RelatedFunctions, &m.RelatedStruct)
+	visited := map[string]map[string]bool{}
+	p.fillRelatedContent(depth, mainFunc, &m.RelatedFunctions, &m.RelatedStruct, visited)
 	return m, mainFunc
 }
 
-func (p *goParser) fillRelatedContent(f *Function, fl *[]SingleFunction, sl *[]SingleStruct) {
+func (p *goParser) fillRelatedContent(depth int, f *Function, fl *[]SingleFunction, sl *[]SingleStruct, visited map[string]map[string]bool) {
+	if depth == 0 {
+		return
+	}
+	if f == nil || (visited[f.PkgPath] != nil && visited[f.PkgPath][f.Name]) {
+		return
+	} else {
+		if visited[f.PkgPath] == nil {
+			visited[f.PkgPath] = map[string]bool{}
+		}
+		visited[f.PkgPath][f.Name] = true
+	}
 	for call, ff := range f.InternalFunctionCalls {
 		s := SingleFunction{
 			CallName: call,
 			Content:  ff.Content,
 		}
 		*fl = append(*fl, s)
-		p.fillRelatedContent(ff, fl, sl)
+		p.fillRelatedContent(depth-1, ff, fl, sl, visited)
 	}
 
 	for call, ff := range f.InternalMethodCalls {
+		content := ff.Content
+		if ff.AssociatedStruct != nil && ff.AssociatedStruct.IsInterface {
+			content = ff.AssociatedStruct.Content
+		}
 		s := SingleFunction{
 			CallName: call,
-			Content:  ff.Content,
+			Content:  content,
 		}
 		*fl = append(*fl, s)
-		p.fillRelatedContent(ff, fl, sl)
+		p.fillRelatedContent(depth-1, ff, fl, sl, visited)
+
 		// for method which has been associated with struct, push the struct
 		if ff.AssociatedStruct != nil && ff.AssociatedStruct.Content != "" {
+			st := ff.AssociatedStruct
+			if visited[st.PkgPath] != nil && visited[st.PkgPath][st.Name] {
+				continue
+			} else if visited[st.PkgPath] == nil {
+				visited[st.PkgPath] = map[string]bool{}
+			}
+			visited[st.PkgPath][st.Name] = true
 			ss := SingleStruct{
-				Name:    ff.PkgPath + "." + ff.AssociatedStruct.Name,
-				Content: ff.AssociatedStruct.Content,
+				Name:    ff.PkgPath + "." + st.Name,
+				Content: st.Content,
 			}
 			*sl = append(*sl, ss)
 		}
-		p.fillRelatedContent(ff, fl, sl)
+
+		p.fillRelatedContent(depth-1, ff, fl, sl, visited)
 	}
 }
 
@@ -255,7 +294,7 @@ func main() {
 	}
 
 	// p.generateStruct()
-	m, _ := p.getMain()
+	m, _ := p.getMain(100)
 	m.Dedup()
 
 	out := bytes.NewBuffer(nil)

src/compress/golang/plugin/go_ast_test.go

@@ -3,8 +3,6 @@ package main
 import (
 	"encoding/json"
 	"testing"
-
-	"github.com/davecgh/go-spew/spew"
 )
 
 func Test_goParser_ParseTilTheEnd(t *testing.T) {
@@ -23,10 +21,10 @@ func Test_goParser_ParseTilTheEnd(t *testing.T) {
 		{
 			name: "test",
 			fields: fields{
-				homePageDir: "../../../../testdata/golang",
+				homePageDir: "/Users/bytedance/GOPATH/work/hertz",
 			},
 			args: args{
-				startDir: "./cmd",
+				startDir: "cmd/hz",
 			},
 		},
 	}
@@ -37,16 +35,16 @@ func Test_goParser_ParseTilTheEnd(t *testing.T) {
 			if err != nil {
 				t.Fatalf("goParser.ParseTilTheEnd() error = %v", err)
 			}
-			spew.Dump(p)
-			out, fun := p.getMain()
+			// spew.Dump(p)
+			out, fun := p.getMain(2)
 			if fun.Name != "main" {
 				t.Fail()
 			}
 			out.Dedup()
 			if out, err := json.MarshalIndent(out, "", "  "); err != nil {
 				t.Fatalf("json.Marshal() error = %v", err)
 			} else {
-				println(string(out))
+				println("size:", len(out))
 			}
 		})
 	}

src/compress/golang/plugin/pkg.go

@@ -34,6 +34,7 @@ func (p *goParser) ParseDir(dir string) (err error) {
 	if p.visited[pkgPath] {
 		return nil
 	}
+	p.visited[pkgPath] = true
 
 	// slow-path: load packages in the dir, including sub pakcages
 	fset := token.NewFileSet()