Generate WeakMaps for private fields in class expressions

Author: Joseph Watts

src/compiler/transformers/esnext.ts

@@ -254,13 +254,9 @@ namespace ts {
             if (some(pendingExpressions)) {
                 statements.push(createExpressionStatement(inlineExpressions(pendingExpressions!)));
             }
+            endPrivateNameEnvironment();
             pendingExpressions = savedPendingExpressions;
 
-            const privateNameExpressions = endPrivateNameEnvironment();
-            if (some(privateNameExpressions)) {
-                statements.push(createExpressionStatement(inlineExpressions(privateNameExpressions)));
-            }
-
             // Emit static property assignment. Because classDeclaration is lexically evaluated,
             // it is safe to emit static property assignment after classDeclaration
             // From ES6 specification:
@@ -277,6 +273,7 @@ namespace ts {
         function visitClassExpression(node: ClassExpression): Expression {
             const savedPendingExpressions = pendingExpressions;
             pendingExpressions = undefined;
+            startPrivateNameEnvironment();
 
             // If this class expression is a transformation of a decorated class declaration,
             // then we want to output the pendingExpressions as statements, not as inlined
@@ -304,7 +301,7 @@ namespace ts {
                 if (isDecoratedClassDeclaration) {
                     Debug.assertDefined(pendingStatements, "Decorated classes transformed by TypeScript are expected to be within a variable declaration.");
 
-                    // Write any pending expressions from elided or moved computed property names
+                    // Write any pending expressions from elided or moved computed property names or private names
                     if (some(pendingExpressions)) {
                         pendingStatements!.push(createExpressionStatement(inlineExpressions(pendingExpressions!)));
                     }
@@ -313,6 +310,7 @@ namespace ts {
                     if (some(staticProperties)) {
                         addInitializedPropertyStatements(pendingStatements!, staticProperties, getInternalName(node));
                     }
+                    endPrivateNameEnvironment();
                     return classExpression;
                 }
                 else {
@@ -337,11 +335,13 @@ namespace ts {
                     expressions.push(startOnNewLine(temp));
 
                     pendingExpressions = savedPendingExpressions;
+                    endPrivateNameEnvironment();
                     return inlineExpressions(expressions);
                 }
             }
 
             pendingExpressions = savedPendingExpressions;
+            endPrivateNameEnvironment();
             return classExpression;
         }
 
@@ -508,36 +508,26 @@ namespace ts {
             return env;
         }
 
-        function endPrivateNameEnvironment(): Expression[] {
-            const env = privateNameEnvironmentStack.pop();
-            Debug.assertDefined(env, "Tried to end private name environment that does not exist.");
-            const initializers: Expression[] = [];
-            forEachEntry(env!, (privateName) => {
-                switch (privateName.type) {
-                    case PrivateNameType.InstanceField: {
-                        hoistVariableDeclaration(privateName.weakMapName);
-                        initializers.push(
-                            createAssignment(
-                                privateName.weakMapName,
-                                createNew(
-                                    createIdentifier("WeakMap"),
-                                    /*typeArguments*/ undefined,
-                                    []
-                                )
-                            )
-                        );
-                        break;
-                    }
-                }
-            });
-            return initializers;
+        function endPrivateNameEnvironment() {
+            privateNameEnvironmentStack.pop();
         }
 
         function addPrivateNameToEnvironment(name: PrivateName) {
             const env = last(privateNameEnvironmentStack);
             const text = getTextOfPropertyName(name) as string;
             const weakMapName = createFileLevelUniqueName("_" + text.substring(1));
+            hoistVariableDeclaration(weakMapName);
             env.set(name.escapedText, { type: PrivateNameType.InstanceField, weakMapName });
+            (pendingExpressions || (pendingExpressions = [])).push(
+                createAssignment(
+                    weakMapName,
+                    createNew(
+                        createIdentifier("WeakMap"),
+                        /*typeArguments*/ undefined,
+                        []
+                    )
+                )
+            );
         }
 
         function enableSubstitutionForClassAliases() {