支持path压缩

costs.go

@@ -8,18 +8,29 @@ import (
 type Costs struct {
 	mu    *sync.Mutex
 	costs map[string]*Cost
+
+	compressPath bool
+	pathDict     map[string]string
 }
 
-func newCosts() *Costs {
+func newCosts(compressPath bool, pathDict map[string]string) *Costs {
+	if pathDict == nil {
+		pathDict = make(map[string]string)
+	}
+
 	return &Costs{
 		mu:    new(sync.Mutex),
 		costs: make(map[string]*Cost),
+
+		compressPath: compressPath,
+		pathDict:     pathDict,
 	}
 }
 
 func (c *Costs) addCost(skip int, cost time.Duration) {
 	path := getStackInfo(skip)
-	c.addCostWithPath(path, cost)
+	text := formatStackInfo(path, c.compressPath, c.pathDict)
+	c.addCostWithPath(text, cost)
 }
 
 func (c *Costs) addCostWithPath(path string, cost time.Duration) {

costwhere.go

@@ -8,8 +8,19 @@ import (
 	"github.com/pkg/errors"
 )
 
-func Init(ctx context.Context) (newCtx context.Context, c *CostWhere) {
-	costs := newCosts()
+func Init(ctx context.Context, opts ...InitOption) (newCtx context.Context, c *CostWhere) {
+	compressPath := true
+	var pathDict map[string]string
+	if len(opts) > 0 {
+		opt := opts[0]
+
+		if opt.CompressPath != nil {
+			compressPath = *opt.CompressPath
+		}
+		pathDict = opt.PathDict
+	}
+
+	costs := newCosts(compressPath, pathDict)
 
 	newCtx = writeThis(ctx, costs)
 
@@ -19,9 +30,12 @@ func Init(ctx context.Context) (newCtx context.Context, c *CostWhere) {
 
 	end := func() {
 		cost := time.Since(startAt)
-		costs.addCostWithPath(path, cost)
+		text := formatStackInfo(path, compressPath, pathDict)
+		costs.addCostWithPath(text, cost)
 	}
-	c = newCostWhere(costs, end, parentPath)
+
+	text := formatStackInfo(parentPath, compressPath, pathDict)
+	c = newCostWhere(costs, end, text)
 
 	return
 }
@@ -47,10 +61,22 @@ func (s *CostWhere) EndWithJSON() (j []byte, err error) {
 	return
 }
 
-func (s *CostWhere) ExportJSON() (j []byte, err error) {
+func (s *CostWhere) Export() (output Output) {
 	stacks := formatCosts(s.costs, s.parentPath)
-	j, err = json.Marshal(stacks)
+
+	output = Output{
+		Stacks: stacks,
+	}
+
+	return
+}
+
+func (s *CostWhere) ExportJSON() (j []byte, err error) {
+	output := s.Export()
+
+	j, err = json.Marshal(output)
 	err = errors.WithStack(err)
+
 	return
 }
 
@@ -74,3 +100,12 @@ func Mark(ctx context.Context) (end func()) {
 
 	return
 }
+
+type InitOption struct {
+	CompressPath *bool // 是否压缩路径，默认为 true
+	PathDict     map[string]string
+}
+
+type Output struct {
+	Stacks []string
+}

format.go

@@ -41,3 +41,18 @@ func isDirectChild(parent string, child string) (is bool) {
 	is = strings.Count(after, ";") == 1
 	return
 }
+
+func formatStackInfo(path []string, compressPath bool, pathDict map[string]string) (text string) {
+	if compressPath {
+		for i, item := range path {
+			newItem, ok := pathDict[item]
+			if ok {
+				path[i] = newItem
+			}
+		}
+	}
+
+	text = strings.Join(path, ";")
+
+	return
+}

runtime.go

@@ -3,10 +3,9 @@ package costwhere
 import (
 	"runtime"
 	"slices"
-	"strings"
 )
 
-func getStackInfo(skip int) (pathText string) {
+func getStackInfo(skip int) (path []string) {
 	const maxDepth = 50
 
 	pcs := make([]uintptr, maxDepth)
@@ -18,14 +17,12 @@ func getStackInfo(skip int) (pathText string) {
 		stacks = append(stacks, f)
 	}
 
-	path := make([]string, 0, len(stacks))
+	path = make([]string, 0, len(stacks))
 	for _, frame := range stacks {
 		path = append(path, frame.Function)
 	}
 
 	slices.Reverse(path)
 
-	pathText = strings.Join(path, ";")
-
 	return
 }