opt: define Import

Author: AsterDY

src/compress/golang/plugin/parse/ctx.go

@@ -267,7 +267,7 @@ type importInfo struct {
 	SysImports        map[string]string
 	ProjectImports    map[string]string
 	ThirdPartyImports map[string][2]string // 0-mod, 1-import
-	Origins           []string
+	Origins           []Import
 }
 
 func (p *GoParser) mockTypes(typ ast.Expr, m map[string]Identity, file []byte, fset *token.FileSet, fpath string, mod string, pkg string, impts *importInfo) (name string, isPointer bool) {

src/compress/golang/plugin/parse/file.go

@@ -237,16 +237,16 @@ func (p *GoParser) parseSelector(ctx *fileContext, expr *ast.SelectorExpr, infos
 						// 	// fmt.Fprintf(os.Stderr, "failed to get type id for %s\n", expr.Name)
 						// 	return false
 						// }
-						*infos.tys = Dedup(*infos.tys, dep)
+						*infos.tys = InsertDependency(*infos.tys, dep)
 						// global var
 					} else if _, ok := v.(*types.Const); ok {
-						*infos.globalVars = Dedup(*infos.globalVars, dep)
+						*infos.globalVars = InsertDependency(*infos.globalVars, dep)
 						// external const
 					} else if _, ok := v.(*types.Var); ok {
-						*infos.globalVars = Dedup(*infos.globalVars, dep)
+						*infos.globalVars = InsertDependency(*infos.globalVars, dep)
 						// external function
 					} else if _, ok := v.(*types.Func); ok {
-						*infos.functionCalls = Dedup(*infos.functionCalls, dep)
+						*infos.functionCalls = InsertDependency(*infos.functionCalls, dep)
 					}
 					return false
 				}
@@ -289,7 +289,7 @@ func (p *GoParser) parseSelector(ctx *fileContext, expr *ast.SelectorExpr, infos
 		if err := p.referCodes(ctx, &id, p.opts.ReferCodeDepth); err != nil {
 			fmt.Fprintf(os.Stderr, "failed to get refer code for %s: %v\n", id.Name, err)
 		}
-		*infos.methodCalls = Dedup(*infos.methodCalls, dep)
+		*infos.methodCalls = InsertDependency(*infos.methodCalls, dep)
 		return false
 	}
 
@@ -380,24 +380,24 @@ func (p *GoParser) parseFunc(ctx *fileContext, funcDecl *ast.FuncDecl) (*Functio
 					// 	// fmt.Fprintf(os.Stderr, "failed to get type id for %s\n", expr.Name)
 					// 	return false
 					// }
-					tys = Dedup(tys, dep)
+					tys = InsertDependency(tys, dep)
 					// global var
 				} else if v, ok := use.(*types.Var); ok {
 					// NOTICE: the Parent of global scope is nil?
 					if isPkgScope(v.Parent()) {
-						globalVars = Dedup(globalVars, dep)
+						globalVars = InsertDependency(globalVars, dep)
 					}
 					// global const
 				} else if c, ok := use.(*types.Const); ok {
 					if isPkgScope(c.Parent()) {
-						globalVars = Dedup(globalVars, dep)
+						globalVars = InsertDependency(globalVars, dep)
 					}
 					return false
 					// function
 				} else if f, ok := use.(*types.Func); ok {
 					// exclude method
 					if f.Type().(*types.Signature).Recv() == nil {
-						functionCalls = Dedup(functionCalls, dep)
+						functionCalls = InsertDependency(functionCalls, dep)
 					}
 				}
 			}

src/compress/golang/plugin/parse/pkg.go

