From ec766cd3176823aad6ee850324d0e0159059134e Mon Sep 17 00:00:00 2001
From: Kevin Pilch <me@pilchie.com>
Date: Thu, 1 Jun 2017 18:39:57 -0700
Subject: [PATCH] Refcount document monikers for source and additional files

Turns out we could end up with the same IVisualStudioHostDocument
for both, if they are both added.

Fixes #19968.
---
 .../ProjectSystem/AbstractProject.cs          | 51 ++++++++++++++++---
 1 file changed, 44 insertions(+), 7 deletions(-)

diff --git a/src/VisualStudio/Core/Def/Implementation/ProjectSystem/AbstractProject.cs b/src/VisualStudio/Core/Def/Implementation/ProjectSystem/AbstractProject.cs
index d4bd5d062e4..9ee7f053638 100644
--- a/src/VisualStudio/Core/Def/Implementation/ProjectSystem/AbstractProject.cs
+++ b/src/VisualStudio/Core/Def/Implementation/ProjectSystem/AbstractProject.cs
@@ -39,7 +39,7 @@ internal abstract partial class AbstractProject : ForegroundThreadAffinitizedObj
         private readonly List<ProjectReference> _projectReferences = new List<ProjectReference>();
         private readonly List<VisualStudioMetadataReference> _metadataReferences = new List<VisualStudioMetadataReference>();
         private readonly Dictionary<DocumentId, IVisualStudioHostDocument> _documents = new Dictionary<DocumentId, IVisualStudioHostDocument>();
-        private readonly Dictionary<string, IVisualStudioHostDocument> _documentMonikers = new Dictionary<string, IVisualStudioHostDocument>(StringComparer.OrdinalIgnoreCase);
+        private readonly Dictionary<string, (IVisualStudioHostDocument document, int refCount)> _documentMonikers = new Dictionary<string, (IVisualStudioHostDocument, int)>(StringComparer.OrdinalIgnoreCase);
         private readonly Dictionary<string, VisualStudioAnalyzer> _analyzers = new Dictionary<string, VisualStudioAnalyzer>(StringComparer.OrdinalIgnoreCase);
         private readonly Dictionary<DocumentId, IVisualStudioHostDocument> _additionalDocuments = new Dictionary<DocumentId, IVisualStudioHostDocument>();
 
@@ -410,8 +410,9 @@ public IVisualStudioHostDocument GetCurrentDocumentFromPath(string filePath)
         {
             lock (_gate)
             {
-                _documentMonikers.TryGetValue(filePath, out var document);
-                return document;
+                return _documentMonikers.TryGetValue(filePath, out var value)
+                    ? value.document
+                    : null;
             }
         }
 
@@ -1015,7 +1016,7 @@ internal void AddDocument(IVisualStudioHostDocument document, bool isCurrentCont
                 lock (_gate)
                 {
                     _documents.Add(document.Id, document);
-                    _documentMonikers.Add(document.Key.Moniker, document);
+                    AddMoniker(document);
                 }
 
                 if (_pushingChangesToWorkspaceHosts)
@@ -1054,7 +1055,7 @@ internal void RemoveDocument(IVisualStudioHostDocument document)
                 lock (_gate)
                 {
                     _documents.Remove(document.Id);
-                    _documentMonikers.Remove(document.Key.Moniker);
+                    RemoveMoniker(document);
                 }
 
                 UninitializeDocument(document);
@@ -1069,7 +1070,7 @@ internal void AddAdditionalDocument(IVisualStudioHostDocument document, bool isC
             lock (_gate)
             {
                 _additionalDocuments.Add(document.Id, document);
-                _documentMonikers.Add(document.Key.Moniker, document);
+                AddMoniker(document);
             }
 
             if (_pushingChangesToWorkspaceHosts)
@@ -1097,12 +1098,48 @@ internal void RemoveAdditionalDocument(IVisualStudioHostDocument document)
             lock (_gate)
             {
                 _additionalDocuments.Remove(document.Id);
-                _documentMonikers.Remove(document.Key.Moniker);
+                RemoveMoniker(document);
             }
 
             UninitializeAdditionalDocument(document);
         }
 
+        private void AddMoniker(IVisualStudioHostDocument document)
+        {
+            var moniker = document.Key.Moniker;
+            if (_documentMonikers.TryGetValue(moniker, out var value))
+            {
+                value.refCount++;
+                _documentMonikers[moniker] = (value.document, value.refCount);
+            }
+            else
+            {
+                _documentMonikers.Add(moniker, (document, 1));
+            }
+        }
+
+        private void RemoveMoniker(IVisualStudioHostDocument document)
+        {
+            var moniker = document.Key.Moniker;
+            if (_documentMonikers.TryGetValue(moniker, out var value))
+            {
+                Debug.Assert(value.document.Equals(document));
+                value.refCount--;
+                if (value.refCount == 0)
+                {
+                    _documentMonikers.Remove(moniker);
+                }
+                else
+                {
+                    _documentMonikers[moniker] = (value.document, value.refCount);
+                }
+            }
+            else
+            {
+                Debug.Fail($"Couldn't find '{moniker}' in {nameof(_documentMonikers)} to remove it.");
+            }
+        }
+
         public virtual void Disconnect()
         {
             AssertIsForeground();
-- 
GitLab