mol-plugin: wip StructureRepresentationManager

src/mol-plugin/index.ts

@@ -25,6 +25,7 @@ export const DefaultPluginSpec: PluginSpec = {
         PluginSpec.Action(StateActions.Volume.DownloadDensity),
         PluginSpec.Action(StateActions.DataFormat.OpenFile),
         PluginSpec.Action(StateActions.Structure.Create3DRepresentationPreset),
+        PluginSpec.Action(StateActions.Structure.Remove3DRepresentationPreset),
         PluginSpec.Action(StateActions.Structure.EnableModelCustomProps),
         PluginSpec.Action(StateActions.Structure.EnableStructureCustomProps),
 

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

@@ -295,6 +295,20 @@ export const Create3DRepresentationPreset = StateAction.build({
     plugin.structureRepresentation.manager.apply(ref, params.type.name, params.type.params);
 });
 
+export const Remove3DRepresentationPreset = StateAction.build({
+    display: { name: 'Remove 3D Representation Preset', description: 'Remove 3D representations.' },
+    from: PluginStateObject.Molecule.Structure,
+    isApplicable(_, t, plugin: PluginContext) { return plugin.structureRepresentation.manager.hasManagedRepresentation(t.ref); },
+    params(a, plugin: PluginContext) {
+        return {
+            type: plugin.structureRepresentation.manager.getOptions(a.data).select
+        };
+    }
+})(({ ref, params }, plugin: PluginContext) => {
+    // TODO: this will be completely handled by the managed and is just for testing purposes
+    plugin.structureRepresentation.manager.remove(params.type, ref);
+});
+
 export const UpdateTrajectory = StateAction.build({
     display: { name: 'Update Trajectory' },
     params: {

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

@@ -6,17 +6,19 @@
 
 import { arrayFind } from '../../../mol-data/util';
 import { Structure } from '../../../mol-model/structure';
-import { StateTransform } from '../../../mol-state';
+import { StateTransform, StateTree, StateSelection } from '../../../mol-state';
 import { Task } from '../../../mol-task';
 import { isProductionMode } from '../../../mol-util/debug';
 import { objectForEach } from '../../../mol-util/object';
 import { ParamDefinition as PD } from '../../../mol-util/param-definition';
 import { PluginContext } from '../../context';
 import { PresetStructureReprentations } from './structure/preset';
-import { StructureRepresentationProvider } from './structure/providers';
+import { StructureRepresentationProvider, RepresentationProviderTags } from './structure/provider';
+import { UniqueArray } from '../../../mol-data/generic';
 
 export class StructureRepresentationManager {
     private providers: StructureRepresentationProvider[] = [];
+    private providerMap: Map<string, StructureRepresentationProvider> = new Map();
 
     readonly defaultProvider = PresetStructureReprentations.default;
 
@@ -40,9 +42,61 @@ export class StructureRepresentationManager {
         return PD.MappedStatic(options[0][0], map, { options });
     }
 
+    hasManagedRepresentation(ref: StateTransform.Ref) {
+        const tree = this.plugin.state.dataState.tree;
+        // TODO: make this state selection function?
+        return StateTree.doPreOrder(tree, tree.transforms.get(ref), { found: false, map: this.providerMap }, (n, _, s) => {
+            if (!n.tags) return;
+            for (const t of n.tags) {
+                if (s.map.has(t)) {
+                    s.found = true;
+                    return false;
+                }
+            }
+        }).found;
+    }
+
+    getManagedRepresentations(ref: StateTransform.Ref) {
+        // TODO: check if Structure etc.
+        const tree = this.plugin.state.dataState.tree;
+        return StateTree.doPreOrder(tree, tree.transforms.get(ref), { found: UniqueArray.create<string, StructureRepresentationProvider>(), map: this.providerMap }, (n, _, s) => {
+            if (!n.tags) return;
+            for (const t of n.tags) {
+                if (s.map.has(t)) UniqueArray.add(s.found, t, s.map.get(t)!);
+            }
+        }).found.array;
+    }
+
     register(provider: StructureRepresentationProvider) {
+        if (this.providerMap.has(provider.id)) {
+            throw new Error(`Repr. provider with id '${provider.id}' already registered.`);
+        }
         // TODO: sort by group
         this.providers.push(provider);
+        this.providerMap.set(provider.id, provider);
+    }
+
+    remove(providerOrId: StructureRepresentationProvider | string, structureRoot?: StateTransform.Ref) {
+        const root = structureRoot || StateTransform.RootRef;
+        const id = typeof providerOrId === 'string' ? providerOrId : providerOrId.id;
+
+        const state = this.plugin.state.dataState;
+        const reprs = StateSelection.findWithAllTags(state.tree, root, new Set([id, RepresentationProviderTags.Representation]));
+
+        const builder = state.build();
+        for (const r of reprs) {
+            builder.delete(r.ref);
+        }
+
+        const tree = builder.currentTree;
+        const selections = StateSelection.findWithAllTags(tree, root, new Set([RepresentationProviderTags.Selection]));
+
+        for (const s of selections) {
+            if (!tree.children.has(s.ref) || tree.children.get(s.ref).size === 0) builder.delete(s.ref);
+        }
+
+        if (builder.editInfo.count === 0) return;
+        return this.plugin.runTask(state.updateTree(builder));
     }
 
     apply<P>(ref: StateTransform.Ref, providerOrId: StructureRepresentationProvider<P> | string, params?: P) {

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

@@ -7,9 +7,11 @@
 import { StateTransforms } from '../../transforms';
 import { StructureComplexElementTypes } from '../../transforms/model';
 import { StructureRepresentation3DHelpers } from '../../transforms/representation';
-import { applyBuiltInSelection } from '../../../util/structure-selection-helper';
+import { StructureSelectionQueries as Q } from '../../../util/structure-selection-helper';
 import { BuiltInStructureRepresentations } from '../../../../mol-repr/structure/registry';
-import { StructureRepresentationProvider } from './providers';
+import { StructureRepresentationProvider, RepresentationProviderTags } from './provider';
+import { StateBuilder } from '../../../../mol-state';
+import { PluginStateObject } from '../../objects';
 
 export const PresetStructureReprentations = {
     default: StructureRepresentationProvider({
@@ -19,34 +21,34 @@ export const PresetStructureReprentations = {
             const root = state.build().to(structureCell.transform.ref);
             const structure = structureCell.obj!.data;
 
-            const tags = this.id;
+            const reprTags = [this.id, RepresentationProviderTags.Representation];
 
-            root.apply(StateTransforms.Model.StructureComplexElement, { type: 'protein-or-nucleic' }, { tags: [this.id, StructureComplexElementTypes['protein-or-nucleic']] })
-                .apply(StateTransforms.Representation.StructureRepresentation3D,
-                    StructureRepresentation3DHelpers.getDefaultParams(plugin, 'cartoon', structure), { tags });
+            applyComplex(root, 'protein-or-nucleic')
+                .applyOrUpdateTagged(reprTags, StateTransforms.Representation.StructureRepresentation3D,
+                    StructureRepresentation3DHelpers.getDefaultParams(plugin, 'cartoon', structure));
 
-            root.apply(StateTransforms.Model.StructureComplexElement, { type: 'ligand' }, { tags: [this.id, StructureComplexElementTypes.ligand] })
-                .apply(StateTransforms.Representation.StructureRepresentation3D,
-                    StructureRepresentation3DHelpers.getDefaultParams(plugin, 'ball-and-stick', structure), { tags });
+            applyComplex(root, 'ligand')
+                .applyOrUpdateTagged(reprTags, StateTransforms.Representation.StructureRepresentation3D,
+                    StructureRepresentation3DHelpers.getDefaultParams(plugin, 'ball-and-stick', structure));
 
-            root.apply(StateTransforms.Model.StructureComplexElement, { type: 'modified' }, { tags: [this.id, StructureComplexElementTypes.modified] })
-                .apply(StateTransforms.Representation.StructureRepresentation3D,
-                    StructureRepresentation3DHelpers.getDefaultParamsWithTheme(plugin, 'ball-and-stick', 'polymer-id', structure, void 0), { tags });
+            applyComplex(root, 'modified')
+                .applyOrUpdateTagged(reprTags, StateTransforms.Representation.StructureRepresentation3D,
+                    StructureRepresentation3DHelpers.getDefaultParamsWithTheme(plugin, 'ball-and-stick', 'polymer-id', structure, void 0));
 
-            const branched = root.apply(StateTransforms.Model.StructureComplexElement, { type: 'branched' }, { tags: [this.id, StructureComplexElementTypes.branched] })
+            const branched = applyComplex(root, 'branched');
 
-            branched.apply(StateTransforms.Representation.StructureRepresentation3D,
-                StructureRepresentation3DHelpers.getDefaultParams(plugin, 'ball-and-stick', structure, { alpha: 0.15 }), { tags });
-            branched.apply(StateTransforms.Representation.StructureRepresentation3D,
-                StructureRepresentation3DHelpers.getDefaultParams(plugin, 'carbohydrate', structure), { tags });
+            branched.applyOrUpdateTagged(reprTags, StateTransforms.Representation.StructureRepresentation3D,
+                StructureRepresentation3DHelpers.getDefaultParams(plugin, 'ball-and-stick', structure, { alpha: 0.15 }));
+            branched.applyOrUpdateTagged(reprTags, StateTransforms.Representation.StructureRepresentation3D,
+                StructureRepresentation3DHelpers.getDefaultParams(plugin, 'carbohydrate', structure));
 
-            root.apply(StateTransforms.Model.StructureComplexElement, { type: 'water' }, { tags: [this.id, StructureComplexElementTypes.water] })
-                .apply(StateTransforms.Representation.StructureRepresentation3D,
-                    StructureRepresentation3DHelpers.getDefaultParams(plugin, 'ball-and-stick', structure, { alpha: 0.51 }), { tags });
+            applyComplex(root, 'water')
+                .applyOrUpdateTagged(reprTags, StateTransforms.Representation.StructureRepresentation3D,
+                    StructureRepresentation3DHelpers.getDefaultParams(plugin, 'ball-and-stick', structure, { alpha: 0.51 }));
 
-            root.apply(StateTransforms.Model.StructureComplexElement, { type: 'coarse' }, { tags: [this.id, StructureComplexElementTypes.coarse] })
-                .apply(StateTransforms.Representation.StructureRepresentation3D,
-                    StructureRepresentation3DHelpers.getDefaultParamsWithTheme(plugin, 'spacefill', 'polymer-id', structure, {}), { tags });
+            applyComplex(root, 'coarse')
+                .applyOrUpdateTagged(reprTags, StateTransforms.Representation.StructureRepresentation3D,
+                    StructureRepresentation3DHelpers.getDefaultParamsWithTheme(plugin, 'spacefill', 'polymer-id', structure, {}));
 
             return state.updateTree(root, { revertIfAborted: true });
         }
@@ -62,10 +64,41 @@ export const PresetStructureReprentations = {
                 repr: [BuiltInStructureRepresentations['gaussian-surface'], () => ({ smoothness: 1 })]
             });
 
-            applyBuiltInSelection(root, 'polymer', this.id)
-                .apply(StateTransforms.Representation.StructureRepresentation3D, params, { tags: this.id });
+            const reprTags = [this.id, RepresentationProviderTags.Representation];
+
+            applySelection(root, 'polymer')
+                .applyOrUpdateTagged(reprTags, StateTransforms.Representation.StructureRepresentation3D, params);
+
+            return state.updateTree(root, { revertIfAborted: true });
+        }
+    }),
+    cartoon: StructureRepresentationProvider({
+        id: 'preset-structure-representation-cartoon',
+        display: { name: 'Cartoon', group: 'Preset' },
+        apply(state, structureCell, _, plugin) {
+            const root = plugin.state.dataState.build().to(structureCell.transform.ref);
+            const structure = structureCell.obj!.data;
+
+            const params = StructureRepresentation3DHelpers.createParams(plugin, structure, {
+                repr: BuiltInStructureRepresentations['cartoon']
+            });
+
+            const reprTags = [this.id, RepresentationProviderTags.Representation];
+
+            applySelection(root, 'polymer')
+                .applyOrUpdateTagged(reprTags, StateTransforms.Representation.StructureRepresentation3D, params);
 
             return state.updateTree(root, { revertIfAborted: true });
         }
     })
-};
+};
+
+function applyComplex(to: StateBuilder.To<PluginStateObject.Molecule.Structure>, type: keyof typeof StructureComplexElementTypes) {
+    return to.applyOrUpdateTagged(type, StateTransforms.Model.StructureComplexElement, { type }, { tags: RepresentationProviderTags.Selection });
+}
+
+function applySelection(to: StateBuilder.To<PluginStateObject.Molecule.Structure>, query: keyof typeof Q) {
+    return to.applyOrUpdateTagged(query, StateTransforms.Model.StructureSelectionFromExpression,
+        { expression: Q[query].expression, label: Q[query].label },
+        { tags: RepresentationProviderTags.Selection });
+}

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

@@ -17,7 +17,14 @@ export interface StructureRepresentationProvider<P = any> {
     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(state: State, structure: StateObjectCell<PluginStateObject.Molecule.Structure>, params: P, plugin: PluginContext): Task<any> | Promise<void> | void,
+    // TODO: Custom remove function for more complicated things
+    // remove?(state: State, ref: string, plugin: PluginContext): void
+}
+
+export const enum RepresentationProviderTags {
+    Representation = 'preset-structure-representation',
+    Selection = 'preset-structure-selection'
 }
 
 export function StructureRepresentationProvider<P>(repr: StructureRepresentationProvider<P>) { return repr; }

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

@@ -17,8 +17,9 @@ export class StateObjectActions extends PluginUIComponent<{ state: State, nodeRe
     }
 
     componentDidMount() {
-        // this.subscribe(this.plugin.state.behavior.currentObject, o => {
-        //     this.forceUpdate();
+        // TODO: handle tree change: some state actions might become invalid
+        // this.subscribe(this.props.state.events.changed, o => {
+        //     this.setState(createStateObjectActionSelectState(this.props));
         // });
 
         this.subscribe(this.plugin.events.state.object.updated, ({ ref, state }) => {
@@ -86,6 +87,11 @@ export class StateObjectActionSelect extends PluginUIComponent<StateObjectAction
     }
 
     componentDidMount() {
+        // TODO: handle tree change: some state actions might become invalid
+        // this.subscribe(this.props.state.events.changed, o => {
+        //     this.setState(createStateObjectActionSelectState(this.props));
+        // });
+
         this.subscribe(this.plugin.events.state.object.updated, ({ ref, state }) => {
             const current = this.current;
             if (current.ref !== ref || current.state !== state) return;

src/mol-state/state/builder.ts

@@ -75,6 +75,7 @@ namespace StateBuilder {
     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>>>
@@ -127,6 +128,32 @@ 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>> {
+            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>>;
+                }
+            }
+
+            const t = tr.apply(this.ref, params, { ...options, tags: tagsUnion(tags, options && options.tags) });
+            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);
+        }
+
         /**
          * A helper to greate a group-like state object and keep the current type.
          */
@@ -180,6 +207,14 @@ namespace StateBuilder {
         //     return this.root;
         // }
 
+        private updateTagged<T extends StateTransformer<any, A, any>>(params: any, tags: string | string[] | undefined) {
+            if (this.state.tree.setParams(this.ref, params) || this.state.tree.setTags(this.ref, tags)) {
+                this.editInfo.count++;
+                this.editInfo.lastUpdate = this.ref;
+                this.state.actions.push({ kind: 'update', ref: this.ref, params });
+            }
+        }
+
         update<T extends StateTransformer<any, A, any>>(transformer: T, params: (old: StateTransformer.Params<T>) => StateTransformer.Params<T>): Root
         update(params: StateTransformer.Params<T> | ((old: StateTransformer.Params<T>) => StateTransformer.Params<T>)): Root
         update<T extends StateTransformer<any, A, any>>(paramsOrTransformer: T | any, provider?: (old: StateTransformer.Params<T>) => StateTransformer.Params<T>) {
@@ -196,10 +231,8 @@ namespace StateBuilder {
             if (this.state.tree.setParams(this.ref, params)) {
                 this.editInfo.count++;
                 this.editInfo.lastUpdate = this.ref;
+                this.state.actions.push({ kind: 'update', ref: this.ref, params });
             }
-
-            this.state.actions.push({ kind: 'update', ref: this.ref, params });
-
             return this.root;
         }
 
@@ -216,4 +249,24 @@ namespace StateBuilder {
             }
         }
     }
+}
+
+function tagsUnion<T>(...arrays: (string[] | string | undefined)[]): string[] | undefined {
+    const set = new Set();
+    const ret = [];
+    for (const xs of arrays) {
+        if (!xs) continue;
+        if (typeof xs === 'string') {
+            if (set.has(xs)) continue;
+            set.add(xs);
+            ret.push(xs);
+        } else {
+            for (const x of xs) {
+                if (set.has(x)) continue;
+                set.add(x);
+                ret.push(x);
+            }
+        }
+    }
+    return ret;
 }

src/mol-state/state/selection.ts

@@ -293,6 +293,28 @@ namespace StateSelection {
         }
         return true;
     }
+
+    export function findWithAllTags<K extends string = string>(tree: StateTree, root: StateTransform.Ref, tags: Set<K>): StateTransform[] {
+        return StateTree.doPreOrder(tree, tree.transforms.get(root), { refs: [], tags }, _findWithAllTags).refs;
+    }
+
+    function _findWithAllTags(n: StateTransform, _: any, s: { refs: StateTransform[], tags: Set<string> }) {
+        if (n.tags) {
+            const len = s.tags.size;
+            let found = 0;
+            for (const t of n.tags) {
+                if (!s.tags.has(t)) continue;
+                found++;
+
+                if (found === len) {
+                    s.refs.push(n);
+                    break;
+                }
+            }
+        } else if (s.tags.size === 0) {
+            s.refs.push(n);
+        }
+    }
 }
 
 export { StateSelection }

src/mol-state/transform.ts

@@ -99,6 +99,8 @@ namespace Transform {
         let tags: string[] | undefined = void 0;
         if (options && options.tags) {
             tags = typeof options.tags === 'string' ? [options.tags] : options.tags;
+            if (tags.length === 0) tags = void 0;
+            else tags.sort();
         }
         return {
             parent,
@@ -121,6 +123,16 @@ namespace Transform {
         return { ...t, state: { ...t.state, ...state } };
     }
 
+    export function withTags(t: Transform, newTags?: string | string[]): Transform {
+        let tags: string[] | undefined = void 0;
+        if (newTags) {
+            tags = typeof newTags === 'string' ? [newTags] : newTags;
+            if (tags.length === 0) tags = void 0;
+            else tags.sort();
+        }
+        return { ...t, tags, version: UUID.create22() };
+    }
+
     export function withParent(t: Transform, parent: Ref): Transform {
         return { ...t, parent, version: UUID.create22() };
     }
@@ -134,6 +146,15 @@ namespace Transform {
         return t.tags.indexOf(tag) >= 0;
     }
 
+    export function hasTags(t: Transform, tags: string | string[]) {
+        if (!t.tags) return typeof tags !== 'string' && tags.length === 0;
+        if (typeof tags === 'string') return hasTag(t, tags);
+        for (const tag of tags) {
+            if (t.tags.indexOf(tag) < 0) return false;
+        }
+        return true;
+    }
+
     export interface Serialized {
         parent: string,
         transformer: string,

src/mol-state/tree/transient.ts

@@ -8,6 +8,7 @@ import { Map as ImmutableMap, OrderedSet } from 'immutable';
 import { StateTransform } from '../transform';
 import { StateTree } from './immutable';
 import { shallowEqual } from '../../mol-util/object';
+import { arrayEqual } from '../../mol-util/array';
 
 export { TransientTree }
 
@@ -181,6 +182,27 @@ class TransientTree implements StateTree {
         return true;
     }
 
+     /** Calls Transform.definition.params.areEqual if available, otherwise uses shallowEqual to check if the params changed */
+     setTags(ref: StateTransform.Ref, tags: string | string[] | undefined) {
+        ensurePresent(this.transforms, ref);
+
+        const transform = this.transforms.get(ref)!;
+
+        const withTags = StateTransform.withParams(transform, tags)
+        // TODO: should this be here?
+        if (arrayEqual(transform.tags, withTags.tags)) {
+            return false;
+        }
+
+        if (!this.changedNodes) {
+            this.changedNodes = true;
+            this.transforms = this.transforms.asMutable();
+        }
+
+        this.transforms.set(transform.ref, withTags);
+        return true;
+    }
+
     assignState(ref: StateTransform.Ref, state?: Partial<StateTransform.State>) {
         ensurePresent(this.transforms, ref);
 

src/mol-util/array.ts

@@ -103,4 +103,16 @@ export function arraySetRemove<T>(xs: T[], x: T) {
     }
     xs.pop();
     return true;
+}
+
+export function arrayEqual<T>(xs?: T[], ys?: T[]) {
+    if (!xs || xs.length === 0) return !ys || ys.length === 0;
+    if (!ys) return false;
+
+    const lenX = xs.length;
+    if (lenX !== ys.length) return false;
+    for (let i = 0; i < lenX; i++) {
+        if (xs[i] !== ys[i]) return false;
+    }
+    return true;
 }