mol-state: undo support

src/mol-plugin-state/manager/structure/component.ts

@@ -113,7 +113,7 @@ class StructureComponentManager extends PluginComponent<StructureComponentManage
             for (const s of structures) {
                 await this.plugin.builders.structure.representation.applyPreset(s.cell, provider, params);
             }
-        });
+        }, { canUndo: true });
     }
 
     clear(structures: ReadonlyArray<StructureRef>) {
@@ -143,7 +143,7 @@ class StructureComponentManager extends PluginComponent<StructureComponentManage
                 if (!selection || selection.elementCount === 0) continue;                
                 this.modifyComponent(b, c, selection, action);
             }
-            await this.dataState.updateTree(b).runInContext(taskCtx);
+            await this.dataState.updateTree(b, { canUndo: true }).runInContext(taskCtx);
         }));
     }
 
@@ -174,7 +174,7 @@ class StructureComponentManager extends PluginComponent<StructureComponentManage
             }
         }
 
-        return this.plugin.managers.structure.hierarchy.remove(toRemove);
+        return this.plugin.managers.structure.hierarchy.remove(toRemove, true);
     }
 
     updateRepresentations(components: ReadonlyArray<StructureComponentRef>, pivot: StructureRepresentationRef, params: StateTransformer.Params<StructureRepresentation3D>) {        
@@ -191,22 +191,24 @@ class StructureComponentManager extends PluginComponent<StructureComponentManage
             update.to(repr.cell).update(params);
         }
 
-        return this.plugin.runTask(this.dataState.updateTree(update));
+        return this.plugin.runTask(this.dataState.updateTree(update, { canUndo: true }));
     }
 
-    async addRepresentation(components: ReadonlyArray<StructureComponentRef>, type: string) {
+    addRepresentation(components: ReadonlyArray<StructureComponentRef>, type: string) {
         if (components.length === 0) return;
 
         const { showHydrogens, visualQuality: quality } = this.state.options;
         const ignoreHydrogens = !showHydrogens;
         const typeParams = { ignoreHydrogens, quality };
 
-        for (const component of components) {
-            await this.plugin.builders.structure.representation.addRepresentation(component.cell, {
-                type: this.plugin.structureRepresentation.registry.get(type),
-                typeParams
-            });
-        }
+        return this.plugin.dataTransaction(async () => {
+            for (const component of components) {
+                await this.plugin.builders.structure.representation.addRepresentation(component.cell, {
+                    type: this.plugin.structureRepresentation.registry.get(type),
+                    typeParams
+                });
+            }
+        }, { canUndo: true });
     }
 
     async add(params: StructureComponentManager.AddParams, structures?: ReadonlyArray<StructureRef>) {
@@ -227,7 +229,7 @@ class StructureComponentManager extends PluginComponent<StructureComponentManage
                     type: this.plugin.structureRepresentation.registry.get(params.representation)
                 });
             }
-        });
+        }, { canUndo: true });
     }
 
     async applyColor(params: StructureComponentManager.ColorParams, structures?: ReadonlyArray<StructureRef>) {
@@ -244,7 +246,7 @@ class StructureComponentManager extends PluginComponent<StructureComponentManage
                     await setStructureOverpaint(this.plugin, s.components, p.color, getLoci, params.representations, p.opacity);
                 }
             }
-        });
+        }, { canUndo: true });
     }
 
     private modifyComponent(builder: StateBuilder.Root, component: StructureComponentRef, by: Structure, action: StructureComponentManager.ModifyAction) {

src/mol-plugin-state/manager/structure/hierarchy.ts

@@ -112,11 +112,11 @@ export class StructureHierarchyManager extends PluginComponent<StructureHierarch
         this.behaviors.current.next({ hierarchy, trajectories, models, structures });
     }
 
