fix & simplify lociApply for instanceGranularity

src/mol-repr/structure/units-visual.ts

@@ -15,7 +15,7 @@ import { createUnitsTransform, includesUnitKind, StructureGroup } from './visual
 import { createRenderObject, GraphicsRenderObject, RenderObjectValues } from '../../mol-gl/render-object';
 import { PickingId } from '../../mol-geo/geometry/picking';
 import { Loci, isEveryLoci, EmptyLoci } from '../../mol-model/loci';
-import { Interval, OrderedSet } from '../../mol-data/int';
+import { Interval } from '../../mol-data/int';
 import { VisualUpdateState } from '../util';
 import { ColorTheme } from '../../mol-theme/color';
 import { createMarkers } from '../../mol-geo/geometry/marker-data';
@@ -267,48 +267,6 @@ export function UnitsVisual<G extends Geometry, P extends StructureParams & Geom
         return false;
     }
 
-    function structureInstanceLoci() {
-        const elements: StructureElement.Loci['elements'][0][] = [];
-        currentStructureGroup.group.units.forEach(unit => {
-            elements.push({
-                unit,
-                indices: OrderedSet.ofSingleton(0 as StructureElement.UnitIndex)
-            });
-        });
-        return StructureElement.Loci(currentStructureGroup.structure, elements);
-    }
-
-    function lociFirstElementOfInstances(loci: Loci) {
-        if (isEveryLoci(loci)) {
-            return structureInstanceLoci();
-        } else if (Structure.isLoci(loci)) {
-            if (!Structure.areRootsEquivalent(loci.structure, currentStructureGroup.structure)) {
-                return EmptyLoci;
-            }
-            return structureInstanceLoci();
-        } else if (StructureElement.Loci.is(loci)) {
-            if (!Structure.areRootsEquivalent(loci.structure, currentStructureGroup.structure)) {
-                return EmptyLoci;
-            }
-            if (StructureElement.Loci.isWholeStructure(loci)) {
-                return structureInstanceLoci();
-            }
-
-            const elements: StructureElement.Loci['elements'][0][] = [];
-            loci.elements.forEach(e => {
-                elements.push({
-                    unit: loci.elements[0].unit,
-                    indices: OrderedSet.ofSingleton(OrderedSet.start(loci.elements[0].indices))
-                });
-            });
-            return StructureElement.Loci(currentStructureGroup.structure, elements);
-        }
-
-        console.log('unexpected loci', loci);
-
-        return loci;
-    }
-
     function eachInstance(loci: Loci, structureGroup: StructureGroup, apply: (interval: Interval) => boolean) {
         let changed = false;
         if (!StructureElement.Loci.is(loci)) return false;
@@ -325,15 +283,18 @@ export function UnitsVisual<G extends Geometry, P extends StructureParams & Geom
     }
 
     function lociApply(loci: Loci, apply: (interval: Interval) => boolean, isMarking: boolean) {
-        if (currentProps.instanceGranularity) {
-            // TODO: change loci to include only the first element of each instance
-            loci = lociFirstElementOfInstances(loci);
-            return eachInstance(loci, currentStructureGroup, apply);
-        }
         if (lociIsSuperset(loci)) {
-            return apply(Interval.ofBounds(0, locationIt.groupCount * locationIt.instanceCount));
+            if (currentProps.instanceGranularity) {
+                return apply(Interval.ofBounds(0, locationIt.instanceCount));
+            } else {
+                return apply(Interval.ofBounds(0, locationIt.groupCount * locationIt.instanceCount));
+            }
         } else {
-            return eachLocation(loci, currentStructureGroup, apply, isMarking);
+            if (currentProps.instanceGranularity) {
+                return eachInstance(loci, currentStructureGroup, apply);
+            } else {
+                return eachLocation(loci, currentStructureGroup, apply, isMarking);
+            }
         }
     }