feat: comand write

Author: AsterDY

README.md

@@ -18,20 +18,17 @@
 
 ![ABCoder](images/ABCoder.png)
 
-ABCoder, an AI-oriented code handling tool, is designed to enhance coding-context for Large-Language-Model (LLM).
-
+ABCoder, an AI-oriented code-processing tool, is designed to enhance coding-context for Large-Language-Model (LLM), simplify AI-assisted-coding process.
 
 ## Features
 
--  Universal Abstract Syntax Tree (UniAST), an language-independent and AI-friendly  coding-context, provides ample and recursive code information for AI or programs.
+-  Universal Abstract Syntax Tree (UniAST), an language-independent and AI-friendly  coding-context AST specfication, providing ample and recursive code information for both AI and hunman.
 
 -  Universal Parser, parses abitary languages to UniAST.
 
 -  Univeral Writer, transforms UniAST back to codes.
 
-- (WIP) Code Understanding and Semantic Querying, which can be used to retrieve codes with natural language for either human or AI.
-  
-Based on these features, ABCoder can help developers to easily implement or enhance many AI-assisted coding applications, such as code reviewer, IDE copilot and so on.
+Based on these features, developers can easily implement or enhance their AI-assisted-coding agent or workflow, such as reviewing, optimizing, translating...
 
 ## Getting Started
 
@@ -41,11 +38,12 @@ go install github.com/cloudwego/abcoder@latest
 ```
 2. Use ABCoder to parse a repository to UniAST (JSON)
 ```bash
-abcoder parse <language> <repo-path> > <AST-path>
+abcoder parse {language} {repo-path} > ast.json
 ```
-3. Use ABCoder as a writer
+3. Do your magic with UniAST...
+4. Use ABCoder to write a UniAST back to codes
 ```bash