@@ -33,14 +33,15 @@ func (p *GoParser) parseImports(fset *token.FileSet, file []byte, mod *Module, i
 	sysImports := make(map[string]string)
 	ret := &importInfo{}
 	for _, imp := range impts {
-		ret.Origins = append(ret.Origins, string(GetRawContent(fset, file, imp, p.opts.CollectComment)))
 		importPath := imp.Path.Value[1 : len(imp.Path.Value)-1] // remove the quotes
 		importAlias := ""
 		// Check if user has defined an alias for current import
 		if imp.Name != nil {
 			importAlias = imp.Name.Name // update the alias
+			ret.Origins = append(ret.Origins, Import{Path: imp.Path.Value, Alias: &importAlias})
 		} else {
 			importAlias = getPackageAlias(importPath)
+			ret.Origins = append(ret.Origins, Import{Path: imp.Path.Value})
 		}
 
 		// Fix: module name may also be like this?

src/lang/collect/export.go

@@ -178,7 +178,7 @@ func (c *Collector) exportSymbol(repo *uniast.Repository, symbol *DocumentSymbol
 					continue
 				}
 				dep := uniast.NewDependency(*tyid, c.fileLine(input.Location))
-				obj.Types = uniast.Dedup(obj.Types, dep)
+				obj.Types = uniast.InsertDependency(obj.Types, dep)
 			}
 		}
 		if info.Inputs != nil {
@@ -190,7 +190,7 @@ func (c *Collector) exportSymbol(repo *uniast.Repository, symbol *DocumentSymbol
 					continue
 				}
 				dep := uniast.NewDependency(*tyid, c.fileLine(input.Location))
-				obj.Params = uniast.Dedup(obj.Params, dep)
+				obj.Params = uniast.InsertDependency(obj.Params, dep)
 			}
 		}
 		if info.Outputs != nil {
@@ -202,7 +202,7 @@ func (c *Collector) exportSymbol(repo *uniast.Repository, symbol *DocumentSymbol
 					continue
 				}
 				dep := uniast.NewDependency(*tyid, c.fileLine(output.Location))
-				obj.Results = uniast.Dedup(obj.Results, dep)
+				obj.Results = uniast.InsertDependency(obj.Results, dep)
 			}
 		}
 		if info.Method != nil && info.Method.Receiver.Symbol != nil {
@@ -253,23 +253,23 @@ func (c *Collector) exportSymbol(repo *uniast.Repository, symbol *DocumentSymbol
 				pdep := uniast.NewDependency(*depid, c.fileLine(dep.Location))
 				switch dep.Symbol.Kind {
 				case lsp.SKFunction:
-					obj.FunctionCalls = uniast.Dedup(obj.FunctionCalls, pdep)
+					obj.FunctionCalls = uniast.InsertDependency(obj.FunctionCalls, pdep)
 				case lsp.SKMethod:
 					if obj.MethodCalls == nil {
 						obj.MethodCalls = make([]uniast.Dependency, 0, len(deps))
 					}
 					// NOTICE: use loc token as key here, to make it more readable
-					obj.MethodCalls = uniast.Dedup(obj.MethodCalls, pdep)
+					obj.MethodCalls = uniast.InsertDependency(obj.MethodCalls, pdep)
 				case lsp.SKVariable, lsp.SKConstant:
 					if obj.GlobalVars == nil {
 						obj.GlobalVars = make([]uniast.Dependency, 0, len(deps))
 					}
-					obj.GlobalVars = uniast.Dedup(obj.GlobalVars, pdep)
+					obj.GlobalVars = uniast.InsertDependency(obj.GlobalVars, pdep)
 				case lsp.SKStruct, lsp.SKTypeParameter, lsp.SKInterface, lsp.SKEnum:
 					if obj.Types == nil {
 						obj.Types = make([]uniast.Dependency, 0, len(deps))
 					}
-					obj.Types = uniast.Dedup(obj.Types, pdep)
+					obj.Types = uniast.InsertDependency(obj.Types, pdep)
 				default:
 					log.Error("dep symbol %s not collected for %v\n", dep.Symbol, id)
 				}

src/lang/golang/writer/ast.go

@@ -0,0 +1,74 @@
+/**
+ * Copyright 2025 ByteDance Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package writer
+
+import (
+	"strings"
+
+	"github.com/cloudwego/abcoder/src/uniast"
+)
+
+func writeImport(sb *strings.Builder, impts []uniast.Import) {
+	if len(impts) == 0 {
+		return
+	}
+	sb.WriteString("import ")
+	if len(impts) == 1 {
+		writeSingleImport(sb, impts[0])
+		return
+	}
+	sb.WriteString("(\n")
+	for i := 0; i < len(impts); i++ {
+		sb.WriteString("\t")
+		writeSingleImport(sb, impts[i])
+	}
+	sb.WriteString(")\n")
+}
+
+func writeSingleImport(sb *strings.Builder, v uniast.Import) {
+	if v.Alias != nil {
+		sb.WriteString(*v.Alias)
+		sb.WriteString(" ")
+	}
+	sb.WriteString(v.Path)
+	sb.WriteString("\n")
+}
+
+// merge the imports of file and nodes, and return the merged imports
+// file is in priority (because it contains alias)
+func mergeImports(priors []uniast.Import, subs []uniast.Import) (ret []uniast.Import) {
+	visited := make(map[string]bool, len(priors)+len(subs))
+	ret = make([]uniast.Import, 0, len(priors)+len(subs))
+	for _, v := range priors {
+
+		if visited[v.Path] {
+			continue
+		} else {
+			visited[v.Path] = true
+			ret = append(ret, v)
+		}
+	}
+	for _, v := range subs {
+		if visited[v.Path] {
+			continue
+		} else {
+			visited[v.Path] = true
+			ret = append(ret, v)
+		}
+	}
+	return
+}

src/lang/golang/writer/write.go

@@ -27,12 +27,14 @@ import (
 	"strconv"
 	"strings"
 
+	"github.com/cloudwego/abcoder/src/lang/utils"
 	"github.com/cloudwego/abcoder/src/uniast"
 )
 
 var _ uniast.Writer = (*Writer)(nil)
 
 type Options struct {
+	RepoDir   string
 	OutDir    string
 	GoVersion string
 }
@@ -44,7 +46,7 @@ type Writer struct {
 
 type fileNode struct {
 	chunks []chunk
-	impts  []string
+	impts  []uniast.Import
 }
 
 type chunk struct {
@@ -101,19 +103,13 @@ func (w *Writer) WriteModule(repo *uniast.Repository, modPath string) error {
 			}
 			sb.WriteString("\n\n")
 
-			var fimpts []string
+			var fimpts []uniast.Import
 			if fi, ok := mod.Files[filepath.Join(mod.Dir, rel, fpath)]; ok && fi.Imports != nil {
 				fimpts = fi.Imports
 			}
-			impts := w.mergeImports(fimpts, f.impts)
+			impts := mergeImports(fimpts, f.impts)
 			if len(impts) > 0 {
-				sb.WriteString("import (\n")
-				for _, v := range impts {
-					sb.WriteString("\t")
-					sb.WriteString(v)
-					sb.WriteString("\n")
-				}
-				sb.WriteString(")\n\n")
+				writeImport(&sb, impts)
 			}
 
 			sort.SliceStable(f.chunks, func(i, j int) bool {
@@ -208,15 +204,15 @@ func (w *Writer) appendNode(node *uniast.Node, pkg string, isMain bool, file str
 	if fs == nil {
 		fs = &fileNode{
 			chunks: make([]chunk, 0, len(node.Dependencies)),
-			impts:  make([]string, 0, len(node.Dependencies)),
+			impts:  make([]uniast.Import, 0, len(node.Dependencies)),
 		}
 		p[fpath] = fs
 	}
 	for _, v := range node.Dependencies {
 		if v.Target.PkgPath == "" || v.Target.PkgPath == pkg {
 			continue
 		}
-		fs.impts = append(fs.impts, strconv.Quote(v.Target.PkgPath))
+		fs.impts = append(fs.impts, uniast.Import{Path: strconv.Quote(v.Target.PkgPath)})
 	}
 
 	// 检查是否有imports
@@ -235,19 +231,19 @@ func (w *Writer) appendNode(node *uniast.Node, pkg string, isMain bool, file str
 }
 
 // receive a piece of golang code, parse it and splits the imports and codes
-func (w Writer) SplitImportsAndCodes(src string) (codes string, imports []string, err error) {
+func (w Writer) SplitImportsAndCodes(src string) (codes string, imports []uniast.Import, err error) {
 	fset := token.NewFileSet()
 	f, err := parser.ParseFile(fset, "", src, parser.SkipObjectResolution)
 	if err != nil {
 		// NOTICE: if parse failed, just return the src
 		return src, nil, nil
 	}
 	for _, imp := range f.Imports {
-		var impt = imp.Path.Value
+		var alias string
 		if imp.Name != nil {
-			impt = fmt.Sprintf("%s %s", imp.Name.Name, impt)
+			alias = imp.Name.Name
 		}
-		imports = append(imports, impt)
+		imports = append(imports, uniast.Import{Path: imp.Path.Value, Alias: &alias})
 	}
 	start := 0
 	for _, s := range f.Decls {
@@ -260,42 +256,58 @@ func (w Writer) SplitImportsAndCodes(src string) (codes string, imports []string
 	return src[start:], imports, nil
 }
 
-func (w *Writer) IdToImport(id uniast.Identity) (string, error) {
-	return strconv.Quote(id.PkgPath), nil
+func (w *Writer) IdToImport(id uniast.Identity) (uniast.Import, error) {
+	return uniast.Import{Path: strconv.Quote(id.PkgPath)}, nil
 }
 
-// merge the imports of file and nodes, and return the merged imports
-// file is in priority (because it contains alias)
-func (w *Writer) mergeImports(priors []string, subs []string) (ret []string) {
-	visited := make(map[string]bool, len(priors)+len(subs))
-	ret = make([]string, 0, len(priors)+len(subs))
-	for _, v := range priors {
-		sp := strings.Split(v, " ")
-		var impt = sp[0]
-		if len(sp) >= 2 {
-			impt = sp[1]
+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)
+	}
+
+	fs := token.NewFileSet()
+	f, err := parser.ParseFile(fs, file.Path, bs, parser.ImportsOnly)
+	if err != nil {
+		return nil, utils.WrapError(err, "fail parse file %s", file.Path)
+	}
+
+	old := make([]uniast.Import, 0, len(f.Imports))
+	for _, imp := range f.Imports {
+		i := uniast.Import{
+			Path: imp.Path.Value,
 		}
-		key, _ := strconv.Unquote(impt)
-		if visited[key] {
-			continue
-		} else {
-			visited[key] = true
-			ret = append(ret, v)
+		if imp.Name != nil {
+			tmp := imp.Name.Name
+			i.Alias = &tmp
 		}
+		old = append(old, i)
+	}
+
+	impts := mergeImports(old, file.Imports)
+	if len(impts) == len(old) {
+		return bs, nil
 	}
-	for _, v := range subs {
-		sp := strings.Split(v, " ")
-		var impt = sp[0]
-		if len(sp) >= 2 {
-			impt = sp[1]
+
+	var sb strings.Builder
+	writeImport(&sb, impts)
+	final := sb.String()
+
+	imptStart := fs.Position(f.Name.End()).Offset + 1
+	if len(f.Imports) > 0 {
+		for imptStart < len(bs) && bs[imptStart] != 'i' {
+			imptStart++
 		}
-		key, _ := strconv.Unquote(impt)
-		if visited[key] {
-			continue
-		} else {
-			visited[key] = true
-			ret = append(ret, v)
+	}
+	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] != ')') {
+			imptEnd++
 		}
+		imptEnd += 2 // for `)`
 	}
-	return
+	r1 := append(bs[:imptStart:imptStart], final...)
+	ret := append(r1, bs[imptEnd:]...)
+	return ret, nil
 }