Recover from mismatched loop terminators with quick fixes

src/DiagnosticMessages.ts

@@ -775,7 +775,24 @@ export let DiagnosticMessages = {
         message: `'${name}' is a reserved builtin and can only be used as a function call (e.g. '${name}(...)')`,
         code: 1147,
         severity: DiagnosticSeverity.Error
-    })
+    }),
+    /**
+     * Emitted when a block recovers from a wrong terminator (e.g. `while ... next` or `for ... end while`).
+     * `expected` lists the legal terminators in preferred-first order; `found` is the actual text.
+     * Quick fixes consume the structured `data` to build "Convert '<found>' to '<expected[i]>'" actions.
+     */
+    mismatchedEndingToken: (expected: string[] = [], found = '') => {
+        const expectedList = Array.isArray(expected) ? expected : [];
+        return {
+            message: `Expected ${expectedList.map(text => `'${text}'`).join(' or ')} but found '${found}'`,
+            code: 1148,
+            data: {
+                expected: expectedList,
+                found: found
+            },
+            severity: DiagnosticSeverity.Error
+        };
+    }
 };
 
 export const DiagnosticCodeMap = {} as Record<keyof (typeof DiagnosticMessages), number>;
@@ -800,3 +817,14 @@ export type DiagnosticMessageType<K extends keyof D, D extends Record<string, (.
     ReturnType<D[K]> &
     //include the missing properties from BsDiagnostic
     Pick<BsDiagnostic, 'range' | 'file' | 'relatedInformation' | 'tags'>;
+
+/**
+ * Refines a diagnostic to its concrete `DiagnosticMessageType<K>` shape (including the typed `data`
+ * payload) when its code matches `DiagnosticCodeMap[key]`.
+ */
+export function isDiagnosticOfType<K extends keyof typeof DiagnosticMessages>(
+    diagnostic: { code?: number | string },
+    key: K
+): diagnostic is DiagnosticMessageType<K> {
+    return diagnostic.code === DiagnosticCodeMap[key];
+}

src/bscPlugin/codeActions/CodeActionsProcessor.spec.ts

@@ -1013,6 +1013,146 @@ describe('CodeActionsProcessor', () => {
         });
     });
 
