Check nextEmbedded: '@pop' in action cases (issue#172061) (microsoft#172193)

Author: dmitriypereverza

src/vs/editor/standalone/common/monarch/monarchCommon.ts

@@ -80,6 +80,7 @@ export interface IAction {
 	// an action is either a group of actions
 	group?: FuzzyAction[];
 
+	hasEmbeddedEndInCases?: boolean;
 	// or a function that returns a fresh action
 	test?: (id: string, matches: string[], state: string, eos: boolean) => FuzzyAction;
 

src/vs/editor/standalone/common/monarch/monarchCompile.ts

@@ -8,6 +8,7 @@
  * into a typed and checked ILexer definition.
  */
 
+import { isString } from '../../../../base/common/types.js';
 import * as monarchCommon from './monarchCommon.js';
 import { IMonarchLanguage, IMonarchLanguageBracket } from './monarchTypes.js';
 
@@ -337,6 +338,7 @@ function compileAction(lexer: monarchCommon.ILexerMin, ruleName: string, action:
 		// build an array of test cases
 		const cases: monarchCommon.IBranch[] = [];
 
+		let hasEmbeddedEndInCases = false;
 		// for each case, push a test function and result value
 		for (const tkey in action.cases) {
 			if (action.cases.hasOwnProperty(tkey)) {
@@ -352,12 +354,17 @@ function compileAction(lexer: monarchCommon.ILexerMin, ruleName: string, action:
 				else {
 					cases.push(createGuard(lexer, ruleName, tkey, val));  // call separate function to avoid local variable capture
 				}
+
+				if (!hasEmbeddedEndInCases) {
+					hasEmbeddedEndInCases = !isString(val) && (val.hasEmbeddedEndInCases || ['@pop', '@popall'].includes(val.nextEmbedded || ''));
+				}
 			}
 		}
 
 		// create a matching function
 		const def = lexer.defaultToken;
 		return {
+			hasEmbeddedEndInCases,
 			test: function (id, matches, state, eos) {
 				for (const _case of cases) {
 					const didmatch = (!_case.test || _case.test(id, matches, state, eos));

src/vs/editor/standalone/common/monarch/monarchLexer.ts

@@ -514,7 +514,7 @@ export class MonarchTokenizer extends Disposable implements languages.ITokenizat
 		let hasEmbeddedPopRule = false;
 
 		for (const rule of rules) {
-			if (!monarchCommon.isIAction(rule.action) || rule.action.nextEmbedded !== '@pop') {
+			if (!monarchCommon.isIAction(rule.action) || !(rule.action.nextEmbedded === '@pop' || rule.action.hasEmbeddedEndInCases)) {
 				continue;
 			}
 			hasEmbeddedPopRule = true;

src/vs/editor/standalone/test/browser/monarch.test.ts

@@ -112,6 +112,93 @@ suite('Monarch', () => {
 		disposables.dispose();
 	});
 
+	test('Test nextEmbedded: "@pop" in cases statement', () => {
+		const disposables = new DisposableStore();
+		const languageService = disposables.add(new LanguageService());
+		const configurationService = new StandaloneConfigurationService(new NullLogService());
+		disposables.add(languageService.registerLanguage({ id: 'sql' }));
+		disposables.add(TokenizationRegistry.register('sql', disposables.add(createMonarchTokenizer(languageService, 'sql', {
+			tokenizer: {
+				root: [
+					[/./, 'token']
+				]
+			}
+		}, configurationService))));
+		const SQL_QUERY_START = '(SELECT|INSERT|UPDATE|DELETE|CREATE|REPLACE|ALTER|WITH)';
+		const tokenizer = disposables.add(createMonarchTokenizer(languageService, 'test1', {
+			tokenizer: {
+				root: [
+					[`(\"\"\")${SQL_QUERY_START}`, [{ 'token': 'string.quote', }, { token: '@rematch', next: '@endStringWithSQL', nextEmbedded: 'sql', },]],
+					[/(""")$/, [{ token: 'string.quote', next: '@maybeStringIsSQL', },]],
+				],
+				maybeStringIsSQL: [
+					[/(.*)/, {
+						cases: {
+							[`${SQL_QUERY_START}\\b.*`]: { token: '@rematch', next: '@endStringWithSQL', nextEmbedded: 'sql', },
+							'@default': { token: '@rematch', switchTo: '@endDblDocString', },
+						}
+					}],
+				],
+				endDblDocString: [
+					['[^\']+', 'string'],
+					['\\\\\'', 'string'],
+					['\'\'\'', 'string', '@popall'],
+					['\'', 'string']
+				],
+				endStringWithSQL: [[/"""/, {
+					cases: {
+						'"""': {
+							cases: {
+								'': { token: 'string.quote', next: '@popall', nextEmbedded: '@pop', }
+							}
+						},
+						'@default': ''
+					}
+				}]],
+			}
+		}, configurationService));
+
+		const lines = [
+			`mysql_query("""SELECT * FROM table_name WHERE ds = '<DATEID>'""")`,
+			`mysql_query("""`,
+			`SELECT *`,
+			`FROM table_name`,
+			`WHERE ds = '<DATEID>'`,
+			`""")`,
+		];
+
+		const actualTokens = getTokens(tokenizer, lines);
+
+		assert.deepStrictEqual(actualTokens, [
+			[
+				new Token(0, 'source.test1', 'test1'),
+				new Token(12, 'string.quote.test1', 'test1'),
+				new Token(15, 'token.sql', 'sql'),
+				new Token(61, 'string.quote.test1', 'test1'),
+				new Token(64, 'source.test1', 'test1')
+			],
+			[
+				new Token(0, 'source.test1', 'test1'),
+				new Token(12, 'string.quote.test1', 'test1')
+			],
+			[
+				new Token(0, 'token.sql', 'sql')
+			],
+			[
+				new Token(0, 'token.sql', 'sql')
+			],
+			[
+				new Token(0, 'token.sql', 'sql')
+			],
+			[
+				new Token(0, 'string.quote.test1', 'test1'),
+				new Token(3, 'source.test1', 'test1')
+			]
+		]);
+		disposables.dispose();
+	});
+
+
 	test('microsoft/monaco-editor#1235: Empty Line Handling', () => {
 		const disposables = new DisposableStore();
 		const configurationService = new StandaloneConfigurationService(new NullLogService());