Exploration to implement reverts using Yjs.

0-yjs-undo-revert-trick.ts

import * as Y from 'yjs';

/**
 * In the current version of Yjs (13.x), there is a no
 * builtin way to create and apply a revert to undo an 
 * arbitrary change.
 * 
 * There are tricks to do it though, and they all revolve 
 * around the idea of using an `UndoManager` to undo a change,
 * and reading the `Uint8Array` update that has been applied by undoing.
 */

export const cloneDoc = (doc: Y.Doc) => {
    const clone = new Y.Doc();
    Y.applyUpdateV2(clone, Y.encodeStateAsUpdateV2(doc));
    return clone;
};

const applyUpdateAndComputeRevert = (doc: Y.Doc, update: Uint8Array) => {
    const clonedDoc = cloneDoc(doc);

    Y.applyUpdateV2(doc, update);

    const um = new Y.UndoManager(clonedDoc.getMap('root'));
    Y.applyUpdateV2(clonedDoc, update);
    const beforeUndo = Y.encodeStateVector(clonedDoc);
    um.undo();
    const revert = Y.encodeStateAsUpdateV2(clonedDoc, beforeUndo);
    um.clear();
    um.destroy();

    return revert;
};

const getRevertForLastUpdate = (doc: Y.Doc) => {
    const clonedDoc = cloneDoc(doc);

    const um = new Y.UndoManager(clonedDoc.getMap('root'));

    const beforeUndo = Y.encodeStateVector(clonedDoc);

    um.undo();

    const revert = Y.encodeStateAsUpdateV2(clonedDoc, beforeUndo);

    um.clear();
    um.destroy();

    return revert;
};

1-yjs-reverts-from-snapshots.ts

import { match } from 'ts-pattern';
import * as Y from 'yjs';

/**
 * # Revert Yjs Updates from a snapshot.
 *
 * - every time we update the document, we create a Y.Snapshot
 * - if broadcasting our update fails because the update is invalid, revert to the previous valid snapshot.
 * - We revert by:
 * ` 1.  creating a document for this snapshot
 *   2.  creating an undo manager
 *   3.  getting the snapshot doc to the current state with applyUpdate
 *   4.  undoing this change
 *   5.  getting the diff between the undo operation and the state before undoing.
 *
 * Pros:
 * -  We only create a doc clone when we need to revert, so performance should be ok for valid updates.
 *
 * Cons:
 * -   We need to disable GC on the local doc (probably fine).
 * -   this will revert all changes since this snapshot, just
 *      not the last one.
 *
 */

const SnapshotSymbol = Symbol.for('@multiplayer/y-snapshot');

type Result = { type: 'ok' } | { type: 'error'; error: object };

// Your application logic to broadcast an update
declare function broadcastUpdate(update: Uint8Array): Promise<Result>;

/**
 * Send updates to other central authority,
 * revert if the update was invalid.
 */
export const bindUpdatesWithReverts = (doc: Y.Doc) => {
    let latestSnapshot: Y.Snapshot;

    // before every transaction,
    // keep a reference to the latest document snapshot.
    doc.on('beforeTransaction', () => {
        latestSnapshot = Y.snapshot(doc);
    });

    doc.on('updateV2', async (update, _origin, _doc, transaction) => {
        if (!transaction.local) {
            return;
        }
        const updateSnapshot = latestSnapshot;
        const result = await broadcastUpdate(update);
        match(result)
            .with({ type: 'ok' }, () => {})
            .with({ type: 'error' }, ({ error }) => {
                console.error('broadcast error', { error });
                /**
                 * We get the revert to get back to the previous
                 * valid state.
                 */
                const revert = getRevertFromSnapshot(doc, updateSnapshot);
                // we apply it
                Y.applyUpdateV2(doc, revert);
            });
    });
};

const SnapshotSymbol = Symbol.for('@multiplayer/y-snapshot');

const LocalRevertOriginSymbol = Symbol.for('@multiplayer/y-snapshot');

/**
 * Returns a YJS update that will undo everything
 * that has happened since a given snapshot.
 *
 * This function does not mutate the doc. To perform
 * the revert, you will need to apply the returned update
 * to the doc with `applyRevert(doc, revertUpdate)`.
 */
