add proper additional document update change handler in incremental solution updater

added unit test

src/VisualStudio/Core/Test.Next/Services/SolutionServiceTests.cs

@@ -191,42 +191,110 @@ public async Task TestAddUpdateRemoveProjects()
             });
         }
 
+        [Fact, Trait(Traits.Feature, Traits.Features.RemoteHost)]
+        public async Task TestAdditionalDocument()
+        {
+            var code = @"class Test { void Method() { } }";
+            using (var workspace = await TestWorkspace.CreateCSharpAsync(code))
+            {
+                var projectId = workspace.CurrentSolution.ProjectIds.First();
+                var additionalDocumentId = DocumentId.CreateNewId(projectId);
+                var additionalDocumentInfo = DocumentInfo.Create(
+                    additionalDocumentId, "additionalFile",
+                    loader: TextLoader.From(TextAndVersion.Create(SourceText.From("test"), VersionStamp.Create())));
+
+                await VerifySolutionUpdate(workspace, s =>
+                {
+                    return s.AddAdditionalDocument(additionalDocumentInfo);
+                });
+
+                workspace.OnAdditionalDocumentAdded(additionalDocumentInfo);
+
+                await VerifySolutionUpdate(workspace, s =>
+                {
+                    return s.WithAdditionalDocumentText(additionalDocumentId, SourceText.From("changed"));
+                });
+
+                await VerifySolutionUpdate(workspace, s =>
+                {
+                    return s.RemoveAdditionalDocument(additionalDocumentId);
+                });
+            }
+        }
+
+        [Fact, Trait(Traits.Feature, Traits.Features.RemoteHost)]
+        public async Task TestDocument()
+        {
+            var code = @"class Test { void Method() { } }";
+
+            using (var workspace = await TestWorkspace.CreateCSharpAsync(code))
+            {
+                var projectId = workspace.CurrentSolution.ProjectIds.First();
+                var documentId = DocumentId.CreateNewId(projectId);
+                var documentInfo = DocumentInfo.Create(
+                    documentId, "sourceFile",
+                    loader: TextLoader.From(TextAndVersion.Create(SourceText.From("class A { }"), VersionStamp.Create())));
+
+                await VerifySolutionUpdate(workspace, s =>
+                {
+                    return s.AddDocument(documentInfo);
+                });
+
+                workspace.OnDocumentAdded(documentInfo);
+
+                await VerifySolutionUpdate(workspace, s =>
+                {
+                    return s.WithDocumentText(documentId, SourceText.From("class Changed { }"));
+                });
+
+                await VerifySolutionUpdate(workspace, s =>
+                {
+                    return s.RemoveDocument(documentId);
+                });
+            }
+        }
+
         private static async Task VerifySolutionUpdate(string code, Func<Solution, Solution> newSolutionGetter)
         {
             using (var workspace = await TestWorkspace.CreateCSharpAsync(code))
             {
-                var map = new Dictionary<Checksum, object>();
+                await VerifySolutionUpdate(workspace, newSolutionGetter);
+            }
+        }
 
