From 26ac303090e7197e725975f1c2c140b902a0bc8b Mon Sep 17 00:00:00 2001
From: Gabriel Musat <gb.mt.me@gmail.com>
Date: Sun, 24 Dec 2023 19:40:48 +0100
Subject: [PATCH] perf: improve export resolving performance using cache

---
 internal/language/exports.go      | 34 ++++++++++++++++++++++---------
 internal/language/exports_test.go |  8 ++++----
 internal/language/parser.go       |  4 ++--
 3 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/internal/language/exports.go b/internal/language/exports.go
index 3717d54..5eb74bd 100644
--- a/internal/language/exports.go
+++ b/internal/language/exports.go
@@ -3,8 +3,10 @@ package language
 import (
 	"context"
 	"errors"
+	"fmt"
 
 	"github.com/elliotchance/orderedmap/v2"
+
 	"github.com/gabotechs/dep-tree/internal/utils"
 )
 
@@ -45,12 +47,28 @@ type ExportsResult struct {
 	Errors []error
 }
 
-type UnwrappedExportsCacheKey string
+type parseExportsKey string
 
 func (p *Parser[F]) ParseExports(
 	ctx context.Context,
 	id string,
 	unwrappedExports bool,
+) (context.Context, *ExportsResult, error) {
+	unwrappedCacheKey := parseExportsKey(fmt.Sprintf("%s-%t", id, unwrappedExports))
+	if cached, ok := ctx.Value(unwrappedCacheKey).(*ExportsResult); ok {
+		return ctx, cached, nil
+	}
+	ctx, result, err := p.uncachedParseExports(ctx, id, unwrappedExports, nil)
+	if err != nil {
+		return ctx, nil, err
+	}
+	return context.WithValue(ctx, unwrappedCacheKey, result), result, nil
+}
+
+func (p *Parser[F]) uncachedParseExports(
+	ctx context.Context,
+	id string,
+	unwrappedExports bool,
 	stack *utils.CallStack,
 ) (context.Context, *ExportsResult, error) {
 	if stack == nil {
@@ -60,11 +78,7 @@ func (p *Parser[F]) ParseExports(
 		return ctx, nil, errors.New("circular export: " + err.Error())
 	}
 
-	unwrappedCacheKey := UnwrappedExportsCacheKey(id)
-	if cached, ok := ctx.Value(unwrappedCacheKey).(*ExportsResult); ok {
-		return ctx, cached, nil
-	}
-	ctx, wrapped, err := p.cachedParseExports(ctx, id)
+	ctx, wrapped, err := p.gatherExportsFromFile(ctx, id)
 	if err != nil {
 		return ctx, nil, err
 	}
@@ -80,7 +94,7 @@ func (p *Parser[F]) ParseExports(
 		}
 
 		var unwrapped *ExportsResult
-		ctx, unwrapped, err = p.ParseExports(ctx, export.Path, unwrappedExports, stack)
+		ctx, unwrapped, err = p.uncachedParseExports(ctx, export.Path, unwrappedExports, stack)
 		if err != nil {
 			exportErrors = append(exportErrors, err)
 			continue
@@ -116,13 +130,13 @@ func (p *Parser[F]) ParseExports(
 	return ctx, &ExportsResult{Exports: exports, Errors: exportErrors}, nil
 }
 
-type ExportsCacheKey string
+type gatherExportsFromFileKey string
 
-func (p *Parser[F]) cachedParseExports(
+func (p *Parser[F]) gatherExportsFromFile(
 	ctx context.Context,
 	filePath string,
 ) (context.Context, *ExportsEntries, error) {
-	cacheKey := ExportsCacheKey(filePath)
+	cacheKey := gatherExportsFromFileKey(filePath)
 	if cached, ok := ctx.Value(cacheKey).(*ExportsEntries); ok {
 		return ctx, cached, nil
 	}
diff --git a/internal/language/exports_test.go b/internal/language/exports_test.go
index 69d459a..4803693 100644
--- a/internal/language/exports_test.go
+++ b/internal/language/exports_test.go
@@ -63,12 +63,12 @@ func TestParser_parseExports_IsCached(t *testing.T) {
 	parser := lang.testParser("1")
 
 	start := time.Now()
-	ctx, _, err := parser.cachedParseExports(ctx, "1")
+	ctx, _, err := parser.gatherExportsFromFile(ctx, "1")
 	a.NoError(err)
 	nonCached := time.Since(start)
 
 	start = time.Now()
-	_, _, err = parser.cachedParseExports(ctx, "1")
+	_, _, err = parser.gatherExportsFromFile(ctx, "1")
 	a.NoError(err)
 	cached := time.Since(start)
 
@@ -229,12 +229,12 @@ func TestParser_CachedUnwrappedParseExports(t *testing.T) {
 			}
 			parser := lang.testParser(tt.Path)
 
-			_, exports, err := parser.ParseExports(context.Background(), "1", true, nil)
+			_, exports, err := parser.ParseExports(context.Background(), "1", true)
 			a.NoError(err)
 
 			a.Equal(tt.ExpectedUnwrapped, exports.Exports)
 
-			_, exports, err = parser.ParseExports(context.Background(), "1", false, nil)
+			_, exports, err = parser.ParseExports(context.Background(), "1", false)
 			a.NoError(err)
 
 			a.Equal(tt.ExpectedWrapped, exports.Exports)
diff --git a/internal/language/parser.go b/internal/language/parser.go
index 167cc16..7a3bc22 100644
--- a/internal/language/parser.go
+++ b/internal/language/parser.go
@@ -136,7 +136,7 @@ func (p *Parser[F]) Deps(ctx context.Context, n *Node) (context.Context, []*Node
 		//  technically is true, but it's not true to say that `foo` is imported from `bar.ts`.
 		//  It's more accurate to say that `bar` is imported from `bar.ts`, even if the alias is `foo`.
 		//  Instead we never unwrap export to avoid this.
-		ctx, exports, err = p.ParseExports(ctx, n.Id, false, nil)
+		ctx, exports, err = p.ParseExports(ctx, n.Id, false)
 		if err != nil {
 			return nil, nil, err
 		}
@@ -161,7 +161,7 @@ func (p *Parser[F]) Deps(ctx context.Context, n *Node) (context.Context, []*Node
 		}
 
 		var exports *ExportsResult
-		ctx, exports, err = p.ParseExports(ctx, importEntry.Path, true, nil)
+		ctx, exports, err = p.ParseExports(ctx, importEntry.Path, true)
 		if err != nil {
 			return ctx, nil, err
 		}