+    describe('mismatched loop terminator', () => {
+        it('offers a single quick fix for `while ... next`', () => {
+            const file = program.setFile('source/main.bs', `
+                sub a()
+                    while n <= 3
+                        n = n + 1
+                    next
+                end sub
+            `);
+            // cursor on the bogus `next` (line 4)
+            testGetCodeActions(file, util.createRange(4, 20, 4, 20), [`Convert 'next' to 'end while'`]);
+        });
+
+        it('replaces only the bogus `next` token, preserving indentation', () => {
+            const file = program.setFile('source/main.bs', `
+                sub a()
+                    while n <= 3
+                        n = n + 1
+                    next
+                end sub
+            `);
+            program.validate();
+            const actions = program.getCodeActions(file.srcPath, util.createRange(4, 20, 4, 20));
+            const fix = actions.find(a => a.title === `Convert 'next' to 'end while'`);
+            const changes = Object.values(fix!.edit!.changes!)[0];
+            expect(changes).to.be.lengthOf(1);
+            expect(changes[0].newText).to.equal('end while');
+            //the replace range is exactly the 4-character `next` token
+            expect(changes[0].range.end.character - changes[0].range.start.character).to.equal(4);
+        });
+
+        it('marks the `end while` quick fix as preferred for `while ... next`', () => {
+            const file = program.setFile('source/main.bs', `
+                sub a()
+                    while n <= 3
+                        n = n + 1
+                    next
+                end sub
+            `);
+            program.validate();
+            const actions = program.getCodeActions(file.srcPath, util.createRange(4, 20, 4, 20));
+            const fix = actions.find(a => a.title === `Convert 'next' to 'end while'`);
+            expect(fix!.isPreferred).to.equal(true);
+        });
+
+        it('offers a fix-all for multiple `while ... next` instances in the same file', () => {
+            const file = program.setFile('source/main.bs', `
+                sub a()
+                    while x
+                        x = false
+                    next
+                    while y
+                        y = false
+                    next
+                end sub
+            `);
+            //cursor on the first bogus `next`
+            testGetCodeActions(file, util.createRange(4, 20, 4, 20), [
+                `Convert 'next' to 'end while'`,
+                `Fix all: Convert 'next' to 'end while'`
+            ]);
+        });
+
+        it('offers two quick fixes (end for preferred, next) for `for ... end while`', () => {
+            const file = program.setFile('source/main.bs', `
+                sub a()
+                    for i = 0 to 3
+                        print i
+                    end while
+                end sub
+            `);
+            testGetCodeActions(file, util.createRange(4, 20, 4, 20), [
+                `Convert 'end while' to 'end for'`,
+                `Convert 'end while' to 'next'`
+            ]);
+        });
+
+        it('marks the `end for` quick fix as preferred and `next` as non-preferred', () => {
+            const file = program.setFile('source/main.bs', `
+                sub a()
+                    for i = 0 to 3
+                        print i
+                    end while
+                end sub
+            `);
+            program.validate();
+            const actions = program.getCodeActions(file.srcPath, util.createRange(4, 20, 4, 20));
+            const endForFix = actions.find(a => a.title === `Convert 'end while' to 'end for'`);
+            const nextFix = actions.find(a => a.title === `Convert 'end while' to 'next'`);
+            expect(endForFix!.isPreferred).to.equal(true);
+            expect(nextFix!.isPreferred).to.not.equal(true);
+        });
+
+        it('offers two quick fixes for `for each ... end while`', () => {
+            const file = program.setFile('source/main.bs', `
+                sub a()
+                    for each x in arr
+                        print x
+                    end while
+                end sub
+            `);
+            testGetCodeActions(file, util.createRange(4, 20, 4, 20), [
+                `Convert 'end while' to 'end for'`,
+                `Convert 'end while' to 'next'`
+            ]);
+        });
+
+        it('offers fix-all variants for both end for and next when there are multiple `for ... end while` instances', () => {
+            const file = program.setFile('source/main.bs', `
+                sub a()
+                    for i = 0 to 3
+                        print i
+                    end while
+                    for j = 0 to 3
+                        print j
+                    end while
+                end sub
+            `);
+            //cursor on the first bogus `end while`
+            testGetCodeActions(file, util.createRange(4, 20, 4, 20), [
+                `Convert 'end while' to 'end for'`,
+                `Convert 'end while' to 'next'`,
+                `Fix all: Convert 'end while' to 'end for'`,
+                `Fix all: Convert 'end while' to 'next'`
+            ]);
+        });
+
+        it('does not offer the quick fix for a valid `for ... next`', () => {
+            const file = program.setFile('source/main.bs', `
+                sub a()
+                    for i = 0 to 3
+                        print i
+                    next
+                end sub
+            `);
+            //cursor on the valid `next`
+            testGetCodeActions(file, util.createRange(4, 20, 4, 20), []);
+        });
+    });
+
     describe('disable diagnostic actions', () => {
         const findLineAction = (actions: CodeAction[], code: string | number) => actions.find(a => a.title.startsWith(`Disable ${code} for this line`));
         const findFileAction = (actions: CodeAction[], code: string | number) => actions.find(a => a.title.startsWith(`Disable ${code} for this file`));

src/bscPlugin/codeActions/CodeActionsProcessor.ts

@@ -3,7 +3,7 @@ import { CodeActionKind } from 'vscode-languageserver';
 import { codeActionUtil } from '../../CodeActionUtil';
 import type { DeleteChange, InsertChange, ReplaceChange } from '../../CodeActionUtil';
 import type { DiagnosticMessageType } from '../../DiagnosticMessages';
-import { DiagnosticCodeMap } from '../../DiagnosticMessages';
+import { DiagnosticCodeMap, isDiagnosticOfType } from '../../DiagnosticMessages';
 import type { BrsFile } from '../../files/BrsFile';
 import type { XmlFile } from '../../files/XmlFile';
 import type { BscFile, BsDiagnostic, OnGetCodeActionsEvent } from '../../interfaces';
@@ -53,6 +53,8 @@ export class CodeActionsProcessor {
                 this.suggestMissingOverrideQuickFixes([diagnostic]);
             } else if (diagnostic.code === DiagnosticCodeMap.cannotUseOverrideKeywordOnConstructorFunction) {
                 this.suggestRemoveOverrideFromConstructorQuickFixes([diagnostic]);
+            } else if (isDiagnosticOfType(diagnostic, 'mismatchedEndingToken')) {
+                this.suggestMismatchedEndingTokenQuickFixes([diagnostic]);
             }
         }
 
@@ -80,6 +82,8 @@ export class CodeActionsProcessor {
                     this.suggestScriptImportCasingQuickFixes(allInFile as DiagnosticMessageType<'scriptImportCaseMismatch'>[]);
                 } else if (code === DiagnosticCodeMap.missingOverrideKeyword) {
                     this.suggestMissingOverrideQuickFixes(allInFile);
+                } else if (code === DiagnosticCodeMap.mismatchedEndingToken) {
+                    this.suggestMismatchedEndingTokenQuickFixes(allInFile as DiagnosticMessageType<'mismatchedEndingToken'>[]);
                 }
             }
         }
@@ -766,6 +770,30 @@ export class CodeActionsProcessor {
         );
     }
 
+    /**
+     * Adds one code action per legal terminator. The first entry of `expected` is marked
+     * `isPreferred`, matching the parser's convention of listing the canonical terminator first.
+     */
+    private suggestMismatchedEndingTokenQuickFixes(diagnostics: DiagnosticMessageType<'mismatchedEndingToken'>[]) {
+        const { expected, found } = diagnostics[0].data;
+        for (let index = 0; index < expected.length; index++) {
+            const replacement = expected[index];
+            const changes = diagnostics.map<ReplaceChange>(diagnostic => ({
+                type: 'replace',
+                filePath: this.event.file.srcPath,
+                range: diagnostic.range,
+                newText: replacement
+            }));
+            this.emitOrFixAll(
+                `Convert '${found}' to '${replacement}'`,
+                `Fix all: Convert '${found}' to '${replacement}'`,
+                changes,
+                diagnostics[0],
+                index === 0
+            );
+        }
+    }
+
     /**
      * Adds code actions to remove the invalid `override` keyword from a constructor method.
      */
