From f91dc5823d6766acd51119b72755d3d9e892fa99 Mon Sep 17 00:00:00 2001
From: Roman Dvornov <rdvornov@gmail.com>
Date: Tue, 10 Sep 2024 19:14:07 +0200
Subject: [PATCH] Fix suggestions for `and`, `or` and `??` operators

---
 CHANGELOG.md                   |  3 ++-
 src/lang/compile.js            |  8 ++++----
 src/lang/nodes/Binary.js       | 24 ++++++++++++++++++------
 src/lang/nodes/Block.js        |  8 +++-----
 src/lang/nodes/Compare.js      | 12 ++++--------
 src/lang/nodes/Conditional.js  | 16 +++++++---------
 src/lang/nodes/Filter.js       | 10 ++++------
 src/lang/nodes/Function.js     |  8 ++------
 src/lang/nodes/Map.js          | 10 ++++------
 src/lang/nodes/MapRecursive.js | 10 ++++------
 src/lang/nodes/Pipeline.js     | 13 ++++---------
 src/lang/nodes/Postfix.js      | 18 +++++-------------
 test/suggestions.js            | 12 ++++++++++++
 13 files changed, 73 insertions(+), 79 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a54f48c..d715679 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,7 +1,8 @@
 ## next
 
-- Improved suggestions for the ternary operator (`?:`)
 - Disallowed duplicate parameter names in function definitions, i.e., `($a, $a) => expr` is now prohibited
+- Improved suggestions for the ternary operator (`?:`)
+- Fixed suggestions for `and`, `or` and `??` operators
 - Fixed a compilation error for the generated code of `function` expressions in stat mode in some cases
 - Fixed a runtime error for a `comparator function` in some scenarios
 - Fixed support for TypedArrays with slice notation and methods such as `numbers()`, `avg()`, `count()`, `sum()`, `median()`, `variance()`, `stdev()` and `percentile()`
