Use trie for removeStringLiteralsMatchedByTemplateLiterals

tsc/internal/checker/checker.go

@@ -25977,28 +25977,82 @@ func (c *Checker) removeRedundantLiteralTypes(types []*Type, includes TypeFlags,
 	return types
 }
 
+// templateLiteralTrieMinComparisons is the number of template-vs-literal comparisons a
+// linear scan would perform at which building the template literal trie starts to pay
+// for its construction cost, as measured by
+// BenchmarkTemplateLiteralMatchingPaths/BenchmarkRemoveStringLiteralsMatchedByTemplateLiterals.
+const templateLiteralTrieMinComparisons = 512
+
 func (c *Checker) removeStringLiteralsMatchedByTemplateLiterals(types []*Type) []*Type {
 	templates := core.Filter(types, c.isPatternLiteralType)
-	if len(templates) != 0 {
-		i := len(types)
-		for i > 0 {
-			i--
-			t := types[i]
-			if t.flags&TypeFlagsStringLiteral != 0 && core.Some(templates, func(template *Type) bool {
-				return c.isTypeMatchedByTemplateLiteralOrStringMapping(t, template)
-			}) {
-				types = slices.Delete(types, i, i+1)
+	if len(templates) == 0 {
+		return types
+	}
+	templateLiterals := core.Filter(templates, func(t *Type) bool {
+		return t.flags&TypeFlagsTemplateLiteral != 0
+	})
+	stringMappings := core.Filter(templates, func(t *Type) bool {
+		return t.flags&TypeFlagsStringMapping != 0
+	})
+	var trie *templateLiteralTrieNode
+	if len(templateLiterals) >= 2 && isTemplateLiteralTrieSelective(templateLiterals) {
+		stringLiteralCount := 0
+		for _, t := range types {
+			if t.flags&TypeFlagsStringLiteral != 0 {
+				stringLiteralCount++
 			}
 		}
+		if len(templateLiterals)*stringLiteralCount >= templateLiteralTrieMinComparisons {
+			trie = c.buildTemplateLiteralTrieFromTypes(templateLiterals)
+		}
+	}
+	i := len(types)
+	for i > 0 {
+		i--
+		t := types[i]
+		if t.flags&TypeFlagsStringLiteral != 0 && c.isStringLiteralMatchedByTemplates(t, trie, templateLiterals, stringMappings) {
+			types = slices.Delete(types, i, i+1)
+		}
 	}
 	return types
 }
 
-func (c *Checker) isTypeMatchedByTemplateLiteralOrStringMapping(t *Type, template *Type) bool {
-	if template.flags&TypeFlagsTemplateLiteral != 0 {
-		return c.isTypeMatchedByTemplateLiteralType(t, template.AsTemplateLiteralType(), c.compareTypesAssignable)
+// isTemplateLiteralTrieSelective reports whether the templates differ in their first or
+// final static text. If they do not, the trie cannot discriminate between them and only
+// adds construction and traversal overhead over the linear scan.
+func isTemplateLiteralTrieSelective(templateLiterals []*Type) bool {
+	firstTexts := templateLiterals[0].AsTemplateLiteralType().texts
+	firstPrefix := firstTexts[0]
+	firstSuffix := firstTexts[len(firstTexts)-1]
+	for _, t := range templateLiterals[1:] {
+		texts := t.AsTemplateLiteralType().texts
+		if texts[0] != firstPrefix || texts[len(texts)-1] != firstSuffix {
+			return true
+		}
+	}
+	return false
+}
+
+func (c *Checker) isStringLiteralMatchedByTemplates(source *Type, trie *templateLiteralTrieNode, templateLiterals []*Type, stringMappings []*Type) bool {
+	if trie != nil {
+		if c.findMatchingTemplateLiteralInTrie(trie, source, c.compareTypesAssignable) != nil {
+			return true
+		}
+	} else if len(templateLiterals) > 0 {
+		if core.Some(templateLiterals, func(tl *Type) bool {
+			return c.isTypeMatchedByTemplateLiteralType(source, tl.AsTemplateLiteralType(), c.compareTypesAssignable)
+		}) {
+			return true
+		}
+	}
+	if len(stringMappings) > 0 {
+		if core.Some(stringMappings, func(sm *Type) bool {
+			return c.isMemberOfStringMapping(source, sm)
+		}) {
+			return true
+		}
 	}
-	return c.isMemberOfStringMapping(t, template)
+	return false
 }
 
 func (c *Checker) removeConstrainedTypeVariables(types []*Type) []*Type {

tsc/internal/checker/relater.go

@@ -1105,6 +1105,125 @@ func (c *Checker) getMatchingUnionConstituentForType(unionType *Type, t *Type) *
 	return c.getConstituentTypeForKeyType(unionType, propType)
 }
 
+func (c *Checker) buildTemplateLiteralTrieFromTypes(templateTypes []*Type) *templateLiteralTrieNode {
+	var arena core.Arena[templateLiteralTrieNode]
+	var suffixArena core.Arena[templateLiteralSuffixTrieNode]
+	root := arena.New()
+	for _, t := range templateTypes {
+		texts := t.AsTemplateLiteralType().texts
+		node := root
+		for i := range len(texts[0]) {
+			node = node.getOrInsertChild(&arena, texts[0][i])
+		}
+		if node.bySuffix == nil {
+			node.bySuffix = suffixArena.New()
+		}
+		// A string literal can only match a template literal if it ends with the template's
+		// final static text, so candidates are indexed by that text reversed.
+		suffixNode := node.bySuffix
+		suffix := texts[len(texts)-1]
+		for i := len(suffix) - 1; i >= 0; i-- {
+			suffixNode = suffixNode.getOrInsertChild(&suffixArena, suffix[i])
+		}
+		suffixNode.types = append(suffixNode.types, t)
+	}
+	return root
+}
+
+func (n *templateLiteralTrieNode) getOrInsertChild(arena *core.Arena[templateLiteralTrieNode], ch byte) *templateLiteralTrieNode {
+	for _, edge := range n.children {
+		if edge.b == ch {
+			return edge.node
+		}
+	}
+	child := arena.New()
+	n.children = append(n.children, templateLiteralTrieEdge{b: ch, node: child})
+	return child
+}
+
+func (n *templateLiteralTrieNode) findChild(ch byte) *templateLiteralTrieNode {
+	for _, edge := range n.children {
+		if edge.b == ch {
+			return edge.node
+		}
+	}
+	return nil
+}
+
+func (n *templateLiteralSuffixTrieNode) getOrInsertChild(arena *core.Arena[templateLiteralSuffixTrieNode], ch byte) *templateLiteralSuffixTrieNode {
+	for _, edge := range n.children {
+		if edge.b == ch {
+			return edge.node
+		}
+	}
+	child := arena.New()
+	n.children = append(n.children, templateLiteralSuffixTrieEdge{b: ch, node: child})
+	return child
+}
+
+func (n *templateLiteralSuffixTrieNode) findChild(ch byte) *templateLiteralSuffixTrieNode {
+	for _, edge := range n.children {
+		if edge.b == ch {
+			return edge.node
+		}
+	}
+	return nil
+}
+
+func (c *Checker) findMatchingTemplateLiteralInTrie(trie *templateLiteralTrieNode, source *Type, compareTypes TypeComparer) *Type {
+	value := source.AsLiteralType().Value().(string)
+	node := trie
+	// Check root candidates (empty-prefix templates like `${number}`)
+	if t := c.findMatchingTemplateLiteralBySuffix(node.bySuffix, value, 0, source, compareTypes); t != nil {
+		return t
+	}
+	for i := range len(value) {
+		node = node.findChild(value[i])
+		if node == nil {
+			return nil
+		}
+		if t := c.findMatchingTemplateLiteralBySuffix(node.bySuffix, value, i+1, source, compareTypes); t != nil {
+			return t
+		}
+	}
+	return nil
+}
+
+func (c *Checker) findMatchingTemplateLiteralBySuffix(trie *templateLiteralSuffixTrieNode, value string, prefixLength int, source *Type, compareTypes TypeComparer) *Type {
+	if trie == nil {
+		return nil
+	}
+	// Candidates with an empty final static text (e.g. `/${string}`) are stored at the root.
+	if t := c.matchTemplateLiteralTrieCandidates(trie.types, value, prefixLength, 0, source, compareTypes); t != nil {
+		return t
+	}
+	node := trie
+	for i := len(value) - 1; i >= 0; i-- {
+		node = node.findChild(value[i])
+		if node == nil {
+			return nil
+		}
+		if t := c.matchTemplateLiteralTrieCandidates(node.types, value, prefixLength, len(value)-i, source, compareTypes); t != nil {
+			return t
+		}
+	}
+	return nil
+}
+
+func (c *Checker) matchTemplateLiteralTrieCandidates(candidates []*Type, value string, prefixLength int, suffixLength int, source *Type, compareTypes TypeComparer) *Type {
+	// The first and last static texts cannot overlap in the source, so a shorter string
+	// can never match (mirrors the guard in inferFromLiteralPartsToTemplateLiteral).
+	if prefixLength+suffixLength > len(value) {
+		return nil
+	}
+	for _, t := range candidates {
+		if c.isTypeMatchedByTemplateLiteralType(source, t.AsTemplateLiteralType(), compareTypes) {
+			return t
+		}
+	}
+	return nil
+}
+
 // Return the name of a discriminant property for which it was possible and feasible to construct a map of
 // constituent types keyed by the literal types of the property by that name in each constituent type. Return
 // an empty string if no such discriminant property exists.