mol-script: added internal.generator.current symbol

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

@@ -9,6 +9,7 @@ import { now } from '../../../mol-util/now';
 import { ElementIndex } from '../model';
 import { Link } from '../structure/unit/links';
 import { LinkType } from '../model/types';
+import { StructureSelection } from './selection';
 
 export interface QueryContextView {
     readonly element: StructureElement.Location;
@@ -33,6 +34,9 @@ export class QueryContext implements QueryContextView {
     /** Current link between atoms */
     readonly atomicLink = QueryContextLinkInfo.empty<Unit.Atomic>();
 
+    /** Supply this from the outside. Used by the internal.generator.current symbol */
+    currentSelection: StructureSelection | undefined = void 0;
+
     setElement(unit: Unit, e: ElementIndex) {
         this.element.unit = unit;
         this.element.element = e;
@@ -88,12 +92,23 @@ export class QueryContext implements QueryContextView {
         }
     }
 
-    constructor(structure: Structure, timeoutMs = 0) {
+    tryGetCurrentSelection() {
+        if (!this.currentSelection) throw new Error('The current selection is not assigned.');
+        return this.currentSelection;
+    }
+
+    constructor(structure: Structure, options?: QueryContextOptions) {
         this.inputStructure = structure;
-        this.timeoutMs = timeoutMs;
+        this.timeoutMs = (options && options.timeoutMs) || 0;
+        this.currentSelection = options && options.currentSelection;
     }
 }
 
+export interface QueryContextOptions {
+    timeoutMs?: number,
+    currentSelection?: StructureSelection
+}
+
 export interface QueryPredicate { (ctx: QueryContext): boolean }
 export interface QueryFn<T = any> { (ctx: QueryContext): T }
 

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

@@ -6,12 +6,12 @@
 
 import { Structure } from '../structure'
 import { StructureSelection } from './selection'
-import { QueryContext, QueryFn } from './context';
+import { QueryContext, QueryFn, QueryContextOptions } from './context';
 
 interface StructureQuery extends QueryFn<StructureSelection> { }
 namespace StructureQuery {
-    export function run(query: StructureQuery, structure: Structure, timeoutMs = 0) {
-        return query(new QueryContext(structure, timeoutMs));
+    export function run(query: StructureQuery, structure: Structure, options?: QueryContextOptions) {
+        return query(new QueryContext(structure, options));
     }
 }
 

src/mol-script/language/symbol-table/internal.ts

@@ -22,6 +22,10 @@ const generator = {
     bundle: symbol(Arguments.Dictionary({
         elements: Argument(Type.Any) // BundleElement[]
     }), Struct.Types.ElementSelectionQuery, 'A selection with single structure containing represented by the bundle.'),
+
+    // Use with caution as this is not "state saveable"
+    // This query should never be used in any State Transform!
+    current: symbol(Arguments.None, Struct.Types.ElementSelectionQuery, 'Current selection provided by the query context. Avoid using this in State Transforms.')
 }
 
 export default {

src/mol-script/runtime/query/table.ts

@@ -327,6 +327,7 @@ const symbols = [
     // Internal
     D(MolScript.internal.generator.bundleElement, (ctx, xs) => bundleElementImpl(xs.groupedUnits(ctx), xs.ranges(ctx), xs.set(ctx))),
     D(MolScript.internal.generator.bundle, (ctx, xs) => bundleGenerator(xs.elements(ctx))),
+    D(MolScript.internal.generator.current, (ctx, xs) => ctx.tryGetCurrentSelection()),
 ];
 
 function atomProp(p: (e: StructureElement.Location) => any): (ctx: QueryContext, _: any) => any {

src/servers/model/server/query.ts

@@ -57,7 +57,7 @@ export async function resolveJob(job: Job): Promise<CifWriter.Encoder<any>> {
         const queries = structures.map(s => job.queryDefinition.query(job.normalizedParams, s));
         const result: Structure[] = [];
         for (let i = 0; i < structures.length; i++) {
-            const s = await StructureSelection.unionStructure(StructureQuery.run(queries[i], structures[i], Config.maxQueryTimeInMs))
+            const s = await StructureSelection.unionStructure(StructureQuery.run(queries[i], structures[i], { timeoutMs: Config.maxQueryTimeInMs }))
             if (s.elementCount > 0) result.push(s);
         }
         perf.end('query');