Address review feedback

src/EditorFeatures/CSharpTest/Workspaces/WorkspaceTests.cs

@@ -887,7 +887,7 @@ public async Task TestAdditionalFile_DocumentChanged()
             }
         }
 
-        [Fact]
+        [Fact, WorkItem(31540, "https://github.com/dotnet/roslyn/issues/31540")]
         public async Task TestAdditionalFile_OpenClose()
         {
             using (var workspace = CreateWorkspace())
@@ -905,12 +905,16 @@ public async Task TestAdditionalFile_OpenClose()
 
                 workspace.OnAdditionalDocumentOpened(additionalDoc.Id, additionalDoc.GetOpenTextContainer());
 
-                // We don't have a GetOpenAdditionalDocumentIds since we don't need it. But make sure additional documents
-                // don't creep into OpenDocumentIds (Bug: 1087470)
-                Assert.Empty(workspace.GetOpenDocumentIds());
+                // Make sure that additional documents are included in GetOpenDocumentIds.
+                var openDocumentIds = workspace.GetOpenDocumentIds();
+                Assert.Equal(1, openDocumentIds.Count());
+                Assert.Equal(additionalDoc.Id, openDocumentIds.Single());
 
                 workspace.OnAdditionalDocumentClosed(additionalDoc.Id, TextLoader.From(TextAndVersion.Create(text, version)));
 
+                // Make sure that closed additional documents are not include in GetOpenDocumentIds.
+                Assert.Empty(workspace.GetOpenDocumentIds());
+
                 // Reopen and close to make sure we are not leaking anything.
                 workspace.OnAdditionalDocumentOpened(additionalDoc.Id, additionalDoc.GetOpenTextContainer());
                 workspace.OnAdditionalDocumentClosed(additionalDoc.Id, TextLoader.From(TextAndVersion.Create(text, version)));
@@ -986,6 +990,27 @@ public void TestAdditionalFile_AddRemove_FromProject()
             }
         }
 
+        [Fact, WorkItem(31540, "https://github.com/dotnet/roslyn/issues/31540")]
+        public void TestAdditionalFile_GetDocumentIdsWithFilePath()
+        {
+            using (var workspace = CreateWorkspace())
+            {
+                const string docFilePath = "filePath1", additionalDocFilePath = "filePath2";
+                var document = new TestHostDocument("public class C { }", filePath: docFilePath);
+                var additionalDoc = new TestHostDocument(@"<setting value = ""goo""", filePath: additionalDocFilePath);
+                var project1 = new TestHostProject(workspace, name: "project1", documents: new[] { document }, additionalDocuments: new[] { additionalDoc });
+                workspace.AddTestProject(project1);
+
+                var documentIdsWithFilePath = workspace.CurrentSolution.GetDocumentIdsWithFilePath(docFilePath);
+                Assert.Equal(1, documentIdsWithFilePath.Count());
+                Assert.Equal(document.Id, documentIdsWithFilePath.Single());
+
+                documentIdsWithFilePath = workspace.CurrentSolution.GetDocumentIdsWithFilePath(additionalDocFilePath);
+                Assert.Equal(1, documentIdsWithFilePath.Count());
+                Assert.Equal(additionalDoc.Id, documentIdsWithFilePath.Single());
+            }
+        }
+
         [Fact, WorkItem(209299, "https://devdiv.visualstudio.com/DevDiv/_workitems?id=209299")]
         public async Task TestLinkedFilesStayInSync()
         {

src/EditorFeatures/Core.Wpf/InlineRename/InlineRenameSession.cs

@@ -179,6 +179,11 @@ private void InitializeOpenBuffers(SnapshotSpan triggerSpan)
                 foreach (var d in _workspace.GetOpenDocumentIds())
                 {
                     var document = _baseSolution.GetDocument(d);
+                    if (document == null)
+                    {
+                        continue;
+                    }
+
                     Contract.ThrowIfFalse(document.TryGetText(out var text));
                     Contract.ThrowIfNull(text);
 

src/VisualStudio/Core/Def/Implementation/EditAndContinue/VsReadOnlyDocumentTracker.cs

@@ -54,9 +54,10 @@ internal void UpdateWorkspaceDocuments()
             foreach (var documentId in _workspace.GetOpenDocumentIds())
             {
                 var document = _workspace.CurrentSolution.GetDocument(documentId);
-                Debug.Assert(document != null);
-
-                SetReadOnly(document);
+                if (document != null)
+                {
+                    SetReadOnly(document);
+                }
             }
         }
 

src/VisualStudio/Core/Def/Implementation/ProjectSystem/VisualStudioProject.cs

@@ -347,7 +347,8 @@ private void OnBatchScopeDisposed()
                 }
 
                 var documentFileNamesAdded = ImmutableArray.CreateBuilder<string>();