export const getRevertFromSnapshot = (
    doc: Y.Doc,
    snapshot: Y.Snapshot,
): Uint8Array => {
    const snapshotDoc = Y.createDocFromSnapshot(doc, snapshot);
    const snapshotVClock = Y.encodeStateVector(snapshot.sv);

    const undoManager = new Y.UndoManager(snapshotDoc.getMap('root'), {
        trackedOrigins: new Set([SnapshotSymbol]),
    });

    const updatesToRevert = Y.encodeStateAsUpdateV2(doc, snapshotVClock);

    // Apply and undo
    Y.applyUpdateV2(snapshotDoc, updatesToRevert, SnapshotSymbol);
    undoManager.undo();

    // this update does and undoes the change.
    // We need to send both to other peers.
    const revert = Y.encodeStateAsUpdateV2(snapshotDoc, snapshotVClock);

    undoManager.clear();
    undoManager.destroy();
    snapshotDoc.destroy();

    return revert;
};

/**
 * Mutate
 */
export const applyRevert = (doc: Y.Doc, revert: Uint8Array) => {
    Y.applyUpdateV2(doc, revert, LocalRevertOriginSymbol);
};

export const isRevertTransaction = (transaction: Y.Transaction) =>
    transaction.origin === LocalRevertOriginSymbol;

2-yjs-reverts-from-undo-manager.ts

import { match } from 'ts-pattern';
import * as Y from 'yjs';

/**
 * # Revert Yjs Updates using an UndoManager and storing a list of reverts in a map.
 *
 * - every time we update the document, we create a revert: Uint8Array.
 * - if broadcasting our update fails because the update is invalid, revert this specific update.
 * - We revert by:
 * ` 1.  creating a document for this snapshot
 *   2.  creating an undo manager
 *   3.  getting the snapshot doc to the current state with applyUpdate
 *   4.  undoing this change
 *   5.  getting the diff between the undo operation and the state before undoing.
 *
 * Pros:
 * -  We create a document clone only once, and reuse it.
 * -  We can revert a specific change instead of rolling back all changes since the invalid update.
 *
 * Cons:
 * -  When making updates `A` and then `B`,  It's unclear if keeping update `B` is safe if
 *    update `A` didn't pass validation.
 */

type Result = { type: 'ok' } | { type: 'error'; error: object };

// Your application logic to broadcast an update
declare function broadcastUpdate(update: Uint8Array): Promise<Result>;

/**
 * Send updates to other central authority,
 * revert if the update was invalid.
 */
export const bindUpdatesWithReverts = (doc: Y.Doc) => {
    const reverter = createReverter(doc);

    onDocUpdate(doc, async (update, updateId, isLocal) => {
        if (!isLocal) {
            return;
        }
        const result = await broadcastUpdate(update);
        match(result)
            .with({ type: 'ok' }, () => {
                reverter.discardRevert(updateId);
            })
            .with({ type: 'error' }, ({ error }) => {
                console.error('broadcast error', { error });
                /**
                 * We get the revert to get back to the previous
                 * valid state.
                 */
                const revert = reverter.getRevert(updateId);
                // we apply it
                if (revert) {
                    Y.applyUpdateV2(doc, revert);
                }
            });
    });
};

const cloneDoc = (doc: Y.Doc) => {
    const clone = new Y.Doc();
    Y.applyUpdateV2(clone, Y.encodeStateAsUpdateV2(doc));
    return clone;
};

const createReverter = (doc: Y.Doc) => {
    const revertDoc = cloneDoc(doc);
    const um = new Y.UndoManager(revertDoc.getMap('root'));

    const getRevert = (update: Uint8Array) => {
        Y.applyUpdateV2(revertDoc, update);
        const beforeUndo = Y.encodeStateVector(revertDoc);
        um.undo();
        const revert = Y.encodeStateAsUpdateV2(revertDoc, beforeUndo);
        um.redo();
        um.clear();
        return revert;
    };

    const revertMap = new Map<string, Uint8Array>();

    const unsubscribe = onDocUpdate(doc, (update, updateId, isLocal) => {
        if (!isLocal) {
            Y.applyUpdateV2(revertDoc, update);
        } else {
            revertMap.set(updateId, getRevert(update));
        }
    });

    const destroy = () => {
        unsubscribe();
        um.destroy();
        revertDoc.destroy();
    };

    return {
        getRevert: (updateId: string) => revertMap.get(updateId),
        discardRevert: (updateId: string) => revertMap.delete(updateId),
        destroy,
    };
};

