Merge pull request #12153 from heejaechang/leakfix2

fixed DifferenceViewer leak

src/EditorFeatures/CSharpTest/CodeActions/Preview/PreviewTests.cs

@@ -14,6 +14,7 @@
 using Roslyn.Test.Utilities;
 using Roslyn.Utilities;
 using Xunit;
+using Microsoft.CodeAnalysis.Editor.Implementation.Preview;
 
 namespace Microsoft.CodeAnalysis.Editor.CSharp.UnitTests.CodeRefactorings
 {
@@ -106,11 +107,11 @@ public async Task TestPickTheRightPreview_NoPreference()
                 var previewObjects = await previews.GetPreviewsAsync();
                 var preview = previewObjects[0];
                 Assert.NotNull(preview);
-                Assert.True(preview is IWpfDifferenceViewer);
-                var diffView = preview as IWpfDifferenceViewer;
-                var text = diffView.RightView.TextBuffer.AsTextContainer().CurrentText.ToString();
+                Assert.True(preview is DifferenceViewerPreview);
+                var diffView = preview as DifferenceViewerPreview;
+                var text = diffView.Viewer.RightView.TextBuffer.AsTextContainer().CurrentText.ToString();
                 Assert.Equal(ChangedDocumentText, text);
-                diffView.Close();
+                diffView.Dispose();
 
                 // Then comes the removed metadata reference.
                 preview = previewObjects[1];

src/EditorFeatures/Core/EditorFeatures.csproj

@@ -279,6 +279,7 @@
     <Compile Include="Implementation\Interactive\IAbstractResetInteractiveCommand.cs" />
     <Compile Include="Implementation\Outlining\AbstractSyntaxOutliner.cs" />
     <Compile Include="Implementation\Outlining\InvalidOutliningRegionException.cs" />
+    <Compile Include="Implementation\Preview\DifferenceViewerPreview.cs" />
     <Compile Include="Implementation\Suggestions\FixMultipleOccurrencesService.cs" />
     <Compile Include="Implementation\Suggestions\FixMultipleSuggestedAction.cs" />
     <Compile Include="Implementation\GoToImplementation\GoToImplementationCommandHandler.cs" />
@@ -799,4 +800,4 @@
   <ImportGroup Label="Targets">
     <Import Project="..\..\..\build\Targets\VSL.Imports.targets" />
   </ImportGroup>
-</Project>
+</Project>
\ No newline at end of file

src/EditorFeatures/Core/Implementation/Preview/DifferenceViewerPreview.cs

+// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
+
+using System;
+using Microsoft.CodeAnalysis.ErrorReporting;
+using Microsoft.VisualStudio.Text.Differencing;
+using Roslyn.Utilities;
+
+namespace Microsoft.CodeAnalysis.Editor.Implementation.Preview
+{
+    internal class DifferenceViewerPreview : IDisposable
+    {
+        private IWpfDifferenceViewer _viewer;
+
+        public DifferenceViewerPreview(IWpfDifferenceViewer viewer)
+        {
+            Contract.ThrowIfNull(viewer);
+            _viewer = viewer;
+        }
+
+        public IWpfDifferenceViewer Viewer => _viewer;
+
+        public void Dispose()
+        {
+            GC.SuppressFinalize(this);
+
+            if (_viewer != null && !_viewer.IsClosed)
+            {
+                _viewer.Close();
+            }
+
+            _viewer = null;
+        }
+
+        ~DifferenceViewerPreview()
+        {
+            // make sure we are not leaking diff viewer
+            // we can't close the view from finalizer thread since it must be same
+            // thread (owner thread) this UI is created.
+            if (Environment.HasShutdownStarted)
+            {
+                return;
+            }
+
+            FatalError.ReportWithoutCrash(new Exception($"Dispose is not called how? viewer state : {_viewer.IsClosed}"));
+        }
+    }
+}
\ No newline at end of file

src/EditorFeatures/Core/Implementation/Preview/PreviewFactoryService.cs

@@ -623,7 +623,7 @@ private ITextBuffer CreateNewPlainTextBuffer(TextDocument document, Cancellation
                 rightWorkspace.EnableDiagnostic();
             }
 
-            return diffViewer;
+            return new DifferenceViewerPreview(diffViewer);
         }
 
         private List<LineSpan> CreateLineSpans(ITextSnapshot textSnapshot, NormalizedSpanCollection allSpans, CancellationToken cancellationToken)

