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 {
+                        groupedUnits.set(e.unit.invariantId, [e.unit.id])
+                    }
                 } else {
-                    elementGroups.set(key, {
-                        units: [e.unit.id],
-                        set: e.set,
-                        ranges: e.ranges
-                    })
+                    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,30 +577,32 @@ 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)
-                if (units.length === 0) continue
-
-                let indices: OrderedSet<UnitIndex>
-                if (e.ranges.length === 0) {
-                    indices = e.set
-                } else if (e.set.length === 0) {
-                    if (e.ranges.length === 2) {
-                        indices = Interval.ofRange(e.ranges[0], e.ranges[1])
+                for (const g of e.groupedUnits) {
+                    const units = getUnitsFromIds(g, parent)
+                    if (units.length === 0) continue
+
+                    let indices: OrderedSet<UnitIndex>
+                    if (e.ranges.length === 0) {
+                        indices = e.set
+                    } else if (e.set.length === 0) {
+                        if (e.ranges.length === 2) {
+                            indices = Interval.ofRange(e.ranges[0], e.ranges[1])
+                        } else {
+                            const _indices = new Int32Array(SortedRanges.size(e.ranges))
+                            SortedRanges.forEach(e.ranges, (v, i) => _indices[i] = v)
+                            indices = SortedArray.ofSortedArray(_indices)
+                        }
                     } else {
-                        const _indices = new Int32Array(SortedRanges.size(e.ranges))
+                        const rangesSize = SortedRanges.size(e.ranges)
+                        const _indices = new Int32Array(e.set.length + rangesSize)
                         SortedRanges.forEach(e.ranges, (v, i) => _indices[i] = v)
-                        indices = SortedArray.ofSortedArray(_indices)
+                        _indices.set(e.set, rangesSize)
+                        indices = SortedArray.ofUnsortedArray(_indices) // requires sort
                     }
-                } else {
-                    const rangesSize = SortedRanges.size(e.ranges)
-                    const _indices = new Int32Array(e.set.length + rangesSize)
-                    SortedRanges.forEach(e.ranges, (v, i) => _indices[i] = v)
-                    _indices.set(e.set, rangesSize)
-                    indices = SortedArray.ofUnsortedArray(_indices) // requires sort
-                }
 
-                for (const unit of units) {
-                    elements.push({ unit, indices })
+                    for (const unit of units) {
+                        elements.push({ unit, indices })
+                    }
                 }
             }
             return Loci(parent, elements)
@@ -602,30 +611,32 @@ namespace StructureElement {
         export function toStructure(query: Query, parent: Structure): Structure {
             const units: Unit[] = []
             for (const e of query.elements) {
-                const _units = getUnitsFromIds(e.units, parent)
-                if (_units.length === 0) continue
-
-                const unit = _units[0] // the elements are grouped by unit.invariantId
-                const rangesSize = SortedRanges.size(e.ranges)
-                const _indices = new Int32Array(e.set.length + rangesSize)
-                let indices: SortedArray<ElementIndex>
-                if (e.ranges.length === 0) {
-                    for (let i = 0, il = e.set.length; i < il; ++i) {
-                        _indices[i] = unit.elements[e.set[i]]
-                    }
-                    indices = SortedArray.ofSortedArray(_indices)
-                } else if (e.set.length === 0) {
-                    SortedRanges.forEach(e.ranges, (v, i) => _indices[i] = unit.elements[v])
-                    indices = SortedArray.ofSortedArray(_indices)
-                } else {
+                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
                     const rangesSize = SortedRanges.size(e.ranges)
-                    SortedRanges.forEach(e.ranges, (v, i) => _indices[i] = unit.elements[v])
-                    SortedRanges.forEach(e.set, (v, i) => _indices[i + rangesSize] = unit.elements[v])
-                    indices = SortedArray.ofUnsortedArray(_indices) // requires sort
-                }
+                    const _indices = new Int32Array(e.set.length + rangesSize)
+                    let indices: SortedArray<ElementIndex>
+                    if (e.ranges.length === 0) {
+                        for (let i = 0, il = e.set.length; i < il; ++i) {
+                            _indices[i] = unit.elements[e.set[i]]
+                        }
+                        indices = SortedArray.ofSortedArray(_indices)
+                    } else if (e.set.length === 0) {
+                        SortedRanges.forEach(e.ranges, (v, i) => _indices[i] = unit.elements[v])
+                        indices = SortedArray.ofSortedArray(_indices)
+                    } else {
+                        const rangesSize = SortedRanges.size(e.ranges)
+                        SortedRanges.forEach(e.ranges, (v, i) => _indices[i] = unit.elements[v])
+                        SortedRanges.forEach(e.set, (v, i) => _indices[i + rangesSize] = unit.elements[v])
+                        indices = SortedArray.ofUnsortedArray(_indices) // requires sort
+                    }
 
-                for (const unit of _units) {
-                    units.push(unit.getChild(indices))
+                    for (const unit of _units) {
+                        units.push(unit.getChild(indices))
+                    }
                 }
             }
             return Structure.create(units, parent)
@@ -634,9 +645,11 @@ namespace StructureElement {
         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
             }