Cache name table on SourceFile

Author: jakebailey

internal/ast/ast.go

@@ -10801,6 +10801,8 @@ type SourceFile struct {
 	tokenFactory     *NodeFactory
 	declarationMapMu sync.Mutex
 	declarationMap   map[string][]*Node
+	nameTableOnce    sync.Once
+	nameTable        map[string]int
 }
 
 func (f *NodeFactory) NewSourceFile(opts SourceFileParseOptions, text string, statements *NodeList, endOfFileToken *TokenNode) *Node {
@@ -10944,6 +10946,64 @@ func (node *SourceFile) ECMALineMap() []core.TextPos {
 	return lineMap
 }
 
+// GetNameTable returns a map of all names in the file to their positions.
+// If the name appears more than once, the value is -1.
+func (file *SourceFile) GetNameTable() map[string]int {
+	file.nameTableOnce.Do(func() {
+		file.nameTable = file.initializeNameTable()
+	})
+	return file.nameTable
+}
+
+func (file *SourceFile) initializeNameTable() map[string]int {
+	nameTable := make(map[string]int)
+
+	isTagName := func(node *Node) bool {
+		return node.Parent != nil && IsJSDocTag(node.Parent) && node.Parent.TagName() == node
+	}
+
+	isArgumentOfElementAccessExpression := func(node *Node) bool {
+		return node != nil && node.Parent != nil &&
+			node.Parent.Kind == KindElementAccessExpression &&
+			node.Parent.AsElementAccessExpression().ArgumentExpression == node
+	}
+
+	// We want to store any numbers/strings if they were a name that could be
+	// related to a declaration.  So, if we have 'import x = require("something")'
+	// then we want 'something' to be in the name table.  Similarly, if we have
+	// "a['propname']" then we want to store "propname" in the name table.
+	literalIsName := func(node *Node) bool {
+		return IsDeclarationName(node) ||
+			node.Parent.Kind == KindExternalModuleReference ||
+			isArgumentOfElementAccessExpression(node) ||
+			IsLiteralComputedPropertyDeclarationName(node)
+	}
+
+	var walk func(node *Node) bool
+	walk = func(node *Node) bool {
+		if IsIdentifier(node) && !isTagName(node) && node.Text() != "" ||
+			IsStringOrNumericLiteralLike(node) && literalIsName(node) ||
+			IsPrivateIdentifier(node) {
+			text := node.Text()
+			if _, ok := nameTable[text]; ok {
+				nameTable[text] = -1
+			} else {
+				nameTable[text] = node.Pos()
+			}
+		}
+
+		node.ForEachChild(walk)
+		jsdocNodes := node.JSDoc(file)
+		for _, jsdoc := range jsdocNodes {
+			jsdoc.ForEachChild(walk)
+		}
+		return false
+	}
+
+	file.AsNode().ForEachChild(walk)
+	return nameTable
+}
+
 func (node *SourceFile) IsBound() bool {
 	return node.isBound.Load()
 }

internal/ls/completions.go

@@ -3582,7 +3582,7 @@ func (l *LanguageService) getJSCompletionEntries(
 	uniqueNames *collections.Set[string],
 	sortedEntries []*lsproto.CompletionItem,
 ) []*lsproto.CompletionItem {
-	nameTable := getNameTable(file)
+	nameTable := file.GetNameTable()
 	for name, pos := range nameTable {
 		// Skip identifiers produced only from the current location
 		if pos == position {

internal/ls/findallreferences.go

@@ -2391,7 +2391,7 @@ func (state *refState) forEachRelatedSymbol(
 
 // Search for all occurrences of an identifier in a source file (and filter out the ones that match).
 func (state *refState) searchForName(sourceFile *ast.SourceFile, search *refSearch) {
-	if _, ok := getNameTable(sourceFile)[search.escapedText]; ok {
+	if _, ok := sourceFile.GetNameTable()[search.escapedText]; ok {
 		state.getReferencesInSourceFile(sourceFile, search, true /*addReferencesHere*/)
 	}
 }

internal/ls/utilities.go

@@ -452,47 +452,6 @@ func isSeparator(node *ast.Node, candidate *ast.Node) bool {
 	return candidate != nil && node.Parent != nil && (candidate.Kind == ast.KindCommaToken || (candidate.Kind == ast.KindSemicolonToken && node.Parent.Kind == ast.KindObjectLiteralExpression))
 }
 
-// Returns a map of all names in the file to their positions.
-// !!! cache this
-func getNameTable(file *ast.SourceFile) map[string]int {
-	nameTable := make(map[string]int)
-	var walk func(node *ast.Node) bool
-
-	walk = func(node *ast.Node) bool {
-		if ast.IsIdentifier(node) && !isTagName(node) && node.Text() != "" ||
-			ast.IsStringOrNumericLiteralLike(node) && literalIsName(node) ||
-			ast.IsPrivateIdentifier(node) {
-			text := node.Text()
-			if _, ok := nameTable[text]; ok {
-				nameTable[text] = -1
-			} else {
-				nameTable[text] = node.Pos()
-			}
-		}
-
-		node.ForEachChild(walk)
-		jsdocNodes := node.JSDoc(file)
-		for _, jsdoc := range jsdocNodes {
-			jsdoc.ForEachChild(walk)
-		}
-		return false
-	}
-
-	file.ForEachChild(walk)
-	return nameTable
-}
-
-// We want to store any numbers/strings if they were a name that could be
-// related to a declaration.  So, if we have 'import x = require("something")'
-// then we want 'something' to be in the name table.  Similarly, if we have
-// "a['propname']" then we want to store "propname" in the name table.
-func literalIsName(node *ast.NumericOrStringLikeLiteral) bool {
-	return ast.IsDeclarationName(node) ||
-		node.Parent.Kind == ast.KindExternalModuleReference ||
-		isArgumentOfElementAccessExpression(node) ||
-		ast.IsLiteralComputedPropertyDeclarationName(node)
-}
-
 func isLiteralNameOfPropertyDeclarationOrIndexAccess(node *ast.Node) bool {
 	// utilities
 	switch node.Parent.Kind {
@@ -582,10 +541,6 @@ func findReferenceInPosition(refs []*ast.FileReference, pos int) *ast.FileRefere
 	return core.Find(refs, func(ref *ast.FileReference) bool { return ref.TextRange.ContainsInclusive(pos) })
 }
 
-func isTagName(node *ast.Node) bool {
-	return node.Parent != nil && ast.IsJSDocTag(node.Parent) && node.Parent.TagName() == node
-}
-
 // Assumes `candidate.pos <= position` holds.
 func positionBelongsToNode(candidate *ast.Node, position int, file *ast.SourceFile) bool {
 	if candidate.Pos() > position {