Added last yielded index

packages/cursorless-engine/src/processTargets/modifiers/scopeHandlers/IteratorInfo.ts

@@ -5,6 +5,7 @@
 interface IteratorInfo<T> {
   iterator: Iterator<T>;
   value: T;
+  index: number;
 }
 
 /**
@@ -19,14 +20,15 @@ interface IteratorInfo<T> {
 export function getInitialIteratorInfos<T>(
   iterators: Iterator<T>[],
 ): IteratorInfo<T>[] {
-  return iterators.flatMap((iterator) => {
+  return iterators.flatMap((iterator, i) => {
     const { value, done } = iterator.next();
     return done
       ? []
       : [
           {
             iterator,
             value,
+            index: i,
           },
         ];
   });
@@ -47,10 +49,10 @@ export function advanceIteratorsUntil<T>(
   criterion: (arg: T) => boolean,
 ): IteratorInfo<T>[] {
   return iteratorInfos.flatMap((iteratorInfo) => {
-    const { iterator } = iteratorInfo;
+    const { iterator, index } = iteratorInfo;
     let { value } = iteratorInfo;
-
     let done: boolean | undefined = false;
+
     while (!done && !criterion(value)) {
       ({ value, done } = iterator.next());
     }
@@ -59,6 +61,6 @@ export function advanceIteratorsUntil<T>(
       return [];
     }
 
-    return [{ iterator, value }];
+    return [{ iterator, value, index }];
   });
 }

packages/cursorless-engine/src/processTargets/modifiers/scopeHandlers/OneOfScopeHandler.ts

@@ -17,6 +17,8 @@ import type {
 
 export class OneOfScopeHandler extends BaseScopeHandler {
   protected isHierarchical = true;
+  private iterationScopeHandler: OneOfScopeHandler | undefined;
+  private lastYieldedIndex: number | undefined;
 
   static create(
     scopeHandlerFactory: ScopeHandlerFactory,
@@ -47,9 +49,8 @@ export class OneOfScopeHandler extends BaseScopeHandler {
     scopeHandlers: ScopeHandler[],
     languageId: string,
   ): ScopeHandler {
-    const iterationScopeType = (): CustomScopeType => ({
-      type: "custom",
-      scopeHandler: new OneOfScopeHandler(
+    const getIterationScopeHandler = () =>
+      new OneOfScopeHandler(
         undefined,
         scopeHandlers.map(
           (scopeHandler) =>
@@ -61,20 +62,29 @@ export class OneOfScopeHandler extends BaseScopeHandler {
         () => {
           throw new Error("Not implemented");
         },
-      ),
-    });
+      );
 
-    return new OneOfScopeHandler(scopeType, scopeHandlers, iterationScopeType);
+    return new OneOfScopeHandler(
+      scopeType,
+      scopeHandlers,
+      getIterationScopeHandler,
+    );
   }
 
   get iterationScopeType(): CustomScopeType {
-    return this.getIterationScopeType();
+    if (this.iterationScopeHandler == null) {
+      this.iterationScopeHandler = this.getIterationScopeHandler();
+    }
+    return {
+      type: "custom",
+      scopeHandler: this.iterationScopeHandler,
+    };
   }
 
   private constructor(
     public readonly scopeType: OneOfScopeType | undefined,
     private scopeHandlers: ScopeHandler[],
-    private getIterationScopeType: () => CustomScopeType,
+    private getIterationScopeHandler: () => OneOfScopeHandler,
   ) {
     super();
   }
@@ -85,6 +95,14 @@ export class OneOfScopeHandler extends BaseScopeHandler {
     direction: Direction,
     hints: ScopeIteratorRequirements,
   ): Iterable<TargetScope> {
+    // If we have used the iteration scope handler we only want to yield from it's handler
+    if (this.iterationScopeHandler?.lastYieldedIndex != null) {
+      const handlerIndex = this.iterationScopeHandler.lastYieldedIndex;
+      const handler = this.scopeHandlers[handlerIndex];
+      yield* handler.generateScopes(editor, position, direction, hints);
+      return;
+    }
+
     const iterators = this.scopeHandlers.map((scopeHandler) =>
       scopeHandler
         .generateScopes(editor, position, direction, hints)
@@ -99,7 +117,9 @@ export class OneOfScopeHandler extends BaseScopeHandler {
       );
 
       // Pick minimum scope according to canonical scope ordering
-      const currentScope = iteratorInfos[0].value;
+      const iteratorInfo = iteratorInfos[0];
+      const currentScope = iteratorInfo.value;
+      this.lastYieldedIndex = iteratorInfo.index;
 
       yield currentScope;