mol-state API tweaks, OpenVolume state action provides input data to getFileFormat

src/mol-plugin/behavior/dynamic/animation.ts

@@ -64,7 +64,7 @@ export const StructureAnimation = PluginBehavior.create<StructureAnimationProps>
         rotate(rad: number) {
             this.updatedUnitTransforms.clear()
             const state = this.ctx.state.dataState
-            const reprs = state.select(q => q.rootsOfType(PluginStateObject.Molecule.Representation3D))
+            const reprs = state.selectQ(q => q.rootsOfType(PluginStateObject.Molecule.Representation3D))
             Mat4.rotate(this.rotMat, this.tmpMat, rad, this.rotVec)
             for (const r of reprs) {
                 if (!SO.isRepresentation3D(r.obj)) return
@@ -112,7 +112,7 @@ export const StructureAnimation = PluginBehavior.create<StructureAnimationProps>
         explode(p: number) {
             this.updatedUnitTransforms.clear()
             const state = this.ctx.state.dataState
-            const reprs = state.select(q => q.rootsOfType(PluginStateObject.Molecule.Representation3D));
+            const reprs = state.selectQ(q => q.rootsOfType(PluginStateObject.Molecule.Representation3D));
             for (const r of reprs) {
                 if (!SO.isRepresentation3D(r.obj)) return
                 const structure = getRootStructure(r, state)
@@ -187,5 +187,5 @@ export const StructureAnimation = PluginBehavior.create<StructureAnimationProps>
 //
 
 function getRootStructure(root: StateObjectCell, state: State) {
-    return state.query(StateSelection.Generators.byValue(root).rootOfType([PluginStateObject.Molecule.Structure]))[0];
+    return state.select(StateSelection.Generators.byValue(root).rootOfType([PluginStateObject.Molecule.Structure]))[0];
 }

src/mol-plugin/behavior/dynamic/labels.ts

@@ -112,7 +112,7 @@ export const SceneLabels = PluginBehavior.create<SceneLabelsProps>({
         /** Update structures to be labeled, returns true if changed */
         private updateStructures(p: SceneLabelsProps) {
             const state = this.ctx.state.dataState
-            const structures = state.select(q => q.rootsOfType(PluginStateObject.Molecule.Structure));
+            const structures = state.selectQ(q => q.rootsOfType(PluginStateObject.Molecule.Structure));
             const rootStructures = new Set<SO.Molecule.Structure>()
             for (const s of structures) {
                 const rootStructure = getRootStructure(s, state)

src/mol-plugin/behavior/static/state.ts

@@ -51,17 +51,17 @@ export function SetCurrentObject(ctx: PluginContext) {
 }
 
 export function Update(ctx: PluginContext) {
-    PluginCommands.State.Update.subscribe(ctx, ({ state, tree }) => ctx.runTask(state.update(tree)));
+    PluginCommands.State.Update.subscribe(ctx, ({ state, tree }) => ctx.runTask(state.updateTree(tree)));
 }
 
 export function ApplyAction(ctx: PluginContext) {
-    PluginCommands.State.ApplyAction.subscribe(ctx, ({ state, action, ref }) => ctx.runTask(state.apply(action.action, action.params, ref)));
+    PluginCommands.State.ApplyAction.subscribe(ctx, ({ state, action, ref }) => ctx.runTask(state.applyAction(action.action, action.params, ref)));
 }
 
 export function RemoveObject(ctx: PluginContext) {
     PluginCommands.State.RemoveObject.subscribe(ctx, ({ state, ref }) => {
         const tree = state.tree.build().delete(ref).getTree();
-        return ctx.runTask(state.update(tree));
+        return ctx.runTask(state.updateTree(tree));
     });
 }
 

src/mol-plugin/command/base.ts

@@ -42,7 +42,7 @@ namespace PluginCommand {
         unsubscribe(): void
     }
 
-    export type Action<T> = (params: T) => void | Promise<void>
+    export type Action<T> = (params: T) => unknown | Promise<unknown>
     type Instance = { cmd: PluginCommand<any>, params: any, resolve: () => void, reject: (e: any) => void }
 
     export class Manager {

src/mol-plugin/context.ts

@@ -152,7 +152,7 @@ export class PluginContext {
             tree.toRoot().apply(b.transformer, b.defaultParams, { ref: b.transformer.id });
         }
 
-        await this.runTask(this.state.behaviorState.update(tree, true));
+        await this.runTask(this.state.behaviorState.updateTree(tree, true));
     }
 
     initDataActions() {
@@ -181,9 +181,6 @@ export class PluginContext {
 
         this.lociLabels = new LociLabelManager(this);
 
-        // TODO: find a better solution for this.
-        setTimeout(() => this.log.message(`Mol* Plugin ${PLUGIN_VERSION} [${PLUGIN_VERSION_DATE.toLocaleString()}]`), 500);
+        this.log.message(`Mol* Plugin ${PLUGIN_VERSION} [${PLUGIN_VERSION_DATE.toLocaleString()}]`);
     }
-
-    // settings = ;
 }

src/mol-plugin/state/actions/basic.ts

@@ -74,7 +74,7 @@ const DownloadStructure = StateAction.build({
     }
 
     const data = b.toRoot().apply(StateTransforms.Data.Download, downloadParams);
-    return state.update(createStructureTree(ctx, data, params.source.params.supportProps));
+    return state.updateTree(createStructureTree(ctx, data, params.source.params.supportProps));
 });
 
 export const OpenStructure = StateAction.build({
@@ -84,7 +84,7 @@ export const OpenStructure = StateAction.build({
 })(({ params, state }, ctx: PluginContext) => {
     const b = state.build();
     const data = b.toRoot().apply(StateTransforms.Data.ReadFile, { file: params.file, isBinary: /\.bcif$/i.test(params.file.name) });
-    return state.update(createStructureTree(ctx, data, false));
+    return state.updateTree(createStructureTree(ctx, data, false));
 });
 
 function createStructureTree(ctx: PluginContext, b: StateTreeBuilder.To<PluginStateObject.Data.Binary | PluginStateObject.Data.String>, supportProps: boolean): StateTree {
@@ -123,7 +123,7 @@ export const CreateComplexRepresentation = StateAction.build({
 })(({ ref, state }, ctx: PluginContext) => {
     const root = state.build().to(ref);
     complexRepresentation(ctx, root);
-    return state.update(root.getTree());
+    return state.updateTree(root.getTree());
 });
 
 export const UpdateTrajectory = StateAction.build({
@@ -133,7 +133,7 @@ export const UpdateTrajectory = StateAction.build({
         by: PD.makeOptional(PD.Numeric(1, { min: -1, max: 1, step: 1 }))
     }
 })(({ params, state }) => {
-    const models = state.select(q => q.rootsOfType(PluginStateObject.Molecule.Model)
+    const models = state.selectQ(q => q.rootsOfType(PluginStateObject.Molecule.Model)
         .filter(c => c.transform.transformer === StateTransforms.Model.ModelFromTrajectory));
 
     const update = state.build();
@@ -157,7 +157,7 @@ export const UpdateTrajectory = StateAction.build({
         }
     }
 
-    return state.update(update);
+    return state.updateTree(update);
 });
 
 //
@@ -175,15 +175,13 @@ function getVolumeData(format: VolumeFormat, b: StateTreeBuilder.To<PluginStateO
 }
 
 function createVolumeTree(format: VolumeFormat, ctx: PluginContext, b: StateTreeBuilder.To<PluginStateObject.Data.Binary | PluginStateObject.Data.String>): StateTree {
-
-    const root = getVolumeData(format, b)
+    return getVolumeData(format, b)
         .apply(StateTransforms.Representation.VolumeRepresentation3D,
-            VolumeRepresentation3DHelpers.getDefaultParamsStatic(ctx, 'isosurface'));
-
-    return root.getTree();
+            VolumeRepresentation3DHelpers.getDefaultParamsStatic(ctx, 'isosurface'))
+        .getTree();
 }
 
-function getFileFormat(format: VolumeFormat | 'auto', file: FileInput): VolumeFormat {
+function getFileFormat(format: VolumeFormat | 'auto', file: FileInput, data?: Uint8Array): VolumeFormat {
     if (format === 'auto') {
         const fileFormat = getFileInfo(file).ext
         if (fileFormat in VolumeFormats) {
@@ -205,11 +203,20 @@ export const OpenVolume = StateAction.build({
             ['auto', 'Automatic'], ['ccp4', 'CCP4'], ['mrc', 'MRC'], ['map', 'MAP'], ['dsn6', 'DSN6'], ['brix', 'BRIX']
         ]),
     }
-})(({ params, state }, ctx: PluginContext) => {
-    const b = state.build();
-    const data = b.toRoot().apply(StateTransforms.Data.ReadFile, { file: params.file, isBinary: true });
-    const format = getFileFormat(params.format, params.file)
-    return state.update(createVolumeTree(format, ctx, data));
+})(async ({ params, state }, ctx: PluginContext) => {
+    const dataTree = state.build().toRoot().apply(StateTransforms.Data.ReadFile, { file: params.file, isBinary: true });
+    const volumeData = await ctx.runTask(state.updateTree(dataTree));
+
+    // Alternative for more complex states where the builder is not a simple StateTreeBuilder.To<>:
+    /*
+    const dataRef = dataTree.ref;
+    await ctx.runTask(state.updateTree(dataTree));
+    const dataCell = state.select(dataRef)[0];
+    */
+
+    const format = getFileFormat(params.format, params.file, volumeData.data as Uint8Array)
+    const volumeTree = state.build().to(dataTree.ref);
+    return state.updateTree(createVolumeTree(format, ctx, volumeTree));
 });
 
 export { DownloadDensity };
@@ -276,5 +283,5 @@ const DownloadDensity = StateAction.build({
     }
 
     const data = b.toRoot().apply(StateTransforms.Data.Download, downloadParams);
-    return state.update(createVolumeTree(format, ctx, data));
+    return state.updateTree(createVolumeTree(format, ctx, data));
 });

src/mol-state/action.ts

@@ -71,7 +71,7 @@ namespace StateAction {
             params: def.params as Transformer.Definition<Transformer.From<T>, any, Transformer.Params<T>>['params'],
             run({ cell, state, params }) {
                 const tree = state.build().to(cell.transform.ref).apply(transformer, params);
-                return state.update(tree);
+                return state.updateTree(tree) as Task<void>;
             }
         })
     }

src/mol-state/object.ts

@@ -105,7 +105,7 @@ export class StateObjectTracker<T extends StateObject> {
     }
 
     update() {
-        const cell = this.state.query(this.query)[0];
+        const cell = this.state.select(this.query)[0];
         const version = cell ? cell.transform.version : void 0;
         const changed = this.cell !== cell || this.version !== version;
         this.cell = cell;

src/mol-state/state.ts

@@ -33,7 +33,7 @@ class State {
     readonly globalContext: unknown = void 0;
     readonly events = {
         cell: {
-            stateUpdated: this.ev<State.ObjectEvent & { cellState: StateObjectCell.State}>(),
+            stateUpdated: this.ev<State.ObjectEvent & { cellState: StateObjectCell.State }>(),
             created: this.ev<State.ObjectEvent & { cell: StateObjectCell }>(),
             removed: this.ev<State.ObjectEvent & { parent: Transform.Ref }>(),
         },
@@ -67,7 +67,7 @@ class State {
 
     setSnapshot(snapshot: State.Snapshot) {
         const tree = StateTree.fromJSON(snapshot.tree);
-        return this.update(tree);
+        return this.updateTree(tree);
     }
 
     setCurrent(ref: Transform.Ref) {
@@ -89,26 +89,28 @@ class State {
     }
 
     /**
-     * Select Cells by ref or a query generated on the fly.
-     * @example state.select('test')
-     * @example state.select(q => q.byRef('test').subtree())
+     * Select Cells using the provided selector.
+     * @example state.query(StateSelection.Generators.byRef('test').ancestorOfType([type]))
+     * @example state.query('test')
      */
-    select(selector: Transform.Ref | ((q: typeof StateSelection.Generators) => StateSelection.Selector)) {
-        if (typeof selector === 'string') return StateSelection.select(selector, this);
-        return StateSelection.select(selector(StateSelection.Generators), this)
+    select(selector: StateSelection.Selector) {
+        return StateSelection.select(selector, this)
     }
 
     /**
-     * Select Cells using the provided selector.
-     * @example state.select('test')
+     * Select Cells by building a query generated on the fly.
      * @example state.select(q => q.byRef('test').subtree())
      */
-    query(selector: StateSelection.Selector) {
-        return StateSelection.select(selector, this)
+    selectQ(selector: (q: typeof StateSelection.Generators) => StateSelection.Selector) {
+        if (typeof selector === 'string') return StateSelection.select(selector, this);
+        return StateSelection.select(selector(StateSelection.Generators), this)
     }
 
-    /** If no ref is specified, apply to root */
-    apply<A extends StateAction>(action: A, params: StateAction.Params<A>, ref: Transform.Ref = Transform.RootRef): Task<void> {
+    /**
+     * Creates a Task that applies the specified StateAction (i.e. must use run* on the result)
+     * If no ref is specified, apply to root.
+     */
+    applyAction<A extends StateAction>(action: A, params: StateAction.Params<A>, ref: Transform.Ref = Transform.RootRef): Task<void> {
         return Task.create('Apply Action', ctx => {
             const cell = this.cells.get(ref);
             if (!cell) throw new Error(`'${ref}' does not exist.`);
@@ -118,41 +120,59 @@ class State {
         });
     }
 
-    update(tree: StateTree | StateTreeBuilder, silent: boolean = false): Task<void> {
-        const _tree = (StateTreeBuilder.is(tree) ? tree.getTree() : tree).asTransient();
+    /**
+     * Reconcialites the existing state tree with the new version.
+     *
+     * If the tree is StateTreeBuilder.To<T>, the corresponding StateObject is returned by the task.
+     * @param tree Tree instance or a tree builder instance
+     * @param doNotReportTiming Indicates whether to log timing of the individual transforms
+     */
+    updateTree<T extends StateObject>(tree: StateTree | StateTreeBuilder | StateTreeBuilder.To<T>, doNotLogTiming?: boolean): Task<T>
+    updateTree(tree: StateTree | StateTreeBuilder, doNotLogTiming?: boolean): Task<void>
+    updateTree(tree: StateTree | StateTreeBuilder, doNotLogTiming: boolean = false): Task<any> {
         return Task.create('Update Tree', async taskCtx => {
             let updated = false;
             try {
-                const oldTree = this._tree;
-                this._tree = _tree;
+                const ctx = this.updateTreeAndCreateCtx(tree, taskCtx, doNotLogTiming);
+                updated = await update(ctx);
+                if (StateTreeBuilder.isTo(tree)) {
+                    const cell = this.select(tree.ref)[0];
+                    return cell && cell.obj;
+                }
+            } finally {
+                if (updated) this.events.changed.next();
+            }
+        });
+    }
+
+    private updateTreeAndCreateCtx(tree: StateTree | StateTreeBuilder, taskCtx: RuntimeContext, doNotLogTiming: boolean) {
+        const _tree = (StateTreeBuilder.is(tree) ? tree.getTree() : tree).asTransient();
+        const oldTree = this._tree;
+        this._tree = _tree;
 
-                const ctx: UpdateContext = {
-                    parent: this,
-                    editInfo: StateTreeBuilder.is(tree) ? tree.editInfo : void 0,
+        const ctx: UpdateContext = {
+            parent: this,
+            editInfo: StateTreeBuilder.is(tree) ? tree.editInfo : void 0,
 
-                    errorFree: this.errorFree,
-                    taskCtx,
-                    oldTree,
-                    tree: _tree,
-                    cells: this.cells as Map<Transform.Ref, StateObjectCell>,
-                    transformCache: this.transformCache,
+            errorFree: this.errorFree,
+            taskCtx,
+            oldTree,
+            tree: _tree,
+            cells: this.cells as Map<Transform.Ref, StateObjectCell>,
+            transformCache: this.transformCache,
 
-                    results: [],
+            results: [],
 
-                    silent,
+            silent: doNotLogTiming,
 
-                    changed: false,
-                    hadError: false,
-                    newCurrent: void 0
-                };
+            changed: false,
+            hadError: false,
+            newCurrent: void 0
+        };
 
-                this.errorFree = true;
-                // TODO: handle "cancelled" error? Or would this be handled automatically?
-                updated = await update(ctx);
-            } finally {
-                if (updated) this.events.changed.next();
-            }
-        });
+        this.errorFree = true;
+
+        return ctx;
     }
 
     constructor(rootObject: StateObject, params?: { globalContext?: unknown }) {
@@ -167,8 +187,8 @@ class State {
             version: root.version,
             errorText: void 0,
             params: {
-                definition: { },
-                values: { }
+                definition: {},
+                values: {}
             }
         });
 
@@ -311,7 +331,7 @@ async function update(ctx: UpdateContext) {
         const current = ctx.parent.current;
         const currentCell = ctx.cells.get(current);
         if (currentCell && (
-                currentCell.obj === StateObject.Null
+            currentCell.obj === StateObject.Null
             || (currentCell.status === 'error' && currentCell.errorText === ParentNullErrorText))) {
             newCurrent = findNewCurrent(ctx.oldTree, current, [], ctx.cells);
             ctx.parent.setCurrent(newCurrent);
@@ -523,7 +543,7 @@ async function updateSubtree(ctx: UpdateContext, root: Ref) {
 
 function resolveParams(ctx: UpdateContext, transform: Transform, src: StateObject) {
     const prms = transform.transformer.definition.params;
-    const definition = prms ? prms(src, ctx.parent.globalContext) : { };
+    const definition = prms ? prms(src, ctx.parent.globalContext) : {};
     const values = transform.params ? transform.params : ParamDefinition.getDefaultValues(definition);
     return { definition, values };
 }

src/mol-state/tree/builder.ts

@@ -33,6 +33,10 @@ namespace StateTreeBuilder {
         return !!obj && typeof (obj as StateTreeBuilder).getTree === 'function';
     }
 
+    export function isTo(obj: any): obj is StateTreeBuilder.To<any> {
+        return !!obj && typeof (obj as StateTreeBuilder).getTree === 'function' && typeof (obj as StateTreeBuilder.To<any>).ref === 'string';
+    }
+
     export class Root implements StateTreeBuilder {
         private state: State;
         get editInfo() { return this.state.editInfo; }