diff --git a/src/mol-base/collections/integer/impl/segmentation.ts b/src/mol-base/collections/integer/impl/segmentation.ts
index 567eee22ded8433c20cdb7d18d8ad3ca0205924b..357252a63c2d0f56f2711e188333bf9039dccbe6 100644
--- a/src/mol-base/collections/integer/impl/segmentation.ts
+++ b/src/mol-base/collections/integer/impl/segmentation.ts
@@ -48,7 +48,6 @@ export class SegmentIterator implements Iterator<Segs.Segment> {
     private segmentMax = 0;
     private setRange = Interval.Empty;
     private value: Segs.Segment = { index: 0, start: 0, end: 0 };
-    private last: number = 0;
 
     hasNext: boolean = false;
 
@@ -65,18 +64,13 @@ export class SegmentIterator implements Iterator<Segs.Segment> {
         return this.value;
     }
 
-    private getSegmentIndex(value: number) {
-        if (value >= this.last) return -1;
-        return SortedArray.findPredecessorIndex(this.segments, value + 1) - 1;
-    }
-
     private updateValue() {
         const segmentEnd = this.segments[this.segmentMin + 1];
+        // TODO: add optimized version for interval and array?
         const setEnd = OrderedSet.findPredecessorIndexInInterval(this.set, segmentEnd, this.setRange);
         this.value.start = Interval.start(this.setRange);
         this.value.end = setEnd;
-        const rEnd = Interval.end(this.setRange);
-        this.setRange = Interval.ofBounds(setEnd, rEnd);
+        this.setRange = Interval.ofBounds(setEnd, Interval.end(this.setRange));
         return setEnd > this.value.start;
     }
 
@@ -87,16 +81,18 @@ export class SegmentIterator implements Iterator<Segs.Segment> {
             return;
         }
 
-        const min = OrderedSet.getAt(this.set, sMin);
-        const max = OrderedSet.getAt(this.set, sMax);
-        this.segmentMin = this.getSegmentIndex(min);
-        this.segmentMax = this.getSegmentIndex(max);
+        this.segmentMin = this.segmentMap[OrderedSet.getAt(this.set, sMin)];
+        this.segmentMax = this.segmentMap[OrderedSet.getAt(this.set, sMax)];
 
-        this.hasNext = this.segmentMax >= this.segmentMin && Interval.size(this.setRange) > 0;
+        this.hasNext = this.segmentMax >= this.segmentMin;
+    }
+
+    setSegment(segment: Segs.Segment) {
+        this.setRange = Interval.ofBounds(segment.start, segment.end);
+        this.updateSegmentRange();
     }
 
-    constructor(private segments: SortedArray, private set: OrderedSet, inputRange: Interval) {
-        this.last = SortedArray.max(segments);
+    constructor(private segments: SortedArray, private segmentMap: Int32Array, private set: OrderedSet, inputRange: Interval) {
         this.setRange = inputRange;
         this.updateSegmentRange();
     }
@@ -104,5 +100,5 @@ export class SegmentIterator implements Iterator<Segs.Segment> {
 
 export function segments(segs: Segmentation, set: OrderedSet, segment?: Segs.Segment) {
     const int = typeof segment !== 'undefined' ? Interval.ofBounds(segment.start, segment.end) : Interval.ofBounds(0, OrderedSet.size(set));
-    return new SegmentIterator(segs.segments, set, int);
+    return new SegmentIterator(segs.segments, segs.segmentMap, set, int);
 }
\ No newline at end of file
diff --git a/src/mol-base/collections/integer/segmentation.ts b/src/mol-base/collections/integer/segmentation.ts
index b986fe401a3e9ee4dc1dc6075d0c7a2a65e47c93..3416c530bf91392db9ebf9444d44627f2d9d4970 100644
--- a/src/mol-base/collections/integer/segmentation.ts
+++ b/src/mol-base/collections/integer/segmentation.ts
@@ -4,7 +4,6 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import Iterator from '../iterator'
 import Interval from './interval'
 import OrderedSet from './ordered-set'
 import * as Impl from './impl/segmentation'
@@ -20,7 +19,7 @@ namespace Segmentation {
     export const projectValue: (segs: Segmentation, set: OrderedSet, value: number) => Interval = Impl.projectValue as any;
 
     // Segment iterator that mutates a single segment object to mark all the segments.
-    export const transientSegments: (segs: Segmentation, set: OrderedSet, segment?: Segment) => Iterator<Segment> = Impl.segments as any;
+    export const transientSegments: (segs: Segmentation, set: OrderedSet, segment?: Segment) => Impl.SegmentIterator = Impl.segments as any;
 }
 
 interface Segmentation {
diff --git a/src/mol-data/structure/query/generators.ts b/src/mol-data/structure/query/generators.ts
index eb6cb540346d750af3826f70c31d76fe3d32078b..9a30d0a3dd643ba7230c8eee198ee13386e50ea5 100644
--- a/src/mol-data/structure/query/generators.ts
+++ b/src/mol-data/structure/query/generators.ts
@@ -74,14 +74,14 @@ function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: A
 
             builder.beginUnit();
             const chainsIt = Segmentation.transientSegments(unit.hierarchy.chainSegments, set);
-            const residues = unit.hierarchy.residueSegments;
+            const residuesIt = Segmentation.transientSegments(unit.hierarchy.residueSegments, set);
             while (chainsIt.hasNext) {
                 const chainSegment = chainsIt.move();
                 l.atom = OrderedSet.getAt(set, chainSegment.start);
                 // test entity and chain
                 if (!entityTest(l) || !chainTest(l)) continue;
 
-                const residuesIt = Segmentation.transientSegments(residues, set, chainSegment);
+                residuesIt.setSegment(chainSegment);
                 while (residuesIt.hasNext) {
                     const residueSegment = residuesIt.move();
                     l.atom = OrderedSet.getAt(set, residueSegment.start);
diff --git a/src/perf-tests/structure.ts b/src/perf-tests/structure.ts
index 2ed618f6517f3923b30ec2bda5a339a0ec66c29c..a1d996834a8da7cdd11d19030045e2ca4d3fbb5a 100644
--- a/src/perf-tests/structure.ts
+++ b/src/perf-tests/structure.ts
@@ -235,8 +235,8 @@ export namespace PropertyAccess {
     // }
 
     export async function run() {
-        //const { structures, models } = await readCIF('./examples/1cbs_full.bcif');
-        const { structures, models } = await readCIF('e:/test/quick/1jj2_full.bcif');
+        const { structures, models } = await readCIF('./examples/1cbs_full.bcif');
+        //const { structures, models } = await readCIF('e:/test/quick/3j3q_full.bcif');
         //const { structures, models } = await readCIF('e:/test/quick/3j3q_updated.cif');
 
         console.log('parsed');