type OnUpdate = (
    update: Uint8Array,
    updateId: string,
    isLocal: boolean,
) => void;

type YJSOnUpdate = (
    update: Uint8Array,
    origin: unknown,
    doc: Y.Doc,
    transaction: Y.Transaction,
) => void;

export const onDocUpdate = (doc: Y.Doc, fn: OnUpdate) => {
    const handler: YJSOnUpdate = (update, _origin, _doc, transaction) => {
        const updateId = crypto.randomUUID();
        fn(update, updateId, transaction.local);
    };
    doc.on('updateV2', handler);
    return () => {
        doc.off('updateV2', handler);
    };
};

3-own-revert-low-level-logic.ts

/* eslint-disable no-underscore-dangle */
import * as Y from 'yjs';

/**
 * Writing our own revert logic.
 *
 * Pros:
 * - Strict minimum amount of logic to perform an update
 * Cons:
 * - Heavily relies on internals, like DeleteSet and StructStore
 * - Clonning DeleteSet and StructStore will likely break due to instanceof mismatches.
 * - This is a lot of brittle code!
 *
 * learned
 * - UndoManager adds a `keep: true` property to the Y.Doc, which disables GC for this item.
 * - This means that keeping an UndoManager on the server effectively disables GC.
 * - This means that our current strategy for storing documents and computing reverts might take too much memory.
 * - We need to kind a way to remove the kept items afterward.
 * -  => There is a way! UndoManager.clear() will remove the `keep: true` flag.
 * -  => we could clone/recreate the doc periodically, for example.
 *
 * , conclusion:
 * - re-implementing a revert would force us to manipulate Y.js' undocumented internals.
 * - The amount of code to copy/paste and adapt is massive.
 */

/**
 * These clases are internal I copy/pasted.z
 * Obviously, this will break instanceof
 * in internal Yjs functions.
 */
class DeleteItem {
    constructor(public clock: number, public len: number) {}
}

class DeleteSet {
    clients = new Map<number, Array<DeleteItem>>();
}

class StructStore {
    clients = new Map<number, Array<Y.GC | Y.Item>>();
    pendingStructs: null | {
        missing: Map<number, number>;
        update: Uint8Array;
    } = null;
    pendingDs: null | Uint8Array = null;
}

export const addToDeleteSet = (
    ds: DeleteSet,
    client: number,
    clock: number,
    length: number,
) => {
    const deletes = ds.clients.get(client) ?? [];
    deletes.push(new DeleteItem(clock, length));
};

class StackItem {
    meta = new Map();
    constructor(public deletions: DeleteSet, public insertions: DeleteSet) {}
}

export const splitItem = (
    transaction: Y.Transaction,
    leftItem: Y.Item,
    diff: number,
) => {
    // create rightItem
    const { client, clock } = leftItem.id;
    const rightItem = new Y.Item(
        Y.createID(client, clock + diff),
        leftItem,
        Y.createID(client, clock + diff - 1),
        leftItem.right,
        leftItem.rightOrigin,
        leftItem.parent,
        leftItem.parentSub,
        leftItem.content.splice(diff),
    );
    if (leftItem.deleted) {
        rightItem.markDeleted();
    }
    if (leftItem.keep) {
        rightItem.keep = true;
    }
    if (leftItem.redone !== null) {
        rightItem.redone = Y.createID(
            leftItem.redone.client,
            leftItem.redone.clock + diff,
        );
    }
    // update left (do not set leftItem.rightOrigin as it will lead to problems when syncing)
    leftItem.right = rightItem;
    // update right
    if (rightItem.right !== null) {
        rightItem.right.left = rightItem;
    }
    // right is more specific.
    transaction._mergeStructs.push(rightItem);
    // update parent._map
    if (rightItem.parentSub !== null && rightItem.right === null) {
        (rightItem.parent as Y.AbstractType<any>)._map.set(
            rightItem.parentSub,
            rightItem,
        );
    }
    leftItem.length = diff;
    return rightItem;
};

