Store selection after undone change with redo events

FIX: When undoing, store the selection after the undone change with the redo
event, so that redoing restores it.

Closes codemirror/dev#1286

src/history.ts

@@ -4,7 +4,7 @@ import {KeyBinding, EditorView} from "@codemirror/view"
 
 const enum BranchName { Done, Undone }
 
-const fromHistory = Annotation.define<{side: BranchName, rest: Branch}>()
+const fromHistory = Annotation.define<{side: BranchName, rest: Branch, selection: EditorSelection}>()
 
 /// Transaction annotation that will prevent that transaction from
 /// being combined with other transactions in the undo history. Given
@@ -47,12 +47,6 @@ const historyConfig = Facet.define<HistoryConfig, Required<HistoryConfig>>({
   }
 })
 
-function changeEnd(changes: ChangeDesc) {
-  let end = 0
-  changes.iterChangedRanges((_, to) => end = to)
-  return end
-}
-
 const historyField_ = StateField.define({
   create() {
     return HistoryState.empty
@@ -63,8 +57,7 @@ const historyField_ = StateField.define({
 
     let fromHist = tr.annotation(fromHistory)
     if (fromHist) {
-      let selection = tr.docChanged ? EditorSelection.single(changeEnd(tr.changes)) : undefined
-      let item = HistEvent.fromTransaction(tr, selection), from = fromHist.side
+      let item = HistEvent.fromTransaction(tr, fromHist.selection), from = fromHist.side
       let other = from == BranchName.Done ? state.undone : state.done
       if (item) other = updateBranch(other, other.length, config.minDepth, item)
       else other = addSelection(other, tr.startState.selection)
@@ -358,14 +351,14 @@ class HistoryState {
                             this.prevTime, this.prevUserEvent)
   }
 
-  pop(side: BranchName, state: EditorState, selection: boolean): Transaction | null {
+  pop(side: BranchName, state: EditorState, onlySelection: boolean): Transaction | null {
     let branch = side == BranchName.Done ? this.done : this.undone
     if (branch.length == 0) return null
-    let event = branch[branch.length - 1]
-    if (selection && event.selectionsAfter.length) {
+    let event = branch[branch.length - 1], selection = event.selectionsAfter[0] || state.selection
+    if (onlySelection && event.selectionsAfter.length) {
       return state.update({
         selection: event.selectionsAfter[event.selectionsAfter.length - 1],
-        annotations: fromHistory.of({side, rest: popSelection(branch)}),
+        annotations: fromHistory.of({side, rest: popSelection(branch), selection}),
         userEvent: side == BranchName.Done ? "select.undo" : "select.redo",
         scrollIntoView: true
       })
@@ -378,7 +371,7 @@ class HistoryState {
         changes: event.changes,
         selection: event.startSelection,
         effects: event.effects,
-        annotations: fromHistory.of({side, rest}),
+        annotations: fromHistory.of({side, rest, selection}),
         filter: false,
         userEvent: side == BranchName.Done ? "undo" : "redo",
         scrollIntoView: true

test/test-history.ts

@@ -63,7 +63,7 @@ describe("history", () => {
 
   it("puts the cursor after the change on redo", () => {
     let state = mkState({}, "one\n\ntwo")
-    state = state.update({changes: {from: 3, insert: "!"}}).state
+    state = state.update({changes: {from: 3, insert: "!"}, selection: {anchor: 4}}).state
     state = state.update({selection: {anchor: state.doc.length}}).state
     state = command(state, undo)
     state = command(state, redo)
@@ -359,6 +359,17 @@ describe("history", () => {
     ist(state.selection.ranges.map(r => r.from).join(","), "0,2,3")
   })
 
+  it("restores selection on redo", () => {
+    let state = mkState({}, "a\nb\nc\n")
+    state = state.update({selection: EditorSelection.create([1, 3, 5].map(n => EditorSelection.cursor(n)))}).state
+    state = state.update(state.replaceSelection("-")).state
+    state = state.update({selection:  {anchor: 0}}).state
+    state = command(state, undo)
+    state = state.update({selection:  {anchor: 0}}).state
+    state = command(state, redo)
+    ist(state.selection.ranges.map(r => r.head).join(","), "2,5,8")
+  })
+
   describe("undoSelection", () => {
     it("allows to undo a change", () => {
       let state = mkState()