mol-model/query added link support to QueryContext

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

@@ -491,6 +491,10 @@ export namespace LinkType {
         // currently at most 16 flags are supported!!
     }
 
+    export function create(flags: Flag): LinkType {
+        return BitFlags.create(flags);
+    }
+
     export function isCovalent(flags: LinkType.Flag) {
         return (flags & LinkType.Flag.Covalent) !== 0;
     }

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

@@ -7,6 +7,7 @@
 import { Structure, StructureElement, Unit } from '../structure';
 import { now } from 'mol-task';
 import { ElementIndex } from '../model';
+import { Link } from '../structure/unit/links';
 
 export interface QueryContextView {
     readonly element: StructureElement;
@@ -15,6 +16,7 @@ export interface QueryContextView {
 
 export class QueryContext implements QueryContextView {
     private currentElementStack: StructureElement[] = [];
+    private currentAtomicLinkStack: Link.Location<Unit.Atomic>[] = [];
     private currentStructureStack: Structure[] = [];
     private inputStructureStack: Structure[] = [];
 
@@ -27,6 +29,9 @@ export class QueryContext implements QueryContextView {
     readonly element: StructureElement = StructureElement.create();
     currentStructure: Structure = void 0 as any;
 
+    /** Current link between atoms */
+    readonly atomicLink: Link.Location<Unit.Atomic> = void 0 as any;
+
     setElement(unit: Unit, e: ElementIndex) {
         this.element.unit = unit;
         this.element.element = e;
@@ -42,6 +47,20 @@ export class QueryContext implements QueryContextView {
         (this.element as StructureElement) = this.currentElementStack.pop()!;
     }
 
+    pushCurrentLink() {
+        if (this.atomicLink) this.currentAtomicLinkStack.push(this.atomicLink);
+        (this.atomicLink as Link.Location<Unit.Atomic>) = Link.Location() as Link.Location<Unit.Atomic>;
+        return this.atomicLink;
+    }
+
+    popCurrentLink() {
+        if (this.currentAtomicLinkStack.length > 0) {
+            (this.atomicLink as Link.Location<Unit.Atomic>) = this.currentAtomicLinkStack.pop()!;
+        } else {
+            (this.atomicLink as any) = void 0;
+        }
+    }
+
     pushCurrentStructure() {
         if (this.currentStructure) this.currentStructureStack.push(this.currentStructure);
     }

src/mol-model/structure/structure/unit/links.ts

@@ -6,18 +6,20 @@
  */
 
 import { Unit, StructureElement } from '../../structure'
+import Structure from '../structure';
+import { LinkType } from '../../model/types';
 
 export * from './links/data'
 export * from './links/intra-compute'
 export * from './links/inter-compute'
 
 namespace Link {
-    export interface Location {
+    export interface Location<U extends Unit = Unit> {
         readonly kind: 'link-location',
-        aUnit: Unit,
+        aUnit: U,
         /** Index into aUnit.elements */
         aIndex: StructureElement.UnitIndex,
-        bUnit: Unit,
+        bUnit: U,
         /** Index into bUnit.elements */
         bIndex: StructureElement.UnitIndex,
     }
@@ -42,6 +44,32 @@ namespace Link {
     export function isLoci(x: any): x is Loci {
         return !!x && x.kind === 'link-loci';
     }
+
+    export function getType(structure: Structure, link: Location<Unit.Atomic>): LinkType {
+        if (link.aUnit === link.bUnit) {
+            const links = link.aUnit.links;
+            const idx = links.getEdgeIndex(link.aIndex, link.bIndex);
+            if (idx < 0) return LinkType.create(LinkType.Flag.None);
+            return LinkType.create(links.edgeProps.flags[idx]);
+        } else {
+            const bond = structure.links.getBondFromLocation(link);
+            if (bond) return LinkType.create(bond.flag);
+            return LinkType.create(LinkType.Flag.None);
+        }
+    }
+
+    export function getOrder(structure: Structure, link: Location<Unit.Atomic>): number {
+        if (link.aUnit === link.bUnit) {
+            const links = link.aUnit.links;
+            const idx = links.getEdgeIndex(link.aIndex, link.bIndex);
+            if (idx < 0) return 0;
+            return links.edgeProps.order[idx];
+        } else {
+            const bond = structure.links.getBondFromLocation(link);
+            if (bond) return bond.order;
+            return 0;
+        }
+    }
 }
 
 export { Link }

src/mol-model/structure/structure/unit/links/data.ts

@@ -9,6 +9,7 @@ import { LinkType } from '../../../model/types'
 import { IntAdjacencyGraph } from 'mol-math/graph';
 import Unit from '../../unit';
 import StructureElement from '../../element';
+import { Link } from '../links';
 
 type IntraUnitLinks = IntAdjacencyGraph<{ readonly order: ArrayLike<number>, readonly flags: ArrayLike<LinkType.Flag> }>
 
@@ -27,17 +28,21 @@ class InterUnitBonds {
     }
 
     /** Index into this.bonds */
-    getBondIndex(indexA: number, unitA: Unit, indexB: number, unitB: Unit): number {
+    getBondIndex(indexA: StructureElement.UnitIndex, unitA: Unit, indexB: StructureElement.UnitIndex, unitB: Unit): number {
         const key = InterUnitBonds.getBondKey(indexA, unitA, indexB, unitB)
         const index = this.bondKeyIndex.get(key)
         return index !== undefined ? index : -1
     }
 
-    getBond(indexA: number, unitA: Unit, indexB: number, unitB: Unit): InterUnitBonds.Bond | undefined {
+    getBond(indexA: StructureElement.UnitIndex, unitA: Unit, indexB: StructureElement.UnitIndex, unitB: Unit): InterUnitBonds.Bond | undefined {
         const index = this.getBondIndex(indexA, unitA, indexB, unitB)
         return index !== -1 ? this.bonds[index] : undefined
     }
 
+    getBondFromLocation(l: Link.Location) {
+        return this.getBond(l.aIndex, l.aUnit, l.bIndex, l.bUnit);
+    }
+
     constructor(private map: Map<number, InterUnitBonds.UnitPairBonds[]>) {
         let count = 0
         const bonds: (InterUnitBonds.Bond)[] = []