component ui related tweaks and fixes

- theme reset should only act on selection
- on-screen ui act similar to controls-ui
- stick to component as name consistently

src/mol-plugin-state/helpers/structure-clipping.ts

@@ -14,17 +14,17 @@ import { StructureComponentRef } from '../manager/structure/hierarchy-state';
 import { EmptyLoci, Loci } from '../../mol-model/loci';
 import { Clipping } from '../../mol-theme/clipping';
 
-type ClippingEachReprCallback = (update: StateBuilder.Root, repr: StateObjectCell<PluginStateObject.Molecule.Structure.Representation3D, StateTransform<typeof StateTransforms.Representation.StructureRepresentation3D>>, clipping?: StateObjectCell<any, StateTransform<typeof StateTransforms.Representation.ClippingStructureRepresentation3DFromBundle>>) => void
+type ClippingEachReprCallback = (update: StateBuilder.Root, repr: StateObjectCell<PluginStateObject.Molecule.Structure.Representation3D, StateTransform<typeof StateTransforms.Representation.StructureRepresentation3D>>, clipping?: StateObjectCell<any, StateTransform<typeof StateTransforms.Representation.ClippingStructureRepresentation3DFromBundle>>) => Promise<void>
 const ClippingManagerTag = 'clipping-controls';
 
