selection manager fixes

- add removed/updated events to substructure-parent-helper
- remap selections

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

@@ -13,13 +13,13 @@ import { EmptyLoci, Loci } from '../../../mol-model/loci';
 import { Structure, StructureElement, StructureSelection } from '../../../mol-model/structure';
 import { Boundary } from '../../../mol-model/structure/structure/util/boundary';
 import { PluginContext } from '../../../mol-plugin/context';
-import { StateObject, StateObjectRef } from '../../../mol-state';
+import { StateObjectRef } from '../../../mol-state';
 import { Task } from '../../../mol-task';
 import { structureElementStatsLabel } from '../../../mol-theme/label';
 import { arrayRemoveAtInPlace } from '../../../mol-util/array';
 import { StatefulPluginComponent } from '../../component';
 import { StructureSelectionQuery } from '../../helpers/structure-selection-query';
-import { PluginStateObject } from '../../objects';
+import { PluginStateObject as PSO } from '../../objects';
 import { UUID } from '../../../mol-util';
 import { StructureRef } from './hierarchy-state';
 
@@ -57,7 +57,8 @@ export class StructureSelectionManager extends StatefulPluginComponent<Structure
     }
 
     private getEntry(s: Structure) {
-        const cell = this.plugin.helpers.substructureParent.get(s);
+        // ignore decorators to get stable ref
+        const cell = this.plugin.helpers.substructureParent.get(s, true);
         if (!cell) return;
         const ref = cell.transform.ref;
         if (!this.entries.has(ref)) {
@@ -216,7 +217,7 @@ export class StructureSelectionManager extends StatefulPluginComponent<Structure
     private clearHistoryForStructure(structure: Structure) {
         const historyEntryToRemove: StructureSelectionHistoryEntry[] = [];
         for (const e of this.state.additionsHistory) {
-            if (e.loci.structure === structure) {
+            if (e.loci.structure.root === structure.root) {
                 historyEntryToRemove.push(e);
             }
         }
@@ -228,7 +229,7 @@ export class StructureSelectionManager extends StatefulPluginComponent<Structure
         }
     }
 
-    private onRemove(ref: string, obj: StateObject | undefined) {
+    private onRemove(ref: string, obj: PSO.Molecule.Structure | undefined) {
         if (this.entries.has(ref)) {
             this.entries.delete(ref);
             if (obj?.data) {
@@ -242,43 +243,49 @@ export class StructureSelectionManager extends StatefulPluginComponent<Structure
         }
     }
 
-    private onUpdate(ref: string, oldObj: StateObject | undefined, obj: StateObject) {
-        if (!PluginStateObject.Molecule.Structure.is(obj)) return;
+    private onUpdate(ref: string, oldObj: PSO.Molecule.Structure | undefined, obj: PSO.Molecule.Structure) {
+        if (!oldObj?.data) return;
 
-        if (this.entries.has(ref)) {
-            if (!PluginStateObject.Molecule.Structure.is(oldObj) || oldObj === obj || oldObj.data === obj.data) return;
+        // no change to structure
+        if (oldObj === obj || oldObj.data === obj.data) return;
 
-            // remap the old selection to be related to the new object if possible.
-            if (Structure.areUnitAndIndicesEqual(oldObj.data, obj.data)) {
-                this.entries.set(ref, remapSelectionEntry(this.entries.get(ref)!, obj.data));
+        // ignore decorators to get stable ref
+        const cell = this.plugin.helpers.substructureParent.get(obj.data, true);
+        if (!cell) return;
 
-                // remap referenceLoci & prevHighlight if needed and possible
-                if (this.referenceLoci?.structure === oldObj.data) {
-                    this.referenceLoci = StructureElement.Loci.remap(this.referenceLoci, obj.data);
-                }
+        ref = cell.transform.ref;
+        if (!this.entries.has(ref)) return;
 
-                // remap history locis if needed and possible
-                let changedHistory = false;
-                for (const e of this.state.additionsHistory) {
-                    if (e.loci.structure === oldObj.data) {
-                        e.loci = StructureElement.Loci.remap(e.loci, obj.data);
-                        changedHistory = true;
-                    }
-                }
-                if (changedHistory) this.events.additionsHistoryUpdated.next();
-            } else {
-                // clear the selection for ref
-                this.entries.set(ref, new SelectionEntry(StructureElement.Loci(obj.data, [])));
+        // remap the old selection to be related to the new object if possible.
+        if (Structure.areUnitAndIndicesEqual(oldObj.data, obj.data)) {
+            this.entries.set(ref, remapSelectionEntry(this.entries.get(ref)!, obj.data));
 
-                if (this.referenceLoci?.structure === oldObj.data) {
-                    this.referenceLoci = undefined;
-                }
+            // remap referenceLoci & prevHighlight if needed and possible
+            if (this.referenceLoci?.structure.root === oldObj.data.root) {
+                this.referenceLoci = StructureElement.Loci.remap(this.referenceLoci, obj.data);
+            }
 
-                this.clearHistoryForStructure(oldObj.data);
+            // remap history locis if needed and possible
+            let changedHistory = false;
+            for (const e of this.state.additionsHistory) {
+                if (e.loci.structure === oldObj.data) {
+                    e.loci = StructureElement.Loci.remap(e.loci, obj.data);
+                    changedHistory = true;
+                }
+            }
+            if (changedHistory) this.events.additionsHistoryUpdated.next();
+        } else {
+            // clear the selection for ref
+            this.entries.set(ref, new SelectionEntry(StructureElement.Loci(obj.data, [])));
 
-                this.state.stats = void 0;
-                this.events.changed.next();
+            if (this.referenceLoci?.structure.root === oldObj.data.root) {
+                this.referenceLoci = undefined;
             }
+
+            this.clearHistoryForStructure(oldObj.data);
+
+            this.state.stats = void 0;
+            this.events.changed.next();
         }
     }
 
@@ -456,11 +463,11 @@ export class StructureSelectionManager extends StatefulPluginComponent<Structure
         }));
     }
 
-    fromSelections(ref: StateObjectRef<PluginStateObject.Molecule.Structure.Selections>) {
+    fromSelections(ref: StateObjectRef<PSO.Molecule.Structure.Selections>) {
         const cell = StateObjectRef.resolveAndCheck(this.plugin.state.data, ref);
         if (!cell || !cell.obj) return;
 
-        if (!PluginStateObject.Molecule.Structure.Selections.is(cell.obj)) {
+        if (!PSO.Molecule.Structure.Selections.is(cell.obj)) {
             console.warn('fromSelections applied to wrong object type.', cell.obj);
             return;
         }
@@ -474,8 +481,9 @@ export class StructureSelectionManager extends StatefulPluginComponent<Structure
     constructor(private plugin: PluginContext) {
         super({ entries: new Map(), additionsHistory: [], stats: SelectionStats() });
 
-        plugin.state.data.events.object.removed.subscribe(e => this.onRemove(e.ref, e.obj));
-        plugin.state.data.events.object.updated.subscribe(e => this.onUpdate(e.ref, e.oldObj, e.obj));
+        // listen to events from substructureParent helper to ensure it is updated
+        plugin.helpers.substructureParent.events.removed.subscribe(e => this.onRemove(e.ref, e.obj));
+        plugin.helpers.substructureParent.events.updated.subscribe(e => this.onUpdate(e.ref, e.oldObj, e.obj));
     }
 }
 

src/mol-plugin/context.ts

@@ -137,6 +137,11 @@ export class PluginContext {
         return this.state.data.build();
     }
 
+    readonly helpers = {
+        substructureParent: new SubstructureParentHelper(this),
+        viewportScreenshot: void 0 as ViewportScreenshotHelper | undefined
+    } as const;
+
     readonly managers = {
         structure: {
             hierarchy: new StructureHierarchyManager(this),
@@ -164,11 +169,6 @@ export class PluginContext {
     readonly customStructureControls = new Map<string, { new(): PluginUIComponent<any, any, any> }>();
     readonly genericRepresentationControls = new Map<string, (selection: StructureHierarchyManager['selection']) => [StructureHierarchyRef[], string]>();
 
-    readonly helpers = {
-        substructureParent: new SubstructureParentHelper(this),
-        viewportScreenshot: void 0 as ViewportScreenshotHelper | undefined
-    } as const;
-
     /**
      * Used to store application specific custom state which is then available
      * to State Actions and similar constructs via the PluginContext.
@@ -243,6 +243,7 @@ export class PluginContext {
         this.state.dispose();
         this.tasks.dispose();
         this.layout.dispose();
+        this.helpers.substructureParent.dispose();
 
         objectForEach(this.managers, m => (m as any)?.dispose?.());
         objectForEach(this.managers.structure, m => (m as any)?.dispose?.());

src/mol-plugin/util/substructure-parent-helper.ts

@@ -8,10 +8,18 @@ import { Structure } from '../../mol-model/structure';
 import { PluginStateObject } from '../../mol-plugin-state/objects';
 import { State, StateObject, StateObjectCell, StateSelection } from '../../mol-state';
 import { PluginContext } from '../context';
+import { RxEventHelper } from '../../mol-util/rx-event-helper';
 
 export { SubstructureParentHelper };
 
 class SubstructureParentHelper {
+    private ev = RxEventHelper.create();
+
+    readonly events = {
+        updated: this.ev<{ ref: string, oldObj: PluginStateObject.Molecule.Structure | undefined, obj: PluginStateObject.Molecule.Structure }>(),
+        removed: this.ev<{ ref: string, obj: PluginStateObject.Molecule.Structure | undefined }>(),
+    }
+
     // private decorators = new Map<string, string[]>();
     private root = new Map<Structure, { ref: string, count: number }>();
     private tracked = new Map<string, Structure>();
@@ -36,7 +44,7 @@ class SubstructureParentHelper {
     }
 
     private addMapping(state: State, ref: string, obj: StateObject) {
-        if (!PluginStateObject.Molecule.Structure.is(obj)) return;
+        if (!PluginStateObject.Molecule.Structure.is(obj)) return false;
 
         this.tracked.set(ref, obj.data);
 
@@ -53,10 +61,11 @@ class SubstructureParentHelper {
                 this.root.set(obj.data, { ref: parent.transform.ref, count: 1 });
             }
         }
+        return true;
     }
 
     private removeMapping(ref: string) {
-        if (!this.tracked.has(ref)) return;
+        if (!this.tracked.has(ref)) return false;
 
         const s = this.tracked.get(ref)!;
         this.tracked.delete(ref);
@@ -68,13 +77,19 @@ class SubstructureParentHelper {
         } else {
             this.root.delete(s);
         }
+        return true;
     }
 
     private updateMapping(state: State, ref: string, oldObj: StateObject | undefined, obj: StateObject) {
-        if (!PluginStateObject.Molecule.Structure.is(obj)) return;
+        if (!PluginStateObject.Molecule.Structure.is(obj)) return false;
 
         this.removeMapping(ref);
         this.addMapping(state, ref, obj);
+        return true;
+    }
+
+    dispose() {
+        this.ev.dispose();
     }
 
     constructor(private plugin: PluginContext) {
@@ -83,11 +98,15 @@ class SubstructureParentHelper {
         });
 
         plugin.state.data.events.object.removed.subscribe(e => {
-            this.removeMapping(e.ref);
+            if (this.removeMapping(e.ref)) {
+                this.events.removed.next({ ref: e.ref, obj: e.obj });
+            }
         });
 
         plugin.state.data.events.object.updated.subscribe(e => {
-            this.updateMapping(e.state, e.ref, e.oldObj, e.obj);
+            if (this.updateMapping(e.state, e.ref, e.oldObj, e.obj)) {
+                this.events.updated.next({ ref: e.ref, oldObj: e.oldObj, obj: e.obj });
+            }
         });
     }
 }

src/mol-repr/structure/units-representation.ts

@@ -51,7 +51,7 @@ export function UnitsRepresentation<P extends StructureParams>(label: string, ct
             if (!_structure && !structure) {
                 throw new Error('missing structure');
             } else if (structure && !_structure) {
-                // console.log(label, 'initial structure')
+                // console.log(label, 'initial structure');
                 // First call with a structure, create visuals for each group.
                 _groups = structure.unitSymmetryGroups;
                 for (let i = 0; i < _groups.length; i++) {
@@ -75,7 +75,7 @@ export function UnitsRepresentation<P extends StructureParams>(label: string, ct
                     const group = _groups[i];
                     const visualGroup = oldVisuals.get(group.hashCode);
                     if (visualGroup) {
-                        // console.log(label, 'found visualGroup to reuse')
+                        // console.log(label, 'found visualGroup to reuse');
                         // console.log('old', visualGroup.group)
                         // console.log('new', group)
                         const { visual } = visualGroup;
@@ -91,7 +91,7 @@ export function UnitsRepresentation<P extends StructureParams>(label: string, ct
                             applyMarkerAction(arr, Interval.ofBounds(0, arr.length), MarkerAction.RemoveHighlight);
                         }
                     } else {
-                        // console.log(label, 'did not find visualGroup to reuse, creating new')
+                        // console.log(label, 'did not find visualGroup to reuse, creating new');
                         // newGroups.push(group)
                         const visual = visualCtor(materialId);
                         const promise = visual.createOrUpdate({ webgl: ctx.webgl, runtime }, _theme, _props, { group, structure });
@@ -118,7 +118,7 @@ export function UnitsRepresentation<P extends StructureParams>(label: string, ct
                 // })
                 // unusedVisuals.forEach(visual => visual.destroy())
             } else if (structure && structure !== _structure && Structure.areEquivalent(structure, _structure)) {
-                // console.log(label, 'structures equivalent but not identical')
+                // console.log(label, 'structures equivalent but not identical');
                 // Expects that for structures with the same hashCode,
                 // the unitSymmetryGroups are the same as well.
                 // Re-uses existing visuals for the groups of the new structure.
@@ -132,14 +132,13 @@ export function UnitsRepresentation<P extends StructureParams>(label: string, ct
                         const promise = visualGroup.visual.createOrUpdate({ webgl: ctx.webgl, runtime }, _theme, _props, { group, structure });
                         if (promise) await promise;
                         visualGroup.group = group;
-
                     } else {
                         throw new Error(`expected to find visual for hashCode ${group.hashCode}`);
                     }
                     if (runtime.shouldUpdate) await runtime.update({ message: 'Creating or updating UnitsVisual', current: i, max: _groups.length });
                 }
             } else {
-                // console.log(label, 'no new structure')
+                // console.log(label, 'no new structure');
                 // No new structure given, just update all visuals with new props.
                 const visualsList: [ UnitsVisual<P>, Unit.SymmetryGroup ][] = []; // TODO avoid allocation
                 visuals.forEach(({ visual, group }) => visualsList.push([ visual, group ]));