-                var documentsToOpen = new List<(DocumentId documentId, SourceTextContainer textContainer, bool isAdditionalDocument)>();
+                var documentsToOpen = new List<(DocumentId documentId, SourceTextContainer textContainer)>();
+                var additionalDocumentsToOpen = new List<(DocumentId documentId, SourceTextContainer textContainer)>();
 
                 _workspace.ApplyBatchChangeToProject(Id, solution =>
                 {
@@ -355,7 +356,6 @@ private void OnBatchScopeDisposed()
                         solution,
                         documentFileNamesAdded,
                         documentsToOpen,
-                        isUpdatingAdditionalDocuments: false,
                         (s, documents) => solution.AddDocuments(documents),
                         (s, id) =>
                         {
@@ -368,8 +368,7 @@ private void OnBatchScopeDisposed()
                     solution = _additionalFiles.UpdateSolutionForBatch(
                         solution,
                         documentFileNamesAdded,
-                        documentsToOpen,
-                        isUpdatingAdditionalDocuments: true,
+                        additionalDocumentsToOpen,
                         (s, documents) =>
                         {
                             foreach (var document in documents)
@@ -472,19 +471,14 @@ private void OnBatchScopeDisposed()
                     return solution;
                 });
 
-                foreach (var (documentId, textContainer, isAdditionalDocument) in documentsToOpen)
+                foreach (var (documentId, textContainer) in documentsToOpen)
                 {
-                    _workspace.ApplyChangeToWorkspace(w =>
-                    {
-                        if (!isAdditionalDocument)
-                        {
-                            w.OnDocumentOpened(documentId, textContainer);
-                        }
-                        else
-                        {
-                            w.OnAdditionalDocumentOpened(documentId, textContainer);
-                        }
-                    });
+                    _workspace.ApplyChangeToWorkspace(w => w.OnDocumentOpened(documentId, textContainer));
+                }
+
+                foreach (var (documentId, textContainer) in additionalDocumentsToOpen)
+                {
+                    _workspace.ApplyChangeToWorkspace(w => w.OnAdditionalDocumentOpened(documentId, textContainer));
                 }
 
                 // Check for those files being opened to start wire-up if necessary
@@ -1444,8 +1438,7 @@ public void ReorderFiles(ImmutableArray<string> filePaths)
             internal Solution UpdateSolutionForBatch(
                 Solution solution,
                 ImmutableArray<string>.Builder documentFileNamesAdded,
-                List<(DocumentId documentId, SourceTextContainer textContainer, bool isAdditionalDocument)> documentsToOpen,
-                bool isUpdatingAdditionalDocuments,
+                List<(DocumentId documentId, SourceTextContainer textContainer)> documentsToOpen,
                 Func<Solution, ImmutableArray<DocumentInfo>, Solution> addDocuments,
                 Func<Solution, DocumentId, Solution> removeDocument)
             {
@@ -1458,7 +1451,7 @@ public void ReorderFiles(ImmutableArray<string> filePaths)
 
                     if (_sourceTextContainersToDocumentIds.TryGetKey(documentInfo.Id, out var textContainer))
                     {
-                        documentsToOpen.Add((documentInfo.Id, textContainer, isUpdatingAdditionalDocuments));
+                        documentsToOpen.Add((documentInfo.Id, textContainer));
                     }
                 }
 

src/VisualStudio/Core/Def/Implementation/ProjectSystem/VisualStudioWorkspaceImpl.OpenFileTracker.cs

@@ -350,7 +350,7 @@ private void TryClosingDocumentsForCookie(uint cookie)
                 var moniker = _runningDocumentTable.GetDocumentMoniker(cookie);
                 _workspace.ApplyChangeToWorkspace(w =>
                 {
-                    var documentIds = _workspace.CurrentSolution.GetDocumentIdsWithFilePath(moniker);
+                    var documentIds = w.CurrentSolution.GetDocumentIdsWithFilePath(moniker);
                     if (documentIds.IsDefaultOrEmpty)
                     {
                         return;
@@ -358,14 +358,15 @@ private void TryClosingDocumentsForCookie(uint cookie)
 
                     foreach (var documentId in documentIds)
                     {
-                        if (!_workspace._documentsNotFromFiles.Contains(documentId))
+                        if (w.IsDocumentOpen(documentId) && !_workspace._documentsNotFromFiles.Contains(documentId))
                         {
-                            if (_workspace.IsDocumentOpen(documentId))
+                            if (w.CurrentSolution.ContainsDocument(documentId))
                             {
                                 w.OnDocumentClosed(documentId, new FileTextLoader(moniker, defaultEncoding: null));
                             }
-                            else if (w.CurrentSolution.ContainsAdditionalDocument(documentId))
+                            else
                             {
+                                Debug.Assert(w.CurrentSolution.ContainsAdditionalDocument(documentId));
                                 w.OnAdditionalDocumentClosed(documentId, new FileTextLoader(moniker, defaultEncoding: null));
                             }
                         }

src/Workspaces/Core/Portable/Workspace/Solution/SolutionState.cs

@@ -485,35 +485,30 @@ public SolutionState AddProject(ProjectInfo projectInfo)
 
         private ImmutableDictionary<string, ImmutableArray<DocumentId>> CreateFilePathToDocumentIdsMapWithAddedProject(ProjectState projectState)
         {
-            var documentIdsAndFilePaths = GetDocumentIdsAndFilePaths(projectState, projectState.DocumentIds, projectState.AdditionalDocumentIds);
-            return CreateFilePathToDocumentIdsMapWithAddedDocuments(documentIdsAndFilePaths);
+            return CreateFilePathToDocumentIdsMapWithAddedDocuments(GetDocumentStates(projectState));
         }
 
-        private ImmutableDictionary<string, ImmutableArray<DocumentId>> CreateFilePathToDocumentIdsMapWithAddedDocuments(IEnumerable<(DocumentId documentId, string filePath)> documentIdsAndFilePaths)
+        private ImmutableDictionary<string, ImmutableArray<DocumentId>> CreateFilePathToDocumentIdsMapWithAddedDocuments(IEnumerable<TextDocumentState> documentStates)
         {
             var builder = _filePathToDocumentIdsMap.ToBuilder();
 
-            foreach (var (documentId, filePath) in documentIdsAndFilePaths)
+            foreach (var documentState in documentStates)
             {
-                if (string.IsNullOrEmpty(filePath))
+                if (string.IsNullOrEmpty(documentState.FilePath))
                 {
                     continue;
                 }
 
-                builder[filePath] = builder.TryGetValue(filePath, out var documentIdsWithPath)
-                    ? documentIdsWithPath.Add(documentId)
-                    : ImmutableArray.Create(documentId);
+                builder[documentState.FilePath] = builder.TryGetValue(documentState.FilePath, out var documentIdsWithPath)
+                    ? documentIdsWithPath.Add(documentState.Id)
+                    : ImmutableArray.Create(documentState.Id);
             }
 
             return builder.ToImmutable();
         }
 
-        private static IEnumerable<(DocumentId documentId, string filePath)> GetDocumentIdsAndFilePaths(ProjectState projectState, IEnumerable<DocumentId> documentIds, IEnumerable<DocumentId> additionalDocumentIds)
-        {
-            var documentIdsAndFilePaths = documentIds.Select(id => (id, projectState.GetDocumentState(id).FilePath));
-            var additionalDocumentIdsAndFilePaths = additionalDocumentIds.Select(id => (id, projectState.GetAdditionalDocumentState(id).FilePath));
-            return documentIdsAndFilePaths.Concat(additionalDocumentIdsAndFilePaths);
-        }
+        private static IEnumerable<TextDocumentState> GetDocumentStates(ProjectState projectState)
+            => projectState.DocumentStates.Values.Concat(projectState.AdditionalDocumentStates.Values);
 
         /// <summary>
         /// Create a new solution instance without the project specified.
@@ -547,33 +542,32 @@ public SolutionState RemoveProject(ProjectId projectId)
 
         private ImmutableDictionary<string, ImmutableArray<DocumentId>> CreateFilePathToDocumentIdsMapWithRemovedProject(ProjectState projectState)
         {
-            var documentIdsAndFilePaths = GetDocumentIdsAndFilePaths(projectState, projectState.DocumentIds, projectState.AdditionalDocumentIds);
-            return CreateFilePathToDocumentIdsMapWithRemovedDocuments(documentIdsAndFilePaths);
+            return CreateFilePathToDocumentIdsMapWithRemovedDocuments(GetDocumentStates(projectState));
         }
 
-        private ImmutableDictionary<string, ImmutableArray<DocumentId>> CreateFilePathToDocumentIdsMapWithRemovedDocuments(IEnumerable<(DocumentId documentId, string filePath)> documentIdsAndFilePaths)
+        private ImmutableDictionary<string, ImmutableArray<DocumentId>> CreateFilePathToDocumentIdsMapWithRemovedDocuments(IEnumerable<TextDocumentState> documentStates)
         {
             var builder = _filePathToDocumentIdsMap.ToBuilder();
 
-            foreach (var (documentId, filePath) in documentIdsAndFilePaths)
+            foreach (var documentState in documentStates)
             {
-                if (string.IsNullOrEmpty(filePath))
+                if (string.IsNullOrEmpty(documentState.FilePath))
                 {
                     continue;
                 }
 
-                if (!builder.TryGetValue(filePath, out var documentIdsWithPath) || !documentIdsWithPath.Contains(documentId))
+                if (!builder.TryGetValue(documentState.FilePath, out var documentIdsWithPath) || !documentIdsWithPath.Contains(documentState.Id))
                 {
                     throw new ArgumentException($"The given documentId was not found in '{nameof(_filePathToDocumentIdsMap)}'.");
                 }
 
                 if (documentIdsWithPath.Length == 1)
                 {
-                    builder.Remove(filePath);
+                    builder.Remove(documentState.FilePath);
                 }
                 else
                 {
-                    builder[filePath] = documentIdsWithPath.Remove(documentId);
+                    builder[documentState.FilePath] = documentIdsWithPath.Remove(documentState.Id);
                 }
             }
 
@@ -1163,7 +1157,7 @@ public SolutionState AddDocuments(ImmutableArray<DocumentInfo> documentInfos)
 
                 newSolutionState = newSolutionState.ForkProject(newProjectState,
                     CompilationTranslationAction.AddDocuments(newDocumentStatesForProject),
-                    newFilePathToDocumentIdsMap: CreateFilePathToDocumentIdsMapWithAddedDocuments(documentInfosInProject.Select(d => (d.Id, newProjectState.GetDocumentState(d.Id).FilePath))));
+                    newFilePathToDocumentIdsMap: CreateFilePathToDocumentIdsMapWithAddedDocuments(documentInfosInProject.Select(d => newProjectState.GetDocumentState(d.Id))));
             }
 
             return newSolutionState;
@@ -1186,10 +1180,10 @@ public SolutionState AddAdditionalDocument(DocumentInfo documentInfo)
                 _solutionServices);
 
             var newProject = oldProject.AddAdditionalDocument(state);
-            var addedDocumentIdAndFilePath = SpecializedCollections.SingletonEnumerable((documentInfo.Id, newProject.GetAdditionalDocumentState(documentInfo.Id).FilePath));
+            var documentStates = SpecializedCollections.SingletonEnumerable(newProject.GetAdditionalDocumentState(documentInfo.Id));
 
             return this.ForkProject(newProject,
-                    newFilePathToDocumentIdsMap: CreateFilePathToDocumentIdsMapWithAddedDocuments(addedDocumentIdAndFilePath));
+                    newFilePathToDocumentIdsMap: CreateFilePathToDocumentIdsMapWithAddedDocuments(documentStates));
         }
 
         /// <summary>
@@ -1202,12 +1196,12 @@ public SolutionState RemoveDocument(DocumentId documentId)
             var oldProject = this.GetProjectState(documentId.ProjectId);
             var oldDocument = oldProject.GetDocumentState(documentId);
             var newProject = oldProject.RemoveDocument(documentId);
-            var removedDocumentIdAndFilePath = SpecializedCollections.SingletonEnumerable((documentId, oldProject.GetDocumentState(documentId).FilePath));
+            var documentStates = SpecializedCollections.SingletonEnumerable(oldProject.GetDocumentState(documentId));
 
             return this.ForkProject(
                 newProject,
                 CompilationTranslationAction.RemoveDocument(oldDocument),
-                newFilePathToDocumentIdsMap: CreateFilePathToDocumentIdsMapWithRemovedDocuments(removedDocumentIdAndFilePath));
+                newFilePathToDocumentIdsMap: CreateFilePathToDocumentIdsMapWithRemovedDocuments(documentStates));
         }
 
         /// <summary>
@@ -1219,10 +1213,10 @@ public SolutionState RemoveAdditionalDocument(DocumentId documentId)
 
             var oldProject = this.GetProjectState(documentId.ProjectId);
             var newProject = oldProject.RemoveAdditionalDocument(documentId);
-            var removedDocumentIdAndFilePath = SpecializedCollections.SingletonEnumerable((documentId, oldProject.GetAdditionalDocumentState(documentId).FilePath));
+            var documentStates = SpecializedCollections.SingletonEnumerable(GetAdditionalDocumentState(documentId));
 
             return this.ForkProject(newProject,
-                newFilePathToDocumentIdsMap: CreateFilePathToDocumentIdsMapWithRemovedDocuments(removedDocumentIdAndFilePath));
+                newFilePathToDocumentIdsMap: CreateFilePathToDocumentIdsMapWithRemovedDocuments(documentStates));
         }
 
         /// <summary>

