Use OrderedSet for Element.Loci

src/mol-data/int/ordered-set.ts

@@ -6,6 +6,7 @@
 
 import * as Base from './impl/ordered-set'
 import Interval from './interval'
+import SortedArray from './sorted-array';
 
 namespace OrderedSet {
     export const Empty: OrderedSet = Base.Empty as any;
@@ -41,6 +42,7 @@ namespace OrderedSet {
     }
 }
 
-interface OrderedSet { '@type': 'int-interval' | 'int-sorted-array' }
+type OrderedSet = Interval | SortedArray
+//{ '@type': 'int-interval' | 'int-sorted-array' }
 
 export default OrderedSet

src/mol-geo/representation/structure/point.ts

@@ -16,7 +16,7 @@ import VertexMap from '../../shape/vertex-map';
 import { SizeTheme } from '../../theme';
 import { createTransforms, createColors, createSizes, markElement } from './utils';
 import { deepEqual, defaults } from 'mol-util';
-import { SortedArray } from 'mol-data/int';
+import { SortedArray, OrderedSet } from 'mol-data/int';
 import { RenderableState, PointValues } from 'mol-gl/renderable';
 import { PickingId } from '../../util/picking';
 import { Loci, EmptyLoci } from 'mol-model/loci';
@@ -161,7 +161,7 @@ export default function PointUnitsRepresentation(): UnitsRepresentation<PointPro
             const { objectId, instanceId, elementId } = pickingId
             if (points.id === objectId) {
                 const unit = currentGroup.units[instanceId]
-                const indices = SortedArray.ofSingleton(elementId)
+                const indices = OrderedSet.ofSingleton(elementId)
                 return Element.Loci([{ unit, indices }])
             }
             return EmptyLoci

src/mol-geo/representation/structure/spacefill.ts

@@ -19,7 +19,7 @@ import { RenderableState, MeshValues } from 'mol-gl/renderable';
 import { getMeshData } from '../../util/mesh-data';
 import { Mesh } from '../../shape/mesh';
 import { PickingId } from '../../util/picking';
-import { SortedArray } from 'mol-data/int';
+import { OrderedSet } from 'mol-data/int';
 import { createMarkers, MarkerAction } from '../../util/marker-data';
 import { Loci, EmptyLoci } from 'mol-model/loci';
 
@@ -149,7 +149,7 @@ export default function SpacefillUnitsRepresentation(): UnitsRepresentation<Spac
             const { objectId, instanceId, elementId } = pickingId
             if (spheres.id === objectId) {
                 const unit = currentGroup.units[instanceId]
-                const indices = SortedArray.ofSingleton(elementId);
+                const indices = OrderedSet.ofSingleton(elementId);
                 return Element.Loci([{ unit, indices }])
             }
             return EmptyLoci

src/mol-geo/representation/structure/utils.ts

@@ -22,6 +22,7 @@ import { MeshBuilder } from '../../shape/mesh-builder';
 import { TextureImage } from 'mol-gl/renderable/util';
 import { applyMarkerAction, MarkerAction } from '../../util/marker-data';
 import { Loci, isEveryLoci } from 'mol-model/loci';
+import { Interval } from 'mol-data/int';
 
 export function createTransforms({ units }: Unit.SymmetryGroup, transforms?: ValueCell<Float32Array>) {
     const unitCount = units.length
@@ -101,11 +102,19 @@ export function markElement(tMarker: ValueCell<TextureImage>, group: Unit.Symmet
         for (const e of loci.elements) {
             const unitIdx = Unit.findUnitById(e.unit.id, group.units)
             if (unitIdx !== -1) {
-                for (let i = 0, il = e.indices.length; i < il; ++i) {
-                    const idx = unitIdx * elementCount + e.indices[i]
-                    if (applyMarkerAction(array, idx, idx + 1, action) && !changed) {
+                if (Interval.is(e.indices)) {
+                    const idxStart = unitIdx * elementCount + Interval.start(e.indices);
+                    const idxEnd = unitIdx * elementCount + Interval.end(e.indices);
+                    if (applyMarkerAction(array, idxStart, idxEnd, action) && !changed) {
                         changed = true
                     }
+                } else {
+                    for (let i = 0, _i = e.indices.length; i < _i; i++) {
+                        const idx = unitIdx * elementCount + e.indices[i];
+                        if (applyMarkerAction(array, idx, idx + 1, action) && !changed) {
+                            changed = true
+                        }
+                    }
                 }
             }
         }

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

@@ -7,7 +7,7 @@
 import { HashSet } from 'mol-data/generic'
 import { Structure, Element, Unit } from '../structure'
 import { structureUnion } from './utils/structure';
-import { SortedArray } from 'mol-data/int';
+import { OrderedSet, SortedArray } from 'mol-data/int';
 
 // A selection is a pair of a Structure and a sequence of unique AtomSets
 type Selection = Selection.Singletons | Selection.Sequence
@@ -36,10 +36,18 @@ namespace Selection {
     }
 
     export function toLoci(sel: Selection): Element.Loci {
-        const loci: { unit: Unit, indices: SortedArray }[] = [];
+        const loci: { unit: Unit, indices: OrderedSet }[] = [];
+        const { unitMap } = sel.source;
 
         for (const unit of unionStructure(sel).units) {
-            loci[loci.length] = { unit, indices: SortedArray.indicesOf(sel.source.unitMap.get(unit.id).elements, unit.elements) }
+            if (unit === unitMap.get(unit.id)) {
+                loci[loci.length] = { unit, indices: OrderedSet.ofBounds(0, unit.elements.length) };
+            } else {
+                loci[loci.length] = {
+                    unit,
+                    indices: OrderedSet.ofSortedArray(SortedArray.indicesOf(sel.source.unitMap.get(unit.id).elements, unit.elements))
+                };
+            }
         }
 
         return Element.Loci(loci);

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

@@ -4,7 +4,7 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import { Tuple, SortedArray } from 'mol-data/int'
+import { Tuple, OrderedSet } from 'mol-data/int'
 import Unit from './unit'
 import Structure from './structure'
 
@@ -55,12 +55,15 @@ namespace Element {
         /** Access i-th element as unit.elements[indices[i]] */
         readonly elements: ReadonlyArray<{
             unit: Unit,
-            /** Indices into the unit.elements array */
-            indices: SortedArray
+            /**
+             * Indices into the unit.elements array.
+             * Can use OrderedSet.forEach to iterate (or OrderedSet.size + OrderedSet.getAt)
+             */
+            indices: OrderedSet
         }>
     }
 
-    export function Loci(elements: ArrayLike<{ unit: Unit, indices: SortedArray }>): Loci {
+    export function Loci(elements: ArrayLike<{ unit: Unit, indices: OrderedSet }>): Loci {
         return { kind: 'element-loci', elements: elements as Loci['elements'] };
     }
 

src/mol-view/label.ts

@@ -7,6 +7,7 @@
 
 import { Unit, Element, Queries } from 'mol-model/structure';
 import { Loci } from 'mol-model/loci';
+import { OrderedSet } from 'mol-data/int';
 
 const elementLocA = Element.Location()
 const elementLocB = Element.Location()
@@ -20,8 +21,8 @@ export function labelFirst(loci: Loci): string {
     switch (loci.kind) {
         case 'element-loci':
             const e = loci.elements[0]
-            if (e && e.indices[0] !== undefined) {
-                return elementLabel(Element.Location(e.unit, e.indices[0]))
+            if (e) {
+                return elementLabel(Element.Location(e.unit, OrderedSet.getAt(e.indices, 0)))
             } else {
                 return 'Unknown'
             }