-export async function setStructureClipping(plugin: PluginContext, components: StructureComponentRef[], groups: Clipping.Groups, lociGetter: (structure: Structure) => StructureElement.Loci | EmptyLoci, types?: string[]) {
-    await eachRepr(plugin, components, (update, repr, clippingCell) => {
+export async function setStructureClipping(plugin: PluginContext, components: StructureComponentRef[], groups: Clipping.Groups, lociGetter: (structure: Structure) => Promise<StructureElement.Loci | EmptyLoci>, types?: string[]) {
+    await eachRepr(plugin, components, async (update, repr, clippingCell) => {
         if (types && types.length > 0 && !types.includes(repr.params!.values.type.name)) return;
 
         const structure = repr.obj!.data.source.data;
         // always use the root structure to get the loci so the clipping
         // stays applicable as long as the root structure does not change
-        const loci = lociGetter(structure.root);
+        const loci = await lociGetter(structure.root);
         if (Loci.isEmpty(loci)) return;
 
         const layer = {
@@ -44,13 +44,13 @@ export async function setStructureClipping(plugin: PluginContext, components: St
     });
 }
 
-function eachRepr(plugin: PluginContext, components: StructureComponentRef[], callback: ClippingEachReprCallback) {
+async function eachRepr(plugin: PluginContext, components: StructureComponentRef[], callback: ClippingEachReprCallback) {
     const state = plugin.state.data;
     const update = state.build();
     for (const c of components) {
         for (const r of c.representations) {
             const clipping = state.select(StateSelection.Generators.ofTransformer(StateTransforms.Representation.ClippingStructureRepresentation3DFromBundle, r.cell.transform.ref).withTag(ClippingManagerTag));
-            callback(update, r.cell, clipping[0]);
+            await callback(update, r.cell, clipping[0]);
         }
     }
 

src/mol-plugin-state/helpers/structure-overpaint.ts

@@ -15,17 +15,17 @@ import { Color } from '../../mol-util/color';
 import { StructureComponentRef } from '../manager/structure/hierarchy-state';
 import { EmptyLoci, Loci } from '../../mol-model/loci';
 
-type OverpaintEachReprCallback = (update: StateBuilder.Root, repr: StateObjectCell<PluginStateObject.Molecule.Structure.Representation3D, StateTransform<typeof StateTransforms.Representation.StructureRepresentation3D>>, overpaint?: StateObjectCell<any, StateTransform<typeof StateTransforms.Representation.OverpaintStructureRepresentation3DFromBundle>>) => void
+type OverpaintEachReprCallback = (update: StateBuilder.Root, repr: StateObjectCell<PluginStateObject.Molecule.Structure.Representation3D, StateTransform<typeof StateTransforms.Representation.StructureRepresentation3D>>, overpaint?: StateObjectCell<any, StateTransform<typeof StateTransforms.Representation.OverpaintStructureRepresentation3DFromBundle>>) => Promise<void>
 const OverpaintManagerTag = 'overpaint-controls';
 
-export async function setStructureOverpaint(plugin: PluginContext, components: StructureComponentRef[], color: Color | -1, lociGetter: (structure: Structure) => StructureElement.Loci | EmptyLoci, types?: string[], alpha = 1) {
-    await eachRepr(plugin, components, (update, repr, overpaintCell) => {
+export async function setStructureOverpaint(plugin: PluginContext, components: StructureComponentRef[], color: Color | -1, lociGetter: (structure: Structure) => Promise<StructureElement.Loci | EmptyLoci>, types?: string[], alpha = 1) {
+    await eachRepr(plugin, components, async (update, repr, overpaintCell) => {
         if (types && types.length > 0 && !types.includes(repr.params!.values.type.name)) return;
 
         const structure = repr.obj!.data.source.data;
         // always use the root structure to get the loci so the overpaint
         // stays applicable as long as the root structure does not change
-        const loci = lociGetter(structure.root);
+        const loci = await lociGetter(structure.root);
         if (Loci.isEmpty(loci)) return;
 
         const layer = {
@@ -47,7 +47,7 @@ export async function setStructureOverpaint(plugin: PluginContext, components: S
 }
 
 export async function clearStructureOverpaint(plugin: PluginContext, components: StructureComponentRef[], types?: string[]) {
-    await eachRepr(plugin, components, (update, repr, overpaintCell) => {
+    await eachRepr(plugin, components, async (update, repr, overpaintCell) => {
         if (types && types.length > 0 && !types.includes(repr.params!.values.type.name)) return;
         if (overpaintCell) {
             update.delete(overpaintCell.transform.ref);
@@ -55,13 +55,13 @@ export async function clearStructureOverpaint(plugin: PluginContext, components:
     });
 }
 
-function eachRepr(plugin: PluginContext, components: StructureComponentRef[], callback: OverpaintEachReprCallback) {
+async function eachRepr(plugin: PluginContext, components: StructureComponentRef[], callback: OverpaintEachReprCallback) {
     const state = plugin.state.data;
     const update = state.build();
     for (const c of components) {
         for (const r of c.representations) {
             const overpaint = state.select(StateSelection.Generators.ofTransformer(StateTransforms.Representation.OverpaintStructureRepresentation3DFromBundle, r.cell.transform.ref).withTag(OverpaintManagerTag));
-            callback(update, r.cell, overpaint[0]);
+            await callback(update, r.cell, overpaint[0]);
         }
     }
 

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

@@ -21,7 +21,7 @@ import { ParamDefinition as PD } from '../../../mol-util/param-definition';
 import { StructureRepresentationPresetProvider } from '../../builder/structure/representation-preset';
 import { StatefulPluginComponent } from '../../component';
 import { StructureComponentParams } from '../../helpers/structure-component';
-import { clearStructureOverpaint, setStructureOverpaint } from '../../helpers/structure-overpaint';
+import { setStructureOverpaint } from '../../helpers/structure-overpaint';
 import { createStructureColorThemeParams, createStructureSizeThemeParams } from '../../helpers/structure-representation-params';
 import { StructureSelectionQueries, StructureSelectionQuery } from '../../helpers/structure-selection-query';
 import { StructureRepresentation3D } from '../../transforms/representation';
@@ -361,14 +361,14 @@ class StructureComponentManager extends StatefulPluginComponent<StructureCompone
     }
 
     async applyTheme(params: StructureComponentManager.ThemeParams, structures?: ReadonlyArray<StructureRef>) {
-        return this.plugin.dataTransaction(async () => {
+        return this.plugin.dataTransaction(async ctx => {
             const xs = structures || this.currentStructures;
             if (xs.length === 0) return;
-            const getLoci = (s: Structure) => this.plugin.managers.structure.selection.getLoci(s);
 
+            const getLoci = async (s: Structure) => StructureSelection.toLociWithSourceUnits(await params.selection.getSelection(this.plugin, ctx, s));
             for (const s of xs) {
                 if (params.action.name === 'reset') {
-                    await clearStructureOverpaint(this.plugin, s.components, params.representations);
+                    await setStructureOverpaint(this.plugin, s.components, -1, getLoci, params.representations, 1);
                 } else if (params.action.name === 'color') {
                     const p = params.action.params;
                     await setStructureOverpaint(this.plugin, s.components, p.color, getLoci, params.representations, p.opacity);
@@ -441,14 +441,13 @@ namespace StructureComponentManager {
     };
     export type Options = PD.Values<typeof OptionsParams>
 
-    export function getAddParams(plugin: PluginContext, params?: { pivot?: StructureRef, allowNone: boolean, hideSelection?: boolean, checkExisting?: boolean, defaultSelection?: StructureSelectionQuery }) {
+    export function getAddParams(plugin: PluginContext, params?: { pivot?: StructureRef, allowNone: boolean, hideSelection?: boolean, checkExisting?: boolean }) {
         const { options } = plugin.query.structure.registry;
         params = {
             pivot: plugin.managers.structure.component.pivotStructure,
             allowNone: true,
             hideSelection: false,
             checkExisting: false,
-            defaultSelection: StructureSelectionQueries.current,
             ...params
         };
         return {
@@ -463,7 +462,9 @@ namespace StructureComponentManager {
     export type AddParams = { selection: StructureSelectionQuery, options: { checkExisting: boolean, label: string }, representation: string }
 
     export function getThemeParams(plugin: PluginContext, pivot: StructureRef | StructureComponentRef | undefined) {
+        const { options } = plugin.query.structure.registry;
         return {
+            selection: PD.Select(options[1][0], options, { isHidden: false }),
             action: PD.MappedStatic('color', {
                 color: PD.Group({
                     color: PD.Color(ColorNames.blue, { isExpanded: true }),

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

@@ -27,7 +27,7 @@ interface StructureComponentControlState extends CollapsableState {
 export class StructureComponentControls extends CollapsableControls<{}, StructureComponentControlState> {
     protected defaultState(): StructureComponentControlState {
         return {
-            header: 'Representation',
+            header: 'Components',
             isCollapsed: false,
             isDisabled: false,
             brand: { accent: 'blue', svg: CubeOutlineSvg }
@@ -147,9 +147,7 @@ interface AddComponentControlsProps {
 
 export class AddComponentControls extends PurePluginUIComponent<AddComponentControlsProps, AddComponentControlsState> {
     createState(): AddComponentControlsState {
-        const params = StructureComponentManager.getAddParams(this.plugin, this.props.forSelection
-            ? { allowNone: false, hideSelection: true, checkExisting: this.props.forSelection }
-            : void 0);
+        const params = StructureComponentManager.getAddParams(this.plugin);
         return { params, values: ParamDefinition.getDefaultValues(params) };
     }
 
@@ -159,8 +157,12 @@ export class AddComponentControls extends PurePluginUIComponent<AddComponentCont
         return this.plugin.managers.structure.component.currentStructures;
     }
 
+    get currentStructures() {
+        return this.plugin.managers.structure.hierarchy.current.structures;
+    }
+
     apply = () => {
-        const structures = this.props.forSelection ? this.plugin.managers.structure.hierarchy.getStructuresWithSelection() : this.selectedStructures;
+        const structures = this.props.forSelection ? this.currentStructures : this.selectedStructures;
         this.props.onApply();
         this.plugin.managers.structure.component.add(this.state.values, structures);
     }

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

@@ -52,7 +52,7 @@ interface StructureSelectionActionsControlsState {
     isBusy: boolean,
     canUndo: boolean,
 
-    action?: StructureSelectionModifier | 'theme' | 'add-repr' | 'help'
+    action?: StructureSelectionModifier | 'theme' | 'add-component' | 'help'
 }
 
 const ActionHeader = new Map<StructureSelectionModifier, string>([
@@ -159,7 +159,7 @@ export class StructureSelectionActionsControls extends PluginUIComponent<{}, Str
     toggleIntersect = this.showAction('intersect')
     toggleSet = this.showAction('set')
     toggleTheme = this.showAction('theme')
-    toggleAddRepr = this.showAction('add-repr')
+    toggleAddComponent = this.showAction('add-component')
     toggleHelp = this.showAction('help')
 
     setGranuality: ParamOnChange = ({ value }) => {
@@ -196,14 +196,14 @@ export class StructureSelectionActionsControls extends PluginUIComponent<{}, Str
                 <ToggleButton icon={SetSvg} title={`${ActionHeader.get('set')}. Hold shift key to keep menu open.`} toggle={this.toggleSet} isSelected={this.state.action === 'set'} disabled={this.isDisabled} />
 
                 <ToggleButton icon={BrushSvg} title='Apply Theme to Selection' toggle={this.toggleTheme} isSelected={this.state.action === 'theme'} disabled={this.isDisabled} style={{ marginLeft: '10px' }}  />
-                <ToggleButton icon={CubeOutlineSvg} title='Create Representation of Selection' toggle={this.toggleAddRepr} isSelected={this.state.action === 'add-repr'} disabled={this.isDisabled} />
-                <IconButton svg={RemoveSvg} title='Subtract Selection from Representations' onClick={this.subtract} disabled={this.isDisabled} />
+                <ToggleButton icon={CubeOutlineSvg} title='Create Component of Selection with Representation' toggle={this.toggleAddComponent} isSelected={this.state.action === 'add-component'} disabled={this.isDisabled} />
+                <IconButton svg={RemoveSvg} title='Remove/subtract Selection from all Components' onClick={this.subtract} disabled={this.isDisabled} />
                 <IconButton svg={RestoreSvg} onClick={this.undo} disabled={!this.state.canUndo || this.isDisabled} title={undoTitle} />
 
                 <ToggleButton icon={HelpOutlineSvg} title='Show/hide help' toggle={this.toggleHelp} style={{ marginLeft: '10px' }} isSelected={this.state.action === 'help'} />
                 <IconButton svg={CancelOutlinedSvg} title='Turn selection mode off' onClick={this.turnOff} />
             </div>
-            {(this.state.action && this.state.action !== 'theme' && this.state.action !== 'add-repr' && this.state.action !== 'help') && <div className='msp-selection-viewport-controls-actions'>
+            {(this.state.action && this.state.action !== 'theme' && this.state.action !== 'add-component' && this.state.action !== 'help') && <div className='msp-selection-viewport-controls-actions'>
                 <ActionMenu header={ActionHeader.get(this.state.action as StructureSelectionModifier)} title='Click to close.' items={this.queries} onSelect={this.selectQuery} noOffset />
             </div>}
             {this.state.action === 'theme' && <div className='msp-selection-viewport-controls-actions'>
@@ -211,9 +211,9 @@ export class StructureSelectionActionsControls extends PluginUIComponent<{}, Str
                     <ApplyThemeControls onApply={this.toggleTheme} />
                 </ControlGroup>
             </div>}
-            {this.state.action === 'add-repr' && <div className='msp-selection-viewport-controls-actions'>
-                <ControlGroup header='Add Representation' title='Click to close.' initialExpanded={true} hideExpander={true} hideOffset={true} onHeaderClick={this.toggleAddRepr} topRightIcon={CloseSvg}>
-                    <AddComponentControls onApply={this.toggleAddRepr} forSelection />
+            {this.state.action === 'add-component' && <div className='msp-selection-viewport-controls-actions'>
+                <ControlGroup header='Add Component' title='Click to close.' initialExpanded={true} hideExpander={true} hideOffset={true} onHeaderClick={this.toggleAddComponent} topRightIcon={CloseSvg}>
+                    <AddComponentControls onApply={this.toggleAddComponent} forSelection />
                 </ControlGroup>
             </div>}
             {this.state.action === 'help' && <div className='msp-selection-viewport-controls-actions'>
@@ -222,7 +222,7 @@ export class StructureSelectionActionsControls extends PluginUIComponent<{}, Str
                         <HelpText>Use <Icon svg={UnionSvg} inline /> <Icon svg={SubtractSvg} inline /> <Icon svg={IntersectSvg} inline /> <Icon svg={SetSvg} inline /> to modify the selection.</HelpText>
                     </HelpGroup>
                     <HelpGroup header='Representation Operations'>
-                        <HelpText>Use <Icon svg={BrushSvg} inline /> <Icon svg={CubeOutlineSvg} inline /> <Icon svg={RemoveSvg} inline /> <Icon svg={RestoreSvg} inline /> to color, create selection/representation, subtract it, or undo actions.</HelpText>
+                        <HelpText>Use <Icon svg={BrushSvg} inline /> <Icon svg={CubeOutlineSvg} inline /> <Icon svg={RemoveSvg} inline /> <Icon svg={RestoreSvg} inline /> to color, create components, remove from components, or undo actions.</HelpText>
                     </HelpGroup>
                     <ViewportHelpContent selectOnly={true} />
                 </ControlGroup>
@@ -304,22 +304,22 @@ export class StructureSelectionStatsControls extends PluginUIComponent<{ hideOnE
     }
 }
 
-interface ApplyColorControlsState {
+interface ApplyThemeControlsState {
     values: StructureComponentManager.ThemeParams
 }
 
-interface ApplyColorControlsProps {
+interface ApplyThemeControlsProps {
     onApply?: () => void
 }
 
-class ApplyThemeControls extends PurePluginUIComponent<ApplyColorControlsProps, ApplyColorControlsState> {
+class ApplyThemeControls extends PurePluginUIComponent<ApplyThemeControlsProps, ApplyThemeControlsState> {
     _params = memoizeLatest((pivot: StructureRef | undefined) => StructureComponentManager.getThemeParams(this.plugin, pivot));
     get params() { return this._params(this.plugin.managers.structure.component.pivotStructure); }
 
     state = { values: ParamDefinition.getDefaultValues(this.params) };
 
     apply = () => {
-        this.plugin.managers.structure.component.applyTheme(this.state.values);
+        this.plugin.managers.structure.component.applyTheme(this.state.values, this.plugin.managers.structure.hierarchy.current.structures);
         this.props.onApply?.();
     }
 

src/mol-plugin/context.ts

@@ -30,7 +30,7 @@ import { Representation } from '../mol-repr/representation';
 import { StructureRepresentationRegistry } from '../mol-repr/structure/registry';
 import { VolumeRepresentationRegistry } from '../mol-repr/volume/registry';
 import { StateTransform } from '../mol-state';
-import { Progress, Task } from '../mol-task';
+import { Progress, Task, RuntimeContext } from '../mol-task';
 import { ColorTheme } from '../mol-theme/color';
 import { SizeTheme } from '../mol-theme/size';
 import { ThemeRegistryContext } from '../mol-theme/theme';
@@ -222,7 +222,7 @@ export class PluginContext {
         this.behaviors.interaction.selectionMode.next(mode);
     }
 
-    dataTransaction(f: () => Promise<void> | void, options?: { canUndo?: string | boolean }) {
+    dataTransaction(f: (ctx: RuntimeContext) => Promise<void> | void, options?: { canUndo?: string | boolean }) {
         return this.runTask(this.state.data.transaction(f, options));
     }
 

src/mol-state/state.ts

@@ -178,7 +178,7 @@ class State {
     private inTransactionError = false;
 
     /** Apply series of updates to the state. If any of them fail, revert to the original state. */
-    transaction(edits: () => Promise<void> | void, options?: { canUndo?: string | boolean }) {
+    transaction(edits: (ctx: RuntimeContext) => Promise<void> | void, options?: { canUndo?: string | boolean }) {
         return Task.create('State Transaction', async ctx => {
             const isNested = this.inTransaction;
 
@@ -191,7 +191,7 @@ class State {
 
                 this.inTransaction = true;
                 this.inTransactionError = false;
-                await edits();
+                await edits(ctx);
 
                 if (this.inTransactionError) {
                     restored = true;