feat: add support for bracket matching/colorization MONGOSH-791 (#10)

Author: addaleax

README.md

@@ -39,11 +39,8 @@ Additionally, it's possible to pass an additional `colorize` property to the opt
 }
 ```
 
-## Known issues
-
-* The implementation in Node.js versions 11 and 12, this module works by monkey-patching the Interface prototype (`readline` module).
-If you use `readline` (or a module that depends on it) somewhere else, you may want to test everything thoroughly. In theory, there should be no side effects.
-* For Node.js versions older than 11, this module does nothing.
+In order to highlighting matching pairs of brackets, a `colorizeMatchingBracket`
+is also available.
 
 ## Credits
 

lib/find-all-matching-brackets.js

@@ -0,0 +1,88 @@
+'use strict';
+module.exports = function(str, ignoreMismatches) {
+  // Find all matching brackets inside a string.
+  // 'stack' maintains a list of all currently open brackets, 'matches' a list
+  // of all closed brackets (i.e. the return value).
+  // If `ignoreMismatches` is true, then e.g. {(x} will be ignored entirely.
+  // If not, then { and } will be matched and the ( discarded.
+  const stack = [];
+  const matches = [];
+  for (let i = 0; i < str.length; i++) {
+    const current = stack.length > 0 ? stack[stack.length - 1] : null;
+    const currentKind = current ? current.kind : '';
+    switch (currentKind) {
+      case '':
+      case '(':
+      case '[':
+      case '{':
+      case '$':
+        switch (str[i]) {
+          case '(':
+          case '[':
+          case '{':
+          case "'":
+          case '"':
+          case '`':
+            stack.push({
+              start: i,
+              end: -1,
+              kind: str[i],
+              parent: current
+            });
+            break;
+          case ')':
+          case ']':
+          case '}':
+            for (let j = stack.length - 1; j >= 0; j--) {
+              const entry = stack[j];
+              if ((entry.kind === '(' && str[i] === ')') ||
+                  (entry.kind === '[' && str[i] === ']') ||
+                  (entry.kind === '{' && str[i] === '}') ||
+                  (entry.kind === '$' && str[i] === '}')) {
+                const isProperMatch = j === stack.length - 1;
+                stack.splice(j); // Unwind the stack in any case.
+                if (!ignoreMismatches || isProperMatch) {
+                  entry.end = i;
+                  matches.push(entry);
+                }
+                break;
+              }
+            }
+            break;
+        }
+        break;
+      case "'":
+      case '"':
+      case '`':
+        switch (str[i]) {
+          case "'":
+          case '"':
+          case '`':
+          case '$': {
+            let j; // Count number of preceding \ characters
+            for (j = 0; j < i && str[i - j - 1] == '\\'; j++);
+            if (j % 2 === 1) {
+              break; // This is an escaped character, so we can ignore it.
+            }
+            if ((currentKind === "'" && str[i] === "'") ||
+                (currentKind === '"' && str[i] === '"') ||
+                (currentKind === '`' && str[i] === '`')) {
+              const entry = stack.pop();
+              entry.end = i;
+              matches.push(entry);
+            } else if (currentKind === '`' && str[i] === '$' && str[i+1] === '{') {
+              stack.push({
+                start: i++,
+                end: -1,
+                kind: '$',
+                parent: current
+              });
+            }
+            break;
+          }
+        }
+        break;
+    }
+  }
+  return matches;
+};

lib/highlight.js

@@ -47,6 +47,9 @@ module.exports = (stream) => {
     if (stream.getColorDepth() >= 8) level = 2;
     if (stream.getColorDepth() >= 24) level = 3;
   }
-  const colorSheet = sheet(new chalk.Instance({ level }));
-  return (s) => emphasize.highlight('js', s, colorSheet).value;
+  const chalkInstance = new chalk.Instance({ level });
+  const colorSheet = sheet(chalkInstance);
+  const highlight = (s) => emphasize.highlight('js', s, colorSheet).value;
+  highlight.colorizeMatchingBracket = (s) => chalkInstance.bgBlue(s);
+  return highlight;
 };

lib/pretty-repl.js

@@ -2,6 +2,7 @@
 const repl = require('repl');
 const highlight = require('./highlight');
 const memoizeStringTransformerMethod = require('./memoize-string-transformer');
+const findAllMatchingBrackets = require('./find-all-matching-brackets');
 const ansiRegex = require('ansi-regex');
 const stripAnsi = require('strip-ansi');
 
@@ -10,6 +11,10 @@ const ansiRegexMatchAll = ansiRegex();
 // Regex that matches ANSI escape sequences only at the beginning of a string.
 const ansiRegexMatchBeginningOnly = new RegExp(`^(${ansiRegexMatchAll.source})`);
 
+// Every open/close pair that should be matched against its counterpart for
+// highlighting.
+const allBrackets = '()[]{}\'"`$';
+
 // Compute the length of the longest common prefix of 'before' and 'after',
 // taking ANSI escape sequences into account. For example:
 // 'abcd', 'abab' -> 2
