feat: update golang writer (#10)

* tmp

* feat:(go_ast) support collect test symbols

* opt:(patcher) support add new node

* feat: support bind package with file

* update

* fix: unmarshal `NodeType`

Author: AsterDY

src/compress/types/types.rs

@@ -78,7 +78,9 @@ pub struct Relation {
     #[serde(rename = "Kind")]
     pub(crate) kind: RelationKind,
     #[serde(rename = "Desc")]
-    pub(crate) desc: String,
+    pub(crate) desc: Option<String>,
+    #[serde(rename = "Codes")]
+    pub(crate) codes: Option<String>,
 }
 
 impl Relation {
@@ -91,7 +93,7 @@ impl Relation {
     }
 }
 
-#[derive(Serialize, Debug, Clone, Default)]
+#[derive(Debug, Clone, Default)]
 pub enum NodeType {
     #[default]
     Unknown,
@@ -111,6 +113,15 @@ impl NodeType {
     }
 }
 
+impl Serialize for NodeType {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        serializer.serialize_str(&self.to_string())
+    }
+}
+
 impl<'de> Deserialize<'de> for NodeType {
     fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
     where

src/lang/collect/collect.go

@@ -32,6 +32,7 @@ type CollectOption struct {
 	LoadExternalSymbol bool
 	NeedStdSymbol      bool
 	NoNeedComment      bool
+	NeedTest           bool
 	Language           lsp.Language
 	Excludes           []string
 }

src/lang/collect/export.go

@@ -49,8 +49,7 @@ func (c *Collector) fileLine(loc Location) uniast.FileLine {
 }
 
 func newModule(name string, dir string) *uniast.Module {
-	ret := uniast.NewModule(name, dir)
-	ret.Language = uniast.Rust
+	ret := uniast.NewModule(name, dir, uniast.Rust)
 	return ret
 }
 

src/lang/go.go

@@ -32,6 +32,9 @@ func callGoParser(ctx context.Context, repoPath string, opts collect.CollectOpti
 	if !opts.NoNeedComment {
 		goopts.CollectComment = true
 	}
+	if opts.NeedTest {
+		goopts.NeedTest = true
+	}
 	goopts.Excludes = opts.Excludes
 	p := parser.NewParser(repoPath, repoPath, goopts)
 	repo, err := p.ParseRepo()

src/lang/golang/parser/ctx.go

@@ -18,6 +18,7 @@ import (
 	"bytes"
 	"fmt"
 	"go/ast"
+	"go/parser"
 	"go/token"
 	"go/types"
 	"os"
@@ -51,9 +52,12 @@ func isExternalID(id *Identity, curmod string) bool {
 		strings.Contains(id.PkgPath, "/kitex_gen/") || strings.Contains(id.PkgPath, "/hertz_gen/")
 }
 
-func newModule(mod string, dir string) *Module {
-	ret := uniast.NewModule(mod, dir)
-	ret.Language = Golang
+const (
+	StdLanguage = "go"
+)
+
+func newModule(mod string, dir string) (ret *Module) {
+	ret = uniast.NewModule(mod, dir, Golang)
 	return ret
 }
 
@@ -82,9 +86,12 @@ func (p *GoParser) referCodes(ctx *fileContext, id *Identity, depth int) (err er
 	if pkg == nil {
 		return fmt.Errorf("cannot find package %s", id.PkgPath)
 	}
-	for i, fpath := range pkg.GoFiles {
-		file := pkg.Syntax[i]
+	for _, fpath := range pkg.GoFiles {
 		bs := p.getFileBytes(fpath)
+		file, err := parser.ParseFile(pkg.Fset, fpath, bs, parser.ParseComments)
+		if err != nil {
+			return err
+		}
 		impts, e := p.parseImports(pkg.Fset, bs, mod, file.Imports)
 		if e != nil {
 			err = e

src/lang/golang/parser/file.go

@@ -29,13 +29,13 @@ import (
 func (p *GoParser) parseFile(ctx *fileContext, f *ast.File) error {
 	cont := true
 	ast.Inspect(f, func(node ast.Node) bool {
-		defer func() {
-			if r := recover(); r != nil {
-				fmt.Fprintf(os.Stderr, "panic: %v in %s:%d\n", r, ctx.filePath, ctx.fset.Position(node.Pos()).Line)
-				cont = false
-				return
-			}
-		}()
+		// defer func() {
+		// 	if r := recover(); r != nil {
+		// 		fmt.Fprintf(os.Stderr, "panic: %v in %s:%d\n", r, ctx.filePath, ctx.fset.Position(node.Pos()).Line)
+		// 		cont = false
+		// 		return
+		// 	}
+		// }()
 		if funcDecl, ok := node.(*ast.FuncDecl); ok {
 			// parse funcs
 			_, ct := p.parseFunc(ctx, funcDecl)

src/lang/golang/parser/option.go

@@ -24,6 +24,7 @@ type Options struct {
 	ReferCodeDepth int
 	Excludes       []string
 	CollectComment bool
+	NeedTest       bool
 }
 
 // type Option func(options *Options)

src/lang/golang/parser/pkg.go

@@ -152,22 +152,26 @@ func (p *GoParser) loadPackages(mod *Module, dir string, pkgPath PkgPath) (err e
 		return nil
 	}
 	fmt.Fprintf(os.Stderr, "[loadPackages] mod: %s, dir: %s, pkgPath: %s", mod.Name, dir, pkgPath)
+	fset := token.NewFileSet()
 	loadCount++
 	// slow-path: load packages in the dir, including sub pakcages
 	opts := packages.NeedFiles | packages.NeedSyntax | packages.NeedTypes | packages.NeedTypesInfo | packages.NeedImports
-	if p.opts.ReferCodeDepth != 0 {
-		opts |= packages.NeedDeps
-	}
-	fset := token.NewFileSet()
-	pkgs, err := packages.Load(&packages.Config{
+	cfg := &packages.Config{
 		Mode: opts,
 		Fset: fset,
 		Dir:  dir,
-	}, pkgPath)
+	}
+	if p.opts.ReferCodeDepth != 0 {
+		opts |= packages.NeedDeps
+	}
+	if p.opts.NeedTest {
+		opts |= packages.NeedForTest
+		cfg.Tests = true
+	}
+	pkgs, err := packages.Load(cfg, pkgPath)
 	if err != nil {
 		return fmt.Errorf("load path '%s' failed: %v", dir, err)
 	}
-
 	for _, pkg := range pkgs {
 		if mm := p.repo.Modules[mod.Name]; mm != nil && (*mm).Packages[pkg.ID] != nil {
 			continue
@@ -207,6 +211,8 @@ func (p *GoParser) loadPackages(mod *Module, dir string, pkgPath PkgPath) (err e
 				f = NewFile(relpath)
 				mod.Files[relpath] = f
 			}
+			pkgid := pkg.ID
+			f.Package = &pkgid
 			f.Imports = imports.Origins
 			if err := p.parseFile(ctx, file); err != nil {
 				return err
@@ -221,7 +227,17 @@ func (p *GoParser) loadPackages(mod *Module, dir string, pkgPath PkgPath) (err e
 			// 	obj.Dependencies = append(obj.Dependencies, imp.ID)
 			// }
 			obj.PkgPath = pkg.ID
+			if strings.HasSuffix(obj.PkgPath, ".test]") {
+				obj.IsTest = true
+			}
+			if strings.HasSuffix(obj.PkgPath, ".test") {
+				delete(mod.Packages, obj.PkgPath)
+			}
 		}
 	}
 	return
 }
+
+func IsTestPackage(pkgPath string) bool {
+	return strings.HasSuffix(pkgPath, ".test") || strings.HasSuffix(pkgPath, ".test]")
+}

src/lang/golang/parser/pkg_test.go

@@ -50,6 +50,7 @@ func Test_goParser_ParseRepo(t *testing.T) {
 			println(abs)
 			p := newGoParser(tt.fields.modName, tt.fields.homePageDir, Options{
 				ReferCodeDepth: 1,
+				NeedTest:       true,
 			})
 			r, err := p.ParseRepo()
 			if err != nil {
@@ -63,15 +64,15 @@ func Test_goParser_ParseRepo(t *testing.T) {
 			}
 			_ = pj
 			_ = os.WriteFile("ast.json", pj, 0644)
-			n, err := p.getNode(NewIdentity("github.com/cloudwego/localsession", "github.com/cloudwego/localsession/backup", "RecoverCtxOndemands"))
+			n, err := p.getNode(NewIdentity("github.com/cloudwego/localsession", "github.com/cloudwego/localsession/backup", "RecoverCtxOnDemands"))
 			if err != nil {
 				t.Fatal(err)
 			}
 			jf, err := json.MarshalIndent(n, "", "  ")
 			if err != nil {
 				t.Fatalf("json.Marshal() error = %v", err)
 			}
-			println(string(jf))
+			os.WriteFile("node.json", jf, 0644)
 		})
 	}
 }

src/lang/golang/writer/write.go

@@ -34,8 +34,8 @@ import (
 var _ uniast.Writer = (*Writer)(nil)
 
 type Options struct {
-	RepoDir   string
-	OutDir    string
+	// RepoDir   string
+	// OutDir    string
 	GoVersion string
 }
 
@@ -61,19 +61,19 @@ func NewWriter(opts Options) *Writer {
 	}
 }
 
-func (w *Writer) WriteRepo(repo *uniast.Repository) error {
+func (w *Writer) WriteRepo(repo *uniast.Repository, outDir string) error {
 	for m, mod := range repo.Modules {
 		if strings.Contains(m, "@") {
 			continue
 		}
-		if err := w.WriteModule(repo, m); err != nil {
+		if err := w.WriteModule(repo, m, outDir); err != nil {
 			return fmt.Errorf("write module %s failed: %v", mod.Name, err)
 		}
 	}
 	return nil
 }
 
-func (w *Writer) WriteModule(repo *uniast.Repository, modPath string) error {
+func (w *Writer) WriteModule(repo *uniast.Repository, modPath string, outDir string) error {
 	mod := repo.Modules[modPath]
 	if mod == nil {
 		return fmt.Errorf("module %s not found", modPath)
@@ -84,7 +84,7 @@ func (w *Writer) WriteModule(repo *uniast.Repository, modPath string) error {
 		}
 	}
 
-	outdir := filepath.Join(w.Options.OutDir, mod.Dir)
+	outdir := filepath.Join(outDir, mod.Dir)
 	for dir, pkg := range w.visited {
 		rel := strings.TrimPrefix(dir, mod.Name)
 		pkgDir := filepath.Join(outdir, rel)
@@ -260,16 +260,12 @@ func (w *Writer) IdToImport(id uniast.Identity) (uniast.Import, error) {
 	return uniast.Import{Path: strconv.Quote(id.PkgPath)}, nil
 }
 
-func (p *Writer) PatchImports(file *uniast.File) ([]byte, error) {
-	bs, err := os.ReadFile(filepath.Join(p.Options.RepoDir, file.Path))
-	if err != nil {
-		return nil, utils.WrapError(err, "fail read file %s", file.Path)
-	}
+func (p *Writer) PatchImports(impts []uniast.Import, file []byte) ([]byte, error) {
 
 	fs := token.NewFileSet()
-	f, err := parser.ParseFile(fs, file.Path, bs, parser.ImportsOnly)
+	f, err := parser.ParseFile(fs, "default.go", file, parser.ImportsOnly)
 	if err != nil {
-		return nil, utils.WrapError(err, "fail parse file %s", file.Path)
+		return nil, utils.WrapError(err, "fail parse file %s", file)
 	}
 
 	old := make([]uniast.Import, 0, len(f.Imports))
@@ -284,9 +280,9 @@ func (p *Writer) PatchImports(file *uniast.File) ([]byte, error) {
 		old = append(old, i)
 	}
 
-	impts := mergeImports(old, file.Imports)
+	impts = mergeImports(old, impts)
 	if len(impts) == len(old) {
-		return bs, nil
+		return file, nil
 	}
 
 	var sb strings.Builder
@@ -295,19 +291,47 @@ func (p *Writer) PatchImports(file *uniast.File) ([]byte, error) {
 
 	imptStart := fs.Position(f.Name.End()).Offset + 1
 	if len(f.Imports) > 0 {
-		for imptStart < len(bs) && bs[imptStart] != 'i' {
+		for imptStart < len(file) && file[imptStart] != 'i' {
 			imptStart++
 		}
 	}
 	imptEnd := imptStart
 	if len(f.Imports) > 1 {
 		imptEnd = fs.Position(f.Imports[len(f.Imports)-1].End()).Offset
-		for len(old) > 1 && imptEnd < len(bs) && (bs[imptEnd] != ')') {
+		for len(old) > 1 && imptEnd < len(file) && (file[imptEnd] != ')') {
 			imptEnd++
 		}
 		imptEnd += 2 // for `)`
 	}
-	r1 := append(bs[:imptStart:imptStart], final...)
-	ret := append(r1, bs[imptEnd:]...)
+	r1 := append(file[:imptStart:imptStart], final...)
+	ret := append(r1, file[imptEnd:]...)
 	return ret, nil
 }
+
+func (p *Writer) CreateFile(fi *uniast.File, mod *uniast.Module) ([]byte, error) {
+	var sb strings.Builder
+	sb.WriteString("package ")
+	pkgName := filepath.Base(filepath.Dir(fi.Path))
+	if fi.Package != nil {
+		pkg := mod.Packages[*fi.Package]
+		if pkg != nil {
+			if pkg.IsMain {
+				pkgName = "main"
+			} else {
+				pkgName = filepath.Base(pkg.PkgPath)
+			}
+		}
+	}
+	if pkgName == "" {
+		return nil, fmt.Errorf("package name is empty")
+	}
+	sb.WriteString(pkgName)
+	sb.WriteString("\n\n")
+
+	if len(fi.Imports) > 0 {
+		writeImport(&sb, fi.Imports)
+	}
+
+	bs := sb.String()
+	return []byte(bs), nil
+}