Unify PluginCommands.Highlight and HighlighAll

Select newly created models and structure for create all models

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

@@ -8,6 +8,7 @@ import { PluginContext } from '../../../mol-plugin/context';
 import { StructureHierarchy, buildStructureHierarchy, ModelRef, StructureComponentRef, StructureRef, HierarchyRef, TrajectoryRef } from './hierarchy-state';
 import { PluginComponent } from '../../component';
 import { SetUtils } from '../../../mol-util/set';
+import { StateTransform } from '../../../mol-state';
 
 interface StructureHierarchyManagerState {
     hierarchy: StructureHierarchy,
@@ -54,7 +55,19 @@ export class StructureHierarchyManager extends PluginComponent<StructureHierarch
         return this.state.current;
     }
 
+    private nextSelection: Set<StateTransform.Ref> = new Set();
     private syncCurrent<T extends HierarchyRef>(hierarchy: StructureHierarchy, current: ReadonlyArray<T>, all: ReadonlyArray<T>): T[] {
+        if (this.nextSelection.size > 0) {
+            const newCurrent: T[] = [];
+            for (const r of all) {
+                if (this.nextSelection.has(r.cell.transform.ref)) {
+                    newCurrent.push(r);
+                }
+            }
+            if (newCurrent.length === 0) return all.length > 0 ? [all[0]] : [];
+            return newCurrent;
+        }
+
         if (current.length === 0) return all.length > 0 ? [all[0]] : [];
 
         const newCurrent: T[] = [];
@@ -81,6 +94,8 @@ export class StructureHierarchyManager extends PluginComponent<StructureHierarch
         const models = this.syncCurrent(hierarchy, this.state.current.models, hierarchy.models);
         const structures = this.syncCurrent(hierarchy, this.state.current.structures, hierarchy.structures);
 
+        this.nextSelection.clear();
+
         this.updateState({ hierarchy, current: { trajectories, models, structures }});
         this.behaviors.current.next({ hierarchy, trajectories, models, structures });
     }
@@ -121,6 +136,9 @@ export class StructureHierarchyManager extends PluginComponent<StructureHierarch
 
     createAllModels(trajectory: TrajectoryRef) {
         return this.plugin.dataTransaction(async () => {
+            this.nextSelection.clear();
+
+            this.nextSelection.add(trajectory.cell.transform.ref);
             if (trajectory.models.length > 0) {
                 await this.clearTrajectory(trajectory);
             }
@@ -129,9 +147,11 @@ export class StructureHierarchyManager extends PluginComponent<StructureHierarch
             for (let i = 0; i < tr.length; i++) {
                 const model = await this.plugin.builders.structure.createModel(trajectory.cell, { modelIndex: i }, { isCollapsed: true });
                 const structure = await this.plugin.builders.structure.createStructure(model, { name: 'deposited', params: { } });
+                this.nextSelection.add(model.ref);
+                this.nextSelection.add(structure.ref);
                 await this.plugin.builders.structure.representation.applyPreset(structure, 'auto', { globalThemeName: 'model-index' });
             }
-        })
+        });
     }
 
     private clearTrajectory(trajectory: TrajectoryRef) {

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

@@ -299,7 +299,7 @@ class StructureComponentGroup extends PurePluginUIComponent<{ group: StructureCo
 
     highlight = (e: React.MouseEvent<HTMLElement>) => {
         e.preventDefault();
-        PluginCommands.State.HighlightMany(this.plugin, this.props.group.map(c => ({ state: c.cell.parent, ref: c.cell.transform.ref })))
+        PluginCommands.State.Highlight(this.plugin, { state: this.props.group[0].cell.parent, ref: this.props.group.map(c => c.cell.transform.ref) });
     }
 
     clearHighlight = (e: React.MouseEvent<HTMLElement>) => {

src/mol-plugin/behavior/static/state.ts

@@ -26,7 +26,6 @@ export function registerDefault(ctx: PluginContext) {
     ToggleExpanded(ctx);
     ToggleVisibility(ctx);
     Highlight(ctx);
-    HighlightMany(ctx);
     ClearHighlights(ctx);
     Snapshots(ctx);
 }
@@ -105,32 +104,28 @@ function setVisibilityVisitor(t: StateTransform, tree: StateTree, ctx: { state:
 }
 
 export function Highlight(ctx: PluginContext) {
-    PluginCommands.State.Highlight.subscribe(ctx, stateRef => highlight(ctx, [stateRef]));
-}
-
-export function HighlightMany(ctx: PluginContext) {
-    PluginCommands.State.HighlightMany.subscribe(ctx, stateRefPairs => highlight(ctx, stateRefPairs));
-}
+    PluginCommands.State.Highlight.subscribe(ctx, ({ state, ref }) => {
+        ctx.managers.interactivity.lociHighlights.clearHighlights();
 
-function highlight(ctx: PluginContext, stateRefPairs: { state: State, ref: StateTransform.Ref }[]) {
-    ctx.managers.interactivity.lociHighlights.clearHighlights();
-    for (const { state, ref } of stateRefPairs) {
-        const cell = state.select(ref)[0];
-        if (!cell) continue;
-        if (SO.Molecule.Structure.is(cell.obj)) {
-            ctx.managers.interactivity.lociHighlights.highlight({ loci: Structure.Loci(cell.obj.data) }, false);
-        } else if (cell && SO.isRepresentation3D(cell.obj)) {
-            const { repr } = cell.obj.data
-            ctx.managers.interactivity.lociHighlights.highlight({ loci: repr.getLoci(), repr }, false);
-        } else if (SO.Molecule.Structure.Selections.is(cell.obj)) {
-            for (const entry of cell.obj.data) {
-                ctx.managers.interactivity.lociHighlights.highlight({ loci: entry.loci }, false);
+        const refs = typeof ref === 'string' ? [ref] : ref;
+        for (const r of refs) {
+            const cell = state.cells.get(r);
+            if (!cell) continue;
+            if (SO.Molecule.Structure.is(cell.obj)) {
+                ctx.managers.interactivity.lociHighlights.highlight({ loci: Structure.Loci(cell.obj.data) }, false);
+            } else if (cell && SO.isRepresentation3D(cell.obj)) {
+                const { repr } = cell.obj.data
+                ctx.managers.interactivity.lociHighlights.highlight({ loci: repr.getLoci(), repr }, false);
+            } else if (SO.Molecule.Structure.Selections.is(cell.obj)) {
+                for (const entry of cell.obj.data) {
+                    ctx.managers.interactivity.lociHighlights.highlight({ loci: entry.loci }, false);
+                }
             }
         }
-    }
 
-    // TODO: highlight volumes?
-    // TODO: select structures of subtree?
+        // TODO: highlight volumes?
+        // TODO: select structures of subtree?
+    });
 }
 
 export function ClearHighlights(ctx: PluginContext) {

src/mol-plugin/commands.ts

@@ -25,8 +25,7 @@ export const PluginCommands = {
 
         ToggleExpanded: PluginCommand<{ state: State, ref: StateTransform.Ref }>(),
         ToggleVisibility: PluginCommand<{ state: State, ref: StateTransform.Ref }>(),
-        Highlight: PluginCommand<{ state: State, ref: StateTransform.Ref }>(),
-        HighlightMany: PluginCommand<{ state: State, ref: StateTransform.Ref }[]>(),
+        Highlight: PluginCommand<{ state: State, ref: StateTransform.Ref | StateTransform.Ref[] }>(),
         ClearHighlights: PluginCommand<{}>(),
 
         Snapshots: {