-                var solution = workspace.CurrentSolution;
-                var service = await GetSolutionServiceAsync(solution, map);
+        private static async Task VerifySolutionUpdate(TestWorkspace workspace, Func<Solution, Solution> newSolutionGetter)
+        {
+            var map = new Dictionary<Checksum, object>();
 
-                var solutionChecksum = await solution.State.GetChecksumAsync(CancellationToken.None);
+            var solution = workspace.CurrentSolution;
+            var service = await GetSolutionServiceAsync(solution, map);
 
-                // update primary workspace
-                await service.UpdatePrimaryWorkspaceAsync(solutionChecksum, CancellationToken.None);
-                var first = await service.GetSolutionAsync(solutionChecksum, CancellationToken.None);
+            var solutionChecksum = await solution.State.GetChecksumAsync(CancellationToken.None);
 
-                Assert.Equal(solutionChecksum, await first.State.GetChecksumAsync(CancellationToken.None));
-                Assert.True(object.ReferenceEquals(PrimaryWorkspace.Workspace.PrimaryBranchId, first.BranchId));
+            // update primary workspace
+            await service.UpdatePrimaryWorkspaceAsync(solutionChecksum, CancellationToken.None);
+            var first = await service.GetSolutionAsync(solutionChecksum, CancellationToken.None);
 
-                // get new solution
-                var newSolution = newSolutionGetter(solution);
-                var newSolutionChecksum = await newSolution.State.GetChecksumAsync(CancellationToken.None);
-                newSolution.AppendAssetMap(map);
+            Assert.Equal(solutionChecksum, await first.State.GetChecksumAsync(CancellationToken.None));
+            Assert.True(object.ReferenceEquals(PrimaryWorkspace.Workspace.PrimaryBranchId, first.BranchId));
 
-                // get solution without updating primary workspace
-                var second = await service.GetSolutionAsync(newSolutionChecksum, CancellationToken.None);
+            // get new solution
+            var newSolution = newSolutionGetter(solution);
+            var newSolutionChecksum = await newSolution.State.GetChecksumAsync(CancellationToken.None);
+            newSolution.AppendAssetMap(map);
 
-                Assert.Equal(newSolutionChecksum, await second.State.GetChecksumAsync(CancellationToken.None));
-                Assert.False(object.ReferenceEquals(PrimaryWorkspace.Workspace.PrimaryBranchId, second.BranchId));
+            // get solution without updating primary workspace
+            var second = await service.GetSolutionAsync(newSolutionChecksum, CancellationToken.None);
 
-                // do same once updating primary workspace
-                await service.UpdatePrimaryWorkspaceAsync(newSolutionChecksum, CancellationToken.None);
-                var third = await service.GetSolutionAsync(newSolutionChecksum, CancellationToken.None);
+            Assert.Equal(newSolutionChecksum, await second.State.GetChecksumAsync(CancellationToken.None));
+            Assert.False(object.ReferenceEquals(PrimaryWorkspace.Workspace.PrimaryBranchId, second.BranchId));
 
-                Assert.Equal(newSolutionChecksum, await third.State.GetChecksumAsync(CancellationToken.None));
-                Assert.True(object.ReferenceEquals(PrimaryWorkspace.Workspace.PrimaryBranchId, third.BranchId));
-            }
+            // do same once updating primary workspace
+            await service.UpdatePrimaryWorkspaceAsync(newSolutionChecksum, CancellationToken.None);
+            var third = await service.GetSolutionAsync(newSolutionChecksum, CancellationToken.None);
+
+            Assert.Equal(newSolutionChecksum, await third.State.GetChecksumAsync(CancellationToken.None));
+            Assert.True(object.ReferenceEquals(PrimaryWorkspace.Workspace.PrimaryBranchId, third.BranchId));
         }
 
         private static async Task<SolutionService> GetSolutionServiceAsync(Solution solution, Dictionary<Checksum, object> map = null)

src/Workspaces/Remote/Core/Services/SolutionCreator.cs

