laurentkempe · sharwell · Sep 12, 2014 · Sep 12, 2014
diff --git a/GitDiffMargin/Core/DiffParseResultEventArgs.cs b/GitDiffMargin/Core/DiffParseResultEventArgs.cs
@@ -1,25 +1,35 @@
 using System;
-using System.Collections.Generic;
+using System.Collections.ObjectModel;
 using GitDiffMargin.Git;
 using Microsoft.VisualStudio.Text;
 
 namespace GitDiffMargin.Core
 {
     public class DiffParseResultEventArgs : ParseResultEventArgs
     {
-        private readonly List<HunkRangeInfo> _diff;
+        private readonly ReadOnlyCollection<HunkRangeInfo> _diffToIndex;
+        private readonly ReadOnlyCollection<HunkRangeInfo> _diffToHead;
 
-        public DiffParseResultEventArgs(ITextSnapshot snapshot, TimeSpan elapsedTime, List<HunkRangeInfo> diff)
+        public DiffParseResultEventArgs(ITextSnapshot snapshot, TimeSpan elapsedTime, ReadOnlyCollection<HunkRangeInfo> diffToIndex, ReadOnlyCollection<HunkRangeInfo> diffToHead)
             : base(snapshot, elapsedTime)
         {
-            _diff = diff;
+            _diffToIndex = diffToIndex;
+            _diffToHead = diffToHead;
         }
 
-        public IEnumerable<HunkRangeInfo> Diff
+        public ReadOnlyCollection<HunkRangeInfo> DiffToIndex
         {
             get
             {
-                return _diff;
+                return _diffToIndex;
+            }
+        }
+
+        public ReadOnlyCollection<HunkRangeInfo> DiffToHead
+        {
+            get
+            {
+                return _diffToHead;
             }
         }
     }

diff --git a/GitDiffMargin/Core/DiffUpdateBackgroundParser.cs b/GitDiffMargin/Core/DiffUpdateBackgroundParser.cs
@@ -1,7 +1,6 @@
 using System;
 using System.Diagnostics;
 using System.IO;
-using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using GitDiffMargin.Git;
@@ -102,7 +101,7 @@ protected override void ReParseImpl()
                 if (!TextDocumentFactoryService.TryGetTextDocument(_documentBuffer, out textDocument)) return;
 
                 var diff = _commands.GetGitDiffFor(textDocument, snapshot);
-                var result = new DiffParseResultEventArgs(snapshot, stopwatch.Elapsed, diff.ToList());
+                var result = new DiffParseResultEventArgs(snapshot, stopwatch.Elapsed, diff.DiffToIndex, diff.DiffToHead);
                 OnParseComplete(result);
             }
             catch (InvalidOperationException)