-    remove(refs: HierarchyRef[]) {
+    remove(refs: HierarchyRef[], canUndo?: boolean) {
         if (refs.length === 0) return;
         const deletes = this.plugin.state.dataState.build();
         for (const r of refs) deletes.delete(r.cell.transform.ref);
-        return this.plugin.runTask(this.plugin.state.dataState.updateTree(deletes));
+        return this.plugin.runTask(this.plugin.state.dataState.updateTree(deletes, { canUndo }));
     }
 
     createAllModels(trajectory: TrajectoryRef) {

src/mol-plugin-ui/structure/components.tsx

@@ -44,13 +44,15 @@ export class StructureComponentControls extends CollapsableControls<{}, Structur
 interface ComponentEditorControlsState {
     action?: 'preset' | 'add' | 'options',
     isEmpty: boolean,
-    isBusy: boolean
+    isBusy: boolean,
+    canUndo: boolean
 }
 
 class ComponentEditorControls extends PurePluginUIComponent<{}, ComponentEditorControlsState> {
     state: ComponentEditorControlsState = {
         isEmpty: true,
-        isBusy: false
+        isBusy: false,
+        canUndo: false
     };
 
     get isDisabled() {
@@ -65,6 +67,9 @@ class ComponentEditorControls extends PurePluginUIComponent<{}, ComponentEditorC
         this.subscribe(this.plugin.behaviors.state.isBusy, v => {
             this.setState({ isBusy: v, action: this.state.action !== 'options' ? void 0 : 'options' })
         });
+        this.subscribe(this.plugin.state.dataState.events.historyUpdated, ({ state }) => {
+            this.setState({ canUndo: state.canUndo });
+        });
     }
 
     private toggleAction(action: ComponentEditorControlsState['action']) {
@@ -103,12 +108,18 @@ class ComponentEditorControls extends PurePluginUIComponent<{}, ComponentEditorC
         else mng.component.applyPreset(structures, item.value as any);
     }
 
+    undo = () => {
+        const task = this.plugin.state.dataState.undo();
+        if (task) this.plugin.runTask(task);
+    }
+
     render() {
         return <>
             <div className='msp-control-row msp-select-row'>
                 <ToggleButton icon='bookmarks' label='Preset' toggle={this.togglePreset} isSelected={this.state.action === 'preset'} disabled={this.isDisabled} />
                 <ToggleButton icon='plus' label='Add' toggle={this.toggleAdd} isSelected={this.state.action === 'add'} disabled={this.isDisabled} />
                 <ToggleButton icon='cog' label='Options' toggle={this.toggleOptions} isSelected={this.state.action === 'options'} disabled={this.isDisabled} />
+                <IconButton customClass='msp-flex-item' style={{ flex: '0 0 40px' }} onClick={this.undo} disabled={!this.state.canUndo} icon='ccw' title='Some mistakes of the past can be undone.' />
             </div>
             {this.state.action === 'preset' && this.presetControls}
             {this.state.action === 'add' && <div className='msp-control-offset'>
@@ -258,7 +269,7 @@ class StructureComponentGroup extends PurePluginUIComponent<{ group: StructureCo
 
     get removeActions(): ActionMenu.Items {
         const ret = [
-            ActionMenu.Item('Remove', 'remove', () => this.plugin.managers.structure.hierarchy.remove(this.props.group))
+            ActionMenu.Item('Remove', 'remove', () => this.plugin.managers.structure.hierarchy.remove(this.props.group, true))
         ];
 
         const reprs = this.pivot.representations;

src/mol-plugin/context.ts

@@ -193,8 +193,8 @@ export class PluginContext {
         return this.tasks.run(task);
     }
 
-    dataTransaction(f: () => Promise<void> | void) {
-        return this.runTask(this.state.dataState.transaction(f));
+    dataTransaction(f: () => Promise<void> | void, options?: { canUndo?: boolean }) {
+        return this.runTask(this.state.dataState.transaction(f, options));
     }
 
     requestTaskAbort(progress: Progress, reason?: string) {

src/mol-state/state.ts

@@ -48,7 +48,8 @@ class State {
         },
         log: this.ev<LogEntry>(),
         changed: this.ev<{ state: State, inTransaction: boolean }>(),
-        isUpdating: this.ev<boolean>()
+        isUpdating: this.ev<boolean>(),
+        historyUpdated: this.ev<{ state: State }>()
     };
 
     readonly behaviors = {
@@ -66,6 +67,35 @@ class State {
     readonly cells: State.Cells = new Map();
     private spine = new StateTreeSpine.Impl(this.cells);
 
+    private historyCapacity = 5;
+    private history: StateTree[] = [];
+
+    private addHistory(tree: StateTree) {
+        if (this.historyCapacity === 0) return;
+
+        this.history.unshift(tree);
+        if (this.history.length > this.historyCapacity) this.history.pop();
+
+        this.events.historyUpdated.next({ state: this });
+    }
+    
+    private clearHistory() {
+        if (this.history.length === 0) return;
+        this.history = [];
+        this.events.historyUpdated.next({ state: this });
+    }
+
+    get canUndo() {
+        return this.history.length > 0;
+    }
+
+    undo() {
+        const tree = this.history.shift();
+        if (!tree) return;
+        this.events.historyUpdated.next({ state: this });
+        return this.updateTree(tree, { canUndo: false });
+    }
+
     getSnapshot(): State.Snapshot {
         return { tree: StateTree.toJSON(this._tree) };
     }
@@ -130,10 +160,12 @@ class State {
     private inTransaction = false;
 
     /** Apply series of updates to the state. If any of them fail, revert to the original state. */
-    transaction(edits: () => Promise<void> | void) {
+    transaction(edits: () => Promise<void> | void, options?: { canUndo?: boolean }) {
         return Task.create('State Transaction', async ctx => {
             const isNested = this.inTransaction;
 
+            // if (!isNested) this.changedInTransaction = false;
+
             const snapshot = this._tree.asImmutable();
             let restored = false;
             try {
@@ -149,6 +181,7 @@ class State {
                 }
             } catch (e) {
                 if (!restored) {
+                    restored = true;
                     await this.updateTree(snapshot).runInContext(ctx);
                     this.events.log.error(e);
                 }
@@ -158,6 +191,11 @@ class State {
                     this.inTransaction = false;
                     this.events.changed.next({ state: this, inTransaction: false });
                     this.events.isUpdating.next(false);
+
+                    if (!restored) {
+                        if (options?.canUndo) this.addHistory(snapshot);
+                        else this.clearHistory();
+                    }
                 }
             }
         });
@@ -178,28 +216,45 @@ class State {
             const removed = await this.updateQueue.enqueue(params);
             if (!removed) return;
 
+            const snapshot = options?.canUndo ? this._tree.asImmutable() : void 0;
+            let reverted = false;
+
             if (!this.inTransaction) this.events.isUpdating.next(true);
             try {
+                this.reverted = false;
                 const ret = options && (options.revertIfAborted || options.revertOnError)
                     ? await this._revertibleTreeUpdate(taskCtx, params, options)
                     : await this._updateTree(taskCtx, params);
+                reverted = this.reverted;
+
                 return ret.cell;
             } finally {
                 this.updateQueue.handled(params);
-                if (!this.inTransaction) this.events.isUpdating.next(false);
+                if (this.inTransaction) return;
+                
+                this.events.isUpdating.next(false);
+                if (!options?.canUndo) {
+                    this.clearHistory();
+                } else if (!reverted) {
+                    this.addHistory(snapshot!);
+                }
             }
         }, () => {
             this.updateQueue.remove(params);
         });
     }
 
+    private reverted = false;
     private updateQueue = new AsyncQueue<UpdateParams>();
 
     private async _revertibleTreeUpdate(taskCtx: RuntimeContext, params: UpdateParams, options: Partial<State.UpdateOptions>) {
         const old = this.tree;
         const ret = await this._updateTree(taskCtx, params);
         let revert = ((ret.ctx.hadError || ret.ctx.wasAborted) && options.revertOnError) || (ret.ctx.wasAborted && options.revertIfAborted);
-        if (revert) return await this._updateTree(taskCtx, { tree: old, options: params.options });
+        if (revert) {
+            this.reverted = true;
+            return await this._updateTree(taskCtx, { tree: old, options: params.options });
+        }
         return ret;
     }
 
@@ -251,11 +306,13 @@ class State {
         return ctx;
     }
 
-    constructor(rootObject: StateObject, params?: { globalContext?: unknown, rootState?: StateTransform.State }) {
+    constructor(rootObject: StateObject, params?: { globalContext?: unknown, rootState?: StateTransform.State, historyCapacity?: number }) {
         this._tree = StateTree.createEmpty(StateTransform.createRoot(params && params.rootState)).asTransient();
         const tree = this._tree;
         const root = tree.root;
 
+        if (params?.historyCapacity !== void 0) this.historyCapacity = params.historyCapacity;
+
         (this.cells as Map<StateTransform.Ref, StateObjectCell>).set(root.ref, {
             parent: this,
             transform: root,
@@ -301,7 +358,8 @@ namespace State {
         doNotLogTiming: boolean,
         doNotUpdateCurrent: boolean,
         revertIfAborted: boolean,
-        revertOnError: boolean
+        revertOnError: boolean,
+        canUndo: boolean
     }
 
     export function create(rootObject: StateObject, params?: { globalContext?: unknown, rootState?: StateTransform.State }) {
@@ -313,7 +371,8 @@ const StateUpdateDefaultOptions: State.UpdateOptions = {
     doNotLogTiming: false,
     doNotUpdateCurrent: true,
     revertIfAborted: false,
-    revertOnError: false
+    revertOnError: false,
+    canUndo: false
 };
 
 type Ref = StateTransform.Ref