src/Workspaces/Core/Portable/Workspace/Workspace_Editor.cs

@@ -495,6 +495,8 @@ protected internal void OnAdditionalDocumentOpened(DocumentId documentId, Source
                 var oldDocument = oldSolution.GetAdditionalDocument(documentId);
                 var oldText = oldDocument.GetTextSynchronously(CancellationToken.None);
 
+                AddToOpenDocumentMap(documentId);
+
                 // keep open document text alive by using PreserveIdentity
                 var newText = textContainer.CurrentText;
                 Solution currentSolution;
@@ -629,6 +631,15 @@ private SourceText GetOpenDocumentText(Solution solution, DocumentId documentId)
             return text;
         }
 
+        private SourceText GetOpenAdditionalDocumentText(Solution solution, DocumentId documentId)
+        {
+            CheckDocumentIsOpen(documentId);
+            var doc = solution.GetAdditionalDocument(documentId);
+            // text should always be preserved, so TryGetText will succeed.
+            doc.TryGetText(out var text);
+            return text;
+        }
+
         /// <summary>
         ///  This method is called during OnSolutionReload.  Override this method if you want to manipulate
         ///  the reloaded solution.
@@ -644,6 +655,10 @@ protected virtual Solution AdjustReloadedSolution(Solution oldSolution, Solution
                 {
                     newSolution = newSolution.WithDocumentText(docId, this.GetOpenDocumentText(oldSolution, docId), PreservationMode.PreserveIdentity);
                 }
+                else if (newSolution.ContainsAdditionalDocument(docId))
+                {
+                    newSolution = newSolution.WithAdditionalDocumentText(docId, this.GetOpenAdditionalDocumentText(oldSolution, docId), PreservationMode.PreserveIdentity);
+                }
             }
 
             return newSolution;
@@ -661,6 +676,10 @@ protected virtual Project AdjustReloadedProject(Project oldProject, Project relo
                 {
                     newSolution = newSolution.WithDocumentText(docId, this.GetOpenDocumentText(oldSolution, docId), PreservationMode.PreserveIdentity);
                 }
+                else if (newSolution.ContainsAdditionalDocument(docId))
+                {
+                    newSolution = newSolution.WithAdditionalDocumentText(docId, this.GetOpenAdditionalDocumentText(oldSolution, docId), PreservationMode.PreserveIdentity);
+                }
             }
 
             return newSolution.GetProject(oldProject.Id);