wip, refactoring mol-model/query

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

@@ -0,0 +1,41 @@
+/**
+ * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { RuntimeContext } from 'mol-task';
+import { Structure, StructureElement } from '../structure';
+
+export interface QueryContextView {
+    readonly taskCtx: RuntimeContext;
+    readonly element: StructureElement;
+}
+
+export class QueryContext implements QueryContextView {
+    private currentStack: StructureElement[] = [];
+
+    readonly structure: Structure;
+    readonly taskCtx: RuntimeContext;
+
+    /** Current element */
+    readonly element: StructureElement = StructureElement.create();
+
+    pushCurrentElement(): StructureElement {
+        this.currentStack[this.currentStack.length] = this.element;
+        (this.element as StructureElement) = StructureElement.create();
+        return this.element;
+    }
+
+    popCurrentElement() {
+        (this.element as StructureElement) = this.currentStack.pop()!;
+    }
+
+    constructor(structure: Structure, taskCtx: RuntimeContext) {
+        this.structure = structure;
+        this.taskCtx = taskCtx;
+    }
+}
+
+export interface QueryPredicate { (ctx: QueryContextView): boolean }
+export interface QueryFn<T = any> { (ctx: QueryContextView): T }

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

@@ -4,49 +4,50 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import Query from './query'
-import Selection from './selection'
-import { StructureElement, Unit, StructureProperties as P } from '../structure'
+import { StructureQuery } from './query'
+import { StructureSelection } from './selection'
+import { Unit, StructureProperties as P } from '../structure'
 import { Segmentation } from 'mol-data/int'
 import { LinearGroupingBuilder } from './utils/builders';
+import { QueryPredicate, QueryFn, QueryContextView } from './context';
 
-export const all: Query.Provider = async (s, ctx) => Selection.Singletons(s, s);
+export const all: StructureQuery = async (ctx) => StructureSelection.Singletons(ctx.structure, ctx.structure);
 
