Add typesafe sync.Map wrapper (#777)

jakebailey · web-flow · commit e772e8f2c0ee · 2025-04-08T23:59:36.000Z
diff --git a/internal/collections/syncmap.go b/internal/collections/syncmap.go
@@ -0,0 +1,59 @@
+package collections
+
+import "sync"
+
+type SyncMap[K comparable, V any] struct {
+	m sync.Map
+}
+
+func (s *SyncMap[K, V]) Load(key K) (value V, ok bool) {
+	val, ok := s.m.Load(key)
+	if !ok {
+		return
+	}
+	return val.(V), true
+}
+
+func (s *SyncMap[K, V]) Store(key K, value V) {
+	s.m.Store(key, value)
+}
+
+func (s *SyncMap[K, V]) LoadOrStore(key K, value V) (actual V, loaded bool) {
+	actualAny, loaded := s.m.LoadOrStore(key, value)
+	return actualAny.(V), loaded
+}
+
+func (s *SyncMap[K, V]) Delete(key K) {
+	s.m.Delete(key)
+}
+
+func (s *SyncMap[K, V]) Clear() {
+	s.m.Clear()
+}
+
+func (s *SyncMap[K, V]) Range(f func(key K, value V) bool) {
+	s.m.Range(func(key, value any) bool {
+		return f(key.(K), value.(V))
+	})
+}
+
+// Size returns the approximate number of items in the map.
+// Note that this is not a precise count, as the map may be modified
+// concurrently while this method is running.
+func (s *SyncMap[K, V]) Size() int {
+	count := 0
+	s.m.Range(func(_, _ any) bool {
+		count++
+		return true
+	})
+	return count
+}
+
+func (s *SyncMap[K, V]) ToMap() map[K]V {
+	m := make(map[K]V, s.Size())
+	s.m.Range(func(key, value any) bool {
+		m[key.(K)] = value.(V)
+		return true
+	})
+	return m
+}
diff --git a/internal/project/documentregistry.go b/internal/project/documentregistry.go
@@ -4,6 +4,7 @@ import (
 	"sync"
 
 	"github.com/microsoft/typescript-go/internal/ast"
+	"github.com/microsoft/typescript-go/internal/collections"
 	"github.com/microsoft/typescript-go/internal/core"
 	"github.com/microsoft/typescript-go/internal/parser"
 	"github.com/microsoft/typescript-go/internal/scanner"
@@ -34,7 +35,7 @@ type registryEntry struct {
 // multiple LanguageService instances.
 type documentRegistry struct {
 	options   tspath.ComparePathsOptions
-	documents sync.Map
+	documents collections.SyncMap[registryKey, *registryEntry]
 }
 
 func newDocumentRegistry(options tspath.ComparePathsOptions) *documentRegistry {
@@ -70,8 +71,7 @@ func (r *documentRegistry) releaseDocument(file *ast.SourceFile, compilerOptions
 }
 
 func (r *documentRegistry) releaseDocumentWithKey(key registryKey) {
-	if entryAny, ok := r.documents.Load(key); ok {
-		entry := entryAny.(*registryEntry)
+	if entry, ok := r.documents.Load(key); ok {
 		entry.mu.Lock()
 		defer entry.mu.Unlock()
 		entry.refCount--
@@ -87,10 +87,9 @@ func (r *documentRegistry) getDocumentWorker(
 	key registryKey,
 ) *ast.SourceFile {
 	scriptTarget := core.IfElse(scriptInfo.scriptKind == core.ScriptKindJSON, core.ScriptTargetJSON, compilerOptions.GetEmitScriptTarget())
-	if entryAny, ok := r.documents.Load(key); ok {
+	if entry, ok := r.documents.Load(key); ok {
 		// We have an entry for this file. However, it may be for a different version of
 		// the script snapshot. If so, update it appropriately.
-		entry := entryAny.(*registryEntry)
 		if entry.sourceFile.Version != scriptInfo.version {
 			sourceFile := parser.ParseSourceFile(scriptInfo.fileName, scriptInfo.path, scriptInfo.text, scriptTarget, scanner.JSDocParsingModeParseAll)
 			sourceFile.Version = scriptInfo.version
@@ -104,11 +103,10 @@ func (r *documentRegistry) getDocumentWorker(
 		// Have never seen this file with these settings. Create a new source file for it.
 		sourceFile := parser.ParseSourceFile(scriptInfo.fileName, scriptInfo.path, scriptInfo.text, scriptTarget, scanner.JSDocParsingModeParseAll)
 		sourceFile.Version = scriptInfo.version
-		entryAny, _ := r.documents.LoadOrStore(key, &registryEntry{
+		entry, _ := r.documents.LoadOrStore(key, &registryEntry{
 			sourceFile: sourceFile,
 			refCount:   0,
 		})
-		entry := entryAny.(*registryEntry)
 		entry.mu.Lock()
 		defer entry.mu.Unlock()
 		entry.refCount++
@@ -118,10 +116,5 @@ func (r *documentRegistry) getDocumentWorker(
 
 // size should only be used for testing.
 func (r *documentRegistry) size() int {
-	count := 0
-	r.documents.Range(func(_, _ any) bool {
-		count++
-		return true
-	})
-	return count
+	return r.documents.Size()
 }
diff --git a/internal/testutil/harnessutil/harnessutil.go b/internal/testutil/harnessutil/harnessutil.go
@@ -429,19 +429,19 @@ type cachedCompilerHost struct {
 	options *core.CompilerOptions
 }
 
-var sourceFileCache sync.Map
+var sourceFileCache collections.SyncMap[sourceFileCacheKey, *ast.SourceFile]
+
+type sourceFileCacheKey struct {
+	core.SourceFileAffectingCompilerOptions
+	fileName        string
+	path            tspath.Path
+	languageVersion core.ScriptTarget
+	text            string
+}
 
 func (h *cachedCompilerHost) GetSourceFile(fileName string, path tspath.Path, languageVersion core.ScriptTarget) *ast.SourceFile {
 	text, _ := h.FS().ReadFile(fileName)
 
-	type sourceFileCacheKey struct {
-		core.SourceFileAffectingCompilerOptions
-		fileName        string
-		path            tspath.Path
-		languageVersion core.ScriptTarget
-		text            string
-	}
-
 	key := sourceFileCacheKey{
 		SourceFileAffectingCompilerOptions: h.options.SourceFileAffecting(),
 		fileName:                           fileName,
@@ -451,7 +451,7 @@ func (h *cachedCompilerHost) GetSourceFile(fileName string, path tspath.Path, la
 	}
 
 	if cached, ok := sourceFileCache.Load(key); ok {
-		return cached.(*ast.SourceFile)
+		return cached
 	}
 
 	// !!! dedupe with compiler.compilerHost
@@ -464,7 +464,7 @@ func (h *cachedCompilerHost) GetSourceFile(fileName string, path tspath.Path, la
 	}
 
 	result, _ := sourceFileCache.LoadOrStore(key, sourceFile)
-	return result.(*ast.SourceFile)
+	return result
 }
 
 func createCompilerHost(fs vfs.FS, defaultLibraryPath string, options *core.CompilerOptions, currentDirectory string) compiler.CompilerHost {
diff --git a/internal/vfs/cachedvfs/cachedvfs.go b/internal/vfs/cachedvfs/cachedvfs.go
@@ -1,19 +1,18 @@
 package cachedvfs
 
 import (
-	"sync"
-
+	"github.com/microsoft/typescript-go/internal/collections"
 	"github.com/microsoft/typescript-go/internal/vfs"
 )
 
 type FS struct {
 	fs vfs.FS
 
-	directoryExistsCache      sync.Map // map[string]bool
-	fileExistsCache           sync.Map // map[string]bool
-	getAccessibleEntriesCache sync.Map // map[string]vfs.Entries
-	realpathCache             sync.Map // map[string]string
-	statCache                 sync.Map // map[string]vfs.FileInfo
+	directoryExistsCache      collections.SyncMap[string, bool]
+	fileExistsCache           collections.SyncMap[string, bool]
+	getAccessibleEntriesCache collections.SyncMap[string, vfs.Entries]
+	realpathCache             collections.SyncMap[string, string]
+	statCache                 collections.SyncMap[string, vfs.FileInfo]
 }
 
 var _ vfs.FS = (*FS)(nil)
@@ -32,7 +31,7 @@ func (fsys *FS) ClearCache() {
 
 func (fsys *FS) DirectoryExists(path string) bool {
 	if ret, ok := fsys.directoryExistsCache.Load(path); ok {
-		return ret.(bool)
+		return ret
 	}
 	ret := fsys.fs.DirectoryExists(path)
 	fsys.directoryExistsCache.Store(path, ret)
@@ -41,7 +40,7 @@ func (fsys *FS) DirectoryExists(path string) bool {
 
 func (fsys *FS) FileExists(path string) bool {
 	if ret, ok := fsys.fileExistsCache.Load(path); ok {
-		return ret.(bool)
+		return ret
 	}
 	ret := fsys.fs.FileExists(path)
 	fsys.fileExistsCache.Store(path, ret)
@@ -50,7 +49,7 @@ func (fsys *FS) FileExists(path string) bool {
 
 func (fsys *FS) GetAccessibleEntries(path string) vfs.Entries {
 	if ret, ok := fsys.getAccessibleEntriesCache.Load(path); ok {
-		return ret.(vfs.Entries)
+		return ret
 	}
 	ret := fsys.fs.GetAccessibleEntries(path)
 	fsys.getAccessibleEntriesCache.Store(path, ret)
@@ -63,7 +62,7 @@ func (fsys *FS) ReadFile(path string) (contents string, ok bool) {
 
 func (fsys *FS) Realpath(path string) string {
 	if ret, ok := fsys.realpathCache.Load(path); ok {
-		return ret.(string)
+		return ret
 	}
 	ret := fsys.fs.Realpath(path)
 	fsys.realpathCache.Store(path, ret)