@@ -52,7 +57,10 @@ class PrettyREPLServer extends repl.REPLServer {
     super(options);
     options.output = options.output || process.stdout;
     this.colorize = (options && options.colorize) || highlight(options.output);
+    this.colorizeMatchingBracket = (options && options.colorizeMatchingBracket) ||
+      highlight(options.output).colorizeMatchingBracket;
     this.lineBeforeInsert = undefined;
+    this.highlightBracketPosition = -1;
 
     // For some reason, tests fail if we don't initialize line to be the empty string.
     // Specifically, `REPLServer.Interface.getCursorPos()` finds itself in a state where `line`
@@ -61,6 +69,30 @@ class PrettyREPLServer extends repl.REPLServer {
     this.__prettyModuleLoaded = __filename;
   }
 
+  // If the cursor is moved onto or off a bracket, refresh the whole line so that
+  // we can mark the matching bracket.
+  _moveCursor (dx) {
+    const cursorWasOnBracket = allBrackets.includes(this.line[this.cursor]);
+    super._moveCursor(dx);
+    const cursorIsOnBracket = allBrackets.includes(this.line[this.cursor]);
+    if (cursorWasOnBracket || cursorIsOnBracket) {
+      this._refreshLine();
+    }
+  }
+
+  // When refreshinng the whole line, find matching brackets and keep the position
+  // of the matching one in mind (if there is any).
+  _refreshLine () {
+    try {
+      if (this.colorizeMatchingBracket && allBrackets.includes(this.line[this.cursor])) {
+        this.highlightBracketPosition = this._findMatchingBracket(this.line, this.cursor);
+      }
+      return super._refreshLine();
+    } finally {
+      this.highlightBracketPosition = -1;
+    }
+  }
+
   _writeToOutput (stringToWrite) {
     // Skip false-y values, and if we print only whitespace or have not yet
     // been fully initialized, just write to output directly.
@@ -141,7 +173,22 @@ class PrettyREPLServer extends repl.REPLServer {
     // In those cases, we split the string into prompt and non-prompt parts,
     // and colorize the full non-prompt part.
     stringToWrite = stringToWrite.substring(this._prompt.length);
-    this.output.write(this._prompt + this._doColorize(stringToWrite));
+    if (this.highlightBracketPosition !== -1) {
+      // If there is a matching bracket, we mark it in the string before
+      // highlighting using BOM characters (because it seems safe to assume
+      // that they are ignored by highlighting) so that we can remember where
+      // the bracket was.
+      stringToWrite =
+        stringToWrite.substring(0, this.highlightBracketPosition) +
+        '\ufeff' + stringToWrite[this.highlightBracketPosition] + '\ufeff' +
+        stringToWrite.substring(this.highlightBracketPosition + 1);
+      stringToWrite = this._doColorize(stringToWrite);
+      // Then remove the BOM characters again and colorize the bracket in between.
+      stringToWrite = stringToWrite.replace(/\ufeff(.+)\ufeff/, (_, bracket) => this.colorizeMatchingBracket(bracket));
+    } else {
+      stringToWrite = this._doColorize(stringToWrite);
+    }
+    this.output.write(this._prompt + stringToWrite);
   }
 
   _insertString (c) {
@@ -157,47 +204,45 @@ class PrettyREPLServer extends repl.REPLServer {
     return this.colorize(str);
   });
 
-  _stripCompleteJSStructures(str) {
+  _stripCompleteJSStructures = memoizeStringTransformerMethod(1000, function(str) {
     // Remove substructures of the JS input string `str` in order to simplify it,
-    // by repeatedly removing matching pairs of quotes and parentheses/brackets.
-    let before;
-    do {
-      before = str;
-      str = this._stripCompleteJSStructuresStep(before);
-    } while (before !== str);
-    return str;
-  }
+    // by removing matching pairs of quotes and parentheses/brackets.
 
-  _stripCompleteJSStructuresStep = memoizeStringTransformerMethod(10000, function(str) {
-    // This regular expression matches:
-    // - matching pairs of (), without any of ()[]{}`'" in between
-    // - matching pairs of [], without any of ()[]{}`'" in between
-    // - matching pairs of {}, without any of ()[]{}`'" in between
-    // - 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
-    // - matching pairs of ``, with only non-`{}\, \\ and \` preceded by an even number of \ in between
-    const re = /\([^\(\)\[\]\{\}`'"]*\)|\[[^\(\)\[\]\{\}`'"]*\]|\{[^\(\)\[\]\{\}`'"]*\}|'([^'\\]|(?<=[^\\](\\\\)*)\\'|\\\\)*'|"([^"\\]|(?<=[^\\](\\\\)*)\\"|\\\\)*"|`([^\{\}`\\]|(?<=[^\\](\\\\)*)\\`|\\\\)*`/g;
-    // 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.
+    // Specifically, 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--;
+    const brackets = this._findAllMatchingBracketsIgnoreMismatches(str);
+    if (brackets.length > 0) {
+      const last = brackets[brackets.length - 1];
+      if (last.kind === '(' && (last.parent == null || last.parent.end === -1))
+        brackets.pop();
     }
-    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);
+    // Remove brackets in reverse order, so that their indices remain valid.
+    for (let i = brackets.length - 1; i >= 0; i--) {
+      str = str.substr(0, brackets[i].start) + str.substr(brackets[i].end + 1);
     }
     return str;
   });
+
+  _findAllMatchingBracketsIgnoreMismatches = memoizeStringTransformerMethod(1000, function(str) {
+    return findAllMatchingBrackets(str, true);
+  });
+  _findAllMatchingBracketsIncludeMismatches = memoizeStringTransformerMethod(1000, function(str) {
+    return findAllMatchingBrackets(str, false);
+  });
+
+  // Find the matching bracket opposite of the one at `position`.
+  _findMatchingBracket(line, position) {
+    const brackets = this._findAllMatchingBracketsIncludeMismatches(line);
+    for (const bracket of brackets) {
+      if (bracket.start === position)
+        return bracket.end;
+      if (bracket.end === position)
+        return bracket.start;
+    }
+    return -1;
+  }
 }
 
 module.exports = {