[ES|QL] Fix functions highlight with ES grammar

packages/kbn-monaco/src/esql/lib/monaco/esql_theme.ts

@@ -77,25 +77,18 @@ export const buildESQlTheme = (): monaco.editor.IStandaloneThemeData => ({
         'row',
         'show',
         'limit',
-        'cidr_match',
         'nulls_ordering_direction',
         'nulls_ordering',
         'null',
-        'boolean_value',
-        'comparison_operator',
         'enrich',
         'on',
         'with',
       ],
       euiThemeVars.euiColorPrimaryText
     ),
 
-    // aggregation functions
-    ...buildRuleGroup(['unary_function'], euiThemeVars.euiColorPrimaryText),
-    // is null functions
-    ...buildRuleGroup(['where_functions'], euiThemeVars.euiColorPrimaryText),
-    // math functions
-    ...buildRuleGroup(['math_function'], euiThemeVars.euiColorPrimaryText),
+    // functions
+    ...buildRuleGroup(['functions'], euiThemeVars.euiColorPrimaryText),
 
     // operators
     ...buildRuleGroup(

packages/kbn-monaco/src/esql/lib/monaco/esql_token_helpers.ts

@@ -0,0 +1,33 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+import { monaco } from '../../../monaco_imports';
+import { nonNullable } from '../ast/ast_helpers';
+import { ESQL_TOKEN_POSTFIX } from '../constants';
+
+export function enrichTokensWithFunctionsMetadata(
+  tokens: monaco.languages.IToken[]
+): monaco.languages.IToken[] {
+  // need to trim spaces as "abs (arg)" is still valid as function
+  const myTokensWithoutSpaces = tokens.filter(
+    ({ scopes }) => scopes !== 'expr_ws' + ESQL_TOKEN_POSTFIX
+  );
+  // find out all unquoted_identifiers index
+  const possiblyFunctions = myTokensWithoutSpaces
+    .map((t, i) => (t.scopes === 'unquoted_identifier' + ESQL_TOKEN_POSTFIX ? i : undefined))
+    .filter(nonNullable);
+
+  // then check if the token next is an opening bracket
+  for (const index of possiblyFunctions) {
+    if (myTokensWithoutSpaces[index + 1]?.scopes === 'lp' + ESQL_TOKEN_POSTFIX) {
+      // set the custom "functions" token (only used in theming)
+      myTokensWithoutSpaces[index].scopes = 'functions' + ESQL_TOKEN_POSTFIX;
+    }
+  }
+  return [...tokens];
+}

packages/kbn-monaco/src/esql/lib/monaco/esql_tokens_provider.ts

@@ -16,6 +16,7 @@ import { ESQLState } from './esql_state';
 
 import { getLexer } from '../antlr_facade';
 import { ESQL_TOKEN_POSTFIX } from '../constants';
+import { enrichTokensWithFunctionsMetadata } from './esql_token_helpers';
 
 const EOF = -1;
 
@@ -69,6 +70,11 @@ export class ESQLTokensProvider implements monaco.languages.TokensProvider {
 
     myTokens.sort((a, b) => a.startIndex - b.startIndex);
 
-    return new ESQLLineTokens(myTokens, prevState.getLineNumber() + 1);
+    // special tratement for functions
+    // the previous custom Kibana grammar baked functions directly as tokens, so highlight was easier
+    // The ES grammar doesn't have the token concept of "function"
+    const tokensWithFunctions = enrichTokensWithFunctionsMetadata(myTokens);
+
+    return new ESQLLineTokens(tokensWithFunctions, prevState.getLineNumber() + 1);
   }
 }