feat(aws-cdk-lib): reduce JavaScript load time, second attempt

packages/@aws-cdk-testing/cli-integ/tests/init-javascript/init-javascript.integtest.ts

@@ -25,8 +25,10 @@ integTest('Test importing CDK from ESM', withTemporaryDirectory(withPackages(asy
 
   // Rewrite some files
   await fs.writeFile(path.join(context.integTestDir, 'new-entrypoint.mjs'), `
-// Test two styles of imports
-import { Stack, aws_sns as sns, aws_sns_subscriptions as subs, aws_sqs as sqs } from 'aws-cdk-lib';
+// Test multiple styles of imports
+import { Stack, aws_sns as sns } from 'aws-cdk-lib';
+import { SqsSubscription } from 'aws-cdk-lib/aws-sns-subscriptions';
+import * as sqs from 'aws-cdk-lib/aws-sqs';
 import * as cdk from 'aws-cdk-lib';
 
 class TestjsStack extends Stack {
@@ -39,7 +41,7 @@ class TestjsStack extends Stack {
 
     const topic = new sns.Topic(this, 'TestjsTopic');
 
-    topic.addSubscription(new subs.SqsSubscription(queue));
+    topic.addSubscription(new SqsSubscription(queue));
   }
 }
 

tools/@aws-cdk/lazify/lib/index.ts

@@ -139,6 +139,8 @@ export function transformFileContents(filename: string, contents: string, progre
 
   file = ts.transform(file, [(ctx: ts.TransformationContext): ts.Transformer<ts.SourceFile> => {
     const factory = ctx.factory;
+    const alreadyEmittedExports = new Set<string>();
+
     const visit: ts.Visitor = node => {
       if (node.parent && ts.isSourceFile(node.parent)
         && ts.isExpressionStatement(node)
@@ -159,8 +161,8 @@ export function transformFileContents(filename: string, contents: string, progre
         const module = require(file);
         const entries = Object.keys(module);
 
-        return entries.map((entry) =>
-          createModuleGetter(factory, entry, requiredModule, (mod) =>
+        return entries.flatMap((entry) =>
+          createModuleGetterOnce(alreadyEmittedExports)(factory, entry, requiredModule, (mod) =>
             factory.createPropertyAccessExpression(mod, entry))
           );
       }
@@ -180,7 +182,7 @@ export function transformFileContents(filename: string, contents: string, progre
 
         const exportName = node.expression.left.name.text;
         const moduleName = node.expression.right.arguments[0].text;
-        return createModuleGetter(factory, exportName, moduleName, (x) => x);
+        return createModuleGetterOnce(alreadyEmittedExports)(factory, exportName, moduleName, (x) => x);
       }
 
       return ts.visitEachChild(node, child => visit(child), ctx);
@@ -212,25 +214,67 @@ function createAssignment(factory: ts.NodeFactory, name: string, expression: ts.
       expression));
 }
 
+/**
+ * Create an lazy getter for a particular value at the module level
+ *
+ * Since Node statically analyzes CommonJS modules to determine its exports
+ * (using the `cjs-module-lexer` module), we need to trick it into recognizing
+ * these exports as legitimate.
+ *
+ * We do that by generating one form it will recognize that doesn't do anything,
+ * in combination with a form that actually works, that doesn't disqualify the
+ * export name.
+ */
 function createModuleGetter(
   factory: ts.NodeFactory,
   exportName: string,
   moduleName: string,
   moduleFormatter: (x: ts.Expression) => ts.Expression,
 ) {
-  return factory.createExpressionStatement(factory.createCallExpression(
-    factory.createPropertyAccessExpression(factory.createIdentifier('Object'), factory.createIdentifier('defineProperty')),
-    undefined,
-    [
-      factory.createIdentifier('exports'),
-      factory.createStringLiteral(exportName),
-      factory.createObjectLiteralExpression([
-        factory.createPropertyAssignment('configurable', factory.createTrue()),
-        factory.createPropertyAssignment('get',
-          factory.createArrowFunction(undefined, undefined, [], undefined, undefined,
-            moduleFormatter(
-              factory.createCallExpression(factory.createIdentifier('require'), undefined, [factory.createStringLiteral(moduleName)])))),
-      ]),
-    ]
-  ));
+  return [
+    // exports.<name> = void 0;
+    factory.createExpressionStatement(factory.createBinaryExpression(
+      factory.createPropertyAccessExpression(
+        factory.createIdentifier('exports'),
+        factory.createIdentifier(exportName)),
+      ts.SyntaxKind.EqualsToken,
+      factory.createVoidZero())),
+    // Object.defineProperty(exports, "<n>" + "<ame>", { get: () =>  });
+    factory.createExpressionStatement(factory.createCallExpression(
+      factory.createPropertyAccessExpression(factory.createIdentifier('Object'), factory.createIdentifier('defineProperty')),
+      undefined,
+      [
+        factory.createIdentifier('exports'),
+        factory.createBinaryExpression(
+          factory.createStringLiteral(exportName.substring(0, 1)),
+          ts.SyntaxKind.PlusToken,
+          factory.createStringLiteral(exportName.substring(1)),
+        ),
+        factory.createObjectLiteralExpression([
+          factory.createPropertyAssignment('enumerable', factory.createTrue()),
+          factory.createPropertyAssignment('configurable', factory.createTrue()),
+          factory.createPropertyAssignment('get',
+            factory.createArrowFunction(undefined, undefined, [], undefined, undefined,
+              moduleFormatter(
+                factory.createCallExpression(factory.createIdentifier('require'), undefined, [factory.createStringLiteral(moduleName)])))),
+        ]),
+      ]
+    )
+  )];
+}
+
+/**
+ * Prevent emitting an export if it has already been emitted before
+ *
+ * This assumes that the symbols have the same definition, and are only duplicated because of
+ * accidental multiple `export *`s.
+ */
+function createModuleGetterOnce(alreadyEmittedExports: Set<string>): typeof createModuleGetter {
+  return (factory, exportName, moduleName, moduleFormatter) => {
+    if (alreadyEmittedExports.has(exportName)) {
+      return [];
+    }
+    alreadyEmittedExports.add(exportName);
+    return createModuleGetter(factory, exportName, moduleName, moduleFormatter);
+  };
 }

tools/@aws-cdk/lazify/package.json

@@ -20,7 +20,8 @@
     "@aws-cdk/cdk-build-tools": "0.0.0",
     "jest": "^29",
     "ts-jest": "^29",
-    "typescript": "^4.5.5"
+    "typescript": "^4.5.5",
+    "cjs-module-lexer": "^1.2.3"
   },
   "dependencies": {
     "esbuild": "^0.19.4",

tools/@aws-cdk/lazify/test/export-star.test.ts

@@ -1,6 +1,7 @@
 import * as fs from 'fs-extra';
 import * as path from 'path';
 import { transformFileContents } from '../lib';
+import { parse } from 'cjs-module-lexer';
 
 // Write a .js file in this directory that will be imported by tests below
 beforeEach(async () => {
@@ -12,21 +13,47 @@ beforeEach(async () => {
 
 test('replace __exportStar with getters', () => {
   const fakeFile = path.join(__dirname, 'index.ts');
-  expect(transformFileContents(fakeFile, [
+
+  const transformed = transformFileContents(fakeFile, [
     '__exportStar(require("./some-module"), exports);'
-  ].join('\n'))).toMatchInlineSnapshot(`
-"Object.defineProperty(exports, "foo", { configurable: true, get: () => require("./some-module").foo });
-Object.defineProperty(exports, "bar", { configurable: true, get: () => require("./some-module").bar });
-"
-`);
+  ].join('\n'));
+
+  expect(parse(transformed).exports).toEqual([
+    'foo',
+    'bar',
+  ]);
+
+  const mod = evalModule(transformed);
+  expect(mod.foo()).toEqual('foo');
+  expect(mod.bar).toEqual(5);
 });
 
 test('replace re-export with getter', () => {
   const fakeFile = path.join(__dirname, 'index.ts');
-  expect(transformFileContents(fakeFile, [
+  const transformed = transformFileContents(fakeFile, [
     'exports.some_module = require("./some-module");',
-  ].join('\n'))).toMatchInlineSnapshot(`
-"Object.defineProperty(exports, "some_module", { configurable: true, get: () => require("./some-module") });
-"
-`);
-});
+  ].join('\n'));
+
+  expect(parse(transformed).exports).toEqual([
+    'some_module',
+  ]);
+
+  const mod = evalModule(transformed);
+  expect(mod.some_module.foo()).toEqual('foo');
+  expect(mod.some_module.bar).toEqual(5);
+});
+
+/**
+ * Fake NodeJS evaluation of a module
+ */
+function evalModule(x: string) {
+  const code = [
+    '(function() {',
+    'const exports = {};',
+    'const module = { exports };',
+    x,
+    'return exports;',
+    '})()',
+  ].join('\n');
+  return eval(code);
+}

yarn.lock

@@ -5468,7 +5468,7 @@ cidr-regex@^3.1.1:
   dependencies:
     ip-regex "^4.1.0"
 
-cjs-module-lexer@^1.0.0:
+cjs-module-lexer@^1.0.0, cjs-module-lexer@^1.2.3:
   version "1.2.3"
   resolved "https://registry.npmjs.org/cjs-module-lexer/-/cjs-module-lexer-1.2.3.tgz#6c370ab19f8a3394e318fe682686ec0ac684d107"
   integrity sha512-0TNiGstbQmCFwt4akjjBg5pLRTSyj/PkWQ1ZoO2zntmg9yLqSRxwEa4iCfQLGjqhiqBfOJa7W/E8wfGrTDmlZQ==