mol-model: fixes to link related queries

src/mol-data/generic/unique-array.ts

@@ -20,6 +20,10 @@ namespace UniqueArray {
         array[array.length] = value;
         return true;
     }
+
+    export function has<K, T>({ keys }: UniqueArray<K, T>, key: K) {
+        return keys.has(key);
+    }
 }
 
 export { UniqueArray }

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

@@ -9,6 +9,7 @@ import { now } from '../../../mol-util/now';
 import { ElementIndex } from '../model';
 import { LinkType } from '../model/types';
 import { StructureSelection } from './selection';
+import { defaultLinkTest } from './queries/internal';
 
 export interface QueryContextView {
     readonly element: StructureElement.Location;
@@ -119,6 +120,35 @@ export class QueryContextLinkInfo<U extends Unit = Unit> {
     type: LinkType = LinkType.Flag.None;
     order: number = 0;
 
+    private testFn: QueryPredicate = defaultLinkTest;
+
+    setTestFn(fn?: QueryPredicate) {
+        this.testFn = fn || defaultLinkTest;
+    }
+
+    test(ctx: QueryContext, trySwap: boolean) {
+        if (this.testFn(ctx)) return true;
+        if (trySwap) {
+            this.swap();
+            return this.testFn(ctx);
+        }
+        return false;
+    }
+
+    private swap() {
+        const idxA = this.aIndex;
+        this.aIndex = this.bIndex;
+        this.bIndex = idxA;
+
+        const unitA = this.a.unit;
+        this.a.unit = this.b.unit;
+        this.b.unit = unitA;
+
+        const eA = this.a.element;
+        this.a.element = this.b.element;
+        this.b.element = eA;
+    }
+
     get length() {
         return StructureElement.Location.distance(this.a, this. b);
     }

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

@@ -15,7 +15,6 @@ import { checkStructureMaxRadiusDistance, checkStructureMinMaxDistance } from '.
 import Structure from '../../structure/structure';
 import StructureElement from '../../structure/element';
 import { SortedArray } from '../../../../mol-data/int';
-import { defaultLinkTest } from './internal';
 
 export function pick(query: StructureQuery, pred: QueryPredicate): StructureQuery {
     return ctx => {
@@ -231,13 +230,11 @@ function withinMinMaxRadius({ queryCtx, selection, target, minRadius, maxRadius,
 interface IsConnectedToCtx {
     queryCtx: QueryContext,
     input: Structure,
-    target: Structure,
-    linkTest: QueryFn<boolean>,
-    tElement: StructureElement.Location
+    target: Structure
 }
 
 function checkConnected(ctx: IsConnectedToCtx, structure: Structure) {
-    const { queryCtx, input, target, linkTest, tElement } = ctx;
+    const { queryCtx, input, target } = ctx;
     const atomicLink = queryCtx.atomicLink;
 
     const interLinks = input.links;
@@ -258,36 +255,41 @@ function checkConnected(ctx: IsConnectedToCtx, structure: Structure) {
         for (let i = 0 as StructureElement.UnitIndex, _i = srcElements.length; i < _i; i++) {
             const inputIndex = SortedArray.indexOf(inputElements, srcElements[i]) as StructureElement.UnitIndex;
 
-            atomicLink.aIndex = inputIndex;
-            atomicLink.a.element = srcElements[i];
+            atomicLink.a.unit = inputUnit;
             atomicLink.b.unit = inputUnit;
 
-            tElement.unit = unit;
+            // tElement.unit = unit;
             for (let l = offset[inputIndex], _l = offset[inputIndex + 1]; l < _l; l++) {
-                tElement.element = inputElements[b[l]];
-                if (!target.hasElement(tElement)) continue;
-                atomicLink.bIndex = b[l] as StructureElement.UnitIndex;
+                // tElement.element = inputElements[b[l]];
                 atomicLink.b.element = inputUnit.elements[b[l]];
+                if (!target.hasElement(atomicLink.b)) continue;
+
+                atomicLink.aIndex = inputIndex;
+                atomicLink.a.element = srcElements[i];
+                atomicLink.bIndex = b[l] as StructureElement.UnitIndex;
                 atomicLink.type = flags[l];
                 atomicLink.order = order[l];
-                if (linkTest(queryCtx)) return true;
+                if (atomicLink.test(queryCtx, true)) return true;
             }
 
             for (let li = 0; li < luCount; li++) {
                 const lu = linkedUnits[li];
-                tElement.unit = lu.unitB;
-                atomicLink.b.unit = lu.unitB;
                 const bElements = lu.unitB.elements;
                 const bonds = lu.getBonds(inputIndex);
                 for (let bi = 0, _bi = bonds.length; bi < _bi; bi++) {
                     const bond = bonds[bi];
-                    tElement.element = bElements[bond.indexB];
-                    if (!target.hasElement(tElement)) continue;
+                    atomicLink.b.unit = lu.unitB;
+                    atomicLink.b.element = bElements[bond.indexB];
+                    if (!target.hasElement(atomicLink.b)) continue;
+
+                    atomicLink.a.unit = inputUnit;
+                    atomicLink.aIndex = inputIndex;
+                    atomicLink.a.element = srcElements[i];
+
                     atomicLink.bIndex = bond.indexB;
-                    atomicLink.b.element = tElement.element;
                     atomicLink.type = bond.flag;
                     atomicLink.order = bond.order;
-                    if (linkTest(queryCtx)) return true;
+                    if (atomicLink.test(queryCtx, true)) return true;
                 }
             }
         }
@@ -314,15 +316,17 @@ export function isConnectedTo({ query, target, disjunct, invert, linkTest }: IsC
         const connCtx: IsConnectedToCtx = {
             queryCtx: ctx,
             input: ctx.inputStructure,
-            target: StructureSelection.unionStructure(targetSel),
-            linkTest: linkTest || defaultLinkTest,
-            tElement: StructureElement.Location.create()
+            target: StructureSelection.unionStructure(targetSel)
         }
 
         const ret = StructureSelection.LinearBuilder(ctx.inputStructure);
         ctx.pushCurrentLink();
+        ctx.atomicLink.setTestFn(linkTest);
+
         StructureSelection.forEach(selection, (s, sI) => {
-            if (checkConnected(connCtx, s)) ret.add(s);
+            if (checkConnected(connCtx, s)) {
+                ret.add(s);
+            }
             if (sI % 5 === 0) ctx.throwIfTimedOut();
         })
         ctx.popCurrentLink();

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

@@ -277,7 +277,6 @@ export function querySelection(selection: StructureQuery, query: StructureQuery,
 }
 
 export function linkedAtomicPairs(linkTest?: QueryPredicate): StructureQuery {
-    const test = linkTest || ((ctx: any) => true);
     return function query_linkedAtomicPairs(ctx) {
         const structure = ctx.inputStructure;
 
@@ -287,6 +286,7 @@ export function linkedAtomicPairs(linkTest?: QueryPredicate): StructureQuery {
 
         ctx.pushCurrentLink();
         const atomicLink = ctx.atomicLink;
+        atomicLink.setTestFn(linkTest);
 
         // Process intra unit links
         for (const unit of structure.units) {
@@ -305,7 +305,8 @@ export function linkedAtomicPairs(linkTest?: QueryPredicate): StructureQuery {
                     atomicLink.b.element = unit.elements[intraLinkB[lI]];
                     atomicLink.type = flags[lI];
                     atomicLink.order = order[lI];
-                    if (test(ctx)) {
+                    // No need to "swap test" because each bond direction will be visited eventually.
+                    if (atomicLink.test(ctx, false)) {
                         const b = structure.subsetBuilder(false);
                         b.beginUnit(unit.id);
                         b.addElement(atomicLink.a.element);
@@ -328,7 +329,8 @@ export function linkedAtomicPairs(linkTest?: QueryPredicate): StructureQuery {
             atomicLink.order = bond.order;
             atomicLink.type = bond.flag;
 
-            if (test(ctx)) {
+            // No need to "swap test" because each bond direction will be visited eventually.
+            if (atomicLink.test(ctx, false)) {
                 const b = structure.subsetBuilder(false);
                 b.addToUnit(atomicLink.a.unit.id, atomicLink.a.element);
                 b.addToUnit(atomicLink.b.unit.id, atomicLink.b.element);

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

@@ -15,7 +15,6 @@ import { structureIntersect, structureSubtract } from '../utils/structure-set';
 import { UniqueArray } from '../../../../mol-data/generic';
 import { StructureSubsetBuilder } from '../../structure/util/subset-builder';
 import StructureElement from '../../structure/element';
-import { defaultLinkTest } from './internal';
 
 function getWholeResidues(ctx: QueryContext, source: Structure, structure: Structure) {
     const builder = source.subsetBuilder(true);
@@ -305,16 +304,16 @@ export interface IncludeConnectedParams {
 }
 
 export function includeConnected({ query, layerCount, wholeResidues, linkTest }: IncludeConnectedParams): StructureQuery {
-    const bt = linkTest || defaultLinkTest;
     const lc = Math.max(layerCount, 0);
     return function query_includeConnected(ctx) {
         const builder = StructureSelection.UniqueBuilder(ctx.inputStructure);
         const src = query(ctx);
         ctx.pushCurrentLink();
+        ctx.atomicLink.setTestFn(linkTest);
         StructureSelection.forEach(src, (s, sI) => {
             let incl = s;
             for (let i = 0; i < lc; i++) {
-                incl = includeConnectedStep(ctx, bt, wholeResidues, incl);
+                incl = includeConnectedStep(ctx, wholeResidues, incl);
             }
             builder.add(incl);
             if (sI % 10 === 0) ctx.throwIfTimedOut();
@@ -325,24 +324,24 @@ export function includeConnected({ query, layerCount, wholeResidues, linkTest }:
     }
 }
 
-function includeConnectedStep(ctx: QueryContext, linkTest: QueryFn<boolean>, wholeResidues: boolean, structure: Structure) {
-    const expanded = expandConnected(ctx, structure, linkTest);
+function includeConnectedStep(ctx: QueryContext, wholeResidues: boolean, structure: Structure) {
+    const expanded = expandConnected(ctx, structure);
     if (wholeResidues) return getWholeResidues(ctx, ctx.inputStructure, expanded);
     return expanded;
 }
 
-function expandConnected(ctx: QueryContext, structure: Structure, linkTest: QueryFn<boolean>) {
+function expandConnected(ctx: QueryContext, structure: Structure) {
     const inputStructure = ctx.inputStructure;
     const interLinks = inputStructure.links;
     const builder = new StructureUniqueSubsetBuilder(inputStructure);
 
     // Note: each link is visited twice so that link.atom-a and link.atom-b both get the "swapped values"
+    const visitedSourceUnits = new Set<number>();
 
     const atomicLink = ctx.atomicLink;
 
     // Process intra unit links
     for (const unit of structure.units) {
-
         if (unit.kind !== Unit.Kind.Atomic) {
             // add the whole unit
             builder.beginUnit(unit.id);
@@ -353,53 +352,71 @@ function expandConnected(ctx: QueryContext, structure: Structure, linkTest: Quer
             continue;
         }
 
-        const inputUnit = inputStructure.unitMap.get(unit.id) as Unit.Atomic;
-        const { offset: intraLinkOffset, b: intraLinkB, edgeProps: { flags, order } } = inputUnit.links;
+        const inputUnitA = inputStructure.unitMap.get(unit.id) as Unit.Atomic;
+        const { offset: intraLinkOffset, b: intraLinkB, edgeProps: { flags, order } } = inputUnitA.links;
 
         // Process intra unit links
-        atomicLink.a.unit = inputUnit;
-        atomicLink.b.unit = inputUnit;
+        atomicLink.a.unit = inputUnitA;
+        atomicLink.b.unit = inputUnitA;
         for (let i = 0, _i = unit.elements.length; i < _i; i++) {
             // add the current element
             builder.addToUnit(unit.id, unit.elements[i]);
 
-            const srcIndex = SortedArray.indexOf(inputUnit.elements, unit.elements[i]);
-            atomicLink.aIndex = srcIndex as StructureElement.UnitIndex;
-            atomicLink.a.element = unit.elements[i];
+            const aIndex = SortedArray.indexOf(inputUnitA.elements, unit.elements[i]) as StructureElement.UnitIndex;
 
             // check intra unit links
-            for (let lI = intraLinkOffset[srcIndex], _lI = intraLinkOffset[srcIndex + 1]; lI < _lI; lI++) {
-                atomicLink.bIndex = intraLinkB[lI] as StructureElement.UnitIndex;
-                atomicLink.b.element = inputUnit.elements[intraLinkB[lI]];
+            for (let lI = intraLinkOffset[aIndex], _lI = intraLinkOffset[aIndex + 1]; lI < _lI; lI++) {
+                const bIndex = intraLinkB[lI] as StructureElement.UnitIndex;
+                const bElement = inputUnitA.elements[bIndex];
+
+                // Check if the element is already present:
+                if (SortedArray.has(unit.elements, bElement) || builder.has(unit.id, bElement)) continue;
+
+                atomicLink.aIndex = aIndex;
+                atomicLink.a.element = unit.elements[i];
+                atomicLink.bIndex = bIndex;
+                atomicLink.b.element = bElement;
                 atomicLink.type = flags[lI];
                 atomicLink.order = order[lI];
-                if (linkTest(ctx)) {
-                    builder.addToUnit(unit.id, inputUnit.elements[intraLinkB[lI]]);
+
+                if (atomicLink.test(ctx, true)) {
+                    builder.addToUnit(unit.id, bElement);
                 }
             }
         }
 
         // Process inter unit links
-        for (const linkedUnit of interLinks.getLinkedUnits(inputUnit)) {
-            atomicLink.b.unit = linkedUnit.unitB;
+        for (const linkedUnit of interLinks.getLinkedUnits(inputUnitA)) {
+            if (visitedSourceUnits.has(linkedUnit.unitB.id)) continue;
+            const currentUnitB = structure.unitMap.get(linkedUnit.unitB.id);
+
             for (const aI of linkedUnit.linkedElementIndices) {
                 // check if the element is in the expanded structure
-                if (!SortedArray.has(unit.elements, inputUnit.elements[aI])) continue;
+                if (!SortedArray.has(unit.elements, inputUnitA.elements[aI])) continue;
 
-                atomicLink.aIndex = aI;
-                atomicLink.a.element = inputUnit.elements[aI];
                 for (const bond of linkedUnit.getBonds(aI)) {
+                    const bElement = linkedUnit.unitB.elements[bond.indexB];
+
+                    // Check if the element is already present:
+                    if ((currentUnitB && SortedArray.has(currentUnitB.elements, bElement)) || builder.has(linkedUnit.unitB.id, bElement)) continue;
+
+                    atomicLink.a.unit = inputUnitA;
+                    atomicLink.aIndex = aI;
+                    atomicLink.a.element = inputUnitA.elements[aI];
+                    atomicLink.b.unit = linkedUnit.unitB;
                     atomicLink.bIndex = bond.indexB;
-                    // TODO: optimize the lookup?
-                    atomicLink.b.element = linkedUnit.unitB.elements[bond.indexB];
+                    atomicLink.b.element = bElement;
                     atomicLink.type = bond.flag;
                     atomicLink.order = bond.order;
-                    if (linkTest(ctx)) {
-                        builder.addToUnit(linkedUnit.unitB.id, linkedUnit.unitB.elements[bond.indexB]);
+
+                    if (atomicLink.test(ctx, true)) {
+                        builder.addToUnit(linkedUnit.unitB.id, bElement);
                     }
                 }
             }
         }
+
+        visitedSourceUnits.add(inputUnitA.id);
     }
 
     return builder.getStructure();

src/mol-model/structure/structure/util/unique-subset-builder.ts

@@ -34,6 +34,12 @@ export class StructureUniqueSubsetBuilder {
         }
     }
 
+    has(parentId: number, e: number) {
+        const unit = this.unitMap.get(parentId);
+        if (!unit) return false;
+        return UniqueArray.has(unit, e);
+    }
+
     beginUnit(parentId: number) {
         this.parentId = parentId;
         if (this.unitMap.has(parentId)) {

src/mol-script/language/symbol-table/structure-query.ts

@@ -94,7 +94,7 @@ const generator = {
     }), Types.ElementSelectionQuery, 'Return all atoms for which the tests are satisfied, grouped into sets.'),
 
     linkedAtomicPairs: symbol(Arguments.Dictionary({
-        0: Argument(Type.Bool, { isOptional: true, defaultValue: true, description: 'Test each link with this predicate. Each link is visited twice with swapped atom order.' }),
+        0: Argument(Type.Bool, { isOptional: true, defaultValue: 'true for covalent links' as any, description: 'Test each link with this predicate. Each link is visited twice with swapped atom order.' }),
         // TODO: shoud we support this or just use queryEach to get similar behavior
         // 'group-by': Argument(Type.Any, { isOptional: true, defaultValue: ``, description: 'Group the links using the privided value' }),
     }), Types.ElementSelectionQuery, 'Return all pairs of atoms for which the test is satisfied.'),