Merge pull request #31558 from mavasani/SyntaxEditorTracking

Add support to Replace and Insert SyntaxEditor APIs to allow tracking…

Author: mavasani

src/Workspaces/Core/Portable/Editing/SyntaxEditor.cs

@@ -13,6 +13,8 @@ public class SyntaxEditor
     {
         private readonly SyntaxGenerator _generator;
         private readonly List<Change> _changes;
+        private bool _allowEditsOnLazilyCreatedTrackedNewNodes;
+        private HashSet<SyntaxNode> _lazyTrackedNewNodesOpt;
 
         /// <summary>
         /// Creates a new <see cref="SyntaxEditor"/> instance.
@@ -36,6 +38,30 @@ internal SyntaxEditor(SyntaxNode root, SyntaxGenerator generator)
             _changes = new List<Change>();
         }
 
+        private SyntaxNode ApplyTrackingToNewNode(SyntaxNode node)
+        {
+            if (node == null)
+            {
+                return null;
+            }
+
+            _lazyTrackedNewNodesOpt = _lazyTrackedNewNodesOpt ?? new HashSet<SyntaxNode>();
+            foreach (var descendant in node.DescendantNodesAndSelf())
+            {
+                _lazyTrackedNewNodesOpt.Add(descendant);
+            }
+
+            return node.TrackNodes(node.DescendantNodesAndSelf());
+        }
+
+        private IEnumerable<SyntaxNode> ApplyTrackingToNewNodes(IEnumerable<SyntaxNode> nodes)
+        {
+            foreach (var node in nodes)
+            {
+                yield return ApplyTrackingToNewNode(node);
+            }
+        }
+
         /// <summary>
         /// The <see cref="SyntaxNode"/> that was specified when the <see cref="SyntaxEditor"/> was constructed.
         /// </summary>
