improved StructureElement.Query

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

@@ -16,7 +16,7 @@ import Structure from './structure';
 import Unit from './unit';
 import { Boundary } from './util/boundary';
 import { StructureProperties } from '../structure';
-import { sortArray, hashFnv32a, hash2, hash3 } from '../../../mol-data/util';
+import { sortArray, hashFnv32a, hash2 } from '../../../mol-data/util';
 import Expression from '../../../mol-script/language/expression';
 import SortedRanges from '../../../mol-data/int/sorted-ranges';
 
@@ -474,8 +474,11 @@ namespace StructureElement {
     }
 
     interface QueryElement {
-        /** Array of `Unit.id`s that share the same `Unit.invariantId` */
-        units: SortedArray<number>,
+        /**
+         * Array (sorted by first element in sub-array) of
+         * arrays of `Unit.id`s that share the same `Unit.invariantId`
+         */
+        groupedUnits: SortedArray<number>[],
         set: SortedArray<UnitIndex>
         ranges: SortedRanges<UnitIndex>
     }
@@ -494,13 +497,17 @@ namespace StructureElement {
             }[] = []
             for (const e of loci.elements) {
                 const { unit, indices } = e
+                if (OrderedSet.size(indices) === 0) continue
 
                 const ranges: UnitIndex[] = [];
                 const set: UnitIndex[] = [];
 
                 if (OrderedSet.isInterval(indices)) {
-                    ranges[ranges.length] = Interval.min(indices)
-                    ranges[ranges.length] = Interval.max(indices)
+                    if (OrderedSet.size(indices) === 1) {
+                        set.push(Interval.min(indices))
+                    } else {
+                        ranges.push(Interval.min(indices), Interval.max(indices))
+                    }
                 } else {
                     let i = 0, len = indices.length;
                     while (i < len) {
@@ -508,13 +515,11 @@ namespace StructureElement {
                         i++;
                         while (i < len && indices[i - 1] + 1 === indices[i]) i++;
                         const end = i;
-                        // TODO: is this a good value?
-                        if (end - start > 12) {
-                            ranges[ranges.length] = indices[start];
-                            ranges[ranges.length] = indices[end - 1];
+                        if (end - start > 2) {
+                            ranges.push(indices[start], indices[end - 1])
                         } else {
                             for (let j = start; j < end; j++) {
-                                set[set.length] = indices[j];
+                                set[set.length] = indices[j]
                             }
                         }
                     }
@@ -528,31 +533,33 @@ namespace StructureElement {
             }
 
             const elementGroups = new Map<number, {
-                units: number[]
+                groupedUnits: Map<number, number[]>
                 set: SortedArray<UnitIndex>
                 ranges: SortedRanges<UnitIndex>
             }>();
             for (let i = 0, il = _elements.length; i < il; ++i) {
                 const e = _elements[i]
-                const key = hash3(hashFnv32a(e.ranges), hashFnv32a(e.set), e.unit.invariantId)
+                const key = hash2(hashFnv32a(e.ranges), hashFnv32a(e.set))
                 if (elementGroups.has(key)) {
-                    elementGroups.get(key)!.units.push(e.unit.id)
+                    const { groupedUnits } = elementGroups.get(key)!
+                    if (groupedUnits.has(e.unit.invariantId)) {
+                        groupedUnits.get(e.unit.invariantId)!.push(e.unit.id)
                     } else {
-                    elementGroups.set(key, {
-                        units: [e.unit.id],
-                        set: e.set,
-                        ranges: e.ranges
-                    })
+                        groupedUnits.set(e.unit.invariantId, [e.unit.id])
+                    }
+                } else {
+                    const groupedUnits = new Map<number, number[]>()
+                    groupedUnits.set(e.unit.invariantId, [e.unit.id])
+                    elementGroups.set(key, { groupedUnits, set: e.set, ranges: e.ranges })
                 }
             }
 
             const elements: QueryElement[] = []
             elementGroups.forEach(e => {
-                elements.push({
-                    units: SortedArray.ofUnsortedArray(e.units),
-                    set: e.set,
-                    ranges: e.ranges
-                })
+                const groupedUnits: SortedArray<number>[] = []
+                e.groupedUnits.forEach(g => groupedUnits.push(SortedArray.ofUnsortedArray(g)))
+                groupedUnits.sort((a, b) => a[0] - b[0]) // sort by first unit id of each group
+                elements.push({ groupedUnits, set: e.set, ranges: e.ranges })
             })
 
             return { elements }
@@ -570,7 +577,8 @@ namespace StructureElement {
         export function toLoci(query: Query, parent: Structure): Loci {
             const elements: Loci['elements'][0][] = []
             for (const e of query.elements) {
-                const units = getUnitsFromIds(e.units, parent)
+                for (const g of e.groupedUnits) {
+                    const units = getUnitsFromIds(g, parent)
                     if (units.length === 0) continue
 
                     let indices: OrderedSet<UnitIndex>
@@ -596,13 +604,15 @@ namespace StructureElement {
                         elements.push({ unit, indices })
                     }
                 }
+            }
             return Loci(parent, elements)
         }
 
         export function toStructure(query: Query, parent: Structure): Structure {
             const units: Unit[] = []
             for (const e of query.elements) {
-                const _units = getUnitsFromIds(e.units, parent)
+                for (const g of e.groupedUnits) {
+                    const _units = getUnitsFromIds(g, parent)
                     if (_units.length === 0) continue
 
                     const unit = _units[0] // the elements are grouped by unit.invariantId
@@ -628,15 +638,18 @@ namespace StructureElement {
                         units.push(unit.getChild(indices))
                     }
                 }
+            }
             return Structure.create(units, parent)
         }
 
         export function areEqual(a: Query, b: Query) {
             if (a.elements.length !== b.elements.length) return false
             for (let i = 0, il = a.elements.length; i < il; ++i) {
-                const elementA = a.elements[i]
-                const elementB = b.elements[i]
-                if (!SortedArray.areEqual(elementA.units, elementB.units)) return false
+                const elementA = a.elements[i], elementB = b.elements[i]
+                if (elementA.groupedUnits.length !== elementB.groupedUnits.length) return false
+                for (let j = 0, jl = elementB.groupedUnits.length; j < jl; ++i) {
+                    if (!SortedArray.areEqual(elementA.groupedUnits[j], elementB.groupedUnits[j])) return false
+                }
                 if (!SortedArray.areEqual(elementA.set, elementB.set)) return false
                 if (!SortedRanges.areEqual(elementA.ranges, elementB.ranges)) return false
             }