diff --git a/src/lang/compile.js b/src/lang/compile.js
index d210268..ec1d3c3 100644
--- a/src/lang/compile.js
+++ b/src/lang/compile.js
@@ -72,7 +72,7 @@ function compileInternal(ast, kind, tolerant = false, suggestions = null) {
 
     function createScope(fn, defCurrent, $ref = '$') {
         const prevScope = ctx.scope;
-        const scopeStart = buffer.length;
+        const scopeStart = buffer.length - 1;
 
         ctx.scope = prevScope.spawn(prevScope.arg1 || kind === 'method', $ref);
 
@@ -88,7 +88,7 @@ function compileInternal(ast, kind, tolerant = false, suggestions = null) {
             if (ctx.scope.firstCurrent) {
                 buffer[ctx.scope.firstCurrent] = stat;
             } else {
-                buffer[scopeStart] = defCurrent(buffer[scopeStart], stat);
+                buffer[scopeStart] += defCurrent(stat);
             }
         }
 
@@ -239,9 +239,9 @@ function compileInternal(ast, kind, tolerant = false, suggestions = null) {
 
     createScope(
         () => walk(ast),
-        (scopeStart, sp) => {
+        (sp) => {
             buffer.push(')');
-            return '(' + sp + ',' + scopeStart;
+            return '(' + sp + ',';
         },
         'data'
     );
diff --git a/src/lang/nodes/Binary.js b/src/lang/nodes/Binary.js
index d64e871..b7c2e6a 100644
--- a/src/lang/nodes/Binary.js
+++ b/src/lang/nodes/Binary.js
@@ -89,9 +89,14 @@ export function compile(node, ctx) {
             ctx.put(`${ctx.buildinFn('bool')}(${tmpVar}=`);
             ctx.node(node.left);
             ctx.put(`)?${tmpVar}:`);
-            ctx.scope.captureCurrent.disabled = true;
-            ctx.node(node.right);
-            ctx.scope.captureCurrent.disabled = false;
+            ctx.createScope(
+                () => ctx.node(node.right),
+                (sp) => {
+                    ctx.put(')');
+                    return '(' + sp + ',';
+                },
+                ctx.scope.$ref
+            );
             break;
         }
 
@@ -100,12 +105,19 @@ export function compile(node, ctx) {
         case '??': {
             const tmpVar = ctx.allocateVar();
 
+            // TODO: replace for Nullish coalescing operator (??) instead of ternary operator once drop support for Node.js below 14.0
+            // https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Operators/Nullish_coalescing
             ctx.put(`(${tmpVar}=`);
             ctx.node(node.left);
             ctx.put(`,${tmpVar}!==null&&${tmpVar}!==undefined)?${tmpVar}:`);
-            ctx.scope.captureCurrent.disabled = true;
-            ctx.node(node.right);
-            ctx.scope.captureCurrent.disabled = false;
+            ctx.createScope(
+                () => ctx.node(node.right),
+                (sp) => {
+                    ctx.put(')');
+                    return '(' + sp + ',';
+                },
+                ctx.scope.$ref
+            );
             break;
         }
 
diff --git a/src/lang/nodes/Block.js b/src/lang/nodes/Block.js
index adeea51..0685113 100644
--- a/src/lang/nodes/Block.js
+++ b/src/lang/nodes/Block.js
@@ -5,23 +5,21 @@ export function suggest(node, ctx) {
 }
 export function compile(node, ctx) {
     if (node.definitions.length) {
+        ctx.put('(()=>{');
         ctx.createScope(
             () => {
                 for (const definition of node.definitions) {
                     ctx.scope.awaitInit.add(definition.declarator.name);
                 }
 
-                ctx.put('(()=>{');
                 ctx.list(node.definitions);
                 ctx.put('return ');
                 ctx.nodeOrCurrent(node.body);
-                ctx.put('})()');
-            },
-            (scopeStart, sp) => {
-                return scopeStart + sp + ';';
             },
+            (sp) => sp + ';',
             ctx.scope.$ref
         );
+        ctx.put('})()');
     } else if (node.body && node.body.type === 'Object') {
         ctx.put('(');
         ctx.nodeOrCurrent(node.body);
diff --git a/src/lang/nodes/Compare.js b/src/lang/nodes/Compare.js
index 4fa3e6c..50cf4d3 100644
--- a/src/lang/nodes/Compare.js
+++ b/src/lang/nodes/Compare.js
@@ -15,16 +15,12 @@ export function compile(node, ctx) {
     }
 
     ctx.put(ctx.buildinFn(cmpFn));
+    ctx.put('((_q=$=>(');
     ctx.createScope(
-        () => {
-            ctx.put('((_q=$=>(');
-            ctx.node(node.query);
-            ctx.put('))(_a),_q(_b))');
-        },
-        (scopeStart, sp) => {
-            return scopeStart + sp + ',';
-        }
+        () => ctx.node(node.query),
+        (sp) => sp + ','
     );
+    ctx.put('))(_a),_q(_b))');
 }
 export function walk(node, ctx) {
     ctx.node(node.query);
diff --git a/src/lang/nodes/Conditional.js b/src/lang/nodes/Conditional.js
index f60d74e..7561ea9 100644
--- a/src/lang/nodes/Conditional.js
+++ b/src/lang/nodes/Conditional.js
@@ -2,20 +2,18 @@ export function compile(node, ctx) {
     ctx.put(ctx.buildinFn('bool'));
     ctx.put('(');
     ctx.nodeOrCurrent(node.test);
+    ctx.put(')?');
     ctx.createScope(
-        () => {
-            ctx.put(')?');
-            ctx.nodeOrCurrent(node.consequent);
-        },
-        (scopeStart, sp) => {
+        () => ctx.nodeOrCurrent(node.consequent),
+        (sp) => {
             ctx.put(')');
-            return scopeStart + '(' + sp + ',';
+            return '(' + sp + ',';
         },
         ctx.scope.$ref
     );
+    ctx.put(':');
     ctx.createScope(
         () => {
-            ctx.put(':');
             if (node.alternate) {
                 if (node.alternate.type === 'Placeholder') {
                     ctx.put('(');
@@ -28,9 +26,9 @@ export function compile(node, ctx) {
                 ctx.put('undefined');
             }
         },
-        (scopeStart, sp) => {
+        (sp) => {
             ctx.put(')');
-            return scopeStart + '(' + sp + ',';
+            return '(' + sp + ',';
         },
         ctx.scope.$ref
     );
diff --git a/src/lang/nodes/Filter.js b/src/lang/nodes/Filter.js
index 1a0124f..9080334 100644
--- a/src/lang/nodes/Filter.js
+++ b/src/lang/nodes/Filter.js
@@ -2,14 +2,12 @@ export function compile(node, ctx) {
     ctx.put(ctx.buildinFn('filter'));
     ctx.put('(');
     ctx.nodeOrCurrent(node.value);
+    ctx.put(',$=>');
     ctx.createScope(
-        () => {
-            ctx.put(',$=>');
-            ctx.node(node.query);
-        },
-        (scopeStart, sp) => {
+        () => ctx.node(node.query),
+        (sp) => {
             ctx.put(')');
-            return scopeStart + '(' + sp + ',';
+            return '(' + sp + ',';
         }
     );
     ctx.put(')');
diff --git a/src/lang/nodes/Function.js b/src/lang/nodes/Function.js
index f9eb036..091c26f 100644
--- a/src/lang/nodes/Function.js
+++ b/src/lang/nodes/Function.js
@@ -12,7 +12,7 @@ export function compile(node, ctx) {
     // but Function doesn't create 2nd argument implicitly to prevent function arity changes
     ctx.put('function(');
     ctx.put(String(args) || '$');
-
+    ctx.put('){return ');
     ctx.createScope(
         () => {
             ctx.scope.arg1 = true;
@@ -22,14 +22,10 @@ export function compile(node, ctx) {
                 ctx.scope.add(arg.name);
             }
 
-            ctx.put('){return ');
             ctx.node(node.body);
         },
-        (scopeStart, sp) => {
-            return scopeStart + sp + ',';
-        }
+        (sp) => sp + ','
     );
-
     ctx.put('}');
 }
 export function walk(node, ctx) {
diff --git a/src/lang/nodes/Map.js b/src/lang/nodes/Map.js
index ddffd9b..262bcde 100644
--- a/src/lang/nodes/Map.js
+++ b/src/lang/nodes/Map.js
@@ -2,14 +2,12 @@ export function compile(node, ctx) {
     ctx.put(ctx.buildinFn('map'));
     ctx.put('(');
     ctx.nodeOrCurrent(node.value);
+    ctx.put(',$=>');
     ctx.createScope(
-        () => {
-            ctx.put(',$=>');
-            ctx.node(node.query);
-        },
-        (scopeStart, sp) => {
+        () => ctx.node(node.query),
+        (sp) => {
             ctx.put(')');
-            return scopeStart + '(' + sp + ',';
+            return '(' + sp + ',';
         }
     );
     ctx.put(')');
diff --git a/src/lang/nodes/MapRecursive.js b/src/lang/nodes/MapRecursive.js
index 9189a1c..789a15a 100644
--- a/src/lang/nodes/MapRecursive.js
+++ b/src/lang/nodes/MapRecursive.js
@@ -2,14 +2,12 @@ export function compile(node, ctx) {
     ctx.put(ctx.buildinFn('mapRecursive'));
     ctx.put('(');
     ctx.nodeOrCurrent(node.value);
+    ctx.put(',$=>');
     ctx.createScope(
-        () => {
-            ctx.put(',$=>');
-            ctx.node(node.query);
-        },
-        (scopeStart, sp) => {
+        () => ctx.node(node.query),
+        (sp) => {
             ctx.put(')');
-            return scopeStart + '(' + sp + ',';
+            return '(' + sp + ',';
         }
     );
     ctx.put(')');
diff --git a/src/lang/nodes/Pipeline.js b/src/lang/nodes/Pipeline.js
index 0462d99..45eeb8c 100644
--- a/src/lang/nodes/Pipeline.js
+++ b/src/lang/nodes/Pipeline.js
@@ -1,15 +1,10 @@
 export function compile(node, ctx) {
+    ctx.put('($=>(');
     ctx.createScope(
-        () => {
-            ctx.put('($=>(');
-            ctx.node(node.right);
-            ctx.put('))');
-        },
-        (scopeStart, sp) => {
-            return scopeStart + sp + ',';
-        }
+        () => ctx.node(node.right),
+        (sp) => sp + ','
     );
-
+    ctx.put('))');
     ctx.put('(');
     ctx.nodeOrCurrent(node.left);
     ctx.put(')');
diff --git a/src/lang/nodes/Postfix.js b/src/lang/nodes/Postfix.js
index 4c78b5f..25eb59a 100644
--- a/src/lang/nodes/Postfix.js
+++ b/src/lang/nodes/Postfix.js
@@ -1,26 +1,18 @@
 export function compile(node, ctx) {
     if (node.operator && node.operator.type) {
+        ctx.put('($=>(');
         ctx.createScope(
-            () => {
-                ctx.put('($=>(');
-                ctx.node(node.operator);
-                ctx.put('))');
-            },
-            (scopeStart, sp) => {
-                return scopeStart + sp + ',';
-            }
+            () => ctx.node(node.operator),
+            (sp) => sp + ','
         );
+        ctx.put('))');
         ctx.put('(');
         ctx.node(node.argument);
         ctx.put(')');
         return;
     }
 
-    switch (node.operator) {
-        default: {
-            ctx.error('Unknown operator "' + node.operator + '"', node);
-        }
-    }
+    ctx.error('Unknown operator "' + node.operator + '"', node);
 }
 export function walk(node, ctx) {
     ctx.node(node.argument);
diff --git a/test/suggestions.js b/test/suggestions.js
index 178d53a..091670d 100644
--- a/test/suggestions.js
+++ b/test/suggestions.js
@@ -548,6 +548,18 @@ describe('query/suggestions', () => {
             suggestion('a', ['"foo":value', '"bar":value', 'foo', 'bar'], 3, 4),
             suggestion('a', ['"foo":value', '"bar":value', 'foo', 'bar'], 3, 4),
             null
+        ],
+        '[{a:1},{a:2,b:3}].((b?true:[]) and |x|)': [
+            suggestion('x', ['a', 'b'], 35, 36),
+            suggestion('x', ['a', 'b'], 35, 36)
+        ],
+        '[{a:1},{a:2,b:3}].((b?true:[]) or |x|)': [
+            suggestion('x', ['a'], 34, 35),
+            suggestion('x', ['a'], 34, 35)
+        ],
+        '[{a:1},{a:2,b:3}].(b ?? |x|)': [
+            suggestion('x', ['a'], 24, 25),
+            suggestion('x', ['a'], 24, 25)
         ]
     });
 });