mol-plugin: cache queries

src/mol-plugin/state/transforms/model.ts

@@ -9,13 +9,13 @@ import { parsePDB } from 'mol-io/reader/pdb/parser';
 import { Vec3 } from 'mol-math/linear-algebra';
 import { trajectoryFromMmCIF } from 'mol-model-formats/structure/mmcif';
 import { trajectoryFromPDB } from 'mol-model-formats/structure/pdb';
-import { Model, ModelSymmetry, Queries, QueryContext, Structure, StructureQuery, StructureSelection as Sel, StructureSymmetry } from 'mol-model/structure';
+import { Model, ModelSymmetry, Queries, QueryContext, Structure, StructureQuery, StructureSelection as Sel, StructureSymmetry, QueryFn } from 'mol-model/structure';
 import { Assembly } from 'mol-model/structure/model/properties/symmetry';
 import { PluginContext } from 'mol-plugin/context';
 import { MolScriptBuilder } from 'mol-script/language/builder';
 import Expression from 'mol-script/language/expression';
 import { compile } from 'mol-script/runtime/query/compiler';
-import { StateObject } from 'mol-state';
+import { StateObject, StateTransformer } from 'mol-state';
 import { RuntimeContext, Task } from 'mol-task';
 import { ParamDefinition as PD } from 'mol-util/param-definition';
 import { stringToWords } from 'mol-util/string';