export const findIndexCleanStart = (
    transaction: Y.Transaction,
    structs: Y.Item[],
    clock: number,
) => {
    const index = Y.findIndexSS(structs, clock);
    const struct = structs[index];
    if (struct.id.clock < clock && struct instanceof Y.Item) {
        structs.splice(
            index + 1,
            0,
            splitItem(transaction, struct, clock - struct.id.clock),
        );
        return index + 1;
    }
    return index;
};

export const getItemCleanStart = (
    transaction: Y.Transaction,
    id: Y.ID,
): Y.Item => {
    const structs: Y.Item[] = transaction.doc.store.clients.get(id.client);
    return structs[findIndexCleanStart(transaction, structs, id.clock)];
};

/**
 * Not sure what this does.
 */
export const followRedone = (store: StructStore, id: Y.ID) => {
    let nextID: Y.ID | null = id;
    let diff = 0;
    let item: Y.Item | undefined;
    do {
        if (diff > 0) {
            nextID = Y.createID(nextID.client, nextID.clock + diff);
        }
        item = Y.getItem(store, nextID);
        diff = nextID.clock - item.id.clock;
        nextID = item.redone;
    } while (nextID !== null && item instanceof Y.Item);
    return {
        item,
        diff,
    };
};

/**
 * Update the "GC'able state" of this item and all of its parent.
 */
export const keepItem = (item: Y.Item, keep: boolean) => {
    let item2: Y.Item | null = item;
    while (item2 !== null && item2.keep !== keep) {
        item2.keep = keep;
        item2 = (item2.parent as Y.AbstractType<any>)._item;
    }
};

export const redoItem = (
    transaction: Y.Transaction,
    item: Y.Item,
    redoitems: Set<Y.Item>,
    itemsToDelete: DeleteSet,
) => {
    const doc = transaction.doc;
    const store = doc.store;
    const ownClientID = doc.clientID;
    const redone = item.redone;
    if (redone !== null) {
        return getItemCleanStart(transaction, redone);
    }
    let parentItem = (item.parent as Y.AbstractType<any>)?._item;
    let left: Y.Item | null = null;
    let right: Y.Item | null;
    // make sure that parent is redone
    if (parentItem !== null && parentItem.deleted === true) {
        // try to undo parent if it will be undone anyway
        if (
            parentItem.redone === null &&
            (!redoitems.has(parentItem) ||
                redoItem(transaction, parentItem, redoitems, itemsToDelete) ===
                    null)
        ) {
            return null;
        }
        while (parentItem.redone !== null) {
            parentItem = getItemCleanStart(transaction, parentItem.redone);
        }
    }
    const parentType: Y.AbstractType<any> | Y.ContentType =
        parentItem === null ? item.parent : parentItem.content.type;

    if (item.parentSub === null) {
        // Is an array item. Insert at the old position
        left = item.left;
        right = item;
        // find next cloned_redo items
        while (left !== null) {
            let leftTrace: Y.Item | null = left;
            // trace redone until parent matches
            while (
                leftTrace !== null &&
                (leftTrace.parent as Y.AbstractType<any>)._item !== parentItem
            ) {
                leftTrace =
                    leftTrace.redone === null
                        ? null
                        : getItemCleanStart(transaction, leftTrace.redone);
            }
            if (
                leftTrace !== null &&
                (leftTrace.parent as Y.AbstractType<any>)._item === parentItem
            ) {
                left = leftTrace;
                break;
            }
            left = left.left;
        }
        while (right !== null) {
            let rightTrace: Y.Item | null = right;
            // trace redone until parent matches
            while (
                rightTrace !== null &&
                (rightTrace.parent as Y.AbstractType<any>)._item !== parentItem
            ) {
                rightTrace =
                    rightTrace.redone === null
                        ? null
                        : getItemCleanStart(transaction, rightTrace.redone);
            }
            if (
                rightTrace !== null &&
                (rightTrace.parent as Y.AbstractType<any>)._item === parentItem
            ) {
                right = rightTrace;
                break;
            }
            right = right.right;
        }
    } else {
        right = null;
        if (item.right) {
            left = item;
            // Iterate right while right is in itemsToDelete
            // If it is intended to delete right while item is redone, we can expect that item should replace right.
            while (
                left !== null &&
                left.right !== null &&
                (left.right.redone || Y.isDeleted(itemsToDelete, left.right.id))
            ) {
                left = left.right;
                // follow redone
                while (left.redone) {
                    left = getItemCleanStart(transaction, left.redone);
                }
            }
            if (left && left.right !== null) {
                // It is not possible to redo this item because it conflicts with a
                // change from another client
                return null;
            }
        } else {
            left =
                (parentType as Y.AbstractType<any>)._map.get(item.parentSub) ||
                null;
        }
    }
    const nextClock = Y.getState(store, ownClientID);
    const nextId = Y.createID(ownClientID, nextClock);
    const redoneItem = new Y.Item(
        nextId,
        left,
        left && left.lastId,
        right,
        right && right.id,
        parentType,
        item.parentSub,
        item.content.copy(),
    );
    item.redone = nextId;
    keepItem(redoneItem, true);
    redoneItem.integrate(transaction, 0);
    return redoneItem;
};

