mol-state: StateObjectSelector

src/apps/viewer/extensions/cellpack/model.ts

@@ -393,7 +393,7 @@ export const LoadCellPackModel = StateAction.build({
                     MS.core.rel.eq([MS.ammp('label_atom_id'), 'P'])
                 ])
             })
-            cellpackTree = cellpackTree.apply(StateTransforms.Model.StructureSelectionFromExpression, { expression }, { state: { isGhost: true } })
+            cellpackTree = cellpackTree.apply(StateTransforms.Model.StructureSelectionFromExpression, { expression }, { state: { isGhost: true } }) as any
         }
         cellpackTree
             .apply(StateTransforms.Representation.StructureRepresentation3D,

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

@@ -72,7 +72,7 @@ export const DscifProvider: DataFormatProvider<any> = {
             const cifStateObject = await state.updateTree(cifBuilder).runInContext(taskCtx)
             const b = state.build().to(cifBuilder.ref);
             const blocks = cifStateObject.data.blocks.slice(1); // zero block contains query meta-data
-            let tree: StateBuilder.To<any>
+            let tree: StateBuilder.To<any, any>
             if (blocks.length === 1) {
                 tree = b
                     .apply(StateTransforms.Volume.VolumeFromDensityServerCif, { blockHeader: blocks[0].header })

src/mol-plugin/state/representation/structure.ts

@@ -23,7 +23,7 @@ export class StructureRepresentationManager {
     private providers: StructureRepresentationProvider[] = [];
     private providerMap: Map<string, StructureRepresentationProvider> = new Map();
 
-    readonly defaultProvider = PresetStructureReprentations.auto;
+    readonly defaultProvider = PresetStructureReprentations.default;
 
     hasProvider(s: Structure) {
         for (const p of this.providers) {
@@ -102,7 +102,7 @@ export class StructureRepresentationManager {
         return this.plugin.runTask(state.updateTree(builder));
     }
 
-    apply<P>(ref: StateTransform.Ref, providerOrId: StructureRepresentationProvider<P> | string, params?: P) {
+    apply<P = any, S = {}>(ref: StateTransform.Ref, providerOrId: StructureRepresentationProvider<P, S> | string, params?: P) {
         const provider = typeof providerOrId === 'string'
             ? arrayFind(this.providers, p => p.id === providerOrId)
             : providerOrId;
@@ -119,10 +119,9 @@ export class StructureRepresentationManager {
             ? PD.getDefaultValues(provider.params(cell.obj.data, this.plugin))
             : {})
 
-        const apply = provider.apply(state, cell, prms, this.plugin);
 
-        if (Task.is(apply)) return this.plugin.runTask(apply);
-        return apply;
+        const task = Task.create<S>(`${provider.display.name}`, ctx => provider.apply(ctx, state, cell, prms, this.plugin) as Promise<S>);
+        return this.plugin.runTask(task);
     }
 
     // init() {

src/mol-plugin/state/representation/structure/preset.ts

@@ -16,20 +16,20 @@ import { PluginStateObject } from '../../objects';
 const auto = StructureRepresentationProvider({
     id: 'preset-structure-representation-auto',
     display: { name: 'Automaic', group: 'Preset' },
-    apply(state, structureCell, _, plugin) {
+    apply(ctx, state, structureCell, _, plugin) {
         const s = structureCell.obj!.data;
 
         // TODO: a way to improve this?
 
         if (s.elementCount < 50000) {
-            return defaultPreset.apply(state, structureCell, void 0, plugin);
+            return defaultPreset.apply(ctx, state, structureCell, void 0, plugin);
         } else if (s.elementCount < 200000) {
-            return proteinAndNucleic.apply(state, structureCell, void 0, plugin);
+            return proteinAndNucleic.apply(ctx, state, structureCell, void 0, plugin);
         } else {
             if (s.unitSymmetryGroups[0].units.length > 10) {
-                return capsid.apply(state, structureCell, void 0, plugin);
+                return capsid.apply(ctx, state, structureCell, void 0, plugin);
             } else {
-                return coarseCapsid.apply(state, structureCell, void 0, plugin);
+                return coarseCapsid.apply(ctx, state, structureCell, void 0, plugin);
             }
         }
     }
@@ -38,7 +38,7 @@ const auto = StructureRepresentationProvider({
 const defaultPreset = StructureRepresentationProvider({
     id: 'preset-structure-representation-default',
     display: { name: 'Default', group: 'Preset' },
-    apply(state, structureCell, _, plugin) {
+    async apply(ctx, state, structureCell, _, plugin) {
         const root = state.build().to(structureCell.transform.ref);
         const structure = structureCell.obj!.data;
 
@@ -48,9 +48,9 @@ const defaultPreset = StructureRepresentationProvider({
             .applyOrUpdateTagged(reprTags, StateTransforms.Representation.StructureRepresentation3D,
                 StructureRepresentation3DHelpers.getDefaultParams(plugin, 'cartoon', structure));
 
-        applyComplex(root, 'ligand')
-            .applyOrUpdateTagged(reprTags, StateTransforms.Representation.StructureRepresentation3D,
-                StructureRepresentation3DHelpers.getDefaultParams(plugin, 'ball-and-stick', structure));
+        const ligand = applyComplex(root, 'ligand');
+        const ligandRepr = ligand.applyOrUpdateTagged(reprTags, StateTransforms.Representation.StructureRepresentation3D,
+            StructureRepresentation3DHelpers.getDefaultParams(plugin, 'ball-and-stick', structure));
 
         applyComplex(root, 'modified')
             .applyOrUpdateTagged(reprTags, StateTransforms.Representation.StructureRepresentation3D,
@@ -71,14 +71,21 @@ const defaultPreset = StructureRepresentationProvider({
             .applyOrUpdateTagged(reprTags, StateTransforms.Representation.StructureRepresentation3D,
                 StructureRepresentation3DHelpers.getDefaultParamsWithTheme(plugin, 'spacefill', 'polymer-id', structure, {}));
 
-        return state.updateTree(root, { revertIfAborted: true });
+        await state.updateTree(root, { revertIfAborted: true }).runInContext(ctx);
+
+        return {
+            ligand: {
+                selection: ligand.selector,
+                repr: ligandRepr.selector
+            }
+        };
     }
 });
 
 const proteinAndNucleic = StructureRepresentationProvider({
     id: 'preset-structure-representation-protein-and-nucleic',
     display: { name: 'Protein & Nucleic', group: 'Preset' },
-    apply(state, structureCell, _, plugin) {
+    async apply(ctx, state, structureCell, _, plugin) {
         const root = plugin.state.dataState.build().to(structureCell.transform.ref);
         const structure = structureCell.obj!.data;
         const reprTags = [this.id, RepresentationProviderTags.Representation];
@@ -91,14 +98,15 @@ const proteinAndNucleic = StructureRepresentationProvider({
             .applyOrUpdateTagged(reprTags, StateTransforms.Representation.StructureRepresentation3D,
                 StructureRepresentation3DHelpers.getDefaultParams(plugin, 'gaussian-surface', structure));
 
-        return state.updateTree(root, { revertIfAborted: true });
+        await state.updateTree(root, { revertIfAborted: true }).runInContext(ctx);
+        return {};
     }
 });
 
 const capsid = StructureRepresentationProvider({
     id: 'preset-structure-representation-capsid',
     display: { name: 'Capsid', group: 'Preset' },
-    apply(state, structureCell, _, plugin) {
+    async apply(ctx, state, structureCell, _, plugin) {
         const root = plugin.state.dataState.build().to(structureCell.transform.ref);
         const structure = structureCell.obj!.data;
 
@@ -111,14 +119,15 @@ const capsid = StructureRepresentationProvider({
         applySelection(root, 'polymer')
             .applyOrUpdateTagged(reprTags, StateTransforms.Representation.StructureRepresentation3D, params);
 
-        return state.updateTree(root, { revertIfAborted: true });
+        await state.updateTree(root, { revertIfAborted: true }).runInContext(ctx);
+        return {};
     }
 });
 
 const coarseCapsid = StructureRepresentationProvider({
     id: 'preset-structure-representation-coarse-capsid',
     display: { name: 'Coarse Capsid', group: 'Preset' },
-    apply(state, structureCell, _, plugin) {
+    async apply(ctx, state, structureCell, _, plugin) {
         const root = plugin.state.dataState.build().to(structureCell.transform.ref);
         const structure = structureCell.obj!.data;
 
@@ -134,14 +143,15 @@ const coarseCapsid = StructureRepresentationProvider({
         applySelection(root, 'trace')
             .applyOrUpdateTagged(reprTags, StateTransforms.Representation.StructureRepresentation3D, params);
 
-        return state.updateTree(root, { revertIfAborted: true });
+        await state.updateTree(root, { revertIfAborted: true }).runInContext(ctx);
+        return {};
     }
 });
 
 const cartoon = StructureRepresentationProvider({
     id: 'preset-structure-representation-cartoon',
     display: { name: 'Cartoon', group: 'Preset' },
-    apply(state, structureCell, _, plugin) {
+    async apply(ctx, state, structureCell, _, plugin) {
         const root = plugin.state.dataState.build().to(structureCell.transform.ref);
         const structure = structureCell.obj!.data;
 
@@ -154,7 +164,8 @@ const cartoon = StructureRepresentationProvider({
         applySelection(root, 'polymer')
             .applyOrUpdateTagged(reprTags, StateTransforms.Representation.StructureRepresentation3D, params);
 
-        return state.updateTree(root, { revertIfAborted: true });
+        await state.updateTree(root, { revertIfAborted: true }).runInContext(ctx);
+        return {};
     }
 });
 

src/mol-plugin/state/representation/structure/provider.ts

@@ -6,18 +6,17 @@
 
 import { PluginContext } from '../../../context';
 import { State, StateObjectCell } from '../../../../mol-state';
-import { Task } from '../../../../mol-task';
+import { RuntimeContext } from '../../../../mol-task';
 import { Structure } from '../../../../mol-model/structure';
 import { ParamDefinition as PD } from '../../../../mol-util/param-definition';
 import { PluginStateObject } from '../../objects';
 
-export interface StructureRepresentationProvider<P = any> {
+export interface StructureRepresentationProvider<P = any, S = {}> {
     id: string,
     display: { name: string, group: string, description?: string },
     isApplicable?(structure: Structure, plugin: PluginContext): boolean,
     params?(structure: Structure | undefined, plugin: PluginContext): PD.Def<P>,
-    // TODO: have create return a "representation structure object" that allows modifications
-    apply(state: State, structure: StateObjectCell<PluginStateObject.Molecule.Structure>, params: P, plugin: PluginContext): Task<any> | Promise<void> | void,
+    apply(ctx: RuntimeContext, state: State, structure: StateObjectCell<PluginStateObject.Molecule.Structure>, params: P, plugin: PluginContext): Promise<S> | S,
     // TODO: Custom remove function for more complicated things
     // remove?(state: State, ref: string, plugin: PluginContext): void
 }
@@ -27,4 +26,4 @@ export const enum RepresentationProviderTags {
     Selection = 'preset-structure-selection'
 }
 
-export function StructureRepresentationProvider<P>(repr: StructureRepresentationProvider<P>) { return repr; }
+export function StructureRepresentationProvider<P, S>(repr: StructureRepresentationProvider<P, S>) { return repr; }

src/mol-plugin/ui/left-panel.tsx

@@ -114,7 +114,7 @@ class DataIcon extends PluginUIComponent<{ set: (tab: LeftPanelTabName) => void
 
     render() {
         return <IconButton
-            icon='flow-tree' toggleState={this.tab === 'data'} onClick={() => this.props.set('data')} title='State Tree' 
+            icon='flow-tree' toggleState={this.tab === 'data'} onClick={() => this.props.set('data')} title='State Tree'
             style={{ position: 'relative' }} extraContent={this.state.changed ? <div className='msp-left-panel-controls-button-data-dirty' /> : void 0} />;
     }
 }

src/mol-state/object.ts

@@ -87,6 +87,7 @@ namespace StateObjectCell {
 
     export type Obj<C extends StateObjectCell> = C extends StateObjectCell<infer T> ? T : never
     export type Transform<C extends StateObjectCell> = C extends StateObjectCell<any, infer T> ? T : never
+    export type Transformer<C extends StateObjectCell> = C extends StateObjectCell<any, StateTransform<infer T>> ? T : never
 }
 
 // TODO: improve the API?
@@ -111,4 +112,27 @@ export class StateObjectTracker<T extends StateObject> {
     }
 
     constructor(private state: State) { }
+}
+
+export class StateObjectSelector<S extends StateObject = StateObject, T extends StateTransformer = StateTransformer> {
+    getObj(): S | undefined {
+        return this.state?.cells.get(this.ref)?.obj as S | undefined;
+    }
+
+    getData(): S['data'] | undefined {
+        return this.getObj()?.data;
+    }
+
+    constructor(public ref: StateTransform.Ref, public state?: State) {
+    }
+}
+
+export namespace StateObjectSelector {
+    export function is(s: any): s is StateObjectSelector {
+        const _s: StateObjectSelector = s;
+        return !!s && !!_s.getData && !!_s.getObj && !!_s.ref;
+    }
+
+    export type Obj<S extends StateObjectSelector> = S extends StateObjectSelector<infer A> ? A : never
+    export type Transformer<S extends StateObjectSelector> = S extends StateObjectSelector<any, infer T> ? T : never
 }

src/mol-state/state.ts

@@ -133,7 +133,7 @@ class State {
      * @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: StateBuilder.To<T>, options?: Partial<State.UpdateOptions>): Task<T>
+    updateTree<T extends StateObject>(tree: StateBuilder.To<T, any>, options?: Partial<State.UpdateOptions>): Task<T>
     updateTree(tree: StateTree | StateBuilder, options?: Partial<State.UpdateOptions>): Task<void>
     updateTree(tree: StateTree | StateBuilder, options?: Partial<State.UpdateOptions>): Task<any> {
         const params: UpdateParams = { tree, options };

src/mol-state/state/builder.ts

@@ -6,7 +6,7 @@
 
 import { StateTree } from '../tree/immutable';
 import { TransientTree } from '../tree/transient';
-import { StateObject, StateObjectCell } from '../object';
+import { StateObject, StateObjectCell, StateObjectSelector } from '../object';
 import { StateTransform } from '../transform';
 import { StateTransformer } from '../transformer';
 import { State } from '../state';
@@ -72,15 +72,22 @@ namespace StateBuilder {
         return !!obj && typeof (obj as StateBuilder).getTree === 'function' && typeof (obj as StateBuilder.To<any>).ref === 'string';
     }
 
+    // type ToFromCell<C extends StateObjectCell> = C extends StateObjectCell<infer A, StateTransform<infer T extends StateTransformer>> ? To<A, any>: never
+
     export class Root implements StateBuilder {
         private state: BuildState;
         get editInfo() { return this.state.editInfo; }
         get currentTree() { return this.state.tree; }
 
         to<A extends StateObject>(ref: StateTransform.Ref): To<A>
-        to<C extends StateObjectCell>(cell: C): To<StateObjectCell.Obj<C>, StateTransform.Transformer<StateObjectCell.Transform<C>>>
-        to(refOrCell: StateTransform.Ref | StateObjectCell) {
-            const ref = typeof refOrCell === 'string' ? refOrCell : refOrCell.transform.ref;
+        to<C extends StateObjectCell>(cell: C): To<StateObjectCell.Obj<C>, StateObjectCell.Transformer<C>>
+        to<S extends StateObjectSelector>(selector: S): To<StateObjectSelector.Obj<S>, StateObjectSelector.Transformer<S>>
+        to(refOrCellOrSelector: StateTransform.Ref | StateObjectCell | StateObjectSelector) {
+            const ref = typeof refOrCellOrSelector === 'string'
+                ? refOrCellOrSelector
+                : StateObjectSelector.is(refOrCellOrSelector)
+                ? refOrCellOrSelector.ref
+                : refOrCellOrSelector.transform.ref;
             return new To<StateObject, StateTransformer>(this.state, ref, this);
         }
         toRoot<A extends StateObject>() { return new To<A>(this.state, this.state.tree.root.ref, this); }
@@ -97,6 +104,7 @@ namespace StateBuilder {
 
     export class To<A extends StateObject, T extends StateTransformer = StateTransformer> implements StateBuilder {
         get editInfo() { return this.state.editInfo; }
+        get selector() { return new StateObjectSelector<A, T>(this.ref, this.state.state); }
 
         readonly ref: StateTransform.Ref;
 
@@ -104,14 +112,13 @@ namespace StateBuilder {
          * Apply the transformed to the parent node
          * If no params are specified (params <- undefined), default params are lazily resolved.
          */
-        apply<T extends StateTransformer<A, any, any>>(tr: T, params?: StateTransformer.Params<T>, options?: Partial<StateTransform.Options>): To<StateTransformer.To<T>> {
+        apply<T extends StateTransformer<A, any, any>>(tr: T, params?: StateTransformer.Params<T>, options?: Partial<StateTransform.Options>): To<StateTransformer.To<T>, T> {
             const t = tr.apply(this.ref, params, options);
             this.state.tree.add(t);
             this.editInfo.count++;
             this.editInfo.lastUpdate = t.ref;
 
             this.state.actions.push({ kind: 'add', transform: t });
-
             return new To(this.state, t.ref, this.root);
         }
 
@@ -119,10 +126,10 @@ namespace StateBuilder {
          * If the ref is present, the transform is applied.
          * Otherwise a transform with the specifed ref is created.
          */
-        applyOrUpdate<T extends StateTransformer<A, any, any>>(ref: StateTransform.Ref, tr: T, params?: StateTransformer.Params<T>, options?: Partial<StateTransform.Options>): To<StateTransformer.To<T>> {
+        applyOrUpdate<T extends StateTransformer<A, any, any>>(ref: StateTransform.Ref, tr: T, params?: StateTransformer.Params<T>, options?: Partial<StateTransform.Options>): To<StateTransformer.To<T>, T> {
             if (this.state.tree.transforms.has(ref)) {
-                this.to(ref).update(params);
-                return this.to(ref) as To<StateTransformer.To<T>>;
+                this.root.to(ref).update(params);
+                return this.root.to(ref);
             } else {
                 return this.apply(tr, params, { ...options, ref });
             }
@@ -132,15 +139,15 @@ namespace StateBuilder {
          * Apply the transformed to the parent node
          * If no params are specified (params <- undefined), default params are lazily resolved.
          */
-        applyOrUpdateTagged<T extends StateTransformer<A, any, any>>(tags: string | string[], tr: T, params?: StateTransformer.Params<T>, options?: Partial<StateTransform.Options>): To<StateTransformer.To<T>> {
+        applyOrUpdateTagged<T extends StateTransformer<A, any, any>>(tags: string | string[], tr: T, params?: StateTransformer.Params<T>, options?: Partial<StateTransform.Options>): To<StateTransformer.To<T>, T> {
             const children = this.state.tree.children.get(this.ref).values();
             while (true) {
                 const child = children.next();
                 if (child.done) break;
                 const tr = this.state.tree.transforms.get(child.value);
                 if (tr && StateTransform.hasTags(tr, tags)) {
-                    this.to(child.value).updateTagged(params, tagsUnion(tr.tags, tags, options && options.tags));
-                    return this.to(child.value) as To<StateTransformer.To<T>>;
+                    this.root.to(child.value).updateTagged(params, tagsUnion(tr.tags, tags, options && options.tags));
+                    return this.root.to(child.value);
                 }
             }
 
@@ -157,14 +164,14 @@ namespace StateBuilder {
         /**
          * A helper to greate a group-like state object and keep the current type.
          */
-        group<T extends StateTransformer<A, any, any>>(tr: T, params?: StateTransformer.Params<T>, options?: Partial<StateTransform.Options>): To<A> {
-            return this.apply(tr, params, options) as To<A>;
+        group<T extends StateTransformer<A, any, any>>(tr: T, params?: StateTransformer.Params<T>, options?: Partial<StateTransform.Options>): To<A, T> {
+            return this.apply(tr, params, options);
         }
 
         /**
          * Inserts a new transform that does not change the object type and move the original children to it.
          */
-        insert<T extends StateTransformer<A, A, any>>(tr: T, params?: StateTransformer.Params<T>, options?: Partial<StateTransform.Options>): To<StateTransformer.To<T>> {
+        insert<T extends StateTransformer<A, A, any>>(tr: T, params?: StateTransformer.Params<T>, options?: Partial<StateTransform.Options>): To<StateTransformer.To<T>, T> {
             // cache the children
             const children = this.state.tree.children.get(this.ref).toArray();
 
@@ -236,7 +243,10 @@ namespace StateBuilder {
             return this.root;
         }
 
-        to<A extends StateObject>(ref: StateTransform.Ref) { return this.root.to<A>(ref); }
+        to<A extends StateObject>(ref: StateTransform.Ref): To<A>
+        to<C extends StateObjectCell>(cell: C): To<StateObjectCell.Obj<C>, StateObjectCell.Transformer<C>>
+        to<S extends StateObjectSelector>(selector: S): To<StateObjectSelector.Obj<S>, StateObjectSelector.Transformer<S>>
+        to(ref: StateTransform.Ref | StateObjectCell | StateObjectSelector) { return  this.root.to(ref as any); }
         toRoot<A extends StateObject>() { return this.root.toRoot<A>(); }
         delete(ref: StateTransform.Ref) { return this.root.delete(ref); }
 
@@ -252,10 +262,11 @@ namespace StateBuilder {
 }
 
 function tagsUnion<T>(...arrays: (string[] | string | undefined)[]): string[] | undefined {
-    const set = new Set();
+    let set: Set<string> | undefined = void 0;
     const ret = [];
     for (const xs of arrays) {
         if (!xs) continue;
+        if (!set) set = new Set();
         if (typeof xs === 'string') {
             if (set.has(xs)) continue;
             set.add(xs);