mol-state: optimized state updates when only one transform changes

src/mol-plugin/ui/plugin.tsx

@@ -37,7 +37,7 @@ export class Plugin extends React.Component<{ plugin: PluginContext }, {}> {
 }
 
 export class _test_CurrentObject extends PluginComponent {
-    init() {
+    componentDidMount() {
         this.subscribe(this.context.behaviors.state.data.currentObject, () => this.forceUpdate());
     }
 

src/mol-plugin/ui/state-tree.tsx

@@ -11,7 +11,7 @@ import { PluginCommands } from 'mol-plugin/command';
 import { PluginComponent } from './base';
 
 export class StateTree extends PluginComponent<{ state: State }, { }> {
-    init() {
+    componentDidMount() {
         this.subscribe(this.props.state.events.changed, () => this.forceUpdate());
     }
 

src/mol-state/state.ts

@@ -21,6 +21,8 @@ export { State }
 class State {
     private _tree: StateTree = StateTree.createEmpty();
     private _current: Transform.Ref = this._tree.root.ref;
+
+    protected errorFree = true;
     private transformCache = new Map<Transform.Ref, unknown>();
 
     private ev = RxEventHelper.create();
@@ -106,13 +108,18 @@ class State {
 
                 const ctx: UpdateContext = {
                     parent: this,
+
+                    errorFree: this.errorFree,
                     taskCtx,
                     oldTree,
                     tree: _tree,
                     cells: this.cells as Map<Transform.Ref, StateObjectCell>,
                     transformCache: this.transformCache,
-                    changed: false
+                    changed: false,
+                    editInfo: StateTreeBuilder.is(tree) ? tree.editInfo : void 0
                 };
+
+                this.errorFree = true;
                 // TODO: handle "cancelled" error? Or would this be handled automatically?
                 updated = await update(ctx);
             } finally {
@@ -162,37 +169,51 @@ type Ref = Transform.Ref
 interface UpdateContext {
     parent: State,
 
+    errorFree: boolean,
     taskCtx: RuntimeContext,
     oldTree: StateTree,
     tree: StateTree,
     cells: Map<Transform.Ref, StateObjectCell>,
     transformCache: Map<Ref, unknown>,
-    changed: boolean
+    changed: boolean,
+
+    editInfo: StateTreeBuilder.EditInfo | undefined
 }
 
 async function update(ctx: UpdateContext) {
-    // 1: find all nodes that will definitely be deleted.
-    // this is done in "post order", meaning that leaves will be deleted first.
-    const deletes = findDeletes(ctx);
-    for (const d of deletes) {
-        const obj = ctx.cells.has(d) ? ctx.cells.get(d)!.obj : void 0;
-        ctx.cells.delete(d);
-        ctx.transformCache.delete(d);
-        ctx.parent.events.object.removed.next({ state: ctx.parent, ref: d, obj });
-        // TODO: handle current object change
-    }
 
-    // 2: Find roots where transform version changed or where nodes will be added.
-    const roots = findUpdateRoots(ctx.cells, ctx.tree);
+    // if only a single node was added/updated, we can skip potentially expensive diffing
+    const fastTrack = !!(ctx.errorFree && ctx.editInfo && ctx.editInfo.count === 1 && ctx.editInfo.lastUpdate && ctx.editInfo.sourceTree === ctx.oldTree);
+
+    let deletes: Transform.Ref[], roots: Transform.Ref[];
 
-    // 3: Init empty cells where not present
+    if (fastTrack) {
+        deletes = [];
+        roots = [ctx.editInfo!.lastUpdate!];
+    } else {
+        // find all nodes that will definitely be deleted.
+        // this is done in "post order", meaning that leaves will be deleted first.
+        deletes = findDeletes(ctx);
+        for (const d of deletes) {
+            const obj = ctx.cells.has(d) ? ctx.cells.get(d)!.obj : void 0;
+            ctx.cells.delete(d);
+            ctx.transformCache.delete(d);
+            ctx.parent.events.object.removed.next({ state: ctx.parent, ref: d, obj });
+            // TODO: handle current object change
+        }
+
+        // Find roots where transform version changed or where nodes will be added.
+        roots = findUpdateRoots(ctx.cells, ctx.tree);
+    }
+
+    // Init empty cells where not present
     // this is done in "pre order", meaning that "parents" will be created 1st.
     initCells(ctx, roots);
 
-    // 4: Set status of cells that will be updated to 'pending'.
+    // Set status of cells that will be updated to 'pending'.
     initCellStatus(ctx, roots);
 
-    // 6: Sequentially update all the subtrees.
+    // Sequentially update all the subtrees.
     for (const root of roots) {
         await updateSubtree(ctx, root);
     }
@@ -266,7 +287,10 @@ function initCells(ctx: UpdateContext, roots: Ref[]) {
 
 /** Set status and error text of the cell. Remove all existing objects in the subtree. */
 function doError(ctx: UpdateContext, ref: Ref, errorText: string | undefined) {
-    if (errorText) setCellStatus(ctx, ref, 'error', errorText);
+    if (errorText) {
+        (ctx.parent as any as { errorFree: boolean }).errorFree = false;
+        setCellStatus(ctx, ref, 'error', errorText);
+    }
 
     const cell = ctx.cells.get(ref)!;
     if (cell.obj) {

src/mol-state/tree/builder.ts

@@ -14,12 +14,20 @@ import { shallowEqual } from 'mol-util';
 export { StateTreeBuilder }
 
 interface StateTreeBuilder {
+    readonly editInfo: StateTreeBuilder.EditInfo,
     getTree(): StateTree
 }
 
 namespace StateTreeBuilder {
+    export interface EditInfo {
+        sourceTree: StateTree,
+        count: number,
+        lastUpdate?: Transform.Ref
+    }
+
     interface State {
-        tree: TransientTree
+        tree: TransientTree,
+        editInfo: EditInfo
     }
 
     export function is(obj: any): obj is StateTreeBuilder {
@@ -28,20 +36,27 @@ namespace StateTreeBuilder {
 
     export class Root implements StateTreeBuilder {
         private state: State;
+        get editInfo() { return this.state.editInfo; }
+
         to<A extends StateObject>(ref: Transform.Ref) { return new To<A>(this.state, ref, this); }
         toRoot<A extends StateObject>() { return new To<A>(this.state, this.state.tree.root.ref, this); }
         delete(ref: Transform.Ref) {
+            this.editInfo.count++;
             this.state.tree.remove(ref);
             return this;
         }
         getTree(): StateTree { return this.state.tree.asImmutable(); }
-        constructor(tree: StateTree) { this.state = { tree: tree.asTransient() } }
+        constructor(tree: StateTree) { this.state = { tree: tree.asTransient(), editInfo: { sourceTree: tree, count: 0, lastUpdate: void 0 } } }
     }
 
     export class To<A extends StateObject> implements StateTreeBuilder {
+        get editInfo() { return this.state.editInfo; }
+
         apply<T extends Transformer<A, any, any>>(tr: T, params?: Transformer.Params<T>, props?: Partial<Transform.Options>): To<Transformer.To<T>> {
             const t = tr.apply(this.ref, params, props);
             this.state.tree.add(t);
+            this.editInfo.count++;
+            this.editInfo.lastUpdate = t.ref;
             return new To(this.state, t.ref, this.root);
         }
 
@@ -64,6 +79,9 @@ namespace StateTreeBuilder {
                 }
             }
 
+            this.editInfo.count++;
+            this.editInfo.lastUpdate = this.ref;
+
             this.state.tree.set(Transform.updateParams(old, params));
             return this.root;
         }