-abcoder write <AST-path>
+abcoder write {language} ast.json
 ```
 
 ## Universal-Abstract-Syntax-Tree Specification
@@ -57,11 +55,12 @@ see [UniAST Specification](docs/uniast-zh.md)
 
 ABCoder currently supports the following languages:
 
-| Language | Parser | Writer |
-| -------- | ------ | ------ |
-| Go       | ✅      | ✅      |
-| Rust     | ✅      | WIP    |
-| C        | WIP    | ❌      |
+| Language | Parser      | Writer      |
+| -------- | ----------- | ----------- |
+| Go       | ✅           | ✅           |
+| Rust     | ✅           | Coming Soon |
+| Kotlin   | Coming Soon | ❌           |
+| C        | WIP         | ❌           |
 
 
 ## Getting Involved

lang/collect/collect.go

@@ -23,13 +23,13 @@ import (
 	"unicode"
 
 	"github.com/cloudwego/abcoder/lang/log"
-	"github.com/cloudwego/abcoder/lang/lsp"
 	. "github.com/cloudwego/abcoder/lang/lsp"
 	"github.com/cloudwego/abcoder/lang/rust"
+	"github.com/cloudwego/abcoder/lang/uniast"
 )
 
 type CollectOption struct {
-	Language           lsp.Language
+	Language           uniast.Language
 	LoadExternalSymbol bool
 	NeedStdSymbol      bool
 	NoNeedComment      bool
@@ -75,9 +75,9 @@ type functionInfo struct {
 	OutputsSorted    []dependency       `json:"-"`
 }
 
-func switchSpec(l lsp.Language) lsp.LanguageSpec {
+func switchSpec(l uniast.Language) LanguageSpec {
 	switch l {
-	case Rust:
+	case uniast.Rust:
 		return &rust.RustSpec{}
 	default:
 		panic(fmt.Sprintf("unsupported language %s", l))
@@ -94,7 +94,7 @@ func NewCollector(repo string, cli *LSPClient) *Collector {
 		deps:  map[*DocumentSymbol][]dependency{},
 		vars:  map[*DocumentSymbol]dependency{},
 	}
-	if cli.Language == Rust {
+	if cli.Language == uniast.Rust {
 		ret.modPatcher = &rust.RustModulePatcher{Root: repo}
 	}
 	return ret
@@ -522,7 +522,7 @@ func (c *Collector) collectImpl(ctx context.Context, sym *DocumentSymbol, depth
 	}
 	var impl string
 	if fn > 0 && fn < len(sym.Tokens) {
-		impl = lsp.ChunkHead(sym.Text, sym.Location.Range.Start, sym.Tokens[fn].Location.Range.Start)
+		impl = ChunkHead(sym.Text, sym.Location.Range.Start, sym.Tokens[fn].Location.Range.Start)
 	}
 	if impl == "" || len(impl) < len(sym.Name) {
 		impl = sym.Name

lang/golang/writer/write.go

@@ -22,11 +22,14 @@ import (
 	"go/parser"
 	"go/token"
 	"os"
+	"os/exec"
 	"path/filepath"
+	"regexp"
 	"sort"
 	"strconv"
 	"strings"
 
+	"github.com/cloudwego/abcoder/lang/log"
 	"github.com/cloudwego/abcoder/lang/uniast"
 	"github.com/cloudwego/abcoder/lang/utils"
 )
@@ -36,7 +39,7 @@ var _ uniast.Writer = (*Writer)(nil)
 type Options struct {
 	// RepoDir   string
 	// OutDir    string
-	GoVersion string
+	CompilerPath string
 }
 
 type Writer struct {
@@ -55,6 +58,9 @@ type chunk struct {
 }
 
 func NewWriter(opts Options) *Writer {
+	if opts.CompilerPath == "" {
+		opts.CompilerPath = "go"
+	}
 	return &Writer{
 		Options: opts,
 		visited: make(map[string]map[string]*fileNode),
@@ -126,12 +132,16 @@ func (w *Writer) WriteModule(repo *uniast.Repository, modPath string, outDir str
 		}
 	}
 
-	// go mod
+	// create go mod
 	var bs strings.Builder
 	bs.WriteString("module ")
 	bs.WriteString(mod.Name)
 	bs.WriteString("\n\ngo ")
-	bs.WriteString(w.Options.GoVersion)
+	goVersion, err := w.GetGoVersion()
+	if err != nil {
+		goVersion = "1.21"
+	}
+	bs.WriteString(goVersion)
 	bs.WriteString("\n\n")
 	if len(mod.Dependencies) > 0 {
 		bs.WriteString("require (\n")
@@ -151,9 +161,31 @@ func (w *Writer) WriteModule(repo *uniast.Repository, modPath string, outDir str
 		return fmt.Errorf("write go.mod failed: %v", err)
 	}
 
+	// go mod tidy
+	cmd := exec.Command(w.Options.CompilerPath, "mod", "tidy")
+	cmd.Dir = outdir
+	if err := cmd.Run(); err != nil {
+		log.Error("go mod tidy failed: %v", err)
+	}
 	return nil
 }
 
+var goVersionRegex = regexp.MustCompile(`go(\d+\.\d+(\.\d+)?)`)
+
+func (w *Writer) GetGoVersion() (string, error) {
+	cmd := exec.Command(w.Options.CompilerPath, "version")
+	out, err := cmd.Output()
+	if err != nil {
+		return "", fmt.Errorf("get go version failed: %v", err)
+	}
+	// extract with regexp
+	matches := goVersionRegex.FindStringSubmatch(string(out))
+	if len(matches) == 0 {
+		return "", fmt.Errorf("get go version failed: %v", err)
+	}
+	return matches[1], nil
+}
+
 func (w *Writer) appendPackage(repo *uniast.Repository, pkg *uniast.Package) error {
 	for _, v := range pkg.Vars {
 		n := repo.GetNode(v.Identity)

lang/golang/writer/write_test.go

@@ -46,7 +46,7 @@ func TestWriter_WriteRepo(t *testing.T) {
 			name: "test",
 			fields: fields{
 				Options: Options{
-					GoVersion: "1.18",
+					CompilerPath: "1.18",
 				},
 			},
 			args:    args{repo: repo},

lang/lsp/client.go

@@ -24,6 +24,7 @@ import (
 	"time"
 
 	"github.com/cloudwego/abcoder/lang/log"
+	"github.com/cloudwego/abcoder/lang/uniast"
 	lsp "github.com/sourcegraph/go-lsp"
 	"github.com/sourcegraph/jsonrpc2"
 )
@@ -39,7 +40,7 @@ type LSPClient struct {
 
 type ClientOptions struct {
 	Server string
-	Language
+	uniast.Language
 	Verbose bool
 }
 

lang/lsp/spec.go

@@ -15,41 +15,9 @@
 package lsp
 
 import (
-	"strings"
-
 	"github.com/cloudwego/abcoder/lang/uniast"
 )
 
-type Language string
-
-const (
-	Rust    Language = "rust"
-	Golang  Language = "golang"
-	Unknown Language = ""
-)
-
-func NewLanguage(l string) Language {
-	switch strings.ToLower(l) {
-	case "rust":
-		return Rust
-	case "go", "golang":
-		return Golang
-	default:
-		return Unknown
-	}
-}
-
-func (l Language) String() string {
-	switch l {
-	case Rust:
-		return "rust"
-	case Golang:
-		return "go"
-	default:
-		return "unknown"
-	}
-}
-
 type LanguageSpec interface {
 	// initialize a root workspace, and return all modules [modulename=>abs-path] inside
 	WorkSpace(root string) (map[string]string, error)

lang/parse.go

@@ -33,8 +33,11 @@ import (
 	"github.com/cloudwego/abcoder/lang/uniast"
 )
 
+// ParseOptions is the options for parsing the repo.
 type ParseOptions struct {
-	LSP     string
+	// LSP sever executable path
+	LSP string
+	// Language of the repo
 	Verbose bool
 	collect.CollectOption
 }
@@ -43,7 +46,10 @@ func Parse(ctx context.Context, uri string, args ParseOptions) ([]byte, error) {
 	if !filepath.IsAbs(uri) {
 		uri, _ = filepath.Abs(uri)
 	}
-	l, lspPath := checkLSP(args.Language, args.LSP)
+	l, lspPath, err := checkLSP(args.Language, args.LSP)
+	if err != nil {
+		return nil, err
+	}
 	openfile, opentime, err := checkRepoPath(uri, l)
 	if err != nil {
 		return nil, err
@@ -80,12 +86,12 @@ func Parse(ctx context.Context, uri string, args ParseOptions) ([]byte, error) {
 	return out, nil
 }
 
-func checkRepoPath(repoPath string, language lsp.Language) (openfile string, wait time.Duration, err error) {
+func checkRepoPath(repoPath string, language uniast.Language) (openfile string, wait time.Duration, err error) {
 	if _, err := os.Stat(repoPath); os.IsNotExist(err) {
 		return "", 0, fmt.Errorf("repository not found: %s", repoPath)
 	}
 	switch language {
-	case lsp.Rust:
+	case uniast.Rust:
 		// NOTICE: open the Cargo.toml file is required for Rust projects
 		openfile, wait = rust.CheckRepo(repoPath)
 	default:
@@ -97,40 +103,29 @@ func checkRepoPath(repoPath string, language lsp.Language) (openfile string, wai
 	return
 }
 
-func checkVerbose(verbose bool, debug bool) {
-	if debug {
-
-	} else if verbose {
-		log.SetLogLevel(log.InfoLevel)
-	} else {
-		log.SetLogLevel(log.ErrorLevel)
-	}
-}
-
-func checkLSP(language lsp.Language, lspPath string) (l lsp.Language, s string) {
+func checkLSP(language uniast.Language, lspPath string) (l uniast.Language, s string, err error) {
 	switch language {
-	case lsp.Rust:
+	case uniast.Rust:
 		l, s = rust.GetDefaultLSP()
-	case lsp.Golang:
-		l = lsp.Golang
+	case uniast.Golang:
+		l = uniast.Golang
 		s = ""
 		if _, err := exec.LookPath("go"); err != nil {
 			if _, err := os.Stat(lspPath); os.IsNotExist(err) {
 				log.Error("Go compiler not found, please make it excutable!\n", lspPath)
-				os.Exit(1)
+				return uniast.Unknown, "", err
 			}
 		}
 		return
 	default:
-		log.Error("Unsupported language: %s\n", language)
-		os.Exit(1)
+		return uniast.Unknown, "", fmt.Errorf("unsupported language: %s", language)
 	}
 	// check if lsp excutable
 	if lspPath != "" {
 		if _, err := exec.LookPath(lspPath); err != nil {
 			if _, err := os.Stat(lspPath); os.IsNotExist(err) {
 				log.Error("Language server %s not found, please make it excutable!\n", lspPath)
-				os.Exit(1)
+				return uniast.Unknown, "", err
 			}
 		}
 		s = lspPath
@@ -140,7 +135,7 @@ func checkLSP(language lsp.Language, lspPath string) (l lsp.Language, s string)
 }
 
 func collectSymbol(ctx context.Context, cli *lsp.LSPClient, repoPath string, opts collect.CollectOption) (repo *uniast.Repository, err error) {
-	if opts.Language == lsp.Golang {
+	if opts.Language == uniast.Golang {
 		repo, err = callGoParser(ctx, repoPath, opts)
 		if err != nil {
 			return nil, err

lang/rust/repo.go

@@ -23,7 +23,7 @@ import (
 	"time"
 
 	"github.com/cloudwego/abcoder/lang/log"
-	"github.com/cloudwego/abcoder/lang/lsp"
+	"github.com/cloudwego/abcoder/lang/uniast"
 	"github.com/cloudwego/abcoder/lang/utils"
 )
 
@@ -73,8 +73,8 @@ next:
 	return openfile, wait
 }
 
-func GetDefaultLSP() (lang lsp.Language, name string) {
-	return lsp.Rust, "rust-analyzer"
+func GetDefaultLSP() (lang uniast.Language, name string) {
+	return uniast.Rust, "rust-analyzer"
 }
 
 func GetLastCommitTime(repo string) time.Time {

lang/uniast/ast.go

@@ -28,9 +28,20 @@ type Language string
 const (
 	Golang  Language = "go"
 	Rust    Language = "rust"
-	Unknown Language = "unknown"
+	Unknown Language = ""
 )
 
+func (l Language) String() string {
+	switch l {
+	case Rust:
+		return "rust"
+	case Golang:
+		return "go"
+	default:
+		return string(l)
+	}
+}
+
 func NewLanguage(lang string) (l Language) {
 	// sp := strings.Split(lang, "@")
 	// if len(sp) > 1 {