Fix input handling for (function() => { ... (#7)

The implementation assumed that in cases like the above input,
everything before the `()` would be highlighted the same regardless
of what follows afterwards, however, for `function`, the parentheses
have actual meaning. Therefore, only remove `()` for code highlighting
input size reduction if they are:
- Not the last match in the code, or
- Followed by something that definitely terminates the input block,
  i.e. a closing bracket, quote, etc.

Fixes: https://jira.mongodb.org/browse/MONGOSH-759

Author: addaleax

lib/pretty-repl.js

@@ -177,7 +177,26 @@ class PrettyREPLServer extends repl.REPLServer {
     // - matching pairs of "", with only non-"\, \\, and \" preceded by an even number of \ in between
     // - matching pairs of ``, with only non-`{}\, \\ and \` preceded by an even number of \ in between
     const re = /\([^\(\)\[\]\{\}`'"]*\)|\[[^\(\)\[\]\{\}`'"]*\]|\{[^\(\)\[\]\{\}`'"]*\}|'([^'\\]|(?<=[^\\](\\\\)*)\\'|\\\\)*'|"([^"\\]|(?<=[^\\](\\\\)*)\\"|\\\\)*"|`([^\{\}`\\]|(?<=[^\\](\\\\)*)\\`|\\\\)*`/g;
-    return str.replace(re, '');
+    // Match the regexp against the input. If there are no matches, we can just return.
+    const matches = [...str.matchAll(re)];
+    if (matches.length === 0) {
+      return str;
+    }
+    // Remove all but the last, non-nested pair of (), because () can affect
+    // whether the previous word is seen as a keyword.
+    // E.g.: When input is `function() {`, do not replace the ().
+    //       When input is `{ foo(); }`, do replace the `()`, then afterwards the `{ ... }`.
+    let startsReplaceIndex = matches.length - 1;
+    const lastMatch = matches[matches.length - 1];
+    if (lastMatch[0].startsWith('(') && !str.substr(lastMatch.index + lastMatch[0].length).match(/[\)\]\}`'"]/)) {
+      startsReplaceIndex--;
+    }
+    for (let i = startsReplaceIndex; i >= 0; i--) {
+      // Replace str1<match>str2 with str1str2. Go backwards so that the match
+      // indices into the string remain valid.
+      str = str.substr(0, matches[i].index) + str.substr(matches[i].index + matches[i][0].length);
+    }
+    return str;
   });
 }
 

test/index.spec.js

@@ -104,7 +104,7 @@ test('memoizeStringTransformerMethod', t => {
 });
 
 test('stripCompleteJSStructures', t => {
-  t.plan(6);
+  t.plan(10);
   const { stdin } = stdio();
   const output = new PassThrough();
   output.isTTY = true;
@@ -119,6 +119,10 @@ test('stripCompleteJSStructures', t => {
   t.equal(prettyRepl._stripCompleteJSStructures(String.raw `"abc\"def"`), '');
   t.equal(prettyRepl._stripCompleteJSStructures(String.raw `"abc\\"def"`), 'def"');
   t.equal(prettyRepl._stripCompleteJSStructures(String.raw `"a\\\\bc\\"def"`), 'def"');
+  t.equal(prettyRepl._stripCompleteJSStructures('(function {}'), '(function ');
+  t.equal(prettyRepl._stripCompleteJSStructures('(function() {'), '(function() {');
+  t.equal(prettyRepl._stripCompleteJSStructures('(function() => {'), '(function() => {');
+  t.equal(prettyRepl._stripCompleteJSStructures('(function() => {}'), '(function => ');
 });
 
 test('full pass-through test', t => {