diff --git a/GitDiffMargin/Core/MarginCore.cs b/GitDiffMargin/Core/MarginCore.cs
@@ -42,6 +42,8 @@ public MarginCore(IWpfTextView textView, ITextDocumentFactoryService textDocumen
             _parser.ParseComplete += HandleParseComplete;
             _parser.RequestParse(false);
 
+            _textView.Options.OptionChanged += HandleOptionChanged;
+
             _textView.Closed += (sender, e) =>
             {
                 _editorFormatMap.FormatMappingChanged -= HandleFormatMappingChanged;
@@ -261,6 +263,18 @@ private static Brush GetBrush(ResourceDictionary properties)
 
             return Brushes.Transparent;
         }
+
+        private void HandleOptionChanged(object sender, EditorOptionChangedEventArgs e)
+        {
+            if (_isDisposed)
+                return;
+
+            if (string.Equals(e.OptionId, GitDiffMarginTextViewOptions.CompareToIndexId.Name, StringComparison.Ordinal))
+            {
+                _parser.RequestParse(true);
+            }
+        }
+
         public bool RollBack(HunkRangeInfo hunkRangeInfo)
         {
             if (hunkRangeInfo.SuppressRollback)
@@ -323,7 +337,12 @@ private void HandleParseComplete(object sender, ParseResultEventArgs e)
             var diffResult = e as DiffParseResultEventArgs;
             if (diffResult == null) return;
 
-            CheckBeginInvokeOnUi(() => OnHunksChanged(diffResult.Diff));
+            CheckBeginInvokeOnUi(
+                () =>
+                {
+                    bool compareToIndex = _textView.Options.GetOptionValue(GitDiffMarginTextViewOptions.CompareToIndexId);
+                    OnHunksChanged(compareToIndex ? diffResult.DiffToIndex : diffResult.DiffToHead);
+                });
         }
 
         private void OnHunksChanged(IEnumerable<HunkRangeInfo> hunkRangeInfos)

diff --git a/GitDiffMargin/DiffMarginBase.cs b/GitDiffMargin/DiffMarginBase.cs
@@ -3,7 +3,6 @@
 using System.Windows.Controls;
 using GitDiffMargin.ViewModel;
 using Microsoft.VisualStudio.Text.Editor;
-using Microsoft.VisualStudio.Text.Editor.OptionsExtensionMethods;
 
 namespace GitDiffMargin
 {
@@ -53,7 +52,7 @@ public bool Enabled
         {
             get
             {
-                return TextView.Options.IsSelectionMarginEnabled();
+                return TextView.Options.GetOptionValue(GitDiffMarginTextViewOptions.DiffMarginEnabledId);
             }
         }
 
@@ -68,7 +67,10 @@ public FrameworkElement VisualElement
 
         private void HandleOptionChanged(object sender, EditorOptionChangedEventArgs e)
         {
-            if (!_isDisposed && e.OptionId == GitDiffMarginTextViewOptions.DiffMarginName)
+            if (_isDisposed)
+                return;
+
+            if (string.Equals(e.OptionId, GitDiffMarginTextViewOptions.DiffMarginEnabledId.Name, StringComparison.Ordinal))
             {
                 UpdateVisibility();
             }

diff --git a/GitDiffMargin/Git/DiffResult.cs b/GitDiffMargin/Git/DiffResult.cs
@@ -0,0 +1,42 @@
+namespace GitDiffMargin.Git
+{
+    using System.Collections.Generic;
+    using System.Collections.ObjectModel;
+    using System.Linq;
+
+    public class DiffResult
+    {
+        private static readonly DiffResult _empty = new DiffResult();
+
+        private DiffResult()
+            : this(Enumerable.Empty<HunkRangeInfo>(), Enumerable.Empty<HunkRangeInfo>())
+        {
+        }
+
+        public DiffResult(IEnumerable<HunkRangeInfo> diffToIndex, IEnumerable<HunkRangeInfo> diffToHead)
+        {
+            DiffToIndex = diffToIndex.ToList().AsReadOnly();
+            DiffToHead = diffToHead.ToList().AsReadOnly();
+        }
+
+        public static DiffResult Empty
+        {
+            get
+            {
+                return _empty;
+            }
+        }
+
+        public ReadOnlyCollection<HunkRangeInfo> DiffToIndex
+        {
+            get;
+            private set;
+        }
+
+        public ReadOnlyCollection<HunkRangeInfo> DiffToHead
+        {
+            get;
+            private set;
+        }
+    }
+}
diff --git a/GitDiffMargin/Git/GitCommands.cs b/GitDiffMargin/Git/GitCommands.cs
@@ -1,6 +1,5 @@
 extern alias vs11;
 using System;
-using System.Collections.Generic;
 using System.ComponentModel.Composition;
 using System.Diagnostics;
 using System.IO;
@@ -28,40 +27,41 @@ public GitCommands(SVsServiceProvider serviceProvider)
 
         private const int ContextLines = 0;
 
-        public IEnumerable<HunkRangeInfo> GetGitDiffFor(ITextDocument textDocument, ITextSnapshot snapshot)
+        public DiffResult GetGitDiffFor(ITextDocument textDocument, ITextSnapshot snapshot)
         {
             var filename = textDocument.FilePath;
             var repositoryPath = GetGitRepository(Path.GetFullPath(filename));
             if (repositoryPath == null)
-                yield break;
+                return DiffResult.Empty;
 
             using (var repo = new Repository(repositoryPath))
             {
                 var workingDirectory = repo.Info.WorkingDirectory;
                 if (workingDirectory == null)
-                    yield break;
+                    return DiffResult.Empty;
 
                 var retrieveStatus = repo.Index.RetrieveStatus(filename);
                 if (retrieveStatus == FileStatus.Nonexistent)
                 {
                     // this occurs if a file within the repository itself (not the working copy) is opened.
-                    yield break;
+                    return DiffResult.Empty;
                 }
 
                 if ((retrieveStatus & FileStatus.Ignored) != 0)
                 {
                     // pointless to show diffs for ignored files
-                    yield break;
+                    return DiffResult.Empty;
                 }
 
                 if (retrieveStatus == FileStatus.Unaltered && !textDocument.IsDirty)
                 {
                     // truly unaltered
-                    yield break;
+                    return DiffResult.Empty;
                 }
 
                 var content = GetCompleteContent(textDocument, snapshot);
-                if (content == null) yield break;
+                if (content == null)
+                    return DiffResult.Empty;
 
                 using (var currentContent = new MemoryStream(content))
                 {
@@ -71,22 +71,24 @@ public IEnumerable<HunkRangeInfo> GetGitDiffFor(ITextDocument textDocument, ITex
 
                     var newBlob = repo.ObjectDatabase.CreateBlob(currentContent, relativeFilepath);
 
-                    bool suppressRollback;
-                    Blob blob;
+                    bool headSuppressRollback;
+                    bool indexSuppressRollback;
+                    Blob headBlob;
+                    Blob indexBlob;
 
                     if ((retrieveStatus & FileStatus.Untracked) != 0 || (retrieveStatus & FileStatus.Added) != 0)
                     {
-                        suppressRollback = true;
+                        headSuppressRollback = true;
 
                         // special handling for added files (would need updating to compare against index)
                         using (var emptyContent = new MemoryStream())
                         {
-                            blob = repo.ObjectDatabase.CreateBlob(emptyContent, relativeFilepath);
+                            headBlob = repo.ObjectDatabase.CreateBlob(emptyContent, relativeFilepath);
                         }
                     }
                     else
                     {
-                        suppressRollback = false;
+                        headSuppressRollback = false;
 
                         Commit from = repo.Head.Tip;
                         TreeEntry fromEntry = from[relativeFilepath];
@@ -97,30 +99,51 @@ public IEnumerable<HunkRangeInfo> GetGitDiffFor(ITextDocument textDocument, ITex
                             foreach (string segment in relativeFilepath.Split(Path.DirectorySeparatorChar))
                             {
                                 if (tree == null)
-                                    yield break;
+                                    return DiffResult.Empty;
 
                                 fromEntry = tree.FirstOrDefault(i => string.Equals(segment, i.Name, StringComparison.OrdinalIgnoreCase));
                                 if (fromEntry == null)
-                                    yield break;
+                                    return DiffResult.Empty;
 
                                 tree = fromEntry.Target as Tree;
                             }
                         }
 
-                        blob = fromEntry.Target as Blob;
-                        if (blob == null)
-                            yield break;
+                        headBlob = fromEntry.Target as Blob;
+                        if (headBlob == null)
+                            return DiffResult.Empty;
                     }
 
-                    var treeChanges = repo.Diff.Compare(blob, newBlob, new CompareOptions { ContextLines = ContextLines, InterhunkLines = 0 });
-
-                    var gitDiffParser = new GitDiffParser(treeChanges.Patch, ContextLines, suppressRollback);
-                    var hunkRangeInfos = gitDiffParser.Parse();
-
-                    foreach (var hunkRangeInfo in hunkRangeInfos)
+                    if ((retrieveStatus & FileStatus.Untracked) != 0)
                     {
-                        yield return hunkRangeInfo;
+                        indexSuppressRollback = true;
+                        indexBlob = headBlob;
                     }
+                    else
+                    {
+                        indexSuppressRollback = false;
+
+                        // the index matches the head unless a specific IndexEntry exists
+                        indexBlob = headBlob;
+                        foreach (var indexEntry in repo.Index)
+                        {
+                            if (string.Equals(indexEntry.Path, relativeFilepath, StringComparison.OrdinalIgnoreCase))
+                            {
+                                indexBlob = repo.Lookup<Blob>(indexEntry.Id);
+                                break;
+                            }
+                        }
+                    }
+
+                    ContentChanges treeChanges = repo.Diff.Compare(headBlob, newBlob, new CompareOptions { ContextLines = ContextLines, InterhunkLines = 0 });
+                    var gitDiffParser = new GitDiffParser(treeChanges.Patch, ContextLines, headSuppressRollback);
+                    var diffToHead = gitDiffParser.Parse();
+
+                    treeChanges = repo.Diff.Compare(indexBlob, newBlob, new CompareOptions { ContextLines = ContextLines, InterhunkLines = 0 });
+                    gitDiffParser = new GitDiffParser(treeChanges.Patch, ContextLines, indexSuppressRollback);
+                    var diffToIndex = gitDiffParser.Parse();
+
+                    return new DiffResult(diffToIndex, diffToHead);
                 }
             }
         }
@@ -154,7 +177,7 @@ private static byte[] GetCompleteContent(ITextDocument textDocument, ITextSnapsh
             "|" + "(?:" + QuotedParameterPattern + "(?:" + UnquotedParameterPattern + QuotedParameterPattern + ")*" + "(?:" + UnquotedParameterPattern + ")?" + ")" +
             ")";
 
-        public void StartExternalDiff(ITextDocument textDocument)
+        public void StartExternalDiff(ITextDocument textDocument, bool compareToIndex)
         {
             if (textDocument == null || string.IsNullOrEmpty(textDocument.FilePath)) return;
 
@@ -174,12 +197,25 @@ public void StartExternalDiff(ITextDocument textDocument)
                 // the name of the object in the database
                 string objectName = Path.GetFileName(filename);
 
+                IndexEntry indexEntry = null;
                 Blob oldBlob = null;
-                var indexEntry = repo.Index[relativePath];
-                if (indexEntry != null)
+                if (compareToIndex)
+                {
+                    indexEntry = repo.Index[relativePath];
+                    if (indexEntry != null)
+                    {
+                        objectName = Path.GetFileName(indexEntry.Path);
+                        oldBlob = repo.Lookup<Blob>(indexEntry.Id);
+                    }
+                }
+                else
                 {
-                    objectName = Path.GetFileName(indexEntry.Path);
-                    oldBlob = repo.Lookup<Blob>(indexEntry.Id);
+                    var headEntry = repo.Head[relativePath];
+                    if (headEntry != null)
+                    {
+                        objectName = Path.GetFileName(headEntry.Path);
+                        oldBlob = repo.Lookup<Blob>(headEntry.Target.Id);
+                    }
                 }
 
                 var tempFileName = Path.GetTempFileName();

diff --git a/GitDiffMargin/Git/IGitCommands.cs b/GitDiffMargin/Git/IGitCommands.cs
@@ -1,6 +1,5 @@
 #region using
 
-using System.Collections.Generic;
 using Microsoft.VisualStudio.Text;
 
 #endregion
@@ -9,9 +8,9 @@ namespace GitDiffMargin.Git
 {
     internal interface IGitCommands
     {
-        IEnumerable<HunkRangeInfo> GetGitDiffFor(ITextDocument textDocument, ITextSnapshot snapshot);
+        DiffResult GetGitDiffFor(ITextDocument textDocument, ITextSnapshot snapshot);
 
-        void StartExternalDiff(ITextDocument textDocument);
+        void StartExternalDiff(ITextDocument textDocument, bool compareToIndex);
 
         /// <summary>
         /// Determines if a file or directory is located within a Git repository.

diff --git a/GitDiffMargin/GitDiffMargin.csproj b/GitDiffMargin/GitDiffMargin.csproj
@@ -197,6 +197,7 @@
     <Compile Include="GitDiffMarginCommandHandler.cs" />
     <Compile Include="GitDiffMarginPackage.cs" />
     <Compile Include="GitDiffMarginTextViewCreationListener.cs" />
+    <Compile Include="Git\DiffResult.cs" />
     <Compile Include="Settings\DiffAdditionEditorFormatDefinition.cs" />
     <Compile Include="DiffMarginBase.cs" />
     <Compile Include="Settings\DiffModificationEditorFormatDefinition.cs" />
@@ -220,6 +221,8 @@
     <Compile Include="ScrollDiffMargin.cs" />
     <Compile Include="ScrollDiffMarginFactory2010.cs" />
     <Compile Include="ScrollDiffMarginFactory2013.cs" />
+    <Compile Include="Settings\CompareToIndex.cs" />
+    <Compile Include="Settings\DiffMarginEnabled.cs" />
     <Compile Include="TextViewCommandFilter.cs" />
     <Compile Include="ViewModel\DiffMarginViewModelBase.cs" />
     <Compile Include="ViewModel\EditorDiffMarginViewModel.cs" />