From cbeb351de47b566a858d1a9990ab7aa282d86877 Mon Sep 17 00:00:00 2001
From: Nathan Sobo <nathan@github.com>
Date: Wed, 4 Apr 2012 18:04:17 -0600
Subject: [PATCH] Restore selections after undo
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This only restores changes made via the CompositeSelection… this makes sense because this is the only way to make changes interactively. Any other changes are made via the api or a command line and should not modify selections when they are undone. Still need to test restoration after redo.
---
 spec/app/editor-spec.coffee        | 21 +++++++++++++++++++++
 spec/app/undo-manager-spec.coffee  |  2 +-
 src/app/buffer.coffee              |  8 ++++----
 src/app/composite-selection.coffee | 22 ++++++++++++----------
 src/app/editor.coffee              |  3 ++-
 5 files changed, 40 insertions(+), 16 deletions(-)

diff --git a/spec/app/editor-spec.coffee b/spec/app/editor-spec.coffee
index 6b417f8bf..e814d78c1 100644
--- a/spec/app/editor-spec.coffee
+++ b/spec/app/editor-spec.coffee
@@ -1632,6 +1632,27 @@ describe "Editor", ->
         expect(buffer.lineForRow(0)).not.toContain "foo"
         expect(buffer.lineForRow(1)).not.toContain "foo"
 
+      it "restores the selected ranges after undo", ->
+        editor.setSelectedBufferRanges([[[1, 6], [1, 10]], [[1, 22], [1, 27]]])
+        editor.delete()
+        editor.delete()
+
+        selections = editor.getSelections()
+        expect(buffer.lineForRow(1)).toBe '  var = function( {'
+        expect(selections[0].getBufferRange()).toEqual [[1, 6], [1, 6]]
+        expect(selections[1].getBufferRange()).toEqual [[1, 17], [1, 17]]
+
+        editor.trigger 'undo'
+        expect(selections[0].getBufferRange()).toEqual [[1, 6], [1, 6]]
+        expect(selections[1].getBufferRange()).toEqual [[1, 18], [1, 18]]
+
+        editor.trigger 'undo'
+        expect(selections[0].getBufferRange()).toEqual [[1, 6], [1, 10]]
+        expect(selections[1].getBufferRange()).toEqual [[1, 22], [1, 27]]
+
+        editor.trigger 'redo'
+        expect(selections[0].getBufferRange()).toEqual [[1, 6], [1, 6]]
+        expect(selections[1].getBufferRange()).toEqual [[1, 18], [1, 18]]
 
     describe "when multiple lines are removed from the buffer (regression)", ->
       it "removes all of them from the dom", ->
diff --git a/spec/app/undo-manager-spec.coffee b/spec/app/undo-manager-spec.coffee
index 812418d74..a2ba830e2 100644
--- a/spec/app/undo-manager-spec.coffee
+++ b/spec/app/undo-manager-spec.coffee
@@ -61,7 +61,7 @@ describe "UndoManager", ->
       expect(buffer.getText()).toContain 'qsport'
 
   describe "startUndoBatch() / endUndoBatch()", ->
-    fit "causes changes in batch to be undone simultaneously and returns an array of ranges to select from undo and redo", ->
+    it "causes changes in batch to be undone simultaneously and returns an array of ranges to select from undo and redo", ->
       buffer.insert([0, 0], "foo")
 
       beforeRanges = [[[1, 2], [1, 2]], [[1, 9], [1, 9]]]
diff --git a/src/app/buffer.coffee b/src/app/buffer.coffee
index b8e3c0ea2..733a4c9a5 100644
--- a/src/app/buffer.coffee
+++ b/src/app/buffer.coffee
@@ -127,11 +127,11 @@ class Buffer
     @lines[oldRange.start.row..oldRange.end.row] = newTextLines
     @trigger 'change', { oldRange, newRange, oldText, newText }
 
-  startUndoBatch: ->
-    @undoManager.startUndoBatch()
+  startUndoBatch: (selectedBufferRanges) ->
+    @undoManager.startUndoBatch(selectedBufferRanges)
 
-  endUndoBatch: ->
-    @undoManager.endUndoBatch()
+  endUndoBatch: (selectedBufferRanges) ->
+    @undoManager.endUndoBatch(selectedBufferRanges)
 
   undo: ->
     @undoManager.undo()
diff --git a/src/app/composite-selection.coffee b/src/app/composite-selection.coffee
index 34b58fe5e..6cfa33faa 100644
--- a/src/app/composite-selection.coffee
+++ b/src/app/composite-selection.coffee
@@ -16,6 +16,9 @@ class CompositeSeleciton
   getSelections: ->
     new Array(@selections...)
 
+  getSelectedBufferRanges: ->
+    selection.getBufferRange() for selection in @getSelections()
+
   getLastSelection: ->
     _.last(@selections)
 
@@ -57,10 +60,12 @@ class CompositeSeleciton
     @getLastSelection().setBufferRange(bufferRange, options)
 
   setBufferRanges: (bufferRanges) ->
-    @clearSelections()
-    @setBufferRange(bufferRanges[0])
-    for bufferRange in bufferRanges[1..]
-      @addSelectionForBufferRange(bufferRange)
+    selections = @getSelections()
+    for bufferRange, i in bufferRanges
+      if selections[i]
+        selections[i].setBufferRange(bufferRange)
+      else
+        @addSelectionForBufferRange(bufferRange)
     @mergeIntersectingSelections()
 
   getBufferRange: (bufferRange) ->
@@ -79,12 +84,9 @@ class CompositeSeleciton
 
   mutateSelectedText: (fn) ->
     selections = @getSelections()
-    if selections.length > 1
-      @editor.buffer.startUndoBatch()
-      fn(selection) for selection in selections
-      @editor.buffer.endUndoBatch()
-    else
-      fn(selections[0])
+    @editor.buffer.startUndoBatch(@getSelectedBufferRanges())
+    fn(selection) for selection in selections
+    @editor.buffer.endUndoBatch()
 
   insertText: (text) ->
     @mutateSelectedText (selection) -> selection.insertText(text)
diff --git a/src/app/editor.coffee b/src/app/editor.coffee
index a9fb0af98..efe8c60e1 100644
--- a/src/app/editor.coffee
+++ b/src/app/editor.coffee
@@ -428,7 +428,8 @@ class Editor extends View
   foldSelection: -> @getSelection().fold()
 
   undo: ->
-    @buffer.undo()
+    if ranges = @buffer.undo()
+      @setSelectedBufferRanges(ranges)
 
   redo: ->
     @buffer.redo()