Add UndoManager.prototype.transact

The `transact` method takes a function and batches all operations within that function as a single transaction to be undone and redone. The edit session now uses generic operations to restore selection state around transactions. The `undo` and `do` methods on operations are now optional. In addition, the undo manager now supports an optional `redo` method on an operation for code that should *only* be run on redo, and not when the operation is initially pushed. This is used by edit session to restore selection state after redo.

src/app/buffer.coffee

@@ -132,10 +132,7 @@ class Buffer
   change: (oldRange, newText) ->
     oldRange = Range.fromObject(oldRange)
     operation = new BufferChangeOperation({buffer: this, oldRange, newText})
-    if @undoManager
-      @undoManager.pushOperation(operation)
-    else
-      operation.do()
+    @pushOperation(operation)
 
   prefixAndSuffixForRange: (range) ->
     prefix: @lines[range.start.row][0...range.start.column]
@@ -145,11 +142,14 @@ class Buffer
     @lines[startRow..endRow] = newLines
     @modified = true
 
-  startUndoBatch: (selectedBufferRanges) ->
-    @undoManager.startUndoBatch(selectedBufferRanges)
+  pushOperation: (operation) ->
+    if @undoManager
+      @undoManager.pushOperation(operation)
+    else
+      operation.do()
 
-  endUndoBatch: (selectedBufferRanges) ->
-    @undoManager.endUndoBatch(selectedBufferRanges)
+  transact: (fn) ->
+    @undoManager.transact(fn)
 
   undo: ->
     @undoManager.undo()

src/app/edit-session.coffee

@@ -177,12 +177,10 @@ class EditSession
     @insertText($native.readFromPasteboard())
 
   undo: ->
-    if ranges = @buffer.undo()
-      @setSelectedBufferRanges(ranges)
+    @buffer.undo()
 
   redo: ->
-    if ranges = @buffer.redo()
-      @setSelectedBufferRanges(ranges)
+    @buffer.redo()
 
   foldSelection: ->
     selection.fold() for selection in @getSelections()
@@ -234,10 +232,15 @@ class EditSession
     @tokenizedBuffer.toggleLineCommentsInRange(range)
 
   mutateSelectedText: (fn) ->
-    selections = @getSelections()
-    @buffer.startUndoBatch(@getSelectedBufferRanges())
-    fn(selection) for selection in selections
-    @buffer.endUndoBatch(@getSelectedBufferRanges())
+    @transact => fn(selection) for selection in @getSelections()
+
+  transact: (fn) ->
+    @buffer.transact =>
+      oldSelectedRanges = @getSelectedBufferRanges()
+      @buffer.pushOperation(undo: => @setSelectedBufferRanges(oldSelectedRanges))
+      fn()
+      newSelectedRanges = @getSelectedBufferRanges()
+      @buffer.pushOperation(redo: => @setSelectedBufferRanges(newSelectedRanges))
 
   getAnchors: ->
     new Array(@anchors...)

src/app/undo-manager.coffee

@@ -5,8 +5,7 @@ module.exports =
 class UndoManager
   undoHistory: null
   redoHistory: null
-  currentBatch: null
-  startBatchCallCount: null
+  currentTransaction: null
 
   constructor: (@buffer) ->
     @startBatchCallCount = 0
@@ -14,36 +13,34 @@ class UndoManager
     @redoHistory = []
 
   pushOperation: (operation) ->
-    if @currentBatch
-      @currentBatch.push(operation)
+    if @currentTransaction
+      @currentTransaction.push(operation)
     else
       @undoHistory.push([operation])
     @redoHistory = []
-    operation.do()
+    operation.do?()
+
+  transact: (fn) ->
+    if @currentTransaction
+      fn()
+    else
+      @currentTransaction = []
+      fn()
+      @undoHistory.push(@currentTransaction) if @currentTransaction.length
+      @currentTransaction = null
 
   undo: ->
     if batch = @undoHistory.pop()
       opsInReverse = new Array(batch...)
       opsInReverse.reverse()
-      op.undo() for op in opsInReverse
+      op.undo?() for op in opsInReverse
       @redoHistory.push batch
       batch.oldSelectionRanges
 
   redo: ->
     if batch = @redoHistory.pop()
-      op.do() for op in batch
+      for op in batch
+        op.do?()
+        op.redo?()
       @undoHistory.push(batch)
       batch.newSelectionRanges
-
-  startUndoBatch: (ranges) ->
-    @startBatchCallCount++
-    return if @startBatchCallCount > 1
-    @currentBatch = []
-    @currentBatch.oldSelectionRanges = ranges
-
-  endUndoBatch: (ranges) ->
-    @startBatchCallCount--
-    return if @startBatchCallCount > 0
-    @currentBatch.newSelectionRanges = ranges
-    @undoHistory.push(@currentBatch) if @currentBatch.length > 0
-    @currentBatch = null