Reduce allocations in template matching tries

src/TemplateEngine/Microsoft.TemplateEngine.Core/Matching/Trie.cs

@@ -6,12 +6,33 @@ namespace Microsoft.TemplateEngine.Core.Matching
     public class Trie<T>
         where T : TerminalBase
     {
+        private byte _singleNextNodeMatch;
+        private TrieNode<T>? _singleNextNode;
+        private Dictionary<byte, TrieNode<T>>? _nextNodes;
+
         public Trie()
         {
-            NextNodes = new Dictionary<byte, TrieNode<T>>();
         }
 
-        public Dictionary<byte, TrieNode<T>> NextNodes { get; }
+        public Dictionary<byte, TrieNode<T>> NextNodes
+        {
+            get
+            {
+                Dictionary<byte, TrieNode<T>>? nextNodes = _nextNodes;
+                if (nextNodes != null)
+                {
+                    return nextNodes;
+                }
+
+                nextNodes = new Dictionary<byte, TrieNode<T>>();
+                if (_singleNextNode != null)
+                {
+                    nextNodes.Add(_singleNextNodeMatch, _singleNextNode);
+                }
+
+                return Interlocked.CompareExchange(ref _nextNodes, nextNodes, null) ?? nextNodes;
+            }
+        }
 
         public int MaxRemainingLength { get; private set; }
 
@@ -23,23 +44,10 @@ public void AddPath(byte[] path, T terminal)
             }
 
             int remainingLength = path.Length - 1;
-            Dictionary<byte, TrieNode<T>>? current = NextNodes;
+            Trie<T> current = this;
             for (int i = 0; i < path.Length; ++i, --remainingLength)
             {
-                if (!current.TryGetValue(path[i], out TrieNode<T> next))
-                {
-                    current[path[i]] = next = new TrieNode<T>(path[i])
-                    {
-                        MaxRemainingLength = remainingLength
-                    };
-                }
-                else
-                {
-                    if (next.MaxRemainingLength < remainingLength)
-                    {
-                        next.MaxRemainingLength = remainingLength;
-                    }
-                }
+                TrieNode<T> next = current.GetOrAddNextNode(path[i], remainingLength);
 
                 if (i == path.Length - 1)
                 {
@@ -58,8 +66,58 @@ public void AddPath(byte[] path, T terminal)
                     }
                 }
 
-                current = next.NextNodes;
+                current = next;
+            }
+        }
+
+        internal bool TryGetNextNode(byte match, out TrieNode<T> next)
+        {
+            if (_nextNodes != null)
+            {
+                return _nextNodes.TryGetValue(match, out next!);
+            }
+
+            if (_singleNextNode != null && _singleNextNodeMatch == match)
+            {
+                next = _singleNextNode;
+                return true;
+            }
+
+            next = null!;
+            return false;
+        }
+
+        private TrieNode<T> GetOrAddNextNode(byte match, int remainingLength)
+        {
+            TrieNode<T> next;
+            if (_nextNodes != null)
+            {
+                if (!_nextNodes.TryGetValue(match, out next!))
+                {
+                    _nextNodes.Add(match, next = new TrieNode<T>(match));
+                }
+            }
+            else if (_singleNextNode == null)
+            {
+                _singleNextNodeMatch = match;
+                _singleNextNode = next = new TrieNode<T>(match);
+            }
+            else if (_singleNextNodeMatch == match)
+            {
+                next = _singleNextNode;
+            }
+            else
+            {
+                next = new TrieNode<T>(match);
+                NextNodes.Add(match, next);
+            }
+
+            if (next.MaxRemainingLength < remainingLength)
+            {
+                next.MaxRemainingLength = remainingLength;
             }
+
+            return next;
         }
     }
 }

src/TemplateEngine/Microsoft.TemplateEngine.Core/Matching/TrieEvaluator.cs

@@ -96,7 +96,7 @@ public bool Accept(byte data, ref int sequenceNumber, out TerminalLocation<T>? t
                 {
                     //If we matched another byte, advance the current node in
                     //  the path and log the encountered terminal (if applicable)
-                    if (path.CurrentNode.NextNodes.TryGetValue(data, out next))
+                    if (path.CurrentNode.TryGetNextNode(data, out next))
                     {
                         path.CurrentNode = next;
 
@@ -121,7 +121,7 @@ public bool Accept(byte data, ref int sequenceNumber, out TerminalLocation<T>? t
             }
 
             //Try to start a new path in the trie
-            if (_trie.NextNodes.TryGetValue(data, out next))
+            if (_trie.TryGetNextNode(data, out next))
             {
                 TriePath<T> path = new TriePath<T>(sequenceNumber)
                 {

test/TemplateEngine/Microsoft.TemplateEngine.Core.UnitTests/TokenTrieTests.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Text;
+using Microsoft.TemplateEngine.Core.Matching;
 using Microsoft.TemplateEngine.Core.Util;
 
 namespace Microsoft.TemplateEngine.Core.UnitTests
@@ -83,6 +84,26 @@ public void VerifyTokenTrieAtBegin()
             Assert.AreEqual(-1, token);
         }
 
+        [TestMethod]
+        public void NextNodesContainsBranchesAndRemainsStable()
+        {
+            TokenTrie trie = new TokenTrie();
+            trie.AddToken("cat"u8.ToArray());
+            trie.AddToken("car"u8.ToArray());
+
+            Dictionary<byte, TrieNode<Token>> rootNodes = trie.NextNodes;
+            Assert.AreSame(rootNodes, trie.NextNodes);
+            Assert.IsTrue(rootNodes.TryGetValue((byte)'c', out TrieNode<Token>? cNode));
+            Assert.IsTrue(cNode.NextNodes.TryGetValue((byte)'a', out TrieNode<Token>? aNode));
+            Assert.HasCount(2, aNode.NextNodes);
+            Assert.IsTrue(aNode.NextNodes.ContainsKey((byte)'r'));
+            Assert.IsTrue(aNode.NextNodes.ContainsKey((byte)'t'));
+
+            trie.AddToken("dog"u8.ToArray());
+            Assert.HasCount(2, rootNodes);
+            Assert.IsTrue(rootNodes.ContainsKey((byte)'d'));
+        }
+
         [TestMethod]
         public void VerifyTokenTrieNotEnoughBufferLeft()
         {