src/EditorFeatures/Core/Implementation/SolutionPreviewResult.cs

 // Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
 
+using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Threading;
@@ -52,25 +53,35 @@ public async Task<IReadOnlyList<object>> GetPreviewsAsync(DocumentId preferredDo
             var result = new List<object>();
             var gotRichPreview = false;
 
-            foreach (var previewItem in _previews)
+            try
             {
-                cancellationToken.ThrowIfCancellationRequested();
-                if (previewItem.Text != null)
+                foreach (var previewItem in _previews)
                 {
-                    result.Add(previewItem.Text);
-                }
-                else if (!gotRichPreview)
-                {
-                    var preview = await previewItem.LazyPreview(cancellationToken).ConfigureAwait(true);
-                    if (preview != null)
+                    cancellationToken.ThrowIfCancellationRequested();
+                    if (previewItem.Text != null)
+                    {
+                        result.Add(previewItem.Text);
+                    }
+                    else if (!gotRichPreview)
                     {
-                        result.Add(preview);
-                        gotRichPreview = true;
+                        var preview = await previewItem.LazyPreview(cancellationToken).ConfigureAwait(true);
+                        if (preview != null)
+                        {
+                            result.Add(preview);
+                            gotRichPreview = true;
+                        }
                     }
                 }
-            }
 
-            return result.Count == 0 ? null : result;
+                return result.Count == 0 ? null : result;
+            }
+            catch (OperationCanceledException)
+            {
+                // make sure we dispose all disposable preview objects before
+                // we let control to exit this method
+                result.OfType<IDisposable>().Do(d => d.Dispose());
+                throw;
+            }
         }
 
         /// <summary>Merge two different previews into one final preview result.  The final preview will

src/EditorFeatures/Core/Implementation/Suggestions/SuggestedAction.cs

