Optimize looking up units by a sequence of words

Calculator/Math/Unit.cs

@@ -1,17 +1,8 @@
-using System.Collections.Generic;
+using System.Collections.Immutable;
-using System.Collections.Immutable;
 
 namespace Calculator.Math;
 
 public sealed record Unit(string ShortName, ImmutableArray<string> LongNames) {
-	internal void AssignNamesTo(Dictionary<string, Unit> nameToUnitDictionary) {
-		nameToUnitDictionary.Add(ShortName, this);
-
-		foreach (string longName in LongNames) {
-			nameToUnitDictionary.Add(longName, this);
-		}
-	}
-
 	public override string ToString() {
 		return ShortName;
 	}

Calculator/Math/UnitUniverse.cs

@@ -10,7 +10,6 @@ using ExtendedNumerics;
 namespace Calculator.Math;
 
 sealed class UnitUniverse(
-	Unit primaryUnit,
 	FrozenDictionary<Unit, Func<Number, Number>> unitToConversionToPrimaryUnit,
 	FrozenDictionary<Unit, Func<Number, Number>> unitToConversionFromPrimaryUnit
 ) {
@@ -61,9 +60,9 @@ sealed class UnitUniverse(
 	}
 
 	internal sealed class Builder {
+		private readonly Unit primaryUnit;
 		private readonly Dictionary<Unit, Func<Number, Number>> unitToConversionToPrimaryUnit = new (ReferenceEqualityComparer.Instance);
 		private readonly Dictionary<Unit, Func<Number, Number>> unitToConversionFromPrimaryUnit = new (ReferenceEqualityComparer.Instance);
-		private readonly Unit primaryUnit;
 
 		public Builder(Unit primaryUnit) {
 			this.primaryUnit = primaryUnit;
@@ -113,7 +112,6 @@ sealed class UnitUniverse(
 
 		public UnitUniverse Build() {
 			return new UnitUniverse(
-				primaryUnit,
 				unitToConversionToPrimaryUnit.ToFrozenDictionary(ReferenceEqualityComparer.Instance),
 				unitToConversionFromPrimaryUnit.ToFrozenDictionary(ReferenceEqualityComparer.Instance)
 			);

Calculator/Math/UnitUniverses.cs

@@ -1,4 +1,5 @@
-using System.Collections.Frozen;
+using System;
+using System.Collections.Frozen;
 using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
 
@@ -6,28 +7,72 @@ namespace Calculator.Math;
 
 sealed class UnitUniverses {
 	private readonly FrozenDictionary<Unit, UnitUniverse> unitToUniverse;
-	private readonly FrozenDictionary<string, Unit> nameToUnit;
+	
+	public WordLookupTrieNode UnitLookupByWords { get; }
 
 	internal UnitUniverses(params UnitUniverse[] universes) {
 		Dictionary<Unit, UnitUniverse> unitToUniverseBuilder = new (ReferenceEqualityComparer.Instance);
-		Dictionary<string, Unit> nameToUnitBuilder = new ();
+		WordLookupTrieNode.Builder unitLookupByWordsBuilder = new ();
-
+		
 		foreach (UnitUniverse universe in universes) {
 			foreach (Unit unit in universe.AllUnits) {
 				unitToUniverseBuilder.Add(unit, universe);
-				unit.AssignNamesTo(nameToUnitBuilder);
+				unitLookupByWordsBuilder.Add(unit);
 			}
 		}
 
 		unitToUniverse = unitToUniverseBuilder.ToFrozenDictionary(ReferenceEqualityComparer.Instance);
-		nameToUnit = nameToUnitBuilder.ToFrozenDictionary();
+		UnitLookupByWords = unitLookupByWordsBuilder.Build();
-	}
-	
-	public bool TryGetUnit(string name, [NotNullWhen(true)] out Unit? unit) {
-		return nameToUnit.TryGetValue(name, out unit);
 	}
 	
 	public bool TryGetUniverse(Unit unit, [NotNullWhen(true)] out UnitUniverse? universe) {
 		return unitToUniverse.TryGetValue(unit, out universe);
 	}
+	
+	public sealed record WordLookupTrieNode(Unit? Unit, FrozenDictionary<string, WordLookupTrieNode> Children) {
+		internal sealed class Builder {
+			private sealed record Node(Unit? Unit, Dictionary<string, Node> Children) {
+				internal static Node Create(Unit? unit = null) {
+					return new Node(unit, new Dictionary<string, Node>());
+				}
+				
+				internal Node Child(string word) {
+					if (!Children.TryGetValue(word, out Node? child)) {
+						Children.Add(word, child = Create());
+					}
+					
+					return child;
+				}
+			}
+			
+			private readonly Node root = Node.Create();
+			
+			public void Add(Unit unit) {
+				Add(unit.ShortName, unit);
+				
+				foreach (string longName in unit.LongNames) {
+					Add(longName, unit);
+				}
+			}
+			
+			private void Add(string name, Unit unit) {
+				Node node = root;
+				string[] words = name.Split(' ');
+
+				foreach (string word in words.AsSpan(..^1)) {
+					node = node.Child(word);
+				}
+				
+				node.Children.Add(words[^1], Node.Create(unit));
+			}
+
+			public WordLookupTrieNode Build() {
+				return Build(root);
+			}
+			
+			private WordLookupTrieNode Build(Node node) {
+				return new WordLookupTrieNode(node.Unit, node.Children.ToFrozenDictionary(static kvp => kvp.Key, kvp => Build(kvp.Value)));
+			}
+		}
+	}
 }

Calculator/Parser/Parser.cs

@@ -1,8 +1,6 @@
 using System;
-using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Diagnostics.CodeAnalysis;
-using System.Linq;
 using Calculator.Math;
 
 namespace Calculator.Parser;
@@ -189,25 +187,23 @@ public sealed class Parser(ImmutableArray<Token> tokens) {
 
 	private bool MatchUnit([NotNullWhen(true)] out Unit? unit) {
 		int position = current;
-		
+
-		List<string> words = [];
+		UnitUniverses.WordLookupTrieNode node = Units.All.UnitLookupByWords;
-		
+
-		while (Match(out Token.Text? text)) {
+		// ReSharper disable once AccessToModifiedClosure
-			words.Add(text.Value);
+		while (Match(token => node.Children.ContainsKey(token.Value), out Token.Text? text)) {
+			node = node.Children[text.Value];
 		}
 
-		for (int i = words.Count; i > 0; i--) {
+		unit = node.Unit;
-			string unitName = string.Join(' ', words.Take(i));
 
-			if (Units.All.TryGetUnit(unitName, out unit)) {
+		if (unit != null) {
-				current = position + i;
+			return true;
-				return true;
+		}
-			}
+		else {
+			current = position;
+			return false;
 		}
-
-		current = position;
-		unit = null;
-		return false;
 	}
 
 	private bool MatchUnitConversionOperator() {