@@ -835,7 +863,8 @@ export class CodeActionsProcessor {
         singleTitle: string,
         fixAllTitle: string,
         changes: Array<InsertChange | DeleteChange | ReplaceChange>,
-        diagnostic: Diagnostic
+        diagnostic: Diagnostic,
+        isPreferred?: boolean
     ) {
         if (changes.length === 0) {
             return;
@@ -845,6 +874,7 @@ export class CodeActionsProcessor {
                 codeActionUtil.createCodeAction({
                     title: singleTitle,
                     diagnostics: [diagnostic],
+                    ...(isPreferred ? { isPreferred: true } : {}),
                     kind: CodeActionKind.QuickFix,
                     changes: changes
                 })
@@ -853,6 +883,7 @@ export class CodeActionsProcessor {
             this.event.codeActions.push(
                 codeActionUtil.createCodeAction({
                     title: fixAllTitle,
+                    ...(isPreferred ? { isPreferred: true } : {}),
                     kind: CodeActionKind.QuickFix,
                     changes: changes
                 })

src/parser/Parser.ts

@@ -1316,18 +1316,26 @@ export class Parser {
 
         this.consumeStatementSeparators();
 
-        const whileBlock = this.block(TokenKind.EndWhile);
+        const whileBlock = this.block(TokenKind.EndWhile, TokenKind.Next);
         let endWhile: Token;
-        if (!whileBlock || this.peek().kind !== TokenKind.EndWhile) {
+        if (whileBlock && this.peek().kind === TokenKind.EndWhile) {
+            endWhile = this.advance();
+        } else if (whileBlock && this.peek().kind === TokenKind.Next) {
+            //recover: a stray `next` is a common mistake when the user means `end while`.
+            //emit a targeted diagnostic and consume the `next` so the rest of the file parses cleanly.
+            this.diagnostics.push({
+                ...DiagnosticMessages.mismatchedEndingToken(['end while'], 'next'),
+                range: this.peek().range
+            });
+            endWhile = this.advance();
+        } else {
             this.diagnostics.push({
                 ...DiagnosticMessages.couldNotFindMatchingEndKeyword('while'),
                 range: this.peek().range
             });
             if (!whileBlock) {
                 throw this.lastDiagnosticAsError();
             }
-        } else {
-            endWhile = this.advance();
         }
 
         return new WhileStatement(
@@ -1363,18 +1371,25 @@ export class Parser {
 
         this.consumeStatementSeparators();
 
-        let body = this.block(TokenKind.EndFor, TokenKind.Next);
+        let body = this.block(TokenKind.EndFor, TokenKind.Next, TokenKind.EndWhile);
         let endForToken: Token;
-        if (!body || !this.checkAny(TokenKind.EndFor, TokenKind.Next)) {
+        if (body && this.checkAny(TokenKind.EndFor, TokenKind.Next)) {
+            endForToken = this.advance();
+        } else if (body && this.peek().kind === TokenKind.EndWhile) {
+            //recover: a stray `end while` is a common mistake when the user means `end for`.
+            this.diagnostics.push({
+                ...DiagnosticMessages.mismatchedEndingToken(['end for', 'next'], 'end while'),
+                range: this.peek().range
+            });
+            endForToken = this.advance();
+        } else {
             this.diagnostics.push({
                 ...DiagnosticMessages.expectedEndForOrNextToTerminateForLoop(),
                 range: this.peek().range
             });
             if (!body) {
                 throw this.lastDiagnosticAsError();
             }
-        } else {
-            endForToken = this.advance();
         }
 
         // WARNING: BrightScript doesn't delete the loop initial value after a for/to loop! It just
@@ -1424,17 +1439,25 @@ export class Parser {
 
         this.consumeStatementSeparators();
 
-        let body = this.block(TokenKind.EndFor, TokenKind.Next);
-        if (!body) {
+        let body = this.block(TokenKind.EndFor, TokenKind.Next, TokenKind.EndWhile);
+        let endFor: Token;
+        if (body && this.checkAny(TokenKind.EndFor, TokenKind.Next)) {
+            endFor = this.advance();
+        } else if (body && this.peek().kind === TokenKind.EndWhile) {
+            //recover: a stray `end while` is a common mistake when the user means `end for`.
+            this.diagnostics.push({
+                ...DiagnosticMessages.mismatchedEndingToken(['end for', 'next'], 'end while'),
+                range: this.peek().range
+            });
+            endFor = this.advance();
+        } else {
             this.diagnostics.push({
                 ...DiagnosticMessages.expectedEndForOrNextToTerminateForLoop(),
                 range: this.peek().range
             });
             throw this.lastDiagnosticAsError();
         }
 
-        let endFor = this.advance();
-
         return new ForEachStatement(
             {
                 forEach: forEach,