@@ -265,8 +265,8 @@ private async Task<Project> UpdateDocumentsAsync(Project project, HashSet<Checks
             {
                 if (!oldMap.ContainsKey(kv.Key))
                 {
-                    // we have new project added
-                    project = AddDocument(project, await CreateDocumentInfoAsync(kv.Value.Checksum).ConfigureAwait(false));
+                    // we have new document added
+                    project = AddDocument(project, await CreateDocumentInfoAsync(kv.Value.Checksum).ConfigureAwait(false), additionalText);
                 }
             }
 
@@ -282,7 +282,8 @@ private async Task<Project> UpdateDocumentsAsync(Project project, HashSet<Checks
                 var newDocumentChecksums = kv.Value;
                 Contract.ThrowIfTrue(oldDocumentChecksums.Checksum == newDocumentChecksums.Checksum);
 
-                project = await UpdateDocumentAsync(project.GetDocument(kv.Key), oldDocumentChecksums, newDocumentChecksums).ConfigureAwait(false);
+                var document = additionalText ? project.GetAdditionalDocument(kv.Key) : project.GetDocument(kv.Key);
+                project = await UpdateDocumentAsync(document, oldDocumentChecksums, newDocumentChecksums, additionalText).ConfigureAwait(false);
             }
 
             // removed project
@@ -290,33 +291,48 @@ private async Task<Project> UpdateDocumentsAsync(Project project, HashSet<Checks
             {
                 if (!newMap.ContainsKey(kv.Key))
                 {
-                    // we have a project removed
-                    project = project.RemoveDocument(kv.Key);
+                    // we have a document removed
+                    if (additionalText)
+                    {
+                        project = project.RemoveAdditionalDocument(kv.Key);
+                    }
+                    else
+                    {
+                        project = project.RemoveDocument(kv.Key);
+                    }
                 }
             }
 
             return project;
         }
 
-        private async Task<Project> UpdateDocumentAsync(Document document, DocumentStateChecksums oldDocumentChecksums, DocumentStateChecksums newDocumentChecksums)
+        private async Task<Project> UpdateDocumentAsync(TextDocument document, DocumentStateChecksums oldDocumentChecksums, DocumentStateChecksums newDocumentChecksums, bool additionalText)
         {
             // changed info
             if (oldDocumentChecksums.Info != newDocumentChecksums.Info)
             {
-                document = await UpdateDocumentInfoAsync(document, newDocumentChecksums.Info).ConfigureAwait(false);
+                document = await UpdateDocumentInfoAsync(document, newDocumentChecksums.Info, additionalText).ConfigureAwait(false);
             }
 
             // changed text
             if (oldDocumentChecksums.Text != newDocumentChecksums.Text)
             {
                 var sourceText = await _assetService.GetAssetAsync<SourceText>(newDocumentChecksums.Text, _cancellationToken).ConfigureAwait(false);
-                document = document.Project.Solution.WithDocumentText(document.Id, sourceText).GetDocument(document.Id);
+
+                if (additionalText)
+                {
+                    document = document.Project.Solution.WithAdditionalDocumentText(document.Id, sourceText).GetAdditionalDocument(document.Id);
+                }
+                else
+                {
+                    document = document.Project.Solution.WithDocumentText(document.Id, sourceText).GetDocument(document.Id);
+                }
             }
 
             return document.Project;
         }
 
-        private async Task<Document> UpdateDocumentInfoAsync(Document document, Checksum infoChecksum)
+        private async Task<TextDocument> UpdateDocumentInfoAsync(TextDocument document, Checksum infoChecksum, bool additionalText)
         {
             var newDocumentInfo = await _assetService.GetAssetAsync<DocumentInfo.DocumentAttributes>(infoChecksum, _cancellationToken).ConfigureAwait(false);
 
@@ -328,11 +344,15 @@ private async Task<Document> UpdateDocumentInfoAsync(Document document, Checksum
 
             if (document.State.Info.Attributes.Folders != newDocumentInfo.Folders)
             {
+                // additional document can't change folder once created
+                Contract.ThrowIfTrue(additionalText);
                 document = document.Project.Solution.WithDocumentFolders(document.Id, newDocumentInfo.Folders).GetDocument(document.Id);
             }
 
             if (document.State.Info.Attributes.SourceCodeKind != newDocumentInfo.SourceCodeKind)
             {
+                // additional document can't change sourcecode kind once created
+                Contract.ThrowIfTrue(additionalText);
                 document = document.Project.Solution.WithDocumentSourceCodeKind(document.Id, newDocumentInfo.SourceCodeKind).GetDocument(document.Id);
             }
 
@@ -493,8 +513,13 @@ private async Task<DocumentInfo> CreateDocumentInfoAsync(Checksum documentChecks
                 documentInfo.IsGenerated);
         }
 
-        private Project AddDocument(Project project, DocumentInfo documentInfo)
+        private Project AddDocument(Project project, DocumentInfo documentInfo, bool additionalText)
         {
+            if (additionalText)
+            {
+                return project.Solution.AddAdditionalDocument(documentInfo).GetProject(project.Id);
+            }
+
             return project.Solution.AddDocument(documentInfo).GetProject(project.Id);
         }
     }