From d2376b7ff90183d303527e11e37e75cffd7838d4 Mon Sep 17 00:00:00 2001
From: Vadim Kibana <82822460+vadimkibana@users.noreply.github.com>
Date: Tue, 22 Oct 2024 20:05:09 +0200
Subject: [PATCH] [ES|QL] Normalize multiplication by one when pretty-printing
 (#197182)

## Summary

Partially addresses https://github.com/elastic/kibana/issues/189258

This change will likely supercede
https://github.com/elastic/kibana/pull/196817

When parsing, currently ES|QL parsers adds extraneous multiply by 1 or
multiply by -1 nodes when parsing arithmetic unary expressions.

For example, `-(x)` is parsed as `-1 * x`.

This change, reverts these when pretty-printing using the
`BasicPrettyPrinter`: `-1 * x` is pretty printed as `-x`.

### Checklist

Delete any items that are not applicable to this PR.

- [x] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios

### For maintainers

- [x] This was checked for breaking API changes and was [labeled
appropriately](https://www.elastic.co/guide/en/kibana/master/contributing.html#_add_your_labels)

(cherry picked from commit 848c9f48dc101744c8fae69951916dc511dfdc8a)
---
 packages/kbn-esql-ast/src/ast/helpers.ts      | 27 ++++++-
 .../__tests__/basic_pretty_printer.test.ts    | 62 ++++++++++++++-
 .../src/pretty_print/basic_pretty_printer.ts  | 79 ++++++++++++++++++-
 .../kbn-esql-ast/src/pretty_print/helpers.ts  |  2 +-
 packages/kbn-esql-ast/src/visitor/utils.ts    |  4 +
 5 files changed, 169 insertions(+), 5 deletions(-)

diff --git a/packages/kbn-esql-ast/src/ast/helpers.ts b/packages/kbn-esql-ast/src/ast/helpers.ts
index 9ca49dcb38822..74a7b5c0991e8 100644
--- a/packages/kbn-esql-ast/src/ast/helpers.ts
+++ b/packages/kbn-esql-ast/src/ast/helpers.ts
@@ -7,11 +7,22 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { ESQLAstNode, ESQLBinaryExpression, ESQLFunction } from '../types';
+import type {
+  ESQLAstNode,
+  ESQLBinaryExpression,
+  ESQLColumn,
+  ESQLFunction,
+  ESQLIntegerLiteral,
+  ESQLLiteral,
+  ESQLProperNode,
+} from '../types';
 import { BinaryExpressionGroup } from './constants';
 
+export const isProperNode = (node: unknown): node is ESQLProperNode =>
+  !!node && typeof node === 'object' && !Array.isArray(node);
+
 export const isFunctionExpression = (node: unknown): node is ESQLFunction =>
-  !!node && typeof node === 'object' && !Array.isArray(node) && (node as any).type === 'function';
+  isProperNode(node) && node.type === 'function';
 
 /**
  * Returns true if the given node is a binary expression, i.e. an operator
@@ -28,6 +39,18 @@ export const isFunctionExpression = (node: unknown): node is ESQLFunction =>
 export const isBinaryExpression = (node: unknown): node is ESQLBinaryExpression =>
   isFunctionExpression(node) && node.subtype === 'binary-expression';
 
+export const isLiteral = (node: unknown): node is ESQLLiteral =>
+  isProperNode(node) && node.type === 'literal';
+
+export const isIntegerLiteral = (node: unknown): node is ESQLIntegerLiteral =>
+  isLiteral(node) && node.literalType === 'integer';
+
+export const isDoubleLiteral = (node: unknown): node is ESQLIntegerLiteral =>
+  isLiteral(node) && node.literalType === 'double';
+
+export const isColumn = (node: unknown): node is ESQLColumn =>
+  isProperNode(node) && node.type === 'column';
+
 /**
  * Returns the group of a binary expression:
  *
diff --git a/packages/kbn-esql-ast/src/pretty_print/__tests__/basic_pretty_printer.test.ts b/packages/kbn-esql-ast/src/pretty_print/__tests__/basic_pretty_printer.test.ts
index 20db9e729f094..9e21c45f75b4b 100644
--- a/packages/kbn-esql-ast/src/pretty_print/__tests__/basic_pretty_printer.test.ts
+++ b/packages/kbn-esql-ast/src/pretty_print/__tests__/basic_pretty_printer.test.ts
@@ -16,7 +16,7 @@ const reprint = (src: string) => {
   const { root } = parse(src);
   const text = BasicPrettyPrinter.print(root);
 
-  // console.log(JSON.stringify(ast, null, 2));
+  // console.log(JSON.stringify(root, null, 2));
 
   return { text };
 };
@@ -194,6 +194,66 @@ describe('single line query', () => {
 
           expect(text).toBe('ROW NOT a');
         });
+
+        test('negative numbers', () => {
+          const { text } = reprint('ROW -1');
+
+          expect(text).toBe('ROW -1');
+        });
+
+        test('negative numbers in brackets', () => {
+          const { text } = reprint('ROW -(1)');
+
+          expect(text).toBe('ROW -1');
+        });
+
+        test('negative column names', () => {
+          const { text } = reprint('ROW -col');
+
+          expect(text).toBe('ROW -col');
+        });
+
+        test('plus unary expression', () => {
+          const { text } = reprint('ROW +(23)');
+
+          expect(text).toBe('ROW 23');
+        });
+
+        test('chained multiple unary expressions', () => {
+          const { text } = reprint('ROW ----+-+(23)');
+
+          expect(text).toBe('ROW -23');
+        });
+
+        test('before another expression', () => {
+          const { text } = reprint('ROW ----+-+(1 + 1)');
+
+          expect(text).toBe('ROW -(1 + 1)');
+        });
+
+        test('negative one from the right side', () => {
+          const { text } = reprint('ROW 2 * -1');
+
+          expect(text).toBe('ROW -2');
+        });
+
+        test('two minuses is plus', () => {
+          const { text } = reprint('ROW --123');
+
+          expect(text).toBe('ROW 123');
+        });
+
+        test('two minuses is plus (float)', () => {
+          const { text } = reprint('ROW --1.23');
+
+          expect(text).toBe('ROW 1.23');
+        });
+
+        test('two minuses is plus (with brackets)', () => {
+          const { text } = reprint('ROW --(123)');
+
+          expect(text).toBe('ROW 123');
+        });
       });
 
       describe('postfix unary expression', () => {
diff --git a/packages/kbn-esql-ast/src/pretty_print/basic_pretty_printer.ts b/packages/kbn-esql-ast/src/pretty_print/basic_pretty_printer.ts
index ec744c65f636e..2f1e3439cd3a3 100644
--- a/packages/kbn-esql-ast/src/pretty_print/basic_pretty_printer.ts
+++ b/packages/kbn-esql-ast/src/pretty_print/basic_pretty_printer.ts
@@ -7,9 +7,18 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { binaryExpressionGroup } from '../ast/helpers';
+import {
+  binaryExpressionGroup,
+  isBinaryExpression,
+  isColumn,
+  isDoubleLiteral,
+  isIntegerLiteral,
+  isLiteral,
+  isProperNode,
+} from '../ast/helpers';
 import { ESQLAstBaseItem, ESQLAstCommand, ESQLAstQueryExpression } from '../types';
 import { ESQLAstExpressionNode, Visitor } from '../visitor';
+import { resolveItem } from '../visitor/utils';
 import { LeafPrinter } from './leaf_printer';
 
 export interface BasicPrettyPrinterOptions {
@@ -152,6 +161,62 @@ export class BasicPrettyPrinter {
     return formatted;
   }
 
+  protected simplifyMultiplicationByOne(
+    node: ESQLAstExpressionNode,
+    minusCount: number = 0
+  ): string | undefined {
+    if (isBinaryExpression(node) && node.name === '*') {
+      let [left, right] = node.args;
+      left = resolveItem(left);
+      right = resolveItem(right);
+
+      if (isProperNode(left) && isProperNode(right)) {
+        if (!!left.formatting || !!right.formatting) {
+          return undefined;
+        }
+        if (isIntegerLiteral(left)) {
+          if (left.value === 1) {
+            return this.simplifyMultiplicationByOne(right, minusCount);
+          } else if (left.value === -1) {
+            return this.simplifyMultiplicationByOne(right, minusCount + 1);
+          }
+        }
+        if (isIntegerLiteral(right)) {
+          if (right.value === 1) {
+            return this.simplifyMultiplicationByOne(left, minusCount);
+          } else if (right.value === -1) {
+            return this.simplifyMultiplicationByOne(left, minusCount + 1);
+          }
+        }
+        return undefined;
+      } else {
+        return undefined;
+      }
+    }
+
+    const isNegative = minusCount % 2 === 1;
+
+    if (isNegative && (isIntegerLiteral(node) || isDoubleLiteral(node)) && node.value < 0) {
+      return BasicPrettyPrinter.expression(
+        {
+          ...node,
+          value: Math.abs(node.value),
+        },
+        this.opts
+      );
+    }
+
+    let expression = BasicPrettyPrinter.expression(node, this.opts);
+    const sign = isNegative ? '-' : '';
+    const needsBrackets = !!sign && !isColumn(node) && !isLiteral(node);
+
+    if (needsBrackets) {
+      expression = `(${expression})`;
+    }
+
+    return sign ? `${sign}${expression}` : expression;
+  }
+
   protected readonly visitor: Visitor<any> = new Visitor()
     .on('visitExpression', (ctx) => {
       return '<EXPRESSION>';
@@ -237,6 +302,18 @@ export class BasicPrettyPrinter {
           const groupLeft = binaryExpressionGroup(left);
           const groupRight = binaryExpressionGroup(right);
 
+          if (
+            node.name === '*' &&
+            ((isIntegerLiteral(left) && Math.abs(left.value) === 1) ||
+              (isIntegerLiteral(right) && Math.abs(right.value) === 1))
+          ) {
+            const formatted = this.simplifyMultiplicationByOne(node);
+
+            if (formatted) {
+              return formatted;
+            }
+          }
+
           let leftFormatted = ctx.visitArgument(0);
           let rightFormatted = ctx.visitArgument(1);
 
diff --git a/packages/kbn-esql-ast/src/pretty_print/helpers.ts b/packages/kbn-esql-ast/src/pretty_print/helpers.ts
index f9d9daac84e7a..1b4a75a119cb2 100644
--- a/packages/kbn-esql-ast/src/pretty_print/helpers.ts
+++ b/packages/kbn-esql-ast/src/pretty_print/helpers.ts
@@ -7,7 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { ESQLAstBaseItem, ESQLProperNode } from '../types';
+import type { ESQLAstBaseItem, ESQLProperNode } from '../types';
 import { Walker } from '../walker';
 
 export interface QueryPrettyPrintStats {
diff --git a/packages/kbn-esql-ast/src/visitor/utils.ts b/packages/kbn-esql-ast/src/visitor/utils.ts
index 0dc95b73cf9d7..da8544ef46c90 100644
--- a/packages/kbn-esql-ast/src/visitor/utils.ts
+++ b/packages/kbn-esql-ast/src/visitor/utils.ts
@@ -36,6 +36,10 @@ export const firstItem = (items: ESQLAstItem[]): ESQLSingleAstItem | undefined =
   }
 };
 
+export const resolveItem = (items: ESQLAstItem | ESQLAstItem[]): ESQLAstItem => {
+  return Array.isArray(items) ? resolveItem(items[0]) : items;
+};
+
 /**
  * Returns the last normalized "single item" from the "item" list.
  *