@@ -261,14 +261,23 @@ const StructureSelection = PluginStateTransform.BuiltIn({
         label: PD.makeOptional(PD.Text('', { isHidden: true }))
     }
 })({
-    apply({ a, params }) {
-        // TODO: use cache, add "update"
+    apply({ a, params, cache }) {
         const compiled = compile<Sel>(params.query);
+        (cache as { compiled: QueryFn<Sel> }).compiled = compiled;
+
         const result = compiled(new QueryContext(a.data));
         const s = Sel.unionStructure(result);
         if (s.elementCount === 0) return StateObject.Null;
         const props = { label: `${params.label || 'Selection'}`, description: structureDesc(s) };
         return new SO.Molecule.Structure(s, props);
+    },
+    update: ({ a, b, oldParams, newParams, cache }) => {
+        if (oldParams.query !== newParams.query) return StateTransformer.UpdateResult.Recreate;
+
+        if (updateStructureFromQuery((cache as { compiled: QueryFn<Sel> }).compiled, a.data, b, newParams.label)) {
+            return StateTransformer.UpdateResult.Updated;
+        }
+        return StateTransformer.UpdateResult.Null;
     }
 });
 
@@ -283,19 +292,41 @@ const UserStructureSelection = PluginStateTransform.BuiltIn({
         label: PD.makeOptional(PD.Text(''))
     }
 })({
-    apply({ a, params }) {
-        // TODO: use cache, add "update"
+    apply({ a, params, cache }) {
         const parsed = parseMolScript(params.query.expression);
         if (parsed.length === 0) throw new Error('No query');
         const query = transpileMolScript(parsed[0]);
         const compiled = compile<Sel>(query);
+        (cache as { compiled: QueryFn<Sel> }).compiled = compiled;
         const result = compiled(new QueryContext(a.data));
         const s = Sel.unionStructure(result);
         const props = { label: `${params.label || 'Selection'}`, description: structureDesc(s) };
         return new SO.Molecule.Structure(s, props);
+    },
+    update: ({ a, b, oldParams, newParams, cache }) => {
+        if (oldParams.query.language !== newParams.query.language || oldParams.query.expression !== newParams.query.expression) {
+            return StateTransformer.UpdateResult.Recreate;
+        }
+
+        updateStructureFromQuery((cache as { compiled: QueryFn<Sel> }).compiled, a.data, b, newParams.label);
+        return StateTransformer.UpdateResult.Updated;
     }
 });
 
+function updateStructureFromQuery(query: QueryFn<Sel>, src: Structure, obj: SO.Molecule.Structure, label?: string) {
+    const result = query(new QueryContext(src));
+    const s = Sel.unionStructure(result);
+    if (s.elementCount === 0) {
+        return false;
+    }
+
+    obj.label = `${label || 'Selection'}`;
+    obj.description = structureDesc(s);
+    obj.data = s;
+    return true;
+}
+
+
 namespace StructureComplexElement {
     export type Types = 'atomic-sequence' | 'water' | 'atomic-het' | 'spheres'
 }

src/mol-state/object.ts

@@ -69,7 +69,9 @@ interface StateObjectCell<T = StateObject> {
     } | undefined;
 
     errorText?: string,
-    obj?: T
+    obj?: T,
+
+    cache: unknown | undefined
 }
 
 namespace StateObjectCell {

src/mol-state/state.ts

@@ -25,7 +25,6 @@ class State {
     private _tree: TransientTree;
 
     protected errorFree = true;
-    private transformCache = new Map<StateTransform.Ref, unknown>();
 
     private ev = RxEventHelper.create();
 
@@ -165,7 +164,6 @@ class State {
             oldTree,
             tree: _tree,
             cells: this.cells as Map<StateTransform.Ref, StateObjectCell>,
-            transformCache: this.transformCache,
 
             results: [],
             stateChanges: [],
@@ -196,7 +194,8 @@ class State {
             params: {
                 definition: {},
                 values: {}
-            }
+            },
+            cache: { }
         });
 
         this.globalContext = params && params.globalContext;
@@ -244,7 +243,6 @@ interface UpdateContext {
     oldTree: StateTree,
     tree: TransientTree,
     cells: Map<StateTransform.Ref, StateObjectCell>,
-    transformCache: Map<Ref, unknown>,
 
     results: UpdateNodeResult[],
     stateChanges: StateTransform.Ref[],
@@ -288,7 +286,6 @@ async function update(ctx: UpdateContext) {
         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);
             deletedObjects.push(obj);
         }
 
@@ -443,7 +440,8 @@ function initCellsVisitor(transform: StateTransform, _: any, { ctx, added }: Ini
         sourceRef: void 0,
         status: 'pending',
         errorText: void 0,
-        params: void 0
+        params: void 0,
+        cache: void 0
     };
     ctx.cells.set(transform.ref, cell);
     added.push(cell);
@@ -516,8 +514,8 @@ function doError(ctx: UpdateContext, ref: Ref, errorText: string | undefined, si
     if (cell.obj) {
         const obj = cell.obj;
         cell.obj = void 0;
+        cell.cache = void 0;
         ctx.parent.events.object.removed.next({ state: ctx.parent, ref, obj });
-        ctx.transformCache.delete(ref);
     }
 
     // remove the objects in the child nodes if they exist
@@ -608,7 +606,7 @@ async function updateNode(ctx: UpdateContext, currentRef: Ref): Promise<UpdateNo
 
     if (!oldTree.transforms.has(currentRef) || !current.params) {
         current.params = params;
-        const obj = await createObject(ctx, currentRef, transform.transformer, parent, params.values);
+        const obj = await createObject(ctx, current, transform.transformer, parent, params.values);
         updateTag(obj, transform);
         current.obj = obj;
 
@@ -619,13 +617,13 @@ async function updateNode(ctx: UpdateContext, currentRef: Ref): Promise<UpdateNo
         current.params = params;
 
         const updateKind = !!current.obj && current.obj !== StateObject.Null
-            ? await updateObject(ctx, currentRef, transform.transformer, parent, current.obj!, oldParams, newParams)
+            ? await updateObject(ctx, current, transform.transformer, parent, current.obj!, oldParams, newParams)
             : StateTransformer.UpdateResult.Recreate;
 
         switch (updateKind) {
             case StateTransformer.UpdateResult.Recreate: {
                 const oldObj = current.obj;
-                const newObj = await createObject(ctx, currentRef, transform.transformer, parent, newParams);
+                const newObj = await createObject(ctx, current, transform.transformer, parent, newParams);
                 updateTag(newObj, transform);
                 current.obj = newObj;
                 return { ref: currentRef, action: 'replaced', oldObj, obj: newObj };
@@ -633,6 +631,10 @@ async function updateNode(ctx: UpdateContext, currentRef: Ref): Promise<UpdateNo
             case StateTransformer.UpdateResult.Updated:
                 updateTag(current.obj, transform);
                 return { ref: currentRef, action: 'updated', obj: current.obj! };
+            case StateTransformer.UpdateResult.Null: {
+                current.obj = StateObject.Null;
+                return { ref: currentRef, action: 'updated', obj: current.obj! };
+            }
             default:
                 return { action: 'none' };
         }
@@ -649,20 +651,15 @@ function runTask<T>(t: T | Task<T>, ctx: RuntimeContext) {
     return t as T;
 }
 
-function createObject(ctx: UpdateContext, ref: Ref, transformer: StateTransformer, a: StateObject, params: any) {
-    const cache = Object.create(null);
-    ctx.transformCache.set(ref, cache);
-    return runTask(transformer.definition.apply({ a, params, cache }, ctx.parent.globalContext), ctx.taskCtx);
+function createObject(ctx: UpdateContext, cell: StateObjectCell, transformer: StateTransformer, a: StateObject, params: any) {
+    if (!cell.cache) cell.cache = Object.create(null);
+    return runTask(transformer.definition.apply({ a, params, cache: cell.cache }, ctx.parent.globalContext), ctx.taskCtx);
 }
 
-async function updateObject(ctx: UpdateContext, ref: Ref, transformer: StateTransformer, a: StateObject, b: StateObject, oldParams: any, newParams: any) {
+async function updateObject(ctx: UpdateContext, cell: StateObjectCell,  transformer: StateTransformer, a: StateObject, b: StateObject, oldParams: any, newParams: any) {
     if (!transformer.definition.update) {
         return StateTransformer.UpdateResult.Recreate;
     }
-    let cache = ctx.transformCache.get(ref);
-    if (!cache) {
-        cache = Object.create(null);
-        ctx.transformCache.set(ref, cache);
-    }
-    return runTask(transformer.definition.update({ a, oldParams, b, newParams, cache }, ctx.parent.globalContext), ctx.taskCtx);
+    if (!cell.cache) cell.cache = Object.create(null);
+    return runTask(transformer.definition.update({ a, oldParams, b, newParams, cache: cell.cache }, ctx.parent.globalContext), ctx.taskCtx);
 }

src/mol-state/transformer.ts

@@ -55,7 +55,7 @@ namespace Transformer {
         newParams: P
     }
 
-    export enum UpdateResult { Unchanged, Updated, Recreate }
+    export enum UpdateResult { Unchanged, Updated, Recreate, Null }
 
     /** Specify default control descriptors for the parameters */
     // export type ParamsDefinition<A extends StateObject = StateObject, P = any> = (a: A, globalCtx: unknown) => { [K in keyof P]: PD.Any }