-export interface AtomQueryParams {
-    entityTest: StructureElement.Predicate,
-    chainTest: StructureElement.Predicate,
-    residueTest: StructureElement.Predicate,
-    atomTest: StructureElement.Predicate,
-    groupBy: StructureElement.Property<any>
+export interface AtomsQueryParams {
+    entityTest: QueryPredicate,
+    chainTest: QueryPredicate,
+    residueTest: QueryPredicate,
+    atomTest: QueryPredicate,
+    groupBy: QueryFn
 }
 
-export interface AtomGroupsQueryParams extends AtomQueryParams {
-    groupBy: StructureElement.Property<any>
-}
+export function residues(params?: Partial<AtomsQueryParams>) { return atoms({ ...params, groupBy: ctx => P.residue.key(ctx.element) }); }
+export function chains(params?: Partial<AtomsQueryParams>) { return atoms({ ...params, groupBy: ctx => P.chain.key(ctx.element) }); }
 
-export function residues(params?: Partial<AtomQueryParams>) { return atoms({ ...params, groupBy: P.residue.key }); }
-export function chains(params?: Partial<AtomQueryParams>) { return atoms({ ...params, groupBy: P.chain.key }); }
+function _true(ctx: QueryContextView) { return true; }
+function _zero(ctx: QueryContextView) { return 0; }
 
-export function atoms(params?: Partial<AtomGroupsQueryParams>): Query.Provider {
+export function atoms(params?: Partial<AtomsQueryParams>): StructureQuery {
     if (!params || (!params.atomTest && !params.residueTest && !params.chainTest && !params.entityTest && !params.groupBy)) return all;
     if (!!params.atomTest && !params.residueTest && !params.chainTest && !params.entityTest && !params.groupBy) return atomGroupsLinear(params.atomTest);
 
-    const normalized: AtomGroupsQueryParams = {
-        entityTest: params.entityTest || P.constant.true,
-        chainTest: params.chainTest || P.constant.true,
-        residueTest: params.residueTest || P.constant.true,
-        atomTest: params.atomTest || P.constant.true,
-        groupBy: params.groupBy || P.constant.zero,
+    const normalized: AtomsQueryParams = {
+        entityTest: params.entityTest || _true,
+        chainTest: params.chainTest || _true,
+        residueTest: params.residueTest || _true,
+        atomTest: params.atomTest || _true,
+        groupBy: params.groupBy || _zero,
     };
 
     if (!params.groupBy) return atomGroupsSegmented(normalized)
     return atomGroupsGrouped(normalized);
 }
 
-function atomGroupsLinear(atomTest: StructureElement.Predicate): Query.Provider {
-    return async (structure, ctx) => {
+function atomGroupsLinear(atomTest: QueryPredicate): StructureQuery {
+    return async (ctx) => {
+        const { structure } = ctx;
         const { units } = structure;
-        const l = StructureElement.create();
+        const l = ctx.pushCurrentElement();
         const builder = structure.subsetBuilder(true);
 
         let progress = 0;
@@ -57,22 +58,23 @@ function atomGroupsLinear(atomTest: StructureElement.Predicate): Query.Provider
             builder.beginUnit(unit.id);
             for (let j = 0, _j = elements.length; j < _j; j++) {
                 l.element = elements[j];
-                if (atomTest(l)) builder.addElement(l.element);
+                if (atomTest(ctx)) builder.addElement(l.element);
             }
             builder.commitUnit();
 
             progress++;
-            if (ctx.shouldUpdate) await ctx.update({ message: 'Atom Groups', current: progress, max: units.length });
+            if (ctx.taskCtx.shouldUpdate) await ctx.taskCtx.update({ message: 'Atom Groups', current: progress, max: units.length });
         }
-
-        return Selection.Singletons(structure, builder.getStructure());
+        ctx.popCurrentElement();
+        return StructureSelection.Singletons(structure, builder.getStructure());
     };
 }
 
-function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: AtomGroupsQueryParams): Query.Provider {
-    return async (structure, ctx) => {
+function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: AtomsQueryParams): StructureQuery {
+    return async (ctx) => {
+        const { structure } = ctx;
         const { units } = structure;
-        const l = StructureElement.create();
+        const l = ctx.pushCurrentElement();
         const builder = structure.subsetBuilder(true);
 
         let progress = 0;
@@ -89,7 +91,7 @@ function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: A
                 const chainSegment = chainsIt.move();
                 l.element = elements[chainSegment.start];
                 // test entity and chain
-                if (!entityTest(l) || !chainTest(l)) continue;
+                if (!entityTest(ctx) || !chainTest(ctx)) continue;
 
                 residuesIt.setSegment(chainSegment);
                 while (residuesIt.hasNext) {
@@ -97,11 +99,11 @@ function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: A
                     l.element = elements[residueSegment.start];
 
                     // test residue
-                    if (!residueTest(l)) continue;
+                    if (!residueTest(ctx)) continue;
 
                     for (let j = residueSegment.start, _j = residueSegment.end; j < _j; j++) {
                         l.element = elements[j];
-                        if (atomTest(l)) {
+                        if (atomTest(ctx)) {
                             builder.addElement(l.element);
                         }
                     }
@@ -110,17 +112,18 @@ function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: A
             builder.commitUnit();
 
             progress++;
-            if (ctx.shouldUpdate) await ctx.update({ message: 'Atom Groups', current: progress, max: units.length });
+            if (ctx.taskCtx.shouldUpdate) await ctx.taskCtx.update({ message: 'Atom Groups', current: progress, max: units.length });
         }
-
-        return Selection.Singletons(structure, builder.getStructure());
+        ctx.popCurrentElement();
+        return StructureSelection.Singletons(structure, builder.getStructure());
     };
 }
 
-function atomGroupsGrouped({ entityTest, chainTest, residueTest, atomTest, groupBy }: AtomGroupsQueryParams): Query.Provider {
-    return async (structure, ctx) => {
+function atomGroupsGrouped({ entityTest, chainTest, residueTest, atomTest, groupBy }: AtomsQueryParams): StructureQuery {
+    return async (ctx) => {
+        const { structure } = ctx;
         const { units } = structure;
-        const l = StructureElement.create();
+        const l = ctx.pushCurrentElement();
         const builder = new LinearGroupingBuilder(structure);
 
         let progress = 0;
@@ -136,7 +139,7 @@ function atomGroupsGrouped({ entityTest, chainTest, residueTest, atomTest, group
                 const chainSegment = chainsIt.move();
                 l.element = elements[chainSegment.start];
                 // test entity and chain
-                if (!entityTest(l) || !chainTest(l)) continue;
+                if (!entityTest(ctx) || !chainTest(ctx)) continue;
 
                 residuesIt.setSegment(chainSegment);
                 while (residuesIt.hasNext) {
@@ -144,19 +147,19 @@ function atomGroupsGrouped({ entityTest, chainTest, residueTest, atomTest, group
                     l.element = elements[residueSegment.start];
 
                     // test residue
-                    if (!residueTest(l)) continue;
+                    if (!residueTest(ctx)) continue;
 
                     for (let j = residueSegment.start, _j = residueSegment.end; j < _j; j++) {
                         l.element = elements[j];
-                        if (atomTest(l)) builder.add(groupBy(l), unit.id, l.element);
+                        if (atomTest(ctx)) builder.add(groupBy(ctx), unit.id, l.element);
                     }
                 }
             }
 
             progress++;
-            if (ctx.shouldUpdate) await ctx.update({ message: 'Atom Groups', current: progress, max: units.length });
+            if (ctx.taskCtx.shouldUpdate) await ctx.taskCtx.update({ message: 'Atom Groups', current: progress, max: units.length });
         }
-
+        ctx.popCurrentElement();
         return builder.getSelection();
     };
 }

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

@@ -7,8 +7,8 @@
 import { Segmentation } from 'mol-data/int';
 import { RuntimeContext } from 'mol-task';
 import { Structure, Unit } from '../structure';
-import Query from './query';
-import Selection from './selection';
+import { StructureQuery } from './query';
+import { StructureSelection } from './selection';
 import { UniqueStructuresBuilder } from './utils/builders';
 import { StructureUniqueSubsetBuilder } from '../structure/util/unique-subset-builder';
 
@@ -37,18 +37,18 @@ function getWholeResidues(ctx: RuntimeContext, source: Structure, structure: Str
     return builder.getStructure();
 }
 
-export function wholeResidues(query: Query.Provider, isFlat: boolean): Query.Provider {
-    return async (structure, ctx) => {
-        const inner = await query(structure, ctx);
-        if (Selection.isSingleton(inner)) {
-            return Selection.Singletons(structure, getWholeResidues(ctx, structure, inner.structure));
+export function wholeResidues(query: StructureQuery, isFlat: boolean): StructureQuery {
+    return async (ctx) => {
+        const inner = await query(ctx);
+        if (StructureSelection.isSingleton(inner)) {
+            return StructureSelection.Singletons(ctx.structure, getWholeResidues(ctx.taskCtx, ctx.structure, inner.structure));
         } else {
-            const builder = new UniqueStructuresBuilder(structure);
+            const builder = new UniqueStructuresBuilder(ctx.structure);
             let progress = 0;
             for (const s of inner.structures) {
-                builder.add(getWholeResidues(ctx, structure, s));
+                builder.add(getWholeResidues(ctx.taskCtx, ctx.structure, s));
                 progress++;
-                if (ctx.shouldUpdate) await ctx.update({ message: 'Whole Residues', current: progress, max: inner.structures.length });
+                if (ctx.taskCtx.shouldUpdate) await ctx.taskCtx.update({ message: 'Whole Residues', current: progress, max: inner.structures.length });
             }
             return builder.getSelection();
         }
@@ -83,17 +83,17 @@ async function getIncludeSurroundings(ctx: RuntimeContext, source: Structure, st
     return !!params.wholeResidues ? getWholeResidues(ctx, source, builder.getStructure()) : builder.getStructure();
 }
 
-export function includeSurroundings(query: Query.Provider, params: IncludeSurroundingsParams): Query.Provider {
-    return async (structure, ctx) => {
-        const inner = await query(structure, ctx);
-        if (Selection.isSingleton(inner)) {
-            const surr = await getIncludeSurroundings(ctx, structure, inner.structure, params);
-            const ret = Selection.Singletons(structure, surr);
+export function includeSurroundings(query: StructureQuery, params: IncludeSurroundingsParams): StructureQuery {
+    return async (ctx) => {
+        const inner = await query(ctx);
+        if (StructureSelection.isSingleton(inner)) {
+            const surr = await getIncludeSurroundings(ctx.taskCtx, ctx.structure, inner.structure, params);
+            const ret = StructureSelection.Singletons(ctx.structure, surr);
             return ret;
         } else {
-            const builder = new UniqueStructuresBuilder(structure);
+            const builder = new UniqueStructuresBuilder(ctx.structure);
             for (const s of inner.structures) {
-                builder.add(await getIncludeSurroundings(ctx, structure, s, params));
+                builder.add(await getIncludeSurroundings(ctx.taskCtx, ctx.structure, s, params));
             }
             return builder.getSelection();
         }

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

@@ -4,32 +4,35 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import { StructureElement, StructureProperties as P } from '../structure'
+import { QueryFn, QueryPredicate, QueryContextView } from './context';
 
 namespace Predicates {
     export interface SetLike<A> { has(v: A): boolean }
     function isSetLike<A>(x: any): x is SetLike<A> { return !!x && !!x.has }
 
-    export function eq<A>(p: StructureElement.Property<A>, value: A): StructureElement.Predicate { return l => p(l) === value; }
-    export function lt<A>(p: StructureElement.Property<A>, value: A): StructureElement.Predicate { return l => p(l) < value; }
-    export function lte<A>(p: StructureElement.Property<A>, value: A): StructureElement.Predicate { return l => p(l) <= value; }
-    export function gt<A>(p: StructureElement.Property<A>, value: A): StructureElement.Predicate { return l => p(l) > value; }
-    export function gte<A>(p: StructureElement.Property<A>, value: A): StructureElement.Predicate { return l => p(l) >= value; }
+    export function eq<A>(p: QueryFn<A>, value: A): QueryPredicate { return l => p(l) === value; }
+    export function lt<A>(p: QueryFn<A>, value: A): QueryPredicate { return l => p(l) < value; }
+    export function lte<A>(p: QueryFn<A>, value: A): QueryPredicate { return l => p(l) <= value; }
+    export function gt<A>(p: QueryFn<A>, value: A): QueryPredicate { return l => p(l) > value; }
+    export function gte<A>(p: QueryFn<A>, value: A): QueryPredicate { return l => p(l) >= value; }
 
-    export function inSet<A>(p: StructureElement.Property<A>, values: SetLike<A> | ArrayLike<A>): StructureElement.Predicate {
+    function _true(ctx: QueryContextView) { return true; }
+    function _false(ctx: QueryContextView) { return false; }
+
+    export function inSet<A>(p: QueryFn<A>, values: SetLike<A> | ArrayLike<A>): QueryPredicate {
         if (isSetLike(values)) {
             return l => values.has(p(l));
         } else {
-            if (values.length === 0) return P.constant.false;
+            if (values.length === 0) return _false;
             const set = new Set<A>();
             for (let i = 0; i < values.length; i++) set.add(values[i]);
             return l => set.has(p(l));
         }
     }
 
-    export function and(...ps: StructureElement.Predicate[]): StructureElement.Predicate {
+    export function and(...ps: QueryPredicate[]): QueryPredicate {
         switch (ps.length) {
-            case 0: return P.constant.true;
+            case 0: return _true;
             case 1: return ps[0];
             case 2: {
                 const a = ps[0], b = ps[1];
@@ -61,9 +64,9 @@ namespace Predicates {
         }
     }
 
-    export function or(...ps: StructureElement.Predicate[]): StructureElement.Predicate {
+    export function or(...ps: QueryPredicate[]): QueryPredicate {
         switch (ps.length) {
-            case 0: return P.constant.false;
+            case 0: return _false;
             case 1: return ps[0];
             case 2: {
                 const a = ps[0], b = ps[1];

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

@@ -1,7 +0,0 @@
-/**
- * Copyright (c) 2017-2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
- *
- * @author David Sehnal <david.sehnal@gmail.com>
- */
-
-// TODO

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

@@ -4,21 +4,16 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import { Task, RuntimeContext } from 'mol-task'
+import { RuntimeContext } from 'mol-task'
 import { Structure } from '../structure'
-import Selection from './selection'
+import { StructureSelection } from './selection'
+import { QueryContext } from './context';
 
-// TODO: Query { (s: Structure): Computation<Selection> }
-
-interface Query { (s: Structure): Task<Selection>, provider: Query.Provider }
-function Query(q: Query.Provider): Query {
-    const ret = (s => Task.create('Query', ctx => q(s, ctx))) as Query;
-    ret.provider = q;
-    return ret;
-}
-
-namespace Query {
-    export interface Provider { (s: Structure, ctx: RuntimeContext): Promise<Selection> }
+interface StructureQuery { (ctx: QueryContext): Promise<StructureSelection> }
+namespace StructureQuery {
+    export function run(query: StructureQuery, structure: Structure, ctx?: RuntimeContext) {
+        return query(new QueryContext(structure, ctx || RuntimeContext.Synchronous))
+    }
 }
 
-export default Query
+export { StructureQuery }

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

@@ -10,32 +10,32 @@ import { structureUnion } from './utils/structure';
 import { OrderedSet, SortedArray } from 'mol-data/int';
 
 // A selection is a pair of a Structure and a sequence of unique AtomSets
-type Selection = Selection.Singletons | Selection.Sequence
+type StructureSelection = StructureSelection.Singletons | StructureSelection.Sequence
 
-namespace Selection {
+namespace StructureSelection {
     // If each element of the selection is a singleton, we can use a more efficient representation.
     export interface Singletons { readonly kind: 'singletons', readonly source: Structure, readonly structure: Structure }
     export interface Sequence { readonly kind: 'sequence', readonly source: Structure, readonly structures: Structure[] }
 
     export function Singletons(source: Structure, structure: Structure): Singletons { return { kind: 'singletons', source, structure } }
     export function Sequence(source: Structure, structures: Structure[]): Sequence { return { kind: 'sequence', source, structures } }
-    export function Empty(source: Structure): Selection { return Singletons(source, Structure.Empty); };
+    export function Empty(source: Structure): StructureSelection { return Singletons(source, Structure.Empty); };
 
-    export function isSingleton(s: Selection): s is Singletons { return s.kind === 'singletons'; }
-    export function isEmpty(s: Selection) { return isSingleton(s) ? s.structure.units.length === 0 : s.structures.length === 0; }
+    export function isSingleton(s: StructureSelection): s is Singletons { return s.kind === 'singletons'; }
+    export function isEmpty(s: StructureSelection) { return isSingleton(s) ? s.structure.units.length === 0 : s.structures.length === 0; }
 
-    export function structureCount(sel: Selection) {
+    export function structureCount(sel: StructureSelection) {
         if (isSingleton(sel)) return sel.structure.elementCount;
         return sel.structures.length;
     }
 
-    export function unionStructure(sel: Selection): Structure {
+    export function unionStructure(sel: StructureSelection): Structure {
         if (isEmpty(sel)) return Structure.Empty;
         if (isSingleton(sel)) return sel.structure;
         return structureUnion(sel.source, sel.structures);
     }
 
-    export function toLoci(sel: Selection): StructureElement.Loci {
+    export function toLoci(sel: StructureSelection): StructureElement.Loci {
         const loci: { unit: Unit, indices: OrderedSet<StructureElement.UnitIndex> }[] = [];
         const { unitMap } = sel.source;
 
@@ -55,7 +55,7 @@ namespace Selection {
 
     export interface Builder {
         add(structure: Structure): void,
-        getSelection(): Selection
+        getSelection(): StructureSelection
     }
 
     function getSelection(source: Structure, structures: Structure[], allSingletons: boolean) {
@@ -104,4 +104,4 @@ namespace Selection {
     // TODO: spatial lookup
 }
 
-export default Selection
+export { StructureSelection }

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

@@ -5,7 +5,7 @@
  */
 
 import { StructureElement, Structure } from '../../structure';
-import Selection from '../selection';
+import { StructureSelection } from '../selection';
 import { HashSet } from 'mol-data/generic';
 import { structureUnion } from './structure';
 import { StructureSubsetBuilder } from '../../structure/util/subset-builder';
@@ -25,8 +25,8 @@ export class UniqueStructuresBuilder {
     }
 
     getSelection() {
-        if (this.allSingletons) return Selection.Singletons(this.source, structureUnion(this.source, this.structures));
-        return Selection.Sequence(this.source, this.structures);
+        if (this.allSingletons) return StructureSelection.Singletons(this.source, structureUnion(this.source, this.structures));
+        return StructureSelection.Sequence(this.source, this.structures);
     }
 
     constructor(private source: Structure) {
@@ -54,14 +54,14 @@ export class LinearGroupingBuilder {
         return true;
     }
 
-    private singletonSelection(): Selection {
+    private singletonSelection(): StructureSelection {
         const builder = this.source.subsetBuilder(true);
         const loc = StructureElement.create();
         for (let i = 0, _i = this.builders.length; i < _i; i++) {
             this.builders[i].setSingletonLocation(loc);
             builder.addToUnit(loc.unit.id, loc.element);
         }
-        return Selection.Singletons(this.source, builder.getStructure());
+        return StructureSelection.Singletons(this.source, builder.getStructure());
     }
 
     private fullSelection() {
@@ -69,12 +69,12 @@ export class LinearGroupingBuilder {
         for (let i = 0, _i = this.builders.length; i < _i; i++) {
             structures[i] = this.builders[i].getStructure();
         }
-        return Selection.Sequence(this.source, structures);
+        return StructureSelection.Sequence(this.source, structures);
     }
 
-    getSelection(): Selection {
+    getSelection(): StructureSelection {
         const len = this.builders.length;
-        if (len === 0) return Selection.Empty(this.source);
+        if (len === 0) return StructureSelection.Empty(this.source);
         if (this.allSingletons()) return this.singletonSelection();
         return this.fullSelection();
     }

src/mol-task/execution/runtime-context.ts

@@ -4,6 +4,8 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
+import { SyncRuntimeContext } from './synchronous';
+
 interface RuntimeContext {
     readonly shouldUpdate: boolean,
     readonly isSynchronous: boolean,
@@ -26,6 +28,8 @@ namespace RuntimeContext {
         max: number,
         canAbort: boolean
     }
+
+    export const Synchronous = SyncRuntimeContext;
 }
 
 export { RuntimeContext }