mol-model: implemented includeConnected query

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

@@ -4,7 +4,7 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import { Segmentation } from '../../../../mol-data/int';
+import { Segmentation, SortedArray } from '../../../../mol-data/int';
 import { Structure, Unit } from '../../structure';
 import { StructureQuery } from '../query';
 import { StructureSelection } from '../selection';
@@ -304,78 +304,99 @@ export interface IncludeConnectedParams {
 }
 
 export function includeConnected({ query, layerCount, wholeResidues, bondTest }: IncludeConnectedParams): StructureQuery {
+    const bt = bondTest || defaultBondTest;
+    const lc = Math.max(layerCount, 0);
     return ctx => {
-        return 0 as any;
+        const builder = StructureSelection.UniqueBuilder(ctx.inputStructure);
+        const src = query(ctx);
+        ctx.pushCurrentLink();
+        StructureSelection.forEach(src, (s, sI) => {
+            let incl = s;
+            for (let i = 0; i < lc; i++) {
+                incl = includeConnectedStep(ctx, bt, wholeResidues, incl);
+            }
+            builder.add(incl);
+            if (sI % 10 === 0) ctx.throwIfTimedOut();
+        });
+        ctx.popCurrentLink();
+        return builder.getSelection();
     }
 }
 
-// function defaultBondTest(ctx: QueryContext) {
-//     return true;
-// }
-
-// interface IncludeConnectedCtx {
-//     queryCtx: QueryContext,
-//     input: Structure,
-//     bondTest: QueryFn<boolean>,
-//     wholeResidues: boolean
-// }
-
-// type FrontierSet = UniqueArray<StructureElement.UnitIndex, StructureElement.UnitIndex>
-// type Frontier = { unitIds: UniqueArray<number>, elements: Map<number /* unit id */, FrontierSet> }
-
-// namespace Frontier {
-//     export function has({ elements }: Frontier, unitId: number, element: StructureElement.UnitIndex) {
-//         if (!elements.has(unitId)) return false;
-//         const xs = elements.get(unitId)!;
-//         return xs.keys.has(element);
-//     }
-
-//     export function create(pivot: Structure, input: Structure) {
-//         const unitIds = UniqueArray.create<number>();
-//         const elements: Frontier['elements'] = new Map();
-//         for (const unit of pivot.units) {
-//             if (!Unit.isAtomic(unit)) continue;
-
-//             UniqueArray.add(unitIds, unit.id, unit.id);
-//             const xs: FrontierSet = UniqueArray.create();
-//             elements.set(unit.id, xs);
-
-//             const pivotElements = unit.elements;
-//             const inputElements = input.unitMap.get(unit.id).elements;
-//             for (let i = 0, _i = pivotElements.length; i < _i; i++) {
-//                 const idx = SortedArray.indexOf(inputElements, pivotElements[i]) as StructureElement.UnitIndex;
-//                 UniqueArray.add(xs, idx, idx);
-//             }
-//         }
-
-//         return { unitIds, elements };
-//     }
-
-//     export function addFrontier(target: Frontier, from: Frontier) {
-//         for (const unitId of from.unitIds.array) {
-//             let xs: FrontierSet;
-//             if (target.elements.has(unitId)) {
-//                 xs = target.elements.get(unitId)!;
-//             } else {
-//                 xs = UniqueArray.create();
-//                 target.elements.set(unitId, xs);
-//                 UniqueArray.add(target.unitIds, unitId, unitId);
-//             }
-
-//             for (const e of from.elements.get(unitId)!.array) {
-//                 UniqueArray.add(xs, e, e);
-//             }
-//         }
-//         return target;
-//     }
-
-//     export function includeWholeResidues(structure: Structure, frontier: Frontier) {
-//         // ...
-//     }
-// }
-
-// function expandFrontier(ctx: IncludeConnectedCtx, currentFrontier: Frontier, result: Frontier): Frontier {
-//     return 0 as any;
-// }
+function includeConnectedStep(ctx: QueryContext, bondTest: QueryFn<boolean>, wholeResidues: boolean, structure: Structure) {
+    const expanded = expandConnected(ctx, structure, bondTest);
+    if (wholeResidues) return getWholeResidues(ctx, ctx.inputStructure, expanded);
+    return expanded;
+}
+
+function expandConnected(ctx: QueryContext, structure: Structure, bondTest: QueryFn<boolean>) {
+    const inputStructure = ctx.inputStructure;
+    const interLinks = inputStructure.links;
+    const builder = new StructureUniqueSubsetBuilder(inputStructure);
+
+    const processedUnits = new Set<number>();
+
+    // Process intra unit links
+    for (const unit of structure.units) {
+        processedUnits.add(unit.id);
+
+        if (unit.kind !== Unit.Kind.Atomic) {
+            // add the whole unit
+            builder.beginUnit(unit.id);
+            for (let i = 0, _i = unit.elements.length; i < _i; i++) {
+                builder.addElement(unit.elements[i]);
+            }
+            builder.commitUnit();
+            continue;
+        }
+
+        const inputUnit = inputStructure.unitMap.get(unit.id) as Unit.Atomic;
+        const { offset: intraLinkOffset, b: intraLinkB } = inputUnit.links;
+
+        // Process intra unit links
+        ctx.atomicLink.aUnit = inputUnit;
+        ctx.atomicLink.bUnit = inputUnit;
+        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]);
+            ctx.atomicLink.aIndex = srcIndex as StructureElement.UnitIndex;
+
+            // check intra unit links
+            for (let lI = intraLinkOffset[srcIndex], _lI = intraLinkOffset[srcIndex + 1]; lI < _lI; lI++) {
+                ctx.atomicLink.bIndex = intraLinkB[lI] as StructureElement.UnitIndex;
+                if (bondTest(ctx)) {
+                    builder.addToUnit(unit.id, inputUnit.elements[intraLinkB[lI]]);
+                }
+            }
+        }
+
+        // Process inter unit links
+        for (const linkedUnit of interLinks.getLinkedUnits(inputUnit)) {
+            if (processedUnits.has(linkedUnit.unitB.id)) continue;
+
+            ctx.atomicLink.bUnit = linkedUnit.unitB;
+            for (const aI of linkedUnit.linkedElementIndices) {
+                // check if the element is in the expanded structure
+                if (!SortedArray.has(unit.elements, inputUnit.elements[aI])) continue;
+
+                ctx.atomicLink.aIndex = aI;
+                for (const bond of linkedUnit.getBonds(aI)) {
+                    ctx.atomicLink.bIndex = bond.indexB;
+                    if (bondTest(ctx)) {
+                        builder.addToUnit(linkedUnit.unitB.id, linkedUnit.unitB.elements[bond.indexB]);
+                    }
+                }
+            }
+        }
+    }
+
+    return builder.getStructure();
+}
+
+function defaultBondTest(ctx: QueryContext) {
+    return true;
+}
 
 // TODO: unionBy (skip this one?), cluster
\ No newline at end of file

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

@@ -231,6 +231,12 @@ const symbols = [
     D(MolScript.structureQuery.modifier.union, (ctx, xs) => Queries.modifiers.union(xs[0] as any)(ctx)),
     D(MolScript.structureQuery.modifier.expandProperty, (ctx, xs) => Queries.modifiers.expandProperty(xs[0] as any, xs['property'])(ctx)),
     D(MolScript.structureQuery.modifier.exceptBy, (ctx, xs) => Queries.modifiers.exceptBy(xs[0] as any, xs['by'] as any)(ctx)),
+    D(MolScript.structureQuery.modifier.includeConnected, (ctx, xs) => Queries.modifiers.includeConnected({
+        query: xs[0] as any,
+        bondTest: xs['bond-test'],
+        wholeResidues: !!(xs['as-whole-residues'] && xs['as-whole-residues'](ctx)),
+        layerCount: (xs['layer-count'] && xs['layer-count'](ctx)) || 1
+    })(ctx)),
 
     // ============= COMBINATORS ================