loader: 完善测试和汇编相关的结构, 并完善测试信息解析

internal/loader/api.go

@@ -28,22 +28,45 @@ type Package struct {
 	Pkg     *types.Package // 类型检查后的包
 	Info    *types.Info    // 包的类型检查信息
 	Files   []*ast.File    // AST语法树
-	WsFiles []WsFile       // 汇编代码
+	WsFiles []*WsFile      // 汇编代码
 
 	SSAPkg   *ssa.Package
-	TestInfo TestInfo
-}
-
-// 汇编代码文件
-type WsFile struct {
-	Name string // 文件名
-	Code string // 汇编代码
+	TestInfo *TestInfo
 }
 
 // 单元测试信息
 type TestInfo struct {
 	Files []string // 测试文件
 	Funcs []string // 测试函数
+
+	Tests    map[string]TestFuncInfo
+	Benchs   map[string]BenchFuncInfo
+	Examples map[string]ExampleFuncInfo
+}
+
+// 测试函数信息
+type TestFuncInfo struct {
+	FuncPos token.Pos // 函数位置
+	Name    string    // 函数名, 不含包路径
+}
+
+// 基准函数信息
+type BenchFuncInfo struct {
+	FuncPos token.Pos // 函数位置
+	Name    string    // 函数名, 不含包路径
+}
+
+// 示例函数信息
+type ExampleFuncInfo struct {
+	FuncPos token.Pos // 函数位置
+	Name    string    // 函数名, 不含包路径
+	Output  string    // 期望输出, 为空表示不验证
+}
+
+// 汇编代码文件
+type WsFile struct {
+	Name string // 文件名
+	Code string // 汇编代码
 }
 
 // 加载程序

internal/loader/loader.go

@@ -180,7 +180,7 @@ func (p *_Loader) Import(pkgpath string) (*types.Package, error) {
 		return nil, err
 	}
 
-	// 解析当前包的 AST
+	// 解析当前包的 AST, 隐含了是否测试模式
 	filenames, pkg.Files, err = p.ParseDir(pkgpath)
 	if err != nil {
 		logger.Tracef(&config.EnableTrace_loader, "err: %v", err)
@@ -253,37 +253,110 @@ func (p *_Loader) Import(pkgpath string) (*types.Package, error) {
 
 	// 提取测试信息
 	if p.cfg.UnitTest {
-		for _, filename := range filenames {
-			if !p.isTestFile(filename) {
-				continue
-			}
-			pkg.TestInfo.Files = append(pkg.TestInfo.Files, filename)
+		pkg.TestInfo, err = p.parseTestInfo(&pkg, filenames)
+		if err != nil {
+			return nil, err
 		}
-		for _, name := range pkg.Pkg.Scope().Names() {
-			if len(name) < len("Test?") {
-				continue
+	}
+
+	logger.Tracef(&config.EnableTrace_loader, "save pkgpath: %v", pkgpath)
+
+	p.prog.Pkgs[pkgpath] = &pkg
+	return pkg.Pkg, nil
+}
+
+func (p *_Loader) parseTestInfo(pkg *Package, filenames []string) (*TestInfo, error) {
+	tInfo := &TestInfo{
+		Tests:    make(map[string]TestFuncInfo),
+		Benchs:   make(map[string]BenchFuncInfo),
+		Examples: make(map[string]ExampleFuncInfo),
+	}
+
+	for i, filename := range filenames {
+		if !p.isTestFile(filename) {
+			continue
+		}
+
+		file := pkg.Files[i]
+		tInfo.Files = append(tInfo.Files, filename)
+
+		// 提取测试/基准/示例函数
+		for _, decl := range file.Decls {
+			if fn, ok := decl.(*ast.FuncDecl); ok {
+				name := fn.Name.Name
+
+				// 函数参数均为空
+				obj := pkg.Pkg.Scope().Lookup(name)
+				{
+					var bValidFuncType = false
+					if fn, ok := obj.(*types.Func); ok {
+						if sig, ok := fn.Type().(*types.Signature); ok {
+							if sig.Recv() == nil && sig.Params().Len() == 0 && sig.Results().Len() == 0 {
+								bValidFuncType = true
+							}
+						}
+					}
+					if !bValidFuncType {
+						continue // skip
+					}
+				}
+
+				switch {
+				case strings.HasPrefix(name, "Test"):
+					tInfo.Funcs = append(tInfo.Funcs, name)
+					tInfo.Tests[name] = TestFuncInfo{
+						FuncPos: obj.Pos(),
+						Name:    name,
+					}
+				case strings.HasPrefix(name, "Bench"):
+					tInfo.Benchs[name] = BenchFuncInfo{
+						FuncPos: obj.Pos(),
+						Name:    name,
+					}
+				case strings.HasPrefix(name, "Example"):
+					tInfo.Examples[name] = ExampleFuncInfo{
+						FuncPos: obj.Pos(),
+						Name:    name,
+						Output:  p.parseExampleOutputComment(file, fn),
+					}
+				}
 			}
-			if !strings.HasPrefix(name, "Test") {
+		}
+	}
+
+	return tInfo, nil
+}
+
+func (p *_Loader) parseExampleOutputComment(f *ast.File, fn *ast.FuncDecl) string {
+	for _, commentGroup := range f.Comments {
+		if commentGroup.Pos() <= fn.Body.Pos() {
+			continue
+		}
+		if commentGroup.End() > fn.Body.End() {
+			break
+		}
+
+		for j, comment := range commentGroup.List {
+			if comment.Text != "// Output:" {
 				continue
 			}
-			obj := pkg.Pkg.Scope().Lookup(name)
-			if fn, ok := obj.(*types.Func); ok {
-				if sig, ok := fn.Type().(*types.Signature); ok {
-					if sig.Recv() == nil && sig.Params().Len() == 0 && sig.Results().Len() == 0 {
-						pkg.TestInfo.Funcs = append(pkg.TestInfo.Funcs, name)
-					}
+
+			var lineTexts []string
+			for _, x := range commentGroup.List[j+1:] {
+				if !strings.HasPrefix(x.Text, "//") {
+					break
 				}
+				lineTexts = append(lineTexts, strings.TrimSpace(x.Text[2:]))
 			}
+
+			return strings.Join(lineTexts, "\n")
 		}
 	}
 
-	logger.Tracef(&config.EnableTrace_loader, "save pkgpath: %v", pkgpath)
-
-	p.prog.Pkgs[pkgpath] = &pkg
-	return pkg.Pkg, nil
+	return ""
 }
 
-func (p *_Loader) ParseDir_wsFiles(pkgpath string) (files []WsFile, err error) {
+func (p *_Loader) ParseDir_wsFiles(pkgpath string) (files []*WsFile, err error) {
 	logger.Tracef(&config.EnableTrace_loader, "pkgpath: %v", pkgpath)
 
 	if p.cfg.WaBackend == "" {
@@ -334,7 +407,7 @@ func (p *_Loader) ParseDir_wsFiles(pkgpath string) (files []WsFile, err error) {
 	}
 
 	for i := 0; i < len(filenames); i++ {
-		files = append(files, WsFile{
+		files = append(files, &WsFile{
 			Name: filenames[i],
 			Code: string(datas[i]),
 		})