Merge pull request #29799 from DustinCampbell/code-model-bug

Fix issue when adding/removing event handlers again and again

src/VisualStudio/Core/Impl/CodeModel/FileCodeModel.cs

@@ -189,14 +189,22 @@ internal void UpdateCodeElementNodeKey(AbstractKeyedCodeElement keyedElement, Sy
                 throw new InvalidOperationException($"Unexpected failure in Code Model while updating node keys {oldNodeKey} -> {newNodeKey}");
             }
 
+            // If we're updating this element with the same node key as an element that's already in the table,
+            // just remove the old element. The old element will continue to function (through its node key), but
+            // the new element will replace it in the cache.
+            if (_codeElementTable.ContainsKey(newNodeKey))
+            {
+                _codeElementTable.Remove(newNodeKey);
+            }
+
             _codeElementTable.Add(newNodeKey, codeElement);
         }
 
         internal void OnCodeElementCreated(SyntaxNodeKey nodeKey, EnvDTE.CodeElement element)
         {
-            // If we're creating an element with the same node key as an element that's already in the table, just remove
-            // the old element. The old element will continue to function but the new element will replace it in the cache.
-
+            // If we're updating this element with the same node key as an element that's already in the table,
+            // just remove the old element. The old element will continue to function (through its node key), but
+            // the new element will replace it in the cache.
             if (_codeElementTable.ContainsKey(nodeKey))
             {
                 _codeElementTable.Remove(nodeKey);
@@ -332,7 +340,7 @@ internal void PerformEdit(Func<Document, Document> action)
             });
         }
 
-        private void ApplyChanges(Microsoft.CodeAnalysis.Workspace workspace, Document document)
+        private void ApplyChanges(Workspace workspace, Document document)
         {
             if (IsBatchOpen)
             {

src/VisualStudio/Core/Test/CodeModel/CSharp/FileCodeModelTests.vb

@@ -4,6 +4,7 @@ Imports System.Threading.Tasks
 Imports Microsoft.CodeAnalysis
 Imports Microsoft.CodeAnalysis.Editor.UnitTests.Workspaces
 Imports Microsoft.CodeAnalysis.Test.Utilities
+Imports Microsoft.CodeAnalysis.Text
 Imports Microsoft.VisualStudio.LanguageServices.Implementation.CodeModel
 Imports Microsoft.VisualStudio.LanguageServices.Implementation.CodeModel.InternalElements
 Imports Microsoft.VisualStudio.LanguageServices.Implementation.Interop
@@ -1167,6 +1168,68 @@ class C
                 End Sub)
         End Sub
 
+        <WorkItem(671189, "https://devdiv.visualstudio.com/DevDiv/_workitems/edit/671189")>
+        <ConditionalWpfFact(GetType(x86)), Trait(Traits.Feature, Traits.Features.CodeModel)>
+        Public Async Function AddShouldNotFailAfterCodeIsDeleted() As Task
+            ' This test attempts to add and remove a method several times via code model,
+            ' verifying a scenario where the WinForms or XAML designer adds an event handler
+            ' and the user later deletes and regenerates the event handler.
+
+            Dim codeBeforeOperationXml =
+<code>
+class C
+{
+}
+</code>
+
+            Dim codeAfterOperationXml =
+<code>
+class C
+{
+    void M(int x, int y)
+    {
+
+    }
+}
+</code>
+
+            Dim codeBeforeOperation = codeBeforeOperationXml.Value.NormalizeLineEndings().Trim()
+            Dim codeAfterOperation = codeAfterOperationXml.Value.NormalizeLineEndings().Trim()
+
+            Using state = CreateCodeModelTestState(GetWorkspaceDefinition(codeBeforeOperationXml))
+                Dim workspace = state.VisualStudioWorkspace
+                Dim fileCodeModel = state.FileCodeModel
+                Assert.NotNull(fileCodeModel)
+
+                For i = 1 To 10
+                    Dim docId = workspace.CurrentSolution.Projects(0).DocumentIds(0)
+                    Dim textBeforeOperation = Await workspace.CurrentSolution.GetDocument(docId).GetTextAsync()
+                    Assert.Equal(codeBeforeOperation, textBeforeOperation.ToString())
+
+                    Dim classC = TryCast(fileCodeModel.CodeElements.Item(1), EnvDTE.CodeClass)
+                    Assert.NotNull(classC)
+                    Assert.Equal("C", classC.Name)
+
+                    fileCodeModel.BeginBatch()
+                    Dim newFunction = classC.AddFunction("M", EnvDTE.vsCMFunction.vsCMFunctionFunction, Type:=EnvDTE.vsCMTypeRef.vsCMTypeRefVoid)
+                    Dim param1 = newFunction.AddParameter("x", EnvDTE.vsCMTypeRef.vsCMTypeRefInt, Position:=-1)
+                    Dim param2 = newFunction.AddParameter("y", EnvDTE.vsCMTypeRef.vsCMTypeRefInt, Position:=-1)
+                    fileCodeModel.EndBatch()
+
+                    Dim solution = workspace.CurrentSolution
+                    Dim textAfterOperation = Await solution.GetDocument(docId).GetTextAsync()
+                    Assert.Equal(codeAfterOperation, textAfterOperation.ToString())
+
+                    Dim newText = textAfterOperation.Replace(
+                        span:=TextSpan.FromBounds(0, textAfterOperation.Length),
+                        newText:=codeBeforeOperation)
+
+                    Dim newSolution = solution.WithDocumentText(docId, newText)
+                    Assert.True(workspace.TryApplyChanges(newSolution))
+                Next
+            End Using
+        End Function
+
         Protected Overrides ReadOnly Property LanguageName As String
             Get
                 Return LanguageNames.CSharp