Implement react-hot-loader transform (#376)

Fixes #228

Some details:
* An eval snippet needed to be added to each class, which I decided to do
  unconditionally for simplicity.
* A previous change already comptued the top-level declared variables, so I
  could just use those.
* There was a bug where parameters in arrow function types were seen as
  top-level variables, so I changed it so types are never considered variable
  declarations.
* In order to register the default export, we need to extract it to a variable,
  which required modifying both import transformers to handle that as a special
  case.
* The ReactHotLoaderTransformer doesn't actually participate in normal
  transform, it just adds the snippets to the start and end.

Cases not handled yet that could be handled in the future:
* Avoid treating `require` statements as top-level declarations.
* Skip react and react-hot-loader files themselves (Sucrase shouldn't be running
  on them anyway).

I tested this end-to-end on a small app to make sure hot reloading works,
including for bound methods.

README.md

@@ -53,8 +53,12 @@ are four main transforms that you may want to enable:
   `const enum`s that need cross-file compilation.
 * **flow**:  Removes Flow type annotations. Does not check types.
 * **imports**: Transforms ES Modules (`import`/`export`) to CommonJS
-  (`require`/`module.exports`) using the same approach as Babel 6 and TypeScript
+  (`require`/`module.exports`) using the same approach as Babel and TypeScript
   with `--esModuleInterop`. Also includes dynamic `import`.
+* **react-hot-loader**: Performs the equivalent of the `react-hot-loader/babel`
+  transform in the [react-hot-loader](https://github.com/gaearon/react-hot-loader)
+  project. This enables advanced hot reloading use cases such as editing of
+  bound methods.
 
 The following proposed JS features are built-in and always transformed:
 * [Class fields](https://github.com/tc39/proposal-class-fields): `class C { x = 1; }`.

src/index.ts

@@ -9,7 +9,7 @@ import RootTransformer from "./transformers/RootTransformer";
 import formatTokens from "./util/formatTokens";
 import getTSImportedNames from "./util/getTSImportedNames";
 
-export type Transform = "jsx" | "typescript" | "flow" | "imports";
+export type Transform = "jsx" | "typescript" | "flow" | "imports" | "react-hot-loader";
 
 export interface SourceMapOptions {
   /**

src/parser/traverser/lval.ts

@@ -34,6 +34,9 @@ export function parseBindingIdentifier(isBlockScope: boolean): void {
 }
 
 export function markPriorBindingIdentifier(isBlockScope: boolean): void {
+  if (state.isType) {
+    return;
+  }
   if (state.scopeDepth === 0) {
     state.tokens[state.tokens.length - 1].identifierRole = IdentifierRole.TopLevelDeclaration;
   } else {

src/transformers/CJSImportTransformer.ts

@@ -1,8 +1,10 @@
 import CJSImportProcessor from "../CJSImportProcessor";
+import NameManager from "../NameManager";
 import {IdentifierRole, isDeclaration, isObjectShorthandDeclaration} from "../parser/tokenizer";
 import {ContextualKeyword} from "../parser/tokenizer/keywords";
 import {TokenType as tt} from "../parser/tokenizer/types";
 import TokenProcessor from "../TokenProcessor";
+import ReactHotLoaderTransformer from "./ReactHotLoaderTransformer";
 import RootTransformer from "./RootTransformer";
 import Transformer from "./Transformer";
 
@@ -18,6 +20,8 @@ export default class CJSImportTransformer extends Transformer {
     readonly rootTransformer: RootTransformer,
     readonly tokens: TokenProcessor,
     readonly importProcessor: CJSImportProcessor,
+    readonly nameManager: NameManager,
+    readonly reactHotLoaderTransformer: ReactHotLoaderTransformer | null,
     readonly enableLegacyBabel5ModuleInterop: boolean,
   ) {
     super();
@@ -296,7 +300,7 @@ export default class CJSImportTransformer extends Transformer {
   private processExportDefault(): void {
     if (
       this.tokens.matches4(tt._export, tt._default, tt._function, tt.name) ||
-      // export default aysnc function
+      // export default async function
       this.tokens.matches5(tt._export, tt._default, tt.name, tt._function, tt.name)
     ) {
       this.tokens.removeInitialToken();
@@ -318,7 +322,17 @@ export default class CJSImportTransformer extends Transformer {
       this.tokens.appendCode(` exports.default = ${name};`);
     } else if (this.tokens.matches3(tt._export, tt._default, tt.at)) {
       throw new Error("Export default statements with decorators are not yet supported.");
+    } else if (this.reactHotLoaderTransformer) {
+      // This is a plain "export default E" statement and we need to assign E to a variable.
+      // Change "export default E" to "let _default; exports.default = _default = E"
+      const defaultVarName = this.nameManager.claimFreeName("_default");
+      this.tokens.replaceToken(`let ${defaultVarName}; exports.`);
+      this.tokens.copyToken();
+      this.tokens.appendCode(` = ${defaultVarName} =`);
+      this.reactHotLoaderTransformer.setExtractedDefaultExportName(defaultVarName);
     } else {
+      // This is a plain "export default E" statement, no additional requirements.
+      // Change "export default E" to "exports.default = E"
       this.tokens.replaceToken("exports.");
       this.tokens.copyToken();
       this.tokens.appendCode(" =");

src/transformers/ESMImportTransformer.ts

@@ -1,7 +1,9 @@
+import NameManager from "../NameManager";
 import {ContextualKeyword} from "../parser/tokenizer/keywords";
 import {TokenType as tt} from "../parser/tokenizer/types";
 import TokenProcessor from "../TokenProcessor";
 import {getNonTypeIdentifiers} from "../util/getNonTypeIdentifiers";
+import ReactHotLoaderTransformer from "./ReactHotLoaderTransformer";
 import Transformer from "./Transformer";
 
 /**
@@ -11,7 +13,12 @@ import Transformer from "./Transformer";
 export default class ESMImportTransformer extends Transformer {
   private nonTypeIdentifiers: Set<string>;
 
-  constructor(readonly tokens: TokenProcessor, readonly isTypeScriptTransformEnabled: boolean) {
+  constructor(
+    readonly tokens: TokenProcessor,
+    readonly nameManager: NameManager,
+    readonly reactHotLoaderTransformer: ReactHotLoaderTransformer | null,
+    readonly isTypeScriptTransformEnabled: boolean,
+  ) {
     super();
     this.nonTypeIdentifiers = isTypeScriptTransformEnabled
       ? getNonTypeIdentifiers(tokens)
@@ -31,6 +38,9 @@ export default class ESMImportTransformer extends Transformer {
     if (this.tokens.matches1(tt._import)) {
       return this.processImport();
     }
+    if (this.tokens.matches2(tt._export, tt._default)) {
+      return this.processExportDefault();
+    }
     return false;
   }
 
@@ -182,4 +192,25 @@ export default class ESMImportTransformer extends Transformer {
   private isTypeName(name: string): boolean {
     return this.isTypeScriptTransformEnabled && !this.nonTypeIdentifiers.has(name);
   }
+
+  private processExportDefault(): boolean {
+    const alreadyHasName =
+      this.tokens.matches4(tt._export, tt._default, tt._function, tt.name) ||
+      // export default async function
+      this.tokens.matches5(tt._export, tt._default, tt.name, tt._function, tt.name) ||
+      this.tokens.matches4(tt._export, tt._default, tt._class, tt.name) ||
+      this.tokens.matches5(tt._export, tt._default, tt._abstract, tt._class, tt.name);
+
+    if (!alreadyHasName && this.reactHotLoaderTransformer) {
+      // This is a plain "export default E" statement and we need to assign E to a variable.
+      // Change "export default E" to "let _default; export default _default = E"
+      const defaultVarName = this.nameManager.claimFreeName("_default");
+      this.tokens.replaceToken(`let ${defaultVarName}; export`);
+      this.tokens.copyToken();
+      this.tokens.appendCode(` ${defaultVarName} =`);
+      this.reactHotLoaderTransformer.setExtractedDefaultExportName(defaultVarName);
+      return true;
+    }
+    return false;
+  }
 }

src/transformers/ReactHotLoaderTransformer.ts

@@ -0,0 +1,66 @@
+import {IdentifierRole} from "../parser/tokenizer";
+import TokenProcessor from "../TokenProcessor";
+import Transformer from "./Transformer";
+
+export default class ReactHotLoaderTransformer extends Transformer {
+  private extractedDefaultExportName: string | null = null;
+
+  constructor(readonly tokens: TokenProcessor, readonly filePath: string) {
+    super();
+  }
+
+  setExtractedDefaultExportName(extractedDefaultExportName: string): void {
+    this.extractedDefaultExportName = extractedDefaultExportName;
+  }
+
+  getPrefixCode(): string {
+    return `
+      (function () {
+        var enterModule = require('react-hot-loader').enterModule;
+        enterModule && enterModule(module);
+      })();`
+      .replace(/\s+/g, " ")
+      .trim();
+  }
+
+  getSuffixCode(): string {
+    const topLevelNames = new Set();
+    for (const token of this.tokens.tokens) {
+      if (
+        token.identifierRole === IdentifierRole.TopLevelDeclaration ||
+        token.identifierRole === IdentifierRole.ObjectShorthandTopLevelDeclaration
+      ) {
+        topLevelNames.add(this.tokens.identifierNameForToken(token));
+      }
+    }
+    const namesToRegister = Array.from(topLevelNames).map((name) => ({
+      variableName: name,
+      uniqueLocalName: name,
+    }));
+    if (this.extractedDefaultExportName) {
+      namesToRegister.push({
+        variableName: this.extractedDefaultExportName,
+        uniqueLocalName: "default",
+      });
+    }
+    return `
+(function () {
+  var reactHotLoader = require('react-hot-loader').default;
+  var leaveModule = require('react-hot-loader').leaveModule;
+  if (!reactHotLoader) {
+    return;
+  }
+${namesToRegister
+      .map(
+        ({variableName, uniqueLocalName}) =>
+          `  reactHotLoader.register(${variableName}, "${uniqueLocalName}", "${this.filePath}");`,
+      )
+      .join("\n")}
+  leaveModule(module);
+})();`;
+  }
+
+  process(): boolean {
+    return false;
+  }
+}

src/transformers/RootTransformer.ts

@@ -10,6 +10,7 @@ import JSXTransformer from "./JSXTransformer";
 import NumericSeparatorTransformer from "./NumericSeparatorTransformer";
 import OptionalCatchBindingTransformer from "./OptionalCatchBindingTransformer";
 import ReactDisplayNameTransformer from "./ReactDisplayNameTransformer";
+import ReactHotLoaderTransformer from "./ReactHotLoaderTransformer";
 import Transformer from "./Transformer";
 import TypeScriptTransformer from "./TypeScriptTransformer";
 
@@ -19,6 +20,7 @@ export default class RootTransformer {
   private tokens: TokenProcessor;
   private generatedVariables: Array<string> = [];
   private isImportsTransformEnabled: boolean;
+  private isReactHotLoaderTransformEnabled: boolean;
 
   constructor(
     sucraseContext: SucraseContext,
@@ -30,6 +32,7 @@ export default class RootTransformer {
     const {tokenProcessor, importProcessor} = sucraseContext;
     this.tokens = tokenProcessor;
     this.isImportsTransformEnabled = transforms.includes("imports");
+    this.isReactHotLoaderTransformEnabled = transforms.includes("react-hot-loader");
 
     this.transformers.push(new NumericSeparatorTransformer(tokenProcessor));
     this.transformers.push(new OptionalCatchBindingTransformer(tokenProcessor, this.nameManager));
@@ -42,6 +45,15 @@ export default class RootTransformer {
       );
     }
 
+    let reactHotLoaderTransformer = null;
+    if (transforms.includes("react-hot-loader")) {
+      if (!options.filePath) {
+        throw new Error("filePath is required when using the react-hot-loader transform.");
+      }
+      reactHotLoaderTransformer = new ReactHotLoaderTransformer(tokenProcessor, options.filePath);
+      this.transformers.push(reactHotLoaderTransformer);
+    }
+
     // Note that we always want to enable the imports transformer, even when the import transform
     // itself isn't enabled, since we need to do type-only import pruning for both Flow and
     // TypeScript.
@@ -54,12 +66,19 @@ export default class RootTransformer {
           this,
           tokenProcessor,
           importProcessor,
+          this.nameManager,
+          reactHotLoaderTransformer,
           enableLegacyBabel5ModuleInterop,
         ),
       );
     } else {
       this.transformers.push(
-        new ESMImportTransformer(tokenProcessor, transforms.includes("typescript")),
+        new ESMImportTransformer(
+          tokenProcessor,
+          this.nameManager,
+          reactHotLoaderTransformer,
+          transforms.includes("typescript"),
+        ),
       );
     }
 
@@ -210,6 +229,11 @@ export default class RootTransformer {
       throw new Error("Expected non-null context ID on class.");
     }
     this.tokens.copyExpectedToken(tt.braceL);
+    if (this.isReactHotLoaderTransformEnabled) {
+      this.tokens.appendCode(
+        "__reactstandin__regenerateByEval(key, code) {this[key] = eval(code);}",
+      );
+    }
 
     const needsConstructorInit =
       constructorInitializerStatements.length + instanceInitializerNames.length > 0;

test/index-test.ts

@@ -4,7 +4,7 @@ import {getFormattedTokens} from "../src";
 
 describe("getFormattedTokens", () => {
   it("formats a simple program", () => {
-    assert.equal(
+    assert.strictEqual(
       getFormattedTokens(
         `\
 if (foo) {

test/prefixes.ts

@@ -7,3 +7,6 @@ if (obj != null) { for (var key in obj) { \
 if (Object.prototype.hasOwnProperty.call(obj, key)) newObj[key] = obj[key]; } } \
 newObj.default = obj; return newObj; } }`;
 export const ESMODULE_PREFIX = 'Object.defineProperty(exports, "__esModule", {value: true});';
+export const RHL_PREFIX = `(function () { \
+var enterModule = require('react-hot-loader').enterModule; enterModule && enterModule(module); \
+})();`;