wip: StructureHierarchyManager

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

@@ -104,7 +104,7 @@ export class RepresentationBuilder {
         }
 
         const tree = builder.currentTree;
-        const selections = StateSelection.findWithAllTags(tree, root, new Set([RepresentationProviderTags.Selection]));
+        const selections = StateSelection.findWithAllTags(tree, root, new Set([RepresentationProviderTags.Component]));
 
         for (const s of selections) {
             if (!tree.children.has(s.ref) || tree.children.get(s.ref).size === 0) builder.delete(s.ref);

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

@@ -166,18 +166,15 @@ export class StructureBuilder {
     }
 
     /** returns undefined if the component is empty/null */
-    async tryCreateComponent(structure: StateObjectRef<SO.Molecule.Structure>, params: StructureComponentParams, tag?: string): Promise<StateObjectRef<SO.Molecule.Structure> | undefined> {
+    async tryCreateComponent(structure: StateObjectRef<SO.Molecule.Structure>, params: StructureComponentParams, key: string, tags?: string[]): Promise<StateObjectRef<SO.Molecule.Structure> | undefined> {
         const state = this.dataState;
 
         const root = state.build().to(structure);
-        let component: StateBuilder.To<SO.Molecule.Structure>;
 
-        if (tag) {
-            const typeTag = `structure-component-${tag}`;
-            component = root.applyOrUpdateTagged(typeTag, StateTransforms.Model.StructureComponent, params, { tags: [StructureBuilderTags.Component, typeTag] });
-        } else {
-            component = root.apply(StateTransforms.Model.StructureComponent, params, { tags: StructureBuilderTags.Component });
-        }
+        const keyTag = `structure-component-${key}`;
+        const component = root.applyOrUpdateTagged(keyTag, StateTransforms.Model.StructureComponent, params, { 
+            tags: tags ? [...tags, StructureBuilderTags.Component, keyTag] : [StructureBuilderTags.Component, keyTag]
+        });
 
         await this.plugin.runTask(this.dataState.updateTree(component));
 

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

@@ -206,7 +206,7 @@ function staticComponent(plugin: PluginContext, structure: StateObjectRef<Plugin
         type: { name: 'static', params: type },
         nullIfEmpty: true,
         label: ''
-    }, `static-${type}`);
+    }, `static-${type}`, [RepresentationProviderTags.Component]);
 }
 
 function selectionComponent(plugin: PluginContext, structure: StateObjectRef<PluginStateObject.Molecule.Structure>, query: keyof typeof Q) {
@@ -214,7 +214,7 @@ function selectionComponent(plugin: PluginContext, structure: StateObjectRef<Plu
         type: { name: 'expression', params: Q[query].expression },
         nullIfEmpty: true,
         label: Q[query].label
-    }, `selection-${query}`);
+    }, `selection-${query}`, [RepresentationProviderTags.Component]);
 }
 
 export const PresetStructureReprentations = {

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

@@ -28,7 +28,7 @@ export namespace StructureRepresentationProvider {
 
 export const enum RepresentationProviderTags {
     Representation = 'preset-structure-representation',
-    Selection = 'preset-structure-selection'
+    Component = 'preset-structure-component'
 }
 
 export function StructureRepresentationProvider<P, S>(repr: StructureRepresentationProvider<P, S>) { return repr; }

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

@@ -1,7 +1,90 @@
 /**
- * Copyright (c) 2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-// TODO: manager that handles representation of models (assembly / symmetry / etc) and structures (components / visuals)
+import { PluginContext } from '../../mol-plugin/context';
+import { StructureHierarchy, buildStructureHierarchy, ModelRef, StructureComponentRef } from './structure/hierarchy';
+import { RxEventHelper } from '../../mol-util/rx-event-helper';
+
+export class StructureHierarchyManager {
+    private ev = RxEventHelper.create();
+
+    readonly behaviors = {
+        hierarchy: this.ev.behavior(StructureHierarchy()),
+        currentModels: this.ev.behavior([] as ReadonlyArray<ModelRef>)
+    }
+
+    private checkCurrent() {
+        const hierarchy = this.behaviors.hierarchy.value;
+        const current = this.behaviors.currentModels.value;
+        if (current.length === 0) {
+            const models = hierarchy.trajectories[0]?.models;
+            if (models) {
+                this.behaviors.currentModels.next(models);
+            }
+            return;
+        }
+
+        const newCurrent: ModelRef[] = [];
+        for (const c of current) {
+            const ref = hierarchy.refs.get(c.cell.transform.ref) as ModelRef;
+            if (!ref) continue;
+            newCurrent.push(ref);
+        }
+
+        this.behaviors.currentModels.next(newCurrent);
+    }
+
+    private sync() {
+        const update = buildStructureHierarchy(this.plugin.state.dataState, this.behaviors.hierarchy.value);
+        if (update.added.length === 0 && update.updated.length === 0 && update.removed.length === 0) {
+            return;
+        }
+
+        this.behaviors.hierarchy.next(update.hierarchy)
+        this.checkCurrent();
+    }
+
+    constructor(private plugin: PluginContext) {
+        plugin.state.dataState.events.changed.subscribe(e => {
+            if (e.inTransaction) return;
+            this.sync();
+        });
+    }
+}
+
+export namespace StructureHierarchyManager {
+    function componentKey(c: StructureComponentRef) {
+        if (!c.cell.transform.tags) return;
+        return [...c.cell.transform.tags].sort().join();
+    }
+
+    export function getCommonComponentPivots(models: ReadonlyArray<ModelRef>) {
+        if (!models[0]?.structures?.length) return [];
+        if (models[0]?.structures?.length === 1) return models[0]?.structures[0]?.components || [];
+
+        const pivots = new Map<string, StructureComponentRef>();
+
+        for (const c of models[0]?.structures[0]?.components) {
+            const key = componentKey(c);
+            if (!key) continue;
+            pivots.set(key, c);
+        }
+
+        for (const m of models) {
+            for (const s of m.structures) {
+                for (const c of s.components) {
+                    const key = componentKey(c);
+                    if (!key) continue;
+                    if (!pivots.has(key)) pivots.delete(key);
+                }
+            }
+        }
+
+        const ret: StructureComponentRef[] = [];
+        pivots.forEach(function (this: StructureComponentRef[], p) { this.push(p) }, ret);
+        return ret;
+    }
+}

src/mol-plugin-state/manager/structure/hierarchy.ts

@@ -0,0 +1,205 @@
+/**
+ * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { PluginStateObject as SO } from '../../objects';
+import { StateObject, StateTransform, State, StateObjectCell, StateTree } from '../../../mol-state';
+import { StructureBuilderTags } from '../../builder/structure';
+import { RepresentationProviderTags } from '../../builder/structure/provider';
+
+export function buildStructureHierarchy(state: State, previous?: StructureHierarchy) {
+    const build = BuildState(state, previous || StructureHierarchy());
+    StateTree.doPreOrder(state.tree, state.tree.root, build, visitCell);
+    if (previous) previous.refs.forEach(isRemoved, build);
+    return { hierarchy: build.hierarchy, added: build.added, updated: build.updated, removed: build.removed };
+}
+
+export interface StructureHierarchy {
+    trajectories: TrajectoryRef[],
+    refs: Map<StateTransform.Ref, HierarchyRef>
+}
+
+export function StructureHierarchy(): StructureHierarchy {
+    return { trajectories: [], refs: new Map() }
+}
+
+interface RefBase<K extends string = string, T extends StateObject = StateObject> {
+    kind: K,
+    cell: StateObjectCell<T>,
+    version: StateTransform['version']
+}
+
+export type HierarchyRef = 
+    | TrajectoryRef
+    | ModelRef | ModelPropertiesRef
+    | StructureRef | StructurePropertiesRef | StructureComponentRef | StructureRepresentationRef
+
+export interface TrajectoryRef extends RefBase<'trajectory', SO.Molecule.Trajectory> {
+    models: ModelRef[]
+}
+
+function TrajectoryRef(cell: StateObjectCell<SO.Molecule.Trajectory>): TrajectoryRef { 
+    return { kind: 'trajectory', cell, version: cell.transform.version, models: [] };
+}
+
+export interface ModelRef extends RefBase<'model', SO.Molecule.Model> {
+    trajectory: TrajectoryRef,
+    properties?: ModelPropertiesRef,
+    structures: StructureRef[]
+}
+
+function ModelRef(cell: StateObjectCell<SO.Molecule.Model>, trajectory: TrajectoryRef): ModelRef {
+    return { kind: 'model', cell, version: cell.transform.version, trajectory, structures: [] };
+}
+
+export interface ModelPropertiesRef extends RefBase<'model-properties', SO.Molecule.Model> {
+    model: ModelRef
+}
+
+function ModelPropertiesRef(cell: StateObjectCell<SO.Molecule.Model>, model: ModelRef): ModelPropertiesRef {
+    return { kind: 'model-properties', cell, version: cell.transform.version, model };
+}
+
+export interface StructureRef extends RefBase<'structure', SO.Molecule.Structure> {
+    model: ModelRef,
+    properties?: StructurePropertiesRef,
+    components: StructureComponentRef[],
+    // representations: StructureRepresentationRef[],
+    currentFocus?: {
+        focus?: StructureComponentRef,
+        surroundings?: StructureComponentRef,
+    },
+    // volumeStreaming?: ....
+}
+
+function StructureRef(cell: StateObjectCell<SO.Molecule.Structure>, model: ModelRef): StructureRef {
+    return { kind: 'structure', cell, version: cell.transform.version, model, components: [] };
+}
+
+export interface StructurePropertiesRef extends RefBase<'structure-properties', SO.Molecule.Structure> {
+    structure: StructureRef
+}
+
+function StructurePropertiesRef(cell: StateObjectCell<SO.Molecule.Structure>, structure: StructureRef): StructurePropertiesRef {
+    return { kind: 'structure-properties', cell, version: cell.transform.version, structure };
+}
+
+export interface StructureComponentRef extends RefBase<'structure-component', SO.Molecule.Structure> {
+    structure: StructureRef,
+    representations: StructureRepresentationRef[],
+}
+
+function StructureComponentRef(cell: StateObjectCell<SO.Molecule.Structure>, structure: StructureRef): StructureComponentRef {
+    return { kind: 'structure-component', cell, version: cell.transform.version, structure, representations: [] };
+}
+
+export interface StructureRepresentationRef extends RefBase<'structure-representation', SO.Molecule.Structure.Representation3D> {
+    component: StructureComponentRef
+}
+
+function StructureRepresentationRef(cell: StateObjectCell<SO.Molecule.Structure.Representation3D>, component: StructureComponentRef): StructureRepresentationRef {
+    return { kind: 'structure-representation', cell, version: cell.transform.version, component };
+}
+
+interface BuildState {
+    state: State,
+    oldHierarchy: StructureHierarchy,
+
+    hierarchy: StructureHierarchy,
+
+    currentTrajectory?: TrajectoryRef,
+    currentModel?: ModelRef,
+    currentStructure?: StructureRef,
+    currentComponent?: StructureComponentRef,
+
+    updated: HierarchyRef[],
+    added: HierarchyRef[],
+    removed: HierarchyRef[]
+}
+
+function BuildState(state: State, oldHierarchy: StructureHierarchy): BuildState {
+    return { state, oldHierarchy, hierarchy: StructureHierarchy(), updated: [], added: [], removed: [] };
+}
+
+function createOrUpdateRefList<R extends HierarchyRef, C extends any[]>(state: BuildState, cell: StateObjectCell, list: R[], ctor: (...args: C) => R, ...args: C) {
+    const ref: R = ctor(...args);
+    list.push(ref);
+    state.hierarchy.refs.set(cell.transform.ref, ref);
+    const old = state.oldHierarchy.refs.get(cell.transform.ref);
+    if (old) {
+        if (old.version !== cell.transform.version) state.updated.push(ref);
+    } else {
+        state.added.push(ref);
+    }
+    return ref;
+}
+
+function createOrUpdateRef<R extends HierarchyRef, C extends any[]>(state: BuildState, cell: StateObjectCell, old: R | undefined, ctor: (...args: C) => R, ...args: C) {
+    const ref: R = ctor(...args);
+    state.hierarchy.refs.set(cell.transform.ref, ref);
+    if (old) {
+        if (old.version !== cell.transform.version) state.updated.push(ref);
+    } else {
+        state.added.push(ref);
+    }
+    return ref;
+}
+
+const tagMap: [string, (state: BuildState, cell: StateObjectCell) => boolean | void][] = [
+    [StructureBuilderTags.Trajectory, (state, cell) => {
+        state.currentTrajectory = createOrUpdateRefList(state, cell, state.hierarchy.trajectories, TrajectoryRef, cell);
+    }],
+    [StructureBuilderTags.Model, (state, cell) => {
+        if (!state.currentTrajectory) return false;
+        state.currentModel = createOrUpdateRefList(state, cell, state.currentTrajectory.models, ModelRef, cell, state.currentTrajectory);
+    }],
+    [StructureBuilderTags.ModelProperties, (state, cell) => {
+        if (!state.currentModel) return false;
+        state.currentModel.properties = createOrUpdateRef(state, cell, state.currentModel.properties, ModelPropertiesRef, cell, state.currentModel);
+    }],
+    [StructureBuilderTags.Structure, (state, cell) => {
+        if (!state.currentModel) return false;
+        state.currentStructure = createOrUpdateRefList(state, cell, state.currentModel.structures, StructureRef, cell, state.currentModel);
+    }],
+    [StructureBuilderTags.StructureProperties, (state, cell) => {
+        if (!state.currentStructure) return false;
+        state.currentStructure.properties = createOrUpdateRef(state, cell, state.currentStructure.properties, StructurePropertiesRef, cell, state.currentStructure);
+    }],
+    [StructureBuilderTags.Component, (state, cell) => {
+        if (!state.currentStructure) return false;
+        state.currentComponent = createOrUpdateRefList(state, cell, state.currentStructure.components, StructureComponentRef, cell, state.currentStructure);
+    }],
+    [RepresentationProviderTags.Representation, (state, cell) => {
+        if (!state.currentComponent) return false;
+        createOrUpdateRefList(state, cell, state.currentComponent.representations, StructureRepresentationRef, cell, state.currentComponent);
+    }]
+]
+
+function isValidCell(cell?: StateObjectCell): cell is StateObjectCell {
+    if (!cell || !cell.parent.cells.has(cell.transform.ref)) return false;
+    const { obj } = cell;
+    if (!obj || obj === StateObject.Null || (cell.status !== 'ok' && cell.status !== 'error')) return false;
+    return true;
+}
+
+function visitCell(t: StateTransform, tree: StateTree, state: BuildState): boolean {
+    const cell = state.state.cells.get(t.ref);
+    if (!isValidCell(cell)) return false;
+
+    for (const [t, f] of tagMap) {
+        if (StateObject.hasTag(cell.obj!, t)) {
+            const stop = f(state, cell);
+            if (stop === false) return false;
+        }
+    }
+
+    return true;
+}
+
+function isRemoved(this: BuildState, ref: HierarchyRef) {
+    const { cell } = ref;
+    if (isValidCell(cell)) return;
+    this.removed.push(ref);
+}

src/mol-plugin-ui/controls.tsx

@@ -20,6 +20,7 @@ import { StructureRepresentationControls } from './structure/representation';
 import { StructureSelectionControls } from './structure/selection';
 import { StructureMeasurementsControls } from './structure/measurements';
 import { Icon } from './controls/icons';
+import { StructureComponentControls } from './structure/components';
 
 export class TrajectoryViewportControls extends PluginUIComponent<{}, { show: boolean, label: string }> {
     state = { show: false, label: '' }
@@ -269,6 +270,7 @@ export class StructureToolsWrapper extends PluginUIComponent {
             <StructureSelectionControls />
             <StructureRepresentationControls />
             <StructureMeasurementsControls />
+            <StructureComponentControls />
         </div>;
     }
 }

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

@@ -27,7 +27,7 @@ export class LeftPanelControls extends PluginUIComponent<{}, { tab: LeftPanelTab
             if (this.state.tab !== tab) this.setState({ tab });
         });
 
-        this.subscribe(this.plugin.state.dataState.events.changed, state => {
+        this.subscribe(this.plugin.state.dataState.events.changed, ({ state }) => {
             if (this.state.tab !== 'data') return;
             if (state.cells.size === 1) this.set('root');
         });

src/mol-plugin-ui/structure/components.tsx

@@ -0,0 +1,41 @@
+/**
+ * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import * as React from 'react';
+import { CollapsableControls, CollapsableState, PurePluginUIComponent } from '../base';
+import { StructureHierarchyManager } from '../../mol-plugin-state/manager/structure';
+import { StructureComponentRef } from '../../mol-plugin-state/manager/structure/hierarchy';
+
+interface StructureComponentControlState extends CollapsableState {
+    isDisabled: boolean
+}
+
+export class StructureComponentControls extends CollapsableControls<{}, StructureComponentControlState> {
+    protected defaultState(): StructureComponentControlState {
+        return { header: 'Structure Components', isCollapsed: false, isDisabled: false };
+    }
+
+    get currentModels() {
+        return this.plugin.managers.structureHierarchy.behaviors.currentModels;
+    }
+
+    componentDidMount() {
+        this.subscribe(this.currentModels, () => this.forceUpdate());
+    }
+
+    renderControls() {
+        const components = StructureHierarchyManager.getCommonComponentPivots(this.currentModels.value)
+        return <>
+            {components.map((c, i) => <StructureComponentEntry key={i} component={c} />)}
+        </>;
+    }
+}
+
+class StructureComponentEntry extends PurePluginUIComponent<{ component: StructureComponentRef }> {
+    render() {
+        return <></>;
+    }
+}

src/mol-plugin/context.ts

@@ -48,6 +48,7 @@ import { CustomProperty } from '../mol-model-props/common/custom-property';
 import { PluginConfigManager } from './config';
 import { DataBuilder } from '../mol-plugin-state/builder/data';
 import { StructureBuilder } from '../mol-plugin-state/builder/structure';
+import { StructureHierarchyManager } from '../mol-plugin-state/manager/structure';
 
 export class PluginContext {
     private disposed = false;
@@ -131,6 +132,10 @@ export class PluginContext {
         representation: void 0 as any as RepresentationBuilder
     };
 
+    readonly managers = {
+        structureHierarchy: new StructureHierarchyManager(this)
+    };
+
     readonly customModelProperties = new CustomProperty.Registry<Model>();
     readonly customStructureProperties = new CustomProperty.Registry<Structure>();
     readonly customParamEditors = new Map<string, StateTransformParameters.Class>();

src/mol-state/object.ts

@@ -47,6 +47,14 @@ namespace StateObject {
         }
     }
 
+    export function hasTag(o: StateObject, t: string) {
+        if (!o.tags) return false;
+        for (const s of o.tags) {
+            if (s === t) return true;
+        }
+        return false;
+    }
+
     /** A special object indicating a transformer result has no value. */
     export const Null: StateObject<any, any> = {
         id: UUID.create22(),

src/mol-state/state.ts

@@ -46,7 +46,7 @@ class State {
             removed: this.ev<State.ObjectEvent & { obj?: StateObject }>()
         },
         log: this.ev<LogEntry>(),
-        changed: this.ev<State>(),
+        changed: this.ev<{ state: State, inTransaction: boolean }>(),
         isUpdating: this.ev<boolean>()
     };
 
@@ -126,12 +126,15 @@ class State {
         });
     }
 
+    private inTransaction = false;
+
     /** Apply series of updates to the state. If any of them fail, revert to the original state. */
     transaction(edits: () => Promise<void> | void) {
         return Task.create('State Transaction', async ctx => {
             const snapshot = this._tree.asImmutable();
             let restored = false;
             try {
+                this.inTransaction = true;
                 await edits();
 
                 let hasError = false;
@@ -145,6 +148,9 @@ class State {
                     await this.updateTree(snapshot).runInContext(ctx);
                     this.events.log.error(e);
                 }
+            } finally {
+                this.inTransaction = false;
+                this.events.changed.next({ state: this, inTransaction: false });
             }
         });
     }
@@ -201,7 +207,7 @@ class State {
         } finally {
             this.spine.current = undefined;
 
-            if (updated) this.events.changed.next(this);
+            if (updated) this.events.changed.next({ state: this, inTransaction: this.inTransaction });
             this.events.isUpdating.next(false);
         }
     }