wip selection model

src/examples/proteopedia-wrapper/helpers.ts

@@ -25,8 +25,6 @@ export namespace ModelInfo {
             const json = JSON.parse(src);
             const data = json && json[id];
 
-            console.log(data);
-
             const assemblies = data[0] && data[0].assemblies;
             if (!assemblies || !assemblies.length) return void 0;
 

src/mol-model/structure/query/queries/combinators.ts

@@ -47,7 +47,7 @@ export function intersect(queries: ArrayLike<StructureQuery>): StructureQuery {
         }
 
         ctx.throwIfTimedOut();
-        const pivotSet = HashSet<Structure>(s => s.hashCode, Structure.areEqual);
+        const pivotSet = HashSet<Structure>(s => s.hashCode, Structure.areUnitAndIndicesEqual);
         StructureSelection.forEach(selections[pivotIndex], s => pivotSet.add(s));
 
         const ret = StructureSelection.UniqueBuilder(ctx.inputStructure);

src/mol-model/structure/query/selection.ts

@@ -84,7 +84,7 @@ namespace StructureSelection {
     class HashBuilderImpl implements Builder {
         private structures: Structure[] = [];
         private allSingletons = true;
-        private uniqueSets = HashSet(Structure.hashCode, Structure.areEqual);
+        private uniqueSets = HashSet(Structure.hashCode, Structure.areUnitAndIndicesEqual);
 
         add(structure: Structure) {
             const atomCount = structure.elementCount;

src/mol-model/structure/query/utils/builders.ts

@@ -12,7 +12,7 @@ import { StructureSubsetBuilder } from '../../structure/util/subset-builder';
 import { ElementIndex } from '../../model';
 
 export class UniqueStructuresBuilder {
-    private set = HashSet(Structure.hashCode, Structure.areEqual);
+    private set = HashSet(Structure.hashCode, Structure.areUnitAndIndicesEqual);
     private structures: Structure[] = [];
     private allSingletons = true;
 

src/mol-model/structure/structure/element.ts

@@ -117,6 +117,13 @@ namespace StructureElement {
             })));
         }
 
+        export function remap(loci: Loci, structure: Structure): Loci {
+            return Loci(structure, loci.elements.map(e => ({
+                unit: structure.unitMap.get(e.unit.id)!,
+                indices: e.indices
+            })));
+        }
+
         export function union(xs: Loci, ys: Loci): Loci {
             if (xs.structure !== ys.structure) throw new Error(`Can't union Loci of different structures.`);
             if (xs.elements.length > ys.elements.length) return union(ys, xs);

src/mol-model/structure/structure/structure.ts

@@ -424,7 +424,7 @@ namespace Structure {
         return hashString(s.units.map(u => Unit.conformationId(u)).join('|'))
     }
 
-    export function areEqual(a: Structure, b: Structure) {
+    export function areUnitAndIndicesEqual(a: Structure, b: Structure) {
         if (a.elementCount !== b.elementCount) return false;
         const len = a.units.length;
         if (len !== b.units.length) return false;

src/mol-plugin/util/canvas3d-identify.ts

@@ -6,7 +6,7 @@
 
 import { PluginContext } from '../context';
 import { PickingId } from 'mol-geo/geometry/picking';
-import { EmptyLoci, Loci } from 'mol-model/loci';
+import { EmptyLoci } from 'mol-model/loci';
 import { Representation } from 'mol-repr/representation';
 import { ModifiersKeys, ButtonsType } from 'mol-util/input/input-observer';
 
@@ -49,7 +49,7 @@ export class Canvas3dIdentifyHelper {
         }
 
         const loci = this.ctx.canvas3d.getLoci(this.id);
-        if (loci.repr !== this.prevLoci.repr || !Loci.areEqual(loci.loci, this.prevLoci.loci)) {
+        if (!Representation.Loci.areEqual(this.prevLoci, loci)) {
             this.ctx.events.canvas3d.highlight.next({ loci, modifiers: this.modifiers });
             this.prevLoci = loci;
         }

src/mol-plugin/util/selection/structure-element.ts → src/mol-plugin/util/selection/structure-loci.ts

@@ -8,27 +8,25 @@ import { StructureElement, Structure } from 'mol-model/structure';
 import { PluginContext } from '../../context';
 import { Loci, EmptyLoci } from 'mol-model/loci';
 import { PluginStateObject } from 'mol-plugin/state/objects';
-import { State } from 'mol-state';
+import { State, StateSelection, StateObject } from 'mol-state';
+import { mapObjectMap } from 'mol-util/object';
 
-export { StructureElementSelectionManager }
+export { StructureLociManager }
 
-class StructureElementSelectionManager {
-    private entries = new Map<string, Entry>();
+class StructureLociManager {
+    private entries = new Map<string, SelectionEntry>();
 
     // maps structure to a parent StateObjectCell
-    private structures = {
+    private mapping = {
         root: new Map<Structure, string>(),
         tracked: new Map<string, Structure>()
-    }
+    };
 
     private getEntry(s: Structure) {
-        if (!this.structures.root.has(s)) return;
-        const key = this.structures.root.get(s)!;
+        if (!this.mapping.root.has(s)) return;
+        const key = this.mapping.root.get(s)!;
         if (!this.entries.has(key)) {
-            const entry: Entry = {
-                selection: StructureElement.Loci(s, []),
-                highlight: StructureElement.Loci(s, []),
-            };
+            const entry = SelectionEntry(s);
             this.entries.set(key, entry);
             return entry;
         }
@@ -39,22 +37,25 @@ class StructureElementSelectionManager {
     add(loci: StructureElement.Loci | Structure.Loci, type: 'selection' | 'highlight'): Loci {
         const entry = this.getEntry(loci.structure);
         if (!entry) return EmptyLoci;
-        entry[type] = Structure.isLoci(loci) ? StructureElement.Loci.all(loci.structure) : StructureElement.Loci.union(entry[type], loci);
-        return entry[type];
+        const xs = entry.elements;
+        xs[type] = Structure.isLoci(loci) ? StructureElement.Loci.all(loci.structure) : StructureElement.Loci.union(xs[type], loci);
+        return xs[type];
     }
 
     remove(loci: StructureElement.Loci | Structure.Loci, type: 'selection' | 'highlight'): Loci {
         const entry = this.getEntry(loci.structure);
         if (!entry) return EmptyLoci;
-        entry[type] = Structure.isLoci(loci) ? StructureElement.Loci(loci.structure, []) : StructureElement.Loci.subtract(entry[type], loci);
-        return entry[type].elements.length === 0 ? EmptyLoci : entry[type];
+        const xs = entry.elements;
+        xs[type] = Structure.isLoci(loci) ? StructureElement.Loci(loci.structure, []) : StructureElement.Loci.subtract(xs[type], loci);
+        return xs[type].elements.length === 0 ? EmptyLoci : xs[type];
     }
 
     set(loci: StructureElement.Loci | Structure.Loci, type: 'selection' | 'highlight'): Loci {
         const entry = this.getEntry(loci.structure);
         if (!entry) return EmptyLoci;
-        entry[type] = Structure.isLoci(loci) ? StructureElement.Loci.all(loci.structure) : loci;
-        return entry[type].elements.length === 0 ? EmptyLoci : entry[type];
+        const xs = entry.elements;
+        xs[type] = Structure.isLoci(loci) ? StructureElement.Loci.all(loci.structure) : loci;
+        return xs[type].elements.length === 0 ? EmptyLoci : xs[type];
     }
 
     private prevHighlight: StructureElement.Loci | undefined = void 0;
@@ -74,42 +75,73 @@ class StructureElementSelectionManager {
         return ret || EmptyLoci;
     }
 
-    private trackCell(state: State, ref: string) {
-        const cell = state.cells.get(ref);
-        if (!cell || !cell.obj || !PluginStateObject.Molecule.Structure.is(cell.obj)) return;
-        const parent = state.selectQ(q => q.byRef(cell.transform.ref).rootOfType([PluginStateObject.Molecule.Structure]))[0];
-        this.structures.tracked.set(cell.transform.ref, cell.obj.data);
+    private addMapping(state: State, ref: string, obj: StateObject) {
+        if (!PluginStateObject.Molecule.Structure.is(obj)) return;
+        const parent = state.select(StateSelection.Generators.byRef(ref).rootOfType([PluginStateObject.Molecule.Structure]))[0];
+        this.mapping.tracked.set(ref, obj.data);
         if (!parent) {
-            this.structures.root.set(cell.obj.data, cell.transform.ref);
+            this.mapping.root.set(obj.data, ref);
         } else {
-            this.structures.root.set(cell.obj.data, parent.transform.ref);
+            this.mapping.root.set(obj.data, parent.transform.ref);
         }
     }
 
-    private untrackCell(state: State, ref: string) {
-        if (!this.structures.tracked.has(ref)) return;
-        const s = this.structures.tracked.get(ref)!;
-        this.structures.tracked.delete(ref);
-        this.structures.root.delete(s);
+    private removeMapping(ref: string) {
+        if (!this.mapping.tracked.has(ref)) return;
+        const s = this.mapping.tracked.get(ref)!;
+        this.mapping.tracked.delete(ref);
+        const root = this.mapping.root.get(s);
+        this.mapping.root.delete(s);
+        if (root === ref) this.entries.delete(ref);
+    }
+
+    private updateMapping(state: State, ref: string, oldObj: StateObject | undefined, obj: StateObject) {
+        if (!PluginStateObject.Molecule.Structure.is(obj)) return;
+
+        if (this.entries.has(ref)) {
+            if (!PluginStateObject.Molecule.Structure.is(oldObj) || 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));
+                return;
+            }
+
+            // clear the selection
+            this.entries.set(ref, SelectionEntry(obj.data));
+        } else {
+            this.removeMapping(ref);
+            this.addMapping(state, ref, obj);
+        }
     }
 
     constructor(plugin: PluginContext) {
         plugin.state.dataState.events.object.created.subscribe(e => {
-            this.trackCell(e.state, e.ref);
+            this.addMapping(e.state, e.ref, e.obj);
         });
 
         plugin.state.dataState.events.object.removed.subscribe(e => {
-            this.untrackCell(e.state, e.ref);
+            this.removeMapping(e.ref);
         });
 
         plugin.state.dataState.events.object.updated.subscribe(e => {
-            this.untrackCell(e.state, e.ref);
-            this.trackCell(e.state, e.ref);
+            this.updateMapping(e.state, e.ref, e.oldObj, e.obj);
         });
     }
 }
 
-interface Entry {
-    selection: StructureElement.Loci,
-    highlight: StructureElement.Loci
+interface SelectionEntry {
+    elements: { [category: string]: StructureElement.Loci }
+}
+
+function SelectionEntry(s: Structure): SelectionEntry {
+    return {
+        elements: { }
+    };
+}
+
+function remapSelectionEntry(e: SelectionEntry, s: Structure): SelectionEntry {
+    return {
+        elements: mapObjectMap(e.elements, (l: StructureElement.Loci) => StructureElement.Loci.remap(l, s))
+    };
 }
\ No newline at end of file

src/mol-util/object.ts

@@ -83,3 +83,11 @@ export function deepClone<T>(source: T): T {
 
     throw new Error(`Can't clone, type "${typeof source}" unsupported`);
 }
+
+export function mapObjectMap<O extends { [k: string]: T }, T, S>(o: O, f: (v: T) => S): { [k: string]: S } {
+    const ret: any = { };
+    for (const k of Object.keys(o)) {
+        ret[k] = f((o as any)[k]);
+    }
+    return ret;
+}
\ No newline at end of file