fix restoration of selection after undo/redo operations

- introduce RecoverableSelection that wraps current selection and allows it to be restored later
- get ySyncPlugin state inside undoManager event hooks

fixes yjs#43, yjs#101, yjs#139

src/plugins/sync-plugin.js

@@ -229,26 +229,13 @@ export const ySyncPlugin = (yXmlFragment, {
 
 /**
  * @param {any} tr
- * @param {any} relSel
+ * @param {RecoverableSelection} recoverableSel
  * @param {ProsemirrorBinding} binding
  */
-const restoreRelativeSelection = (tr, relSel, binding) => {
-  if (relSel !== null && relSel.anchor !== null && relSel.head !== null) {
-    const anchor = relativePositionToAbsolutePosition(
-      binding.doc,
-      binding.type,
-      relSel.anchor,
-      binding.mapping
-    )
-    const head = relativePositionToAbsolutePosition(
-      binding.doc,
-      binding.type,
-      relSel.head,
-      binding.mapping
-    )
-    if (anchor !== null && head !== null) {
-      tr = tr.setSelection(TextSelection.create(tr.doc, anchor, head))
-    }
+const restoreRelativeSelection = (tr, recoverableSel, binding) => {
+  if (recoverableSel !== null && recoverableSel.valid()) {
+    const selection = recoverableSel.restore(binding, tr.doc)
+    tr = tr.setSelection(selection)
   }
 }
 
@@ -265,6 +252,63 @@ export const getRelativeSelection = (pmbinding, state) => ({
   )
 })
 
+export const createRecoverableSelection = (pmbinding, state) => {
+  const sel = new RecoverableSelection(pmbinding, state.selection)
+  state.selection.map(state.doc, sel)
+  return sel
+}
+
+export class RecoverableSelection {
+  constructor (pmbinding, selection, recoverMode = false) {
+    this.records = []
+    this.pmbinding = pmbinding
+    this.selection = selection
+    this.recoverMode = recoverMode
+  }
+
+  restore (pmbinding, doc) {
+    return this.selection.map(doc, new RecoveryMapping(pmbinding, this.records))
+  }
+
+  valid () {
+    return !!this.records.length && this.records.every(r => r.relPos)
+  }
+
+  map (pos) {
+    const relPos = absolutePositionToRelativePosition(pos, this.pmbinding.type, this.pmbinding.mapping)
+    this.records.push({ pos, relPos })
+    return pos
+  }
+
+  mapResult (pos) {
+    return { deleted: false, pos: this.map(pos) }
+  }
+}
+
+export class RecoveryMapping {
+  constructor (pmbinding, records) {
+    this.pmbinding = pmbinding
+    this.records = records
+  }
+
+  map (pos) {
+    return this.mapResult(pos).pos
+  }
+
+  mapResult (pos) {
+    for (const rec of this.records) {
+      if (rec.pos === pos) {
+        const mappedPos = relativePositionToAbsolutePosition(this.pmbinding.doc, this.pmbinding.type, rec.relPos, this.pmbinding.mapping)
+        if (mappedPos === null) {
+          return { deleted: true, pos }
+        }
+        return { deleted: false, pos: mappedPos }
+      }
+    }
+    throw new Error('not recorded')
+  }
+}
+
 /**
  * Binding for prosemirror.
  *
@@ -290,13 +334,18 @@ export class ProsemirrorBinding {
      */
     // @ts-ignore
     this.doc = yXmlFragment.doc
+    /**
+     * last selection as relative positions in the Yjs model
+     */
+    this.beforePatchSelection = null
+    this.lastProsemirrorState = prosemirrorView.state
     /**
      * current selection as relative positions in the Yjs model
      */
     this.beforeTransactionSelection = null
     this.beforeAllTransactions = () => {
       if (this.beforeTransactionSelection === null) {
-        this.beforeTransactionSelection = getRelativeSelection(
+        this.beforeTransactionSelection = createRecoverableSelection(
           this,
           prosemirrorView.state
         )
@@ -550,8 +599,10 @@ export class ProsemirrorBinding {
 
   _prosemirrorChanged (doc) {
     this.doc.transact(() => {
+      this.beforePatchSelection = createRecoverableSelection(this, this.lastProsemirrorState)
+      this.lastProsemirrorState = this.prosemirrorView.state
       updateYFragment(this.doc, this.type, doc, this.mapping)
-      this.beforeTransactionSelection = getRelativeSelection(
+      this.beforeTransactionSelection = createRecoverableSelection(
         this,
         this.prosemirrorView.state
       )

src/plugins/undo-plugin.js

@@ -1,6 +1,6 @@
 import { Plugin } from 'prosemirror-state' // eslint-disable-line
 
-import { getRelativeSelection } from './sync-plugin.js'
+import { createRecoverableSelection } from './sync-plugin.js'
 import { UndoManager, Item, ContentType, XmlElement, Text } from 'yjs'
 import { yUndoPluginKey, ySyncPluginKey } from './keys.js'
 
@@ -57,7 +57,7 @@ export const yUndoPlugin = ({ protectedNodes = defaultProtectedNodes, trackedOri
       if (binding) {
         return {
           undoManager,
-          prevSel: getRelativeSelection(binding, oldState),
+          prevSel: createRecoverableSelection(binding, oldState),
           hasUndoOps,
           hasRedoOps
         }
@@ -74,15 +74,16 @@ export const yUndoPlugin = ({ protectedNodes = defaultProtectedNodes, trackedOri
     }
   },
   view: view => {
-    const ystate = ySyncPluginKey.getState(view.state)
     const undoManager = yUndoPluginKey.getState(view.state).undoManager
     undoManager.on('stack-item-added', ({ stackItem }) => {
+      const ystate = ySyncPluginKey.getState(view.state)
       const binding = ystate.binding
       if (binding) {
-        stackItem.meta.set(binding, yUndoPluginKey.getState(view.state).prevSel)
+        stackItem.meta.set(binding, binding.beforePatchSelection)
       }
     })
-    undoManager.on('stack-item-popped', ({ stackItem }) => {
+    undoManager.on('stack-item-pop', ({ stackItem }) => {
+      const ystate = ySyncPluginKey.getState(view.state)
       const binding = ystate.binding
       if (binding) {
         binding.beforeTransactionSelection = stackItem.meta.get(binding) || binding.beforeTransactionSelection