From 75763735c424df99feea72acd7a875ac180752d9 Mon Sep 17 00:00:00 2001
From: oldme <45782393+oldme-git@users.noreply.github.com>
Date: Mon, 15 Apr 2024 20:16:44 +0800
Subject: [PATCH] enhance: improve command `gf gen ctrl` using AST for parsing
 `DstFolder` (#3478)

---
 .../internal/cmd/genctrl/genctrl_calculate.go | 44 ++++++++++++++-----
 .../cmd/genctrl/genctrl_generate_ctrl.go      |  2 +-
 .../genctrl/genctrl_generate_ctrl_clear.go    |  1 +
 .../cmd/genctrl/genctrl_generate_sdk.go       |  2 +
 4 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/cmd/gf/internal/cmd/genctrl/genctrl_calculate.go b/cmd/gf/internal/cmd/genctrl/genctrl_calculate.go
index 21021d798..f794559c9 100644
--- a/cmd/gf/internal/cmd/genctrl/genctrl_calculate.go
+++ b/cmd/gf/internal/cmd/genctrl/genctrl_calculate.go
@@ -69,26 +69,22 @@ func (c CGenCtrl) getApiItemsInDst(dstFolder string) (items []apiItem, err error
 		Path  string
 		Alias string
 	}
-	var fileContent string
 	filePaths, err := gfile.ScanDir(dstFolder, "*.go", true)
 	if err != nil {
 		return nil, err
 	}
 	for _, filePath := range filePaths {
-		fileContent = gfile.GetContents(filePath)
-		match, err := gregex.MatchString(`import\s+\(([\s\S]+?)\)`, fileContent)
-		if err != nil {
-			return nil, err
-		}
-		if len(match) < 2 {
-			continue
-		}
 		var (
 			array       []string
 			importItems []importItem
-			importLines = gstr.SplitAndTrim(match[1], "\n")
+			importLines []string
 			module      = gfile.Basename(gfile.Dir(filePath))
 		)
+		importLines, err = c.getImportsInDst(filePath)
+		if err != nil {
+			return nil, err
+		}
+
 		// retrieve all imports.
 		for _, importLine := range importLines {
 			array = gstr.SplitAndTrim(importLine, " ")
@@ -104,11 +100,15 @@ func (c CGenCtrl) getApiItemsInDst(dstFolder string) (items []apiItem, err error
 			}
 		}
 		// retrieve all api usages.
+		// retrieve it without using AST, but use regular expressions to retrieve.
+		// It's because the api definition is simple and regular.
+		// Use regular expressions to get better performance.
+		fileContent := gfile.GetContents(filePath)
 		matches, err := gregex.MatchAllString(PatternCtrlDefinition, fileContent)
 		if err != nil {
 			return nil, err
 		}
-		for _, match = range matches {
+		for _, match := range matches {
 			// try to find the import path of the api.
 			var (
 				importPath string
@@ -177,3 +177,25 @@ func (c CGenCtrl) getStructsNameInSrc(filePath string) (structsName []string, er
 
 	return
 }
+
+// getImportsInDst retrieves all import paths in the file.
+func (c CGenCtrl) getImportsInDst(filePath string) (imports []string, err error) {
+	var (
+		fileContent = gfile.GetContents(filePath)
+		fileSet     = token.NewFileSet()
+	)
+
+	node, err := parser.ParseFile(fileSet, "", fileContent, parser.ParseComments)
+	if err != nil {
+		return
+	}
+
+	ast.Inspect(node, func(n ast.Node) bool {
+		if imp, ok := n.(*ast.ImportSpec); ok {
+			imports = append(imports, imp.Path.Value)
+		}
+		return true
+	})
+
+	return
+}
diff --git a/cmd/gf/internal/cmd/genctrl/genctrl_generate_ctrl.go b/cmd/gf/internal/cmd/genctrl/genctrl_generate_ctrl.go
index 19d3dca15..5a71c2972 100644
--- a/cmd/gf/internal/cmd/genctrl/genctrl_generate_ctrl.go
+++ b/cmd/gf/internal/cmd/genctrl/genctrl_generate_ctrl.go
@@ -141,7 +141,7 @@ func (c *controllerGenerator) doGenerateCtrlItem(dstModuleFolderPath string, ite
 			"{MethodName}": item.MethodName,
 		})
 
-		if gstr.Contains(gfile.GetContents(methodFilePath), fmt.Sprintf(`func (c *%v) %v`, ctrlName, item.MethodName)) {
+		if gstr.Contains(gfile.GetContents(methodFilePath), fmt.Sprintf(`func (c *%v) %v(`, ctrlName, item.MethodName)) {
 			return
 		}
 		if err = gfile.PutContentsAppend(methodFilePath, gstr.TrimLeft(content)); err != nil {
diff --git a/cmd/gf/internal/cmd/genctrl/genctrl_generate_ctrl_clear.go b/cmd/gf/internal/cmd/genctrl/genctrl_generate_ctrl_clear.go
index ce7b248e9..b419399e6 100644
--- a/cmd/gf/internal/cmd/genctrl/genctrl_generate_ctrl_clear.go
+++ b/cmd/gf/internal/cmd/genctrl/genctrl_generate_ctrl_clear.go
@@ -38,6 +38,7 @@ func (c *controllerClearer) doClear(dstModuleFolderPath string, item apiItem) (e
 		))
 		fileContent = gstr.Trim(gfile.GetContents(methodFilePath))
 	)
+	// retrieve it without using AST, because it's simple.
 	match, err := gregex.MatchString(`.+?Req.+?Res.+?{([\s\S]+?)}`, fileContent)
 	if err != nil {
 		return err
diff --git a/cmd/gf/internal/cmd/genctrl/genctrl_generate_sdk.go b/cmd/gf/internal/cmd/genctrl/genctrl_generate_sdk.go
index 25a628e77..e045ae0d9 100644
--- a/cmd/gf/internal/cmd/genctrl/genctrl_generate_sdk.go
+++ b/cmd/gf/internal/cmd/genctrl/genctrl_generate_sdk.go
@@ -103,6 +103,7 @@ func (c *apiSdkGenerator) doGenerateSdkIClient(
 	// append the import path to current import paths.
 	if !gstr.Contains(fileContent, moduleImportPath) {
 		isDirty = true
+		// It is without using AST, because it is from a template.
 		fileContent, err = gregex.ReplaceString(
 			`(import \([\s\S]*?)\)`,
 			fmt.Sprintf("$1\t%s\n)", moduleImportPath),
@@ -116,6 +117,7 @@ func (c *apiSdkGenerator) doGenerateSdkIClient(
 	// append the function definition to interface definition.
 	if !gstr.Contains(fileContent, interfaceFuncDefinition) {
 		isDirty = true
+		// It is without using AST, because it is from a template.
 		fileContent, err = gregex.ReplaceString(
 			`(type IClient interface {[\s\S]*?)}`,
 			fmt.Sprintf("$1\t%s\n}", interfaceFuncDefinition),