Cache name table on SourceFile (#2363)

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,39 @@ 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() {
+		nameTable := make(map[string]int, file.IdentifierCount)
+
+		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.ForEachChild(walk)
+
+		file.nameTable = nameTable
+	})
+	return file.nameTable
+}
+
 func (node *SourceFile) IsBound() bool {
 	return node.isBound.Load()
 }

internal/ast/utilities.go

@@ -4159,3 +4159,24 @@ func GetRestParameterElementType(node *ParameterDeclarationNode) *Node {
 	}
 	return nil
 }
+
+func isTagName(node *Node) bool {
+	return node.Parent != nil && IsJSDocTag(node.Parent) && node.Parent.TagName() == 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.
+func literalIsName(node *Node) bool {
+	return IsDeclarationName(node) ||
+		node.Parent.Kind == KindExternalModuleReference ||
+		isArgumentOfElementAccessExpression(node) ||
+		IsLiteralComputedPropertyDeclarationName(node)
+}
+
+func isArgumentOfElementAccessExpression(node *Node) bool {
+	return node != nil && node.Parent != nil &&
+		node.Parent.Kind == KindElementAccessExpression &&
+		node.Parent.AsElementAccessExpression().ArgumentExpression == node
+}

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 {
@@ -524,12 +483,6 @@ func isObjectBindingElementWithoutPropertyName(bindingElement *ast.Node) bool {
 		bindingElement.PropertyName() == nil
 }
 
-func isArgumentOfElementAccessExpression(node *ast.Node) bool {
-	return node != nil && node.Parent != nil &&
-		node.Parent.Kind == ast.KindElementAccessExpression &&
-		node.Parent.AsElementAccessExpression().ArgumentExpression == node
-}
-
 func isRightSideOfPropertyAccess(node *ast.Node) bool {
 	return node.Parent.Kind == ast.KindPropertyAccessExpression && node.Parent.Name() == node
 }
@@ -582,10 +535,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 {