@@ -51,7 +77,8 @@ internal SyntaxEditor(SyntaxNode root, SyntaxGenerator generator)
         /// </summary>
         public SyntaxNode GetChangedRoot()
         {
-            var nodes = Enumerable.Distinct(_changes.Select(c => c.Node));
+            var nodes = Enumerable.Distinct(_changes.Where(c => OriginalRoot.Contains(c.Node))
+                                                    .Select(c => c.Node));
             var newRoot = OriginalRoot.TrackNodes(nodes);
 
             foreach (var change in _changes)
@@ -67,7 +94,7 @@ public SyntaxNode GetChangedRoot()
         /// </summary>
         public void TrackNode(SyntaxNode node)
         {
-            CheckNodeInTree(node);
+            CheckNodeInOriginalTreeOrTracked(node);
             _changes.Add(new NoChange(node));
         }
 
@@ -87,7 +114,7 @@ public void RemoveNode(SyntaxNode node)
         /// <param name="options">Options that affect how node removal works.</param>
         public void RemoveNode(SyntaxNode node, SyntaxRemoveOptions options)
         {
-            CheckNodeInTree(node);
+            CheckNodeInOriginalTreeOrTracked(node);
             _changes.Add(new RemoveChange(node, options));
         }
 
@@ -99,14 +126,26 @@ public void RemoveNode(SyntaxNode node, SyntaxRemoveOptions options)
         /// The node passed into the compute function includes changes from prior edits. It will not appear as a descendant of the original root.</param>
         public void ReplaceNode(SyntaxNode node, Func<SyntaxNode, SyntaxGenerator, SyntaxNode> computeReplacement)
         {
-            CheckNodeInTree(node);
-            _changes.Add(new ReplaceChange(node, computeReplacement));
+            CheckNodeInOriginalTreeOrTracked(node);
+            if (computeReplacement == null)
+            {
+                throw new ArgumentNullException(nameof(computeReplacement));
+            }
+
+            _allowEditsOnLazilyCreatedTrackedNewNodes = true;
+            _changes.Add(new ReplaceChange(node, computeReplacement, this));
         }
 
         internal void ReplaceNode<TArgument>(SyntaxNode node, Func<SyntaxNode, SyntaxGenerator, TArgument, SyntaxNode> computeReplacement, TArgument argument)
         {
-            CheckNodeInTree(node);
-            _changes.Add(new ReplaceChange<TArgument>(node, computeReplacement, argument));
+            CheckNodeInOriginalTreeOrTracked(node);
+            if (computeReplacement == null)
+            {
+                throw new ArgumentNullException(nameof(computeReplacement));
+            }
+
+            _allowEditsOnLazilyCreatedTrackedNewNodes = true;
+            _changes.Add(new ReplaceChange<TArgument>(node, computeReplacement, argument, this));
         }
 
         /// <summary>
@@ -116,13 +155,14 @@ internal void ReplaceNode<TArgument>(SyntaxNode node, Func<SyntaxNode, SyntaxGen
         /// <param name="newNode">The new node that will be placed into the tree in the existing node's location.</param>
         public void ReplaceNode(SyntaxNode node, SyntaxNode newNode)
         {
-            CheckNodeInTree(node);
+            CheckNodeInOriginalTreeOrTracked(node);
             if (node == newNode)
             {
                 return;
             }
 
-            this.ReplaceNode(node, (n, g) => newNode);
+            newNode = ApplyTrackingToNewNode(newNode);
+            _changes.Add(new ReplaceChange(node, (n, g) => newNode, this));
         }
 
         /// <summary>
@@ -132,7 +172,13 @@ public void ReplaceNode(SyntaxNode node, SyntaxNode newNode)
         /// <param name="newNodes">The nodes to place before the existing node. These nodes must be of a compatible type to be placed in the same list containing the existing node.</param>
         public void InsertBefore(SyntaxNode node, IEnumerable<SyntaxNode> newNodes)
         {
-            CheckNodeInTree(node);
+            CheckNodeInOriginalTreeOrTracked(node);
+            if (newNodes == null)
+            {
+                throw new ArgumentNullException(nameof(newNodes));
+            }
+
+            newNodes = ApplyTrackingToNewNodes(newNodes);
             _changes.Add(new InsertChange(node, newNodes, isBefore: true));
         }
 
@@ -142,10 +188,7 @@ public void InsertBefore(SyntaxNode node, IEnumerable<SyntaxNode> newNodes)
         /// <param name="node">The node already existing in the tree that the new nodes will be placed before. This must be a node this is contained within a syntax list.</param>
         /// <param name="newNode">The node to place before the existing node. This node must be of a compatible type to be placed in the same list containing the existing node.</param>
         public void InsertBefore(SyntaxNode node, SyntaxNode newNode)
-        {
-            CheckNodeInTree(node);
-            this.InsertBefore(node, new[] { newNode });
-        }
+            => InsertBefore(node, new[] { newNode });
 
         /// <summary>
         /// Insert the new nodes after the specified node already existing in the tree.
@@ -154,7 +197,13 @@ public void InsertBefore(SyntaxNode node, SyntaxNode newNode)
         /// <param name="newNodes">The nodes to place after the existing node. These nodes must be of a compatible type to be placed in the same list containing the existing node.</param>
         public void InsertAfter(SyntaxNode node, IEnumerable<SyntaxNode> newNodes)
         {
-            CheckNodeInTree(node);
+            CheckNodeInOriginalTreeOrTracked(node);
+            if (newNodes == null)
+            {
+                throw new ArgumentNullException(nameof(newNodes));
+            }
+
+            newNodes = ApplyTrackingToNewNodes(newNodes);
             _changes.Add(new InsertChange(node, newNodes, isBefore: false));
         }
 
@@ -164,17 +213,37 @@ public void InsertAfter(SyntaxNode node, IEnumerable<SyntaxNode> newNodes)
         /// <param name="node">The node already existing in the tree that the new nodes will be placed after. This must be a node this is contained within a syntax list.</param>
         /// <param name="newNode">The node to place after the existing node. This node must be of a compatible type to be placed in the same list containing the existing node.</param>
         public void InsertAfter(SyntaxNode node, SyntaxNode newNode)
-        {
-            CheckNodeInTree(node);
-            this.InsertAfter(node, new[] { newNode });
-        }
+            => this.InsertAfter(node, new[] { newNode });
 
-        private void CheckNodeInTree(SyntaxNode node)
+        private void CheckNodeInOriginalTreeOrTracked(SyntaxNode node)
         {
-            if (!OriginalRoot.Contains(node))
+            if (node == null)
+            {
+                throw new ArgumentNullException(nameof(node));
+            }
+
+            if (OriginalRoot.Contains(node))
             {
-                throw new ArgumentException(Microsoft.CodeAnalysis.WorkspacesResources.The_node_is_not_part_of_the_tree, nameof(node));
+                // Node is contained in the original tree.
+                return;
+            }
+
+            if (_allowEditsOnLazilyCreatedTrackedNewNodes)
+            {
+                // This could be a node that is handed to us lazily from one of the prior edits
+                // which support lazy replacement nodes, we conservatively avoid throwing here.
+                // If this was indeed an unsupported node, syntax editor will throw an exception later
+                // when attempting to compute changed root.
+                return;
             }
+            
+            if(_lazyTrackedNewNodesOpt?.Contains(node) == true)
+            {
+                // Node is one of the new nodes, which is already tracked and supported.
+                return;
+            }
+
+            throw new ArgumentException(WorkspacesResources.The_node_is_not_part_of_the_tree, nameof(node));
         }
 
         private abstract class Change
@@ -221,17 +290,23 @@ public override SyntaxNode Apply(SyntaxNode root, SyntaxGenerator generator)
         private class ReplaceChange : Change
         {
             private readonly Func<SyntaxNode, SyntaxGenerator, SyntaxNode> _modifier;
+            private readonly SyntaxEditor _editor;
 
-            public ReplaceChange(SyntaxNode node, Func<SyntaxNode, SyntaxGenerator, SyntaxNode> modifier)
+            public ReplaceChange(
+                SyntaxNode node,
+                Func<SyntaxNode, SyntaxGenerator, SyntaxNode> modifier,
+                SyntaxEditor editor)
                 : base(node)
             {
                 _modifier = modifier;
+                _editor = editor;
             }
 
             public override SyntaxNode Apply(SyntaxNode root, SyntaxGenerator generator)
             {
                 var current = root.GetCurrentNode(this.Node);
                 var newNode = _modifier(current, generator);
+                newNode = _editor.ApplyTrackingToNewNode(newNode);
                 return generator.ReplaceNode(root, current, newNode);
             }
         }
@@ -240,21 +315,25 @@ private class ReplaceChange<TArgument> : Change
         {
             private readonly Func<SyntaxNode, SyntaxGenerator, TArgument, SyntaxNode> _modifier;
             private readonly TArgument _argument;
+            private readonly SyntaxEditor _editor;
 
             public ReplaceChange(
                 SyntaxNode node,
                 Func<SyntaxNode, SyntaxGenerator, TArgument, SyntaxNode> modifier,
-                TArgument argument)
+                TArgument argument,
+                SyntaxEditor editor)
                 : base(node)
             {
                 _modifier = modifier;
                 _argument = argument;
+                _editor = editor;
             }
 
             public override SyntaxNode Apply(SyntaxNode root, SyntaxGenerator generator)
             {
                 var current = root.GetCurrentNode(this.Node);
                 var newNode = _modifier(current, generator, _argument);
+                newNode = _editor.ApplyTrackingToNewNode(newNode);
                 return generator.ReplaceNode(root, current, newNode);
             }
         }