Remove descendent node tracking from SyntaxEditor

src/Features/Core/Portable/PullMemberUp/MembersPuller.cs

@@ -32,7 +32,8 @@ internal static class MembersPuller
         /// Annotation used to mark imports that we move over, so that we can remove these imports if they are unnecessary
         /// (and so we don't remove any other unnecessary imports)
         /// </summary>
-        private static readonly SyntaxAnnotation s_annotation = new("PullMemberRemovableImport");
+        private static readonly SyntaxAnnotation s_removableImportAnnotation = new("PullMemberRemovableImport");
+        private static readonly SyntaxAnnotation s_destinationNodeAnnotation = new("DestinationNode");
 
         public static CodeAction TryComputeCodeAction(
             Document document,
@@ -312,7 +313,9 @@ private static async Task<Solution> PullMembersIntoClassAsync(
             var options = await destinationEditor.OriginalDocument.GetCleanCodeGenerationOptionsAsync(fallbackOptions, cancellationToken).ConfigureAwait(false);
             var info = codeGenerationService.GetInfo(context, options.GenerationOptions, destinationEditor.OriginalDocument.Project.ParseOptions);
 
-            var newDestination = codeGenerationService.AddMembers(destinationSyntaxNode, pullUpMembersSymbols, info, cancellationToken);
+            var newDestination = codeGenerationService
+                .AddMembers(destinationSyntaxNode, pullUpMembersSymbols, info, cancellationToken)
+                .WithAdditionalAnnotations(s_destinationNodeAnnotation);
             using var _ = PooledHashSet<SyntaxNode>.GetInstance(out var sourceImports);
 
             var syntaxFacts = destinationEditor.OriginalDocument.GetRequiredLanguageService<ISyntaxFactsService>();
@@ -329,7 +332,7 @@ private static async Task<Solution> PullMembersIntoClassAsync(
                     sourceImports.Add(
                         destinationEditor.Generator.NamespaceImportDeclaration(
                             resultNamespace.ToDisplayString(SymbolDisplayFormats.NameFormat))
-                        .WithAdditionalAnnotations(s_annotation));
+                        .WithAdditionalAnnotations(s_removableImportAnnotation));
                 }
 
                 foreach (var syntax in symbolToDeclarations[analysisResult.Member])
@@ -342,7 +345,7 @@ private static async Task<Solution> PullMembersIntoClassAsync(
                         .Select(import => import
                             .WithoutLeadingTrivia()
                             .WithTrailingTrivia(originalMemberEditor.Generator.ElasticCarriageReturnLineFeed)
-                            .WithAdditionalAnnotations(s_annotation)));
+                            .WithAdditionalAnnotations(s_removableImportAnnotation)));
 
                     if (!analysisResult.MakeMemberDeclarationAbstract || analysisResult.Member.IsAbstract)
                     {
@@ -378,7 +381,7 @@ private static async Task<Solution> PullMembersIntoClassAsync(
             destinationEditor.ReplaceNode(destinationEditor.OriginalRoot, (node, generator) => addImportsService.AddImports(
                 destinationEditor.SemanticModel.Compilation,
                 node,
-                node.GetCurrentNode(newDestination),
+                node.GetAnnotatedNodes(s_destinationNodeAnnotation).FirstOrDefault(),
                 sourceImports,
                 generator,
                 options.CleanupOptions.AddImportOptions,
@@ -387,12 +390,12 @@ private static async Task<Solution> PullMembersIntoClassAsync(
             var removeImportsService = destinationEditor.OriginalDocument.GetRequiredLanguageService<IRemoveUnnecessaryImportsService>();
             var destinationDocument = await removeImportsService.RemoveUnnecessaryImportsAsync(
                 destinationEditor.GetChangedDocument(),
-                node => node.HasAnnotation(s_annotation),
+                node => node.HasAnnotation(s_removableImportAnnotation),
                 options.CleanupOptions.FormattingOptions,
                 cancellationToken).ConfigureAwait(false);
 
             // Format whitespace trivia within the import statements we pull up
-            destinationDocument = await Formatter.FormatAsync(destinationDocument, s_annotation, options.CleanupOptions.FormattingOptions, cancellationToken).ConfigureAwait(false);
+            destinationDocument = await Formatter.FormatAsync(destinationDocument, s_removableImportAnnotation, options.CleanupOptions.FormattingOptions, cancellationToken).ConfigureAwait(false);
 
             var destinationRoot = AddLeadingTriviaBeforeFirstMember(
                 await destinationDocument.GetSyntaxRootAsync(cancellationToken).ConfigureAwait(false),

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

@@ -4,22 +4,59 @@
 
 using System;
 using System.Collections.Generic;
-using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using Microsoft.CodeAnalysis.Host;
 using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.Editing
 {
     /// <summary>
-    /// An editor for making changes to a syntax tree. 
+    /// An editor for making changes to a syntax tree. The editor works by giving a list of changes to perform to a
+    /// particular tree <em>in order</em>.  Changes are given a <see cref="SyntaxNode"/> they will apply to in the
+    /// original tree the editor is created for.  The semantics of application are as follows:
+    /// 
+    /// <list type="number">
+    /// <item>
+    /// The original root provided is used as the 'current' root for all operations.  This 'current' root will
+    /// continually be updated, becoming the new 'current' root.  The original root is never changed.
+    /// </item>
+    /// <item>
+    /// Each change has its given <see cref="SyntaxNode"/> tracked, using a <see cref="SyntaxAnnotation"/>, producing a
+    /// 'current' root that tracks all of them.  This allows that same node to be found after prior changes are applied
+    /// which mutate the tree.
+    /// </item>
+    /// <item>
+    /// Each change is then applied in order it was added to the editor.
+    /// </item>
+    /// <item>
+    /// A change first attempts to find its <see cref="SyntaxNode"/> in the 'current' root.  If that node cannot be
+    /// found, the operation will fail with an <see cref="ArgumentException"/>.
+    /// </item>
+    /// <item>
+    /// The particular change will run on that node, removing, replacing, or inserting around it according to the
+    /// change.  If the change is passed a delegate as its 'compute' argument, it will be given the <see
+    /// cref="SyntaxNode"/> found in the current root.  The 'current' root will then be updated by replacing the current
+    /// node with the new computed node.
+    /// </item>
+    /// <item>
+    /// The 'current' root is then returned.
+    /// </item>
+    /// </list>
     /// </summary>
+    /// <remarks>
+    /// The above editing strategy makes it an error for a client of the editor to add a change that updates a parent
+    /// node and then adds a change that updates a child node (unless the parent change is certain to contain the
+    /// child), and attempting this will throw at runtime.  If a client ever needs to update both a child and a parent,
+    /// it <em>should</em> add the child change first, and then the parent change.  And the parent change should pass an
+    /// appropriate 'compute' callback so it will see the results of the child change.
+    /// <para/> If a client wants to make a replacement, then find the <em>value</em> <see cref="SyntaxNode"/> put into
+    /// the tree, that can be done by adding a dedicated annotation to that node and then looking it back up in the
+    /// 'current' node passed to a 'compute' callback.
+    /// </remarks>
     public class SyntaxEditor
     {
         private readonly SyntaxGenerator _generator;
         private readonly List<Change> _changes = new();
-        private bool _allowEditsOnLazilyCreatedTrackedNewNodes;
-        private HashSet<SyntaxNode>? _lazyTrackedNewNodesOpt;
 
         /// <summary>
         /// Creates a new <see cref="SyntaxEditor"/> instance.
@@ -53,32 +90,6 @@ internal SyntaxEditor(SyntaxNode root, SyntaxGenerator generator)
             _generator = generator;
         }
 
-        [return: NotNullIfNotNull(nameof(node))]
-        private SyntaxNode? ApplyTrackingToNewNode(SyntaxNode? node)
-        {
-            if (node == null)
-            {
-                return null;
-            }
-
-            _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)
-            {
-                var result = ApplyTrackingToNewNode(node);
-                yield return result;
-            }
-        }
-
         /// <summary>
         /// The <see cref="SyntaxNode"/> that was specified when the <see cref="SyntaxEditor"/> was constructed.
         /// </summary>
@@ -99,9 +110,7 @@ public SyntaxNode GetChangedRoot()
             var newRoot = OriginalRoot.TrackNodes(nodes);
 
             foreach (var change in _changes)
-            {
                 newRoot = change.Apply(newRoot, _generator);
-            }
 
             return newRoot;
         }
@@ -143,36 +152,27 @@ public void ReplaceNode(SyntaxNode node, Func<SyntaxNode, SyntaxGenerator, Synta
         {
             CheckNodeInOriginalTreeOrTracked(node);
             if (computeReplacement == null)
-            {
                 throw new ArgumentNullException(nameof(computeReplacement));
-            }
 
-            _allowEditsOnLazilyCreatedTrackedNewNodes = true;
-            _changes.Add(new ReplaceChange(node, computeReplacement, this));
+            _changes.Add(new ReplaceChange(node, computeReplacement));
         }
 
         internal void ReplaceNode(SyntaxNode node, Func<SyntaxNode, SyntaxGenerator, IEnumerable<SyntaxNode>> computeReplacement)
         {
             CheckNodeInOriginalTreeOrTracked(node);
             if (computeReplacement == null)
-            {
                 throw new ArgumentNullException(nameof(computeReplacement));
-            }
 
-            _allowEditsOnLazilyCreatedTrackedNewNodes = true;
-            _changes.Add(new ReplaceWithCollectionChange(node, computeReplacement, this));
+            _changes.Add(new ReplaceWithCollectionChange(node, computeReplacement));
         }
 
         internal void ReplaceNode<TArgument>(SyntaxNode node, Func<SyntaxNode, SyntaxGenerator, TArgument, SyntaxNode> computeReplacement, TArgument argument)
         {
             CheckNodeInOriginalTreeOrTracked(node);
             if (computeReplacement == null)
-            {
                 throw new ArgumentNullException(nameof(computeReplacement));
-            }
 
-            _allowEditsOnLazilyCreatedTrackedNewNodes = true;
-            _changes.Add(new ReplaceChange<TArgument>(node, computeReplacement, argument, this));
+            _changes.Add(new ReplaceChange<TArgument>(node, computeReplacement, argument));
         }
 
         /// <summary>
@@ -184,12 +184,9 @@ public void ReplaceNode(SyntaxNode node, SyntaxNode newNode)
         {
             CheckNodeInOriginalTreeOrTracked(node);
             if (node == newNode)
-            {
                 return;
-            }
 
-            newNode = ApplyTrackingToNewNode(newNode);
-            _changes.Add(new ReplaceChange(node, (n, g) => newNode, this));
+            _changes.Add(new ReplaceChange(node, (n, g) => newNode));
         }
 
         /// <summary>
@@ -201,11 +198,8 @@ public void InsertBefore(SyntaxNode node, IEnumerable<SyntaxNode> newNodes)
         {
             CheckNodeInOriginalTreeOrTracked(node);
             if (newNodes == null)
-            {
                 throw new ArgumentNullException(nameof(newNodes));
-            }
 
-            newNodes = ApplyTrackingToNewNodes(newNodes);
             _changes.Add(new InsertChange(node, newNodes, isBefore: true));
         }
 
@@ -226,11 +220,8 @@ public void InsertAfter(SyntaxNode node, IEnumerable<SyntaxNode> newNodes)
         {
             CheckNodeInOriginalTreeOrTracked(node);
             if (newNodes == null)
-            {
                 throw new ArgumentNullException(nameof(newNodes));
-            }
 
-            newNodes = ApplyTrackingToNewNodes(newNodes);
             _changes.Add(new InsertChange(node, newNodes, isBefore: false));
         }
 
@@ -245,30 +236,10 @@ public void InsertAfter(SyntaxNode node, SyntaxNode newNode)
         private void CheckNodeInOriginalTreeOrTracked(SyntaxNode node)
         {
             if (node == null)
-            {
                 throw new ArgumentNullException(nameof(node));
-            }
 
             if (OriginalRoot.Contains(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));
         }
@@ -284,9 +255,7 @@ public SyntaxNode Apply(SyntaxNode root, SyntaxGenerator generator)
             {
                 var currentNode = root.GetCurrentNode(OriginalNode);
                 if (currentNode is null)
-                {
                     Contract.Fail($"GetCurrentNode returned null with the following node: {OriginalNode}");
-                }
 
                 return Apply(root, currentNode, generator);
             }
@@ -325,78 +294,53 @@ protected override SyntaxNode Apply(SyntaxNode root, SyntaxNode currentNode, Syn
         private sealed class ReplaceChange : Change
         {
             private readonly Func<SyntaxNode, SyntaxGenerator, SyntaxNode?> _modifier;
-            private readonly SyntaxEditor _editor;
 
             public ReplaceChange(
                 SyntaxNode node,
-                Func<SyntaxNode, SyntaxGenerator, SyntaxNode?> modifier,
-                SyntaxEditor editor)
+                Func<SyntaxNode, SyntaxGenerator, SyntaxNode?> modifier)
                 : base(node)
             {
                 Contract.ThrowIfNull(node);
                 _modifier = modifier;
-                _editor = editor;
             }
 
             protected override SyntaxNode Apply(SyntaxNode root, SyntaxNode currentNode, SyntaxGenerator generator)
-            {
-                var newNode = _modifier(currentNode, generator);
-                newNode = _editor.ApplyTrackingToNewNode(newNode);
-                return ValidateNewRoot(generator.ReplaceNode(root, currentNode, newNode));
-            }
+                => ValidateNewRoot(generator.ReplaceNode(root, currentNode, _modifier(currentNode, generator)));
         }
 
         private sealed class ReplaceWithCollectionChange : Change
         {
             private readonly Func<SyntaxNode, SyntaxGenerator, IEnumerable<SyntaxNode>> _modifier;
-            private readonly SyntaxEditor _editor;
 
             public ReplaceWithCollectionChange(
                 SyntaxNode node,
-                Func<SyntaxNode, SyntaxGenerator, IEnumerable<SyntaxNode>> modifier,
-                SyntaxEditor editor)
+                Func<SyntaxNode, SyntaxGenerator, IEnumerable<SyntaxNode>> modifier)
                 : base(node)
             {
                 _modifier = modifier;
-                _editor = editor;
             }
 
             protected override SyntaxNode Apply(SyntaxNode root, SyntaxNode currentNode, SyntaxGenerator generator)
-            {
-                var newNodes = _modifier(currentNode, generator).ToList();
-                for (var i = 0; i < newNodes.Count; i++)
-                {
-                    newNodes[i] = _editor.ApplyTrackingToNewNode(newNodes[i]);
-                }
-
-                return SyntaxGenerator.ReplaceNode(root, currentNode, newNodes);
-            }
+                => SyntaxGenerator.ReplaceNode(root, currentNode, _modifier(currentNode, generator));
         }
 
         private sealed 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,
-                SyntaxEditor editor)
+                TArgument argument)
                 : base(node)
             {
                 _modifier = modifier;
                 _argument = argument;
-                _editor = editor;
             }
 
             protected override SyntaxNode Apply(SyntaxNode root, SyntaxNode currentNode, SyntaxGenerator generator)
-            {
-                var newNode = _modifier(currentNode, generator, _argument);
-                newNode = _editor.ApplyTrackingToNewNode(newNode);
-                return ValidateNewRoot(generator.ReplaceNode(root, currentNode, newNode));
-            }
+                => ValidateNewRoot(generator.ReplaceNode(root, currentNode, _modifier(currentNode, generator, _argument)));
         }
 
         private sealed class InsertChange : Change