feat: add repl top level await support

Author: ejose19

package-lock.json

@@ -162,6 +162,8 @@
     "@tsconfig/node12": "^1.0.7",
     "@tsconfig/node14": "^1.0.0",
     "@tsconfig/node16": "^1.0.1",
+    "acorn": "^8.4.1",
+    "acorn-walk": "^8.1.1",
     "arg": "^4.1.0",
     "create-require": "^1.1.0",
     "diff": "^4.0.1",

package.json

@@ -18,6 +18,7 @@ import {
 } from './util';
 import { readConfig } from './configuration';
 import type { TSCommon, TSInternal } from './ts-compiler-types';
+import { processTopLevelAwait } from './repl-top-level-await';
 
 export { TSCommon };
 export { createRepl, CreateReplOptions, ReplService } from './repl';
@@ -1177,6 +1178,17 @@ export function create(rawOptions: CreateOptions = {}): Service {
   // Create a simple TypeScript compiler proxy.
   function compile(code: string, fileName: string, lineOffset = 0) {
     const normalizedFileName = normalizeSlashes(fileName);
+
+    if (
+      /* ts-node flag or node --experimental-repl-await flag &&*/
+      code.includes('await')
+    ) {
+      const wrappedCode = processTopLevelAwait(code) || code;
+      if (code !== wrappedCode) {
+        code = wrappedCode;
+      }
+    }
+
     const [value, sourceMap] = getOutput(code, normalizedFileName);
     const output = updateOutput(
       value,

src/index.ts

@@ -0,0 +1,238 @@
+// Based on https://github.com/nodejs/node/blob/975bbbc443c47df777d460fadaada3c6d265b321/lib/internal/repl/await.js
+import { Node, Parser } from 'acorn';
+import { base, recursive, RecursiveWalkerFn, WalkerCallback } from 'acorn-walk';
+import { Recoverable } from 'repl';
+
+const walk = {
+  base,
+  recursive,
+};
+
+const noop = () => {};
+const visitorsWithoutAncestors: VisitorsWithoutAncestors = {
+  ClassDeclaration(node, state, c) {
+    if (state.ancestors[state.ancestors.length - 2] === state.body) {
+      state.prepend(node, `${node.id.name}=`);
+    }
+    walk.base.ClassDeclaration(node, state, c);
+  },
+  ForOfStatement(node, state, c) {
+    if (node.await === true) {
+      state.containsAwait = true;
+    }
+    walk.base.ForOfStatement(node, state, c);
+  },
+  FunctionDeclaration(node, state, c) {
+    state.prepend(node, `${node.id.name}=`);
+  },
+  FunctionExpression: noop,
+  ArrowFunctionExpression: noop,
+  MethodDefinition: noop,
+  AwaitExpression(node, state, c) {
+    state.containsAwait = true;
+    walk.base.AwaitExpression(node, state, c);
+  },
+  ReturnStatement(node, state, c) {
+    state.containsReturn = true;
+    walk.base.ReturnStatement(node, state, c);
+  },
+  VariableDeclaration(node, state, c) {
+    if (
+      node.kind === 'var' ||
+      state.ancestors[state.ancestors.length - 2] === state.body
+    ) {
+      if (node.declarations.length === 1) {
+        state.replace(node.start, node.start + node.kind.length, 'void');
+      } else {
+        state.replace(node.start, node.start + node.kind.length, 'void (');
+      }
+
+      node.declarations.forEach((decl) => {
+        state.prepend(decl, '(');
+        state.append(decl, decl.init ? ')' : '=undefined)');
+      });
+
+      if (node.declarations.length !== 1) {
+        state.append(node.declarations[node.declarations.length - 1], ')');
+      }
+    }
+
+    walk.base.VariableDeclaration(node, state, c);
+  },
+};
+
+const visitors: Record<string, RecursiveWalkerFn<State>> = {};
+for (const nodeType of Object.keys(walk.base)) {
+  const callback =
+    nodeType in visitorsWithoutAncestors
+      ? visitorsWithoutAncestors[nodeType as keyof VisitorsWithoutAncestors]
+      : walk.base[nodeType];
+
+  visitors[nodeType] = (node, state, c) => {
+    const isNew = node !== state.ancestors[state.ancestors.length - 1];
+    if (isNew) {
+      state.ancestors.push(node);
+    }
+    callback(node as CustomRecursiveWalkerNode, state, c);
+    if (isNew) {
+      state.ancestors.pop();
+    }
+  };
+}
+
+export function processTopLevelAwait(src: string) {
+  const wrapPrefix = '(async () => { ';
+  const wrapped = `${wrapPrefix}${src} })()`;
+  const wrappedArray = Array.from(wrapped);
+  let root;
+  try {
+    root = Parser.parse(wrapped, { ecmaVersion: 'latest' }) as RootNode;
+  } catch (e) {
+    if (e.message.startsWith('Unterminated ')) throw new Recoverable(e);
+    // If the parse error is before the first "await", then use the execution
+    // error. Otherwise we must emit this parse error, making it look like a
+    // proper syntax error.
+    const awaitPos = src.indexOf('await');
+    const errPos = e.pos - wrapPrefix.length;
+    if (awaitPos > errPos) return null;
+    // Convert keyword parse errors on await into their original errors when
+    // possible.
+    if (
+      errPos === awaitPos + 6 &&
+      e.message.includes('Expecting Unicode escape sequence')
+    )
+      return null;
+    if (errPos === awaitPos + 7 && e.message.includes('Unexpected token'))
+      return null;
+    const line = e.loc.line;
+    const column = line === 1 ? e.loc.column - wrapPrefix.length : e.loc.column;
+    let message =
+      '\n' +
+      src.split('\n')[line - 1] +
+      '\n' +
+      ' '.repeat(column) +
+      '^\n\n' +
+      e.message.replace(/ \([^)]+\)/, '');
+    // V8 unexpected token errors include the token string.
+    if (message.endsWith('Unexpected token'))
+      message += " '" + src[e.pos - wrapPrefix.length] + "'";
+    // eslint-disable-next-line no-restricted-syntax
+    throw new SyntaxError(message);
+  }
+  const body = root.body[0].expression.callee.body;
+  const state: State = {
+    body,
+    ancestors: [],
+    replace(from: number, to: number, str: string) {
+      for (let i = from; i < to; i++) {
+        wrappedArray[i] = '';
+      }
+      if (from === to) str += wrappedArray[from];
+      wrappedArray[from] = str;
+    },
+    prepend(node: Node, str: string) {
+      wrappedArray[node.start] = str + wrappedArray[node.start];
+    },
+    append(node: Node, str: string) {
+      wrappedArray[node.end - 1] += str;
+    },
+    containsAwait: false,
+    containsReturn: false,
+  };
+
+  walk.recursive(body, state, visitors);
+
+  // Do not transform if
+  // 1. False alarm: there isn't actually an await expression.
+  // 2. There is a top-level return, which is not allowed.
+  if (!state.containsAwait || state.containsReturn) {
+    return null;
+  }
+
+  const last = body.body[body.body.length - 1];
+  if (last.type === 'ExpressionStatement') {
+    // For an expression statement of the form
+    // ( expr ) ;
+    // ^^^^^^^^^^   // last
+    //   ^^^^       // last.expression
+    //
+    // We do not want the left parenthesis before the `return` keyword;
+    // therefore we prepend the `return (` to `last`.
+    //
+    // On the other hand, we do not want the right parenthesis after the
+    // semicolon. Since there can only be more right parentheses between
+    // last.expression.end and the semicolon, appending one more to
+    // last.expression should be fine.
+    state.prepend(last, 'return (');
+    state.append(last.expression, ')');
+  }
+
+  return wrappedArray.join('');
+}
+
+type CustomNode<T> = Node & T;
+type RootNode = CustomNode<{
+  body: Array<
+    CustomNode<{
+      expression: CustomNode<{
+        callee: CustomNode<{
+          body: CustomNode<{
+            body: Array<CustomNode<{ expression: Node }>>;
+          }>;
+        }>;
+      }>;
+    }>
+  >;
+}>;
+type ClassDeclarationNode = CustomNode<{ id: CustomNode<{ name: string }> }>;
+type ForOfStatementNode = CustomNode<{ await: boolean }>;
+type VariableDeclarationNode = CustomNode<{
+  kind: string;
+  declarations: Array<
+    CustomNode<{
+      init: boolean;
+    }>
+  >;
+}>;
+
+interface State {
+  body: Node;
+  ancestors: Node[];
+  replace: (from: number, to: number, str: string) => void;
+  prepend: (node: Node, str: string) => void;
+  append: (from: Node, str: string) => void;
+  containsAwait: boolean;
+  containsReturn: boolean;
+}
+
+type NOOP = () => void;
+
+type VisitorsWithoutAncestors = {
+  ClassDeclaration: CustomRecursiveWalkerFn<ClassDeclarationNode>;
+  ForOfStatement: CustomRecursiveWalkerFn<ForOfStatementNode>;
+  FunctionDeclaration: CustomRecursiveWalkerFn<ClassDeclarationNode>;
+  FunctionExpression: NOOP;
+  ArrowFunctionExpression: NOOP;
+  MethodDefinition: NOOP;
+  AwaitExpression: CustomRecursiveWalkerFn<ClassDeclarationNode>;
+  ReturnStatement: CustomRecursiveWalkerFn<ClassDeclarationNode>;
+  VariableDeclaration: CustomRecursiveWalkerFn<VariableDeclarationNode>;
+};
+
+type UnionToIntersection<U> = (U extends any ? (k: U) => void : never) extends (
+  k: infer I
+) => void
+  ? I
+  : never;
+type CustomRecursiveWalkerNode = UnionToIntersection<
+  Exclude<
+    Parameters<VisitorsWithoutAncestors[keyof VisitorsWithoutAncestors]>[0],
+    undefined
+  >
+>;
+
+type CustomRecursiveWalkerFn<N extends Node> = (
+  node: N,
+  state: State,
+  c: WalkerCallback<State>
+) => void;

src/repl-top-level-await.ts

@@ -109,7 +109,7 @@ export function createRepl(options: CreateReplOptions = {}) {
     return _eval(service!, state, code);
   }
 
-  function nodeEval(
+  async function nodeEval(
     code: string,
     _context: any,
     _filename: string,
@@ -125,7 +125,7 @@ export function createRepl(options: CreateReplOptions = {}) {
     }
 
     try {
-      result = evalCode(code);
+      result = await evalCode(code);
     } catch (error) {
       if (error instanceof TSError) {
         // Support recoverable compilations using >= node 6.
@@ -208,7 +208,7 @@ export function createEvalAwarePartialHost(
 /**
  * Evaluate the code snippet.
  */
-function _eval(service: Service, state: EvalState, input: string) {
+async function _eval(service: Service, state: EvalState, input: string) {
   const lines = state.lines;
   const isCompletion = !/\n$/.test(input);
   const undo = appendEval(state, input);