@@ -236,11 +236,19 @@ public virtual async Task<object> GetPreviewAsync(CancellationToken cancellation
             // Light bulb will always invoke this function on the UI thread.
             AssertIsForeground();
 
+            var previewPaneService = Workspace.Services.GetService<IPreviewPaneService>();
+            if (previewPaneService == null)
+            {
+                return null;
+            }
+
+            // after this point, this method should only return at GetPreviewPane. otherwise, DifferenceViewer will leak
+            // since there is no one to close the viewer
             var preferredDocumentId = Workspace.GetDocumentIdInCurrentContext(SubjectBuffer.AsTextContainer());
             var preferredProjectId = preferredDocumentId?.ProjectId;
 
             var extensionManager = this.Workspace.Services.GetService<IExtensionManager>();
-            var previewContent = await extensionManager.PerformFunctionAsync(Provider, async () =>
+            var previewContents = await extensionManager.PerformFunctionAsync(Provider, async () =>
             {
                 // We need to stay on UI thread after GetPreviewResultAsync() so that TakeNextPreviewAsync()
                 // below can execute on UI thread. We use ConfigureAwait(true) to stay on the UI thread.
@@ -259,14 +267,6 @@ public virtual async Task<object> GetPreviewAsync(CancellationToken cancellation
                 // GetPreviewPane() below needs to run on UI thread. We use ConfigureAwait(true) to stay on the UI thread.
             }, defaultValue: null).ConfigureAwait(true);
 
-            var previewPaneService = Workspace.Services.GetService<IPreviewPaneService>();
-            if (previewPaneService == null)
-            {
-                return null;
-            }
-
-            cancellationToken.ThrowIfCancellationRequested();
-
             // GetPreviewPane() needs to run on the UI thread.
             AssertIsForeground();
 
@@ -274,7 +274,7 @@ public virtual async Task<object> GetPreviewAsync(CancellationToken cancellation
             string projectType;
             Workspace.GetLanguageAndProjectType(preferredProjectId, out language, out projectType);
 
-            return previewPaneService.GetPreviewPane(GetDiagnostic(), language, projectType, previewContent);
+            return previewPaneService.GetPreviewPane(GetDiagnostic(), language, projectType, previewContents);
         }
 
         protected virtual DiagnosticData GetDiagnostic()

src/EditorFeatures/Test/CodeActions/AbstractCodeActionTest.cs

 // Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
 
-using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Threading;
@@ -8,14 +7,10 @@
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.CodeActions;
 using Microsoft.CodeAnalysis.CodeRefactorings;
+using Microsoft.CodeAnalysis.Editor.Implementation.Preview;
 using Microsoft.CodeAnalysis.Editor.UnitTests.Workspaces;
-using Microsoft.CodeAnalysis.Options;
-using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Shared.Utilities;
 using Microsoft.CodeAnalysis.Text;
-using Microsoft.CodeAnalysis.UnitTests;
-using Microsoft.VisualStudio.Text.Differencing;
-using Microsoft.VisualStudio.Text.Editor;
 using Roslyn.Test.Utilities;
 using Roslyn.Utilities;
 using Xunit;
@@ -93,10 +88,10 @@ private static async Task VerifyPreviewContents(TestWorkspace workspace, string
             {
                 var editHandler = workspace.ExportProvider.GetExportedValue<ICodeActionEditHandlerService>();
                 var content = (await editHandler.GetPreviews(workspace, operations, CancellationToken.None).GetPreviewsAsync())[0];
-                var diffView = content as IWpfDifferenceViewer;
-                Assert.NotNull(diffView);
-                var previewContents = diffView.RightView.TextBuffer.AsTextContainer().CurrentText.ToString();
-                diffView.Close();
+                var diffView = content as DifferenceViewerPreview;
+                Assert.NotNull(diffView.Viewer);
+                var previewContents = diffView.Viewer.RightView.TextBuffer.AsTextContainer().CurrentText.ToString();
+                diffView.Dispose();
 
                 Assert.Equal(expectedPreviewContents, previewContents);
             }

src/EditorFeatures/Test/Diagnostics/AbstractUserDiagnosticTest.cs

@@ -23,6 +23,7 @@
 using Roslyn.Test.Utilities;
 using Roslyn.Utilities;
 using Xunit;
+using Microsoft.CodeAnalysis.Editor.Implementation.Preview;
 
 namespace Microsoft.CodeAnalysis.Editor.UnitTests.Diagnostics
 {
@@ -374,9 +375,9 @@ protected async Task TestEquivalenceKeyAsync(string initialMarkup, string equiva
             {
                 // If there is just one document change then we expect the preview to be a WpfTextView
                 var content = (await editHandler.GetPreviews(workspace, operations, CancellationToken.None).GetPreviewsAsync())[0];
-                var diffView = content as IWpfDifferenceViewer;
-                Assert.NotNull(diffView);
-                diffView.Close();
+                var diffView = content as DifferenceViewerPreview;
+                Assert.NotNull(diffView.Viewer);
+                diffView.Dispose();
             }
             else
             {
@@ -390,11 +391,11 @@ protected async Task TestEquivalenceKeyAsync(string initialMarkup, string equiva
                     {
                         if (preview != null)
                         {
-                            var diffView = preview as IWpfDifferenceViewer;
-                            if (diffView != null)
+                            var diffView = preview as DifferenceViewerPreview;
+                            if (diffView?.Viewer != null)
                             {
                                 hasPreview = true;
-                                diffView.Close();
+                                diffView.Dispose();
                                 break;
                             }
                         }

src/EditorFeatures/Test/EditorServicesTest.csproj

@@ -222,6 +222,7 @@
     <Compile Include="Outlining\AbstractSyntaxOutlinerTests.cs" />
     <Compile Include="Outlining\AbstractSyntaxTriviaOutlinerTests.cs" />
     <Compile Include="Outlining\OutliningServiceTests.cs" />
+    <Compile Include="Preview\MockPreviewPaneService.cs" />
     <Compile Include="Preview\TestOnly_CompilerDiagnosticAnalyzerProviderService.cs" />
     <Compile Include="DocCommentFormatting\DocCommentFormattingTests.cs" />
     <Compile Include="DocumentationComments\AbstractDocumentationCommentTests.cs" />
@@ -340,4 +341,4 @@
     <Import Project="..\..\..\build\Targets\VSL.Imports.targets" />
     <Import Project="..\..\..\build\Targets\Roslyn.Toolsets.Xunit.targets" />
   </ImportGroup>
-</Project>
+</Project>
\ No newline at end of file

src/EditorFeatures/Test/Preview/MockPreviewPaneService.cs

+// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
+
+using System;
+using System.Collections.Generic;
+using System.Composition;
+using System.Linq;
+using System.Windows.Controls;
+using Microsoft.CodeAnalysis.Diagnostics;
+using Microsoft.CodeAnalysis.Editor.Host;
+using Microsoft.CodeAnalysis.Host.Mef;
+using Roslyn.Utilities;
+
+namespace Microsoft.CodeAnalysis.Editor.UnitTests.Preview
+{
+    [ExportWorkspaceService(typeof(IPreviewPaneService), ServiceLayer.Host), Shared]
+    internal class MockPreviewPaneService : IPreviewPaneService
+    {
+        public object GetPreviewPane(DiagnosticData diagnostic, string language, string projectType, IReadOnlyList<object> previewContents)
+        {
+            var contents = previewContents ?? SpecializedCollections.EmptyEnumerable<object>();
+
+            foreach (var content in contents.OfType<IDisposable>())
+            {
+                content.Dispose();
+            }
+
+            // test only mock object
+            return new Grid();
+        }
+    }
+}

src/EditorFeatures/Test/Preview/PreviewWorkspaceTests.cs

@@ -26,6 +26,7 @@
 using Roslyn.Test.Utilities;
 using Roslyn.Utilities;
 using Xunit;
+using Microsoft.CodeAnalysis.Editor.Implementation.Preview;
 
 namespace Microsoft.CodeAnalysis.Editor.UnitTests.Preview
 {
@@ -230,48 +231,49 @@ public async Task TestPreviewDiagnosticTaggerInPreviewPane()
                 var newDocument = oldDocument.WithText(oldText.WithChanges(new TextChange(new TextSpan(0, oldText.Length), "class C { }")));
 
                 // create a diff view
-                WpfTestCase.RequireWpfFact($"{nameof(TestPreviewDiagnosticTaggerInPreviewPane)} creates a {nameof(IWpfDifferenceViewer)}");
+                WpfTestCase.RequireWpfFact($"{nameof(TestPreviewDiagnosticTaggerInPreviewPane)} creates a {nameof(DifferenceViewerPreview)}");
 
                 var previewFactoryService = workspace.ExportProvider.GetExportedValue<IPreviewFactoryService>();
-                var diffView = (IWpfDifferenceViewer)(await previewFactoryService.CreateChangedDocumentPreviewViewAsync(oldDocument, newDocument, CancellationToken.None));
-
-                var foregroundService = workspace.GetService<IForegroundNotificationService>();
-                var optionsService = workspace.Services.GetService<IOptionService>();
+                using (var diffView = (DifferenceViewerPreview)(await previewFactoryService.CreateChangedDocumentPreviewViewAsync(oldDocument, newDocument, CancellationToken.None)))
+                {
+                    var foregroundService = workspace.GetService<IForegroundNotificationService>();
+                    var optionsService = workspace.Services.GetService<IOptionService>();
 
-                var waiter = new ErrorSquiggleWaiter();
-                var listeners = AsynchronousOperationListener.CreateListeners(FeatureAttribute.ErrorSquiggles, waiter);
+                    var waiter = new ErrorSquiggleWaiter();
+                    var listeners = AsynchronousOperationListener.CreateListeners(FeatureAttribute.ErrorSquiggles, waiter);
 
-                // set up tagger for both buffers
-                var leftBuffer = diffView.LeftView.BufferGraph.GetTextBuffers(t => t.ContentType.IsOfType(ContentTypeNames.CSharpContentType)).First();
-                var leftProvider = new DiagnosticsSquiggleTaggerProvider(optionsService, diagnosticService, foregroundService, listeners);
-                var leftTagger = leftProvider.CreateTagger<IErrorTag>(leftBuffer);
-                using (var leftDisposable = leftTagger as IDisposable)
-                {
-                    var rightBuffer = diffView.RightView.BufferGraph.GetTextBuffers(t => t.ContentType.IsOfType(ContentTypeNames.CSharpContentType)).First();
-                    var rightProvider = new DiagnosticsSquiggleTaggerProvider(optionsService, diagnosticService, foregroundService, listeners);
-                    var rightTagger = rightProvider.CreateTagger<IErrorTag>(rightBuffer);
-                    using (var rightDisposable = rightTagger as IDisposable)
+                    // set up tagger for both buffers
+                    var leftBuffer = diffView.Viewer.LeftView.BufferGraph.GetTextBuffers(t => t.ContentType.IsOfType(ContentTypeNames.CSharpContentType)).First();
+                    var leftProvider = new DiagnosticsSquiggleTaggerProvider(optionsService, diagnosticService, foregroundService, listeners);
+                    var leftTagger = leftProvider.CreateTagger<IErrorTag>(leftBuffer);
+                    using (var leftDisposable = leftTagger as IDisposable)
                     {
-                        // wait up to 20 seconds for diagnostics
-                        taskSource.Task.Wait(20000);
-                        if (!taskSource.Task.IsCompleted)
+                        var rightBuffer = diffView.Viewer.RightView.BufferGraph.GetTextBuffers(t => t.ContentType.IsOfType(ContentTypeNames.CSharpContentType)).First();
+                        var rightProvider = new DiagnosticsSquiggleTaggerProvider(optionsService, diagnosticService, foregroundService, listeners);
+                        var rightTagger = rightProvider.CreateTagger<IErrorTag>(rightBuffer);
+                        using (var rightDisposable = rightTagger as IDisposable)
                         {
-                            // something is wrong
-                            FatalError.Report(new System.Exception("not finished after 20 seconds"));
+                            // wait up to 20 seconds for diagnostics
+                            taskSource.Task.Wait(20000);
+                            if (!taskSource.Task.IsCompleted)
+                            {
+                                // something is wrong
+                                FatalError.Report(new System.Exception("not finished after 20 seconds"));
+                            }
+
+                            // wait taggers
+                            await waiter.CreateWaitTask();
+
+                            // check left buffer
+                            var leftSnapshot = leftBuffer.CurrentSnapshot;
+                            var leftSpans = leftTagger.GetTags(leftSnapshot.GetSnapshotSpanCollection()).ToList();
+                            Assert.Equal(1, leftSpans.Count);
+
+                            // check right buffer
+                            var rightSnapshot = rightBuffer.CurrentSnapshot;
+                            var rightSpans = rightTagger.GetTags(rightSnapshot.GetSnapshotSpanCollection()).ToList();
+                            Assert.Equal(0, rightSpans.Count);
                         }
-
-                        // wait taggers
-                        await waiter.CreateWaitTask();
-
-                        // check left buffer
-                        var leftSnapshot = leftBuffer.CurrentSnapshot;
-                        var leftSpans = leftTagger.GetTags(leftSnapshot.GetSnapshotSpanCollection()).ToList();
-                        Assert.Equal(1, leftSpans.Count);
-
-                        // check right buffer
-                        var rightSnapshot = rightBuffer.CurrentSnapshot;
-                        var rightSpans = rightTagger.GetTags(rightSnapshot.GetSnapshotSpanCollection()).ToList();
-                        Assert.Equal(0, rightSpans.Count);
                     }
                 }
             }

src/EditorFeatures/VisualBasicTest/CodeActions/Preview/PreviewTests.vb

@@ -5,6 +5,7 @@ Imports System.Threading.Tasks
 Imports System.Windows.Controls
 Imports Microsoft.CodeAnalysis.CodeActions
 Imports Microsoft.CodeAnalysis.CodeRefactorings
+Imports Microsoft.CodeAnalysis.Editor.Implementation.Preview
 Imports Microsoft.CodeAnalysis.Editor.UnitTests.Workspaces
 Imports Microsoft.CodeAnalysis.Text
 Imports Microsoft.VisualStudio.Text.Differencing
@@ -87,11 +88,11 @@ Namespace Microsoft.CodeAnalysis.Editor.VisualBasic.UnitTests.CodeRefactorings
                 Dim previewObjects = Await previews.GetPreviewsAsync()
                 Dim preview = previewObjects(0)
                 Assert.NotNull(preview)
-                Assert.True(TypeOf preview Is IWpfDifferenceViewer)
-                Dim diffView = DirectCast(preview, IWpfDifferenceViewer)
-                Dim text = diffView.RightView.TextBuffer.AsTextContainer().CurrentText.ToString()
+                Assert.True(TypeOf preview Is DifferenceViewerPreview)
+                Dim diffView = DirectCast(preview, DifferenceViewerPreview)
+                Dim text = diffView.Viewer.RightView.TextBuffer.AsTextContainer().CurrentText.ToString()
                 Assert.Equal(s_changedDocumentText, text)
-                diffView.Close()
+                diffView.Dispose()
 
                 ' Then comes the removed metadata reference.
                 preview = previewObjects(1)

src/VisualStudio/Core/Def/Implementation/PreviewPane/PreviewPane.xaml.cs

@@ -9,8 +9,8 @@
 using System.Windows.Navigation;
 using Microsoft.CodeAnalysis.Diagnostics.Log;
 using Microsoft.VisualStudio.LanguageServices.Implementation.Utilities;
-using Microsoft.VisualStudio.Text.Differencing;
 using Roslyn.Utilities;
+using Microsoft.CodeAnalysis.Editor.Implementation.Preview;
 
 namespace Microsoft.VisualStudio.LanguageServices.Implementation.PreviewPane
 {
@@ -26,7 +26,7 @@ internal partial class PreviewPane : UserControl, IDisposable
 
         private bool _isExpanded;
         private double _heightForThreeLineTitle;
-        private IWpfDifferenceViewer _previewDiffViewer;
+        private DifferenceViewerPreview _differenceViewerPreview;
 
         public PreviewPane(
             Image severityIcon,
@@ -108,40 +108,87 @@ private FrameworkElement CreatePreviewElement(IReadOnlyList<object> previewItems
 
             var grid = new Grid();
 
-            for (var i = 0; i < previewItems.Count; i++)
+            foreach (var previewItem in previewItems)
             {
-                var previewItem = previewItems[i];
+                var previewElement = GetPreviewElement(previewItem);
 
-                FrameworkElement previewElement = null;
-                if (previewItem is IWpfDifferenceViewer)
+                // no preview element
+                if (previewElement == null)
                 {
-                    _previewDiffViewer = (IWpfDifferenceViewer)previewItem;
-                    previewElement = _previewDiffViewer.VisualElement;
-                    PreviewDockPanel.Background = _previewDiffViewer.InlineView.Background;
+                    continue;
                 }
-                else if (previewItem is string)
+
+                // the very first preview
+                if (grid.RowDefinitions.Count == 0)
                 {
-                    previewElement = GetPreviewForString((string)previewItem);
+                    grid.RowDefinitions.Add(new RowDefinition());
                 }
-                else if (previewItem is FrameworkElement)
+                else
                 {
-                    previewElement = (FrameworkElement)previewItem;
+                    grid.RowDefinitions.Add(new RowDefinition() { Height = GridLength.Auto });
                 }
 
-                var rowDefinition = i == 0 ? new RowDefinition() : new RowDefinition() { Height = GridLength.Auto };
-                grid.RowDefinitions.Add(rowDefinition);
+                // set row position of the element
+                Grid.SetRow(previewElement, grid.RowDefinitions.Count - 1);
 
-                Grid.SetRow(previewElement, grid.RowDefinitions.IndexOf(rowDefinition));
+                // add the element to the grid
                 grid.Children.Add(previewElement);
 
-                if (i == 0)
+                // set width of the grid same as the first element
+                if (grid.Children.Count == 1)
                 {
                     grid.Width = previewElement.Width;
                 }
             }
 
-            var preview = grid.Children.Count == 0 ? (FrameworkElement)grid.Children[0] : grid;
-            return preview;
+            if (grid.Children.Count == 0)
+            {
+                // no preview
+                return null;
+            }
+
+            // if there is only 1 item, just take preview element as it is without grid
+            if (grid.Children.Count == 1)
+            {
+                var preview = grid.Children[0];
+
+                // we need to take it out from visual tree
+                grid.Children.Clear();
+
+                return (FrameworkElement)preview;
+            }
+
+            return grid;
+        }
+
+        private FrameworkElement GetPreviewElement(object previewItem)
+        {
+            if (previewItem is DifferenceViewerPreview)
+            {
+                // Contract is there should be only 1 diff viewer, otherwise we leak.
+                Contract.ThrowIfFalse(_differenceViewerPreview == null);
+
+                // cache the diff viewer so that we can close it when panel goes away.
+                // this is a bit wierd since we are mutating state here.
+                _differenceViewerPreview = (DifferenceViewerPreview)previewItem;
+                PreviewDockPanel.Background = _differenceViewerPreview.Viewer.InlineView.Background;
+
+                var previewElement = _differenceViewerPreview.Viewer.VisualElement;
+                return previewElement;
+            }
+
+            if (previewItem is string)
+            {
+                return GetPreviewForString((string)previewItem);
+            }
+
+            if (previewItem is FrameworkElement)
+            {
+                return (FrameworkElement)previewItem;
+            }
+
+            // preview item we don't know how to show to users
+            return null;
         }
 
         private void InitializeHyperlinks(Uri helpLink, string helpLinkToolTipText)
@@ -235,28 +282,12 @@ private bool HasDescription
             }
         }
 
-        #region IDisposable Implementation
-        private bool _disposedValue = false;
-
-        // VS editor will call Dispose at which point we should Close() the embedded IWpfDifferenceViewer.
-        protected virtual void Dispose(bool disposing)
-        {
-            if (!_disposedValue)
-            {
-                if (disposing && (_previewDiffViewer != null) && !_previewDiffViewer.IsClosed)
-                {
-                    _previewDiffViewer.Close();
-                }
-            }
-
-            _disposedValue = true;
-        }
-
         void IDisposable.Dispose()
         {
-            Dispose(true);
+            // VS editor will call Dispose at which point we should Close() the embedded IWpfDifferenceViewer.
+            _differenceViewerPreview?.Dispose();
+            _differenceViewerPreview = null;
         }
-        #endregion
 
         private void LearnMoreHyperlink_RequestNavigate(object sender, RequestNavigateEventArgs e)
         {