const createRevertStackItemRaw = (transaction: Y.Transaction) => {
    const insertions = new DeleteSet();
    transaction.afterState.forEach((endClock, client) => {
        const startClock = transaction.beforeState.get(client) || 0;
        const len = endClock - startClock;
        if (len > 0) {
            addToDeleteSet(insertions, client, startClock, len);
        }
    });
    const revertItem = new StackItem(transaction.deleteSet, insertions);

    // This part mutates the Y.Doc to add a keep: true boolean
    // to tell Y.js not to GC the item.
    // Does it stays there forever?
    Y.iterateDeletedStructs(transaction, transaction.deleteSet, (item2) => {
        if (item2 instanceof Y.Item) {
            keepItem(item2, true);
        }
    });

    return revertItem;
};

const applyRevertStackItemRaw = (doc: Y.Doc, stackItem: StackItem) => {
    Y.transact(doc, (transaction) => {
        const store = doc.store;

        /**
         * Collect things that have been inserted and should be deleted
         */
        const itemsToDelete: Y.Item[] = [];
        Y.iterateDeletedStructs(transaction, stackItem.insertions, (struct) => {
            if (struct instanceof Y.GC) {
                return;
            }
            if (struct instanceof Y.Item) {
                // TODO understand why this is here
                if (struct.redone !== null) {
                    let { item, diff } = followRedone(store, struct.id);
                    if (diff > 0) {
                        item = getItemCleanStart(
                            transaction,
                            Y.createID(item.id.client, item.id.clock + diff),
                        );
                    }
                    struct = item;
                }
                if (!struct.deleted) {
                    itemsToDelete.push(struct as Y.Item);
                }
            }
        });

        /**
         * Collect things that have been deleted and should be re applied
         */
        const itemsToRedo = new Set<Y.Item>();
        Y.iterateDeletedStructs(transaction, stackItem.deletions, (struct) => {
            if (
                struct instanceof Y.Item &&
                // Never redo structs in stackItem.insertions because they were created and deleted in the same capture interval.
                !Y.isDeleted(stackItem.insertions, struct.id)
            ) {
                itemsToRedo.add(struct);
            }
        });

        /**
         * Apply redo inserts
         */
        itemsToRedo.forEach((struct) => {
            redoItem(transaction, struct, itemsToRedo, stackItem.insertions);
        });

        /**
         * Apply deletes
         */
        // We want to delete in reverse order so that children are deleted before
        // parents, so we have more information available when items are filtered.
        for (let i = itemsToDelete.length - 1; i >= 0; i--) {
            const item = itemsToDelete[i];
            item.delete(transaction);
        }
    });
};

const revert = (doc: Y.Doc, transaction: Y.Transaction) => {
    // Create the revert and store it
    const revertItem = createRevertStackItemRaw(transaction);
    // apply the revert.
    applyRevertStackItemRaw(doc, revertItem);
};