From 2769b68192a8d72e34e3fef81b8a2558e9aded44 Mon Sep 17 00:00:00 2001
From: Tao_ta <745103130@qq.com>
Date: Mon, 23 Sep 2024 15:31:35 +0800
Subject: [PATCH] feat: add comment into AST, usage in definition can get
 document.

---
 tools/compiler/internal/ast.go        | 19 +++++++++++++++++++
 tools/compiler/internal/base.go       |  3 ++-
 tools/compiler/internal/definition.go | 18 ++++++++++++++++--
 tools/compiler/internal/setups.go     |  2 +-
 4 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/tools/compiler/internal/ast.go b/tools/compiler/internal/ast.go
index 55ed517dd..730fe63c3 100644
--- a/tools/compiler/internal/ast.go
+++ b/tools/compiler/internal/ast.go
@@ -221,3 +221,22 @@ func getCodeSelectorExprList(file *ast.File) ([]*ast.SelectorExpr, error) {
 
 	return sv.selectorExprList, nil
 }
+
+type functionExprVisitor struct {
+	funcList []*ast.FuncDecl
+}
+
+func (v *functionExprVisitor) Visit(node ast.Node) ast.Visitor {
+	if fnExpr, ok := node.(*ast.FuncDecl); ok {
+		v.funcList = append(v.funcList, fnExpr)
+	}
+	return v
+}
+
+func getCodeFuncDecl(file *ast.File) ([]*ast.FuncDecl, error) {
+	fv := &functionExprVisitor{}
+
+	ast.Walk(fv, file)
+
+	return fv.funcList, nil
+}
diff --git a/tools/compiler/internal/base.go b/tools/compiler/internal/base.go
index aa28a5799..61bce3508 100644
--- a/tools/compiler/internal/base.go
+++ b/tools/compiler/internal/base.go
@@ -42,7 +42,8 @@ func GetDefinition(fileName string, fileMap map[string]string) (interface{}, err
 	if err != nil {
 		fmt.Println("Internal error: ", err)
 	}
-	definitionList := getDefinitionList(info)
+	fnList, _ := getCodeFuncDecl(pkg[PKG].Files[fileName])
+	definitionList := getDefinitionList(info, fnList)
 	definitionList.Position(fset)
 	filter := definitions{}
 	for _, item := range definitionList {
diff --git a/tools/compiler/internal/definition.go b/tools/compiler/internal/definition.go
index a9172d7dc..42f71dd22 100644
--- a/tools/compiler/internal/definition.go
+++ b/tools/compiler/internal/definition.go
@@ -33,6 +33,7 @@ type usage struct {
 	InsertText  string  `json:"insertText"`
 	Params      []param `json:"params"`
 	Type        string  `json:"type"`
+	Document    string  `json:"document"`
 }
 
 type param struct {
@@ -127,7 +128,7 @@ func (d definitions) Position(fset *token.FileSet) {
 		def.From.EndPosition = Position(fset.Position(token.Pos(def.From.EndPos)))
 	}
 }
-func getDefinitionList(info *typesutil.Info) definitions {
+func getDefinitionList(info *typesutil.Info, fnList []*ast.FuncDecl) definitions {
 	var definitionList definitions
 
 	defs := info.Defs
@@ -138,7 +139,11 @@ func getDefinitionList(info *typesutil.Info) definitions {
 			continue
 		}
 		definition := createDefinitionItem(ident, obj)
-		definition.Usages = createUsages(obj, ident.Name)
+		definitionUsage := createUsages(obj, ident.Name)
+		if commentList := checkFnExist(fnList, ident); commentList != nil {
+			definitionUsage[0].Document = commentList.Text()
+		}
+		definition.Usages = definitionUsage
 		definitionList = append(definitionList, definition)
 	}
 
@@ -350,3 +355,12 @@ func extractParams(signature *types.Signature) (signList []string, sampleList []
 	}
 	return
 }
+
+func checkFnExist(fnList []*ast.FuncDecl, ident *ast.Ident) *ast.CommentGroup {
+	for _, fnDecl := range fnList {
+		if fnDecl.Name == ident && fnDecl.Doc != nil {
+			return fnDecl.Doc
+		}
+	}
+	return nil
+}
diff --git a/tools/compiler/internal/setups.go b/tools/compiler/internal/setups.go
index 44d693b22..2c8f77415 100644
--- a/tools/compiler/internal/setups.go
+++ b/tools/compiler/internal/setups.go
@@ -28,7 +28,7 @@ func initSPXMod() *gopmod.Module {
 func initSPXParserConf() parser.Config {
 	return parser.Config{
 		ClassKind: gopbuild.ClassKind,
-		Mode:      parser.DeclarationErrors,
+		Mode:      parser.AllErrors | parser.ParseComments,
 	}
 }