Add support for iterating over enums

Timmmm · Tim Hutt · commit 5b2b6c5a76dd · 2021-01-26T12:23:12.000Z
I'm fairly confident about the codegen stuff, but not about the type checker. Iterating over `const enum`s is not supported. This includes a simple test case.
diff --git a/src/compiler/checker.ts b/src/compiler/checker.ts
@@ -35078,6 +35078,17 @@ namespace ts {
                 return anyIterationTypes;
             }
 
+            const symbolType = getDeclaredTypeOfSymbol(type.symbol);
+
+            if (symbolType.symbol.flags & SymbolFlags.RegularEnum) {
+                return createIterationTypes(
+                    createTupleType([
+                        stringType,
+                        symbolType,
+                    ]),
+                );
+            }
+
             if (use & IterationUse.AllowsAsyncIterablesFlag) {
                 const iterationTypes =
                     getIterationTypesOfIterableCached(type, asyncIterationTypesResolver) ||
diff --git a/src/compiler/transformers/ts.ts b/src/compiler/transformers/ts.ts
@@ -2425,7 +2425,20 @@ namespace ts {
 
             const statements: Statement[] = [];
             startLexicalEnvironment();
-            const members = map(node.members, transformEnumMember);
+            const membersAndExpressions = node.members.map(member => {
+                const [nameExpression, valueExpression] = getEnumMemberExpressions(member);
+                return {
+                    member,
+                    nameExpression,
+                    valueExpression,
+                };
+            });
+            const members = membersAndExpressions.map(m => transformEnumMember(
+                m.member,
+                m.nameExpression,
+                m.valueExpression,
+            ));
+            members.push(transformEnumMembersToIterator(membersAndExpressions));
             insertStatementsAfterStandardPrologue(statements, endLexicalEnvironment());
             addRange(statements, members);
 
@@ -2437,20 +2450,37 @@ namespace ts {
         }
 
         /**
-         * Transforms an enum member into a statement.
+         * Get an expression for the name and value of an enum member.
+         * For example `enum Maybe { Yes = 1, No = 2 }` would give
+         * [`"Yes"`, `1`] and [`"No"`, `2`] as the expressions.
          *
          * @param member The enum member node.
          */
-        function transformEnumMember(member: EnumMember): Statement {
+        function getEnumMemberExpressions(member: EnumMember): [Expression, Expression] {
             // enums don't support computed properties
             // we pass false as 'generateNameForComputedPropertyName' for a backward compatibility purposes
             // old emitter always generate 'expression' part of the name as-is.
             const name = getExpressionForPropertyName(member, /*generateNameForComputedPropertyName*/ false);
             const valueExpression = transformEnumMemberDeclarationValue(member);
+            return [name, valueExpression];
+        }
+
+        /**
+         * Transforms an enum member into a statement.
+         *
+         * @param member The enum member node.
+         * @param nameExpression A string literal expression containing the member name
+         * @param valueExpression An expression containing the member value.
+         */
+        function transformEnumMember(
+            member: EnumMember,
+            nameExpression: Expression,
+            valueExpression: Expression,
+        ): Statement {
             const innerAssignment = factory.createAssignment(
                 factory.createElementAccessExpression(
                     currentNamespaceContainerName,
-                    name
+                    nameExpression
                 ),
                 valueExpression
             );
@@ -2461,7 +2491,7 @@ namespace ts {
                         currentNamespaceContainerName,
                         innerAssignment
                     ),
-                    name
+                    nameExpression
                 );
             return setTextRange(
                 factory.createExpressionStatement(
@@ -2474,6 +2504,56 @@ namespace ts {
             );
         }
 
+        /**
+         * Transforms all enum members into a generator function that yields
+         * the enum members.
+         *
+         * @param members The enum member nodes.
+         */
+        function transformEnumMembersToIterator(
+            memberExpressions: {
+                nameExpression: Expression,
+                valueExpression: Expression,
+            }[],
+        ): Statement {
+            const symbolIterator = factory.createPropertyAccessExpression(
+                factory.createIdentifier("Symbol"),
+                "iterator",
+            );
+            const iteratorMember = factory.createElementAccessExpression(
+                currentNamespaceContainerName,
+                symbolIterator,
+            );
+
+            const yieldExpression = factory.createYieldExpression(
+                factory.createToken(SyntaxKind.AsteriskToken),
+                factory.createArrayLiteralExpression(memberExpressions.map(
+                    expressions => factory.createArrayLiteralExpression([
+                        expressions.nameExpression,
+                        expressions.valueExpression,
+                    ]),
+                )),
+            );
+            const generatorBody = factory.createBlock([
+                factory.createExpressionStatement(yieldExpression)
+            ]);
+
+            const iteratorGenerator = factory.createFunctionExpression(
+                /*modifiers*/ undefined,
+                factory.createToken(SyntaxKind.AsteriskToken),
+                /*name*/ undefined,
+                [],
+                [],
+                /*type*/ undefined,
+                generatorBody,
+            );
+            const outerAssignment = factory.createAssignment(
+                iteratorMember,
+                iteratorGenerator,
+            );
+            return factory.createExpressionStatement(outerAssignment);
+        }
+
         /**
          * Transforms the value of an enum member.
          *
diff --git a/tests/baselines/reference/enumIteratorBasics1.js b/tests/baselines/reference/enumIteratorBasics1.js
@@ -0,0 +1,35 @@
+//// [enumIteratorBasics1.ts]
+enum Test {
+  Foo,
+  Bar,
+  Baz,
+}
+
+const AlsoTest = Test;
+
+for (const member of AlsoTest) {
+  const x: string = member[0];
+  const y: Test = member[1];
+  console.log(x, y);
+}
+
+const members: [string, Test][] = [...AlsoTest];
+console.log(members);
+
+
+//// [enumIteratorBasics1.js]
+var Test;
+(function (Test) {
+    Test[Test["Foo"] = 0] = "Foo";
+    Test[Test["Bar"] = 1] = "Bar";
+    Test[Test["Baz"] = 2] = "Baz";
+    Test[Symbol.iterator] = function* () { yield* [["Foo", 0], ["Bar", 1], ["Baz", 2]]; };
+})(Test || (Test = {}));
+const AlsoTest = Test;
+for (const member of AlsoTest) {
+    const x = member[0];
+    const y = member[1];
+    console.log(x, y);
+}
+const members = [...AlsoTest];
+console.log(members);
diff --git a/tests/baselines/reference/enumIteratorBasics1.symbols b/tests/baselines/reference/enumIteratorBasics1.symbols
@@ -0,0 +1,52 @@
+=== tests/cases/compiler/enumIteratorBasics1.ts ===
+enum Test {
+>Test : Symbol(Test, Decl(enumIteratorBasics1.ts, 0, 0))
+
+  Foo,
+>Foo : Symbol(Test.Foo, Decl(enumIteratorBasics1.ts, 0, 11))
+
+  Bar,
+>Bar : Symbol(Test.Bar, Decl(enumIteratorBasics1.ts, 1, 6))
+
+  Baz,
+>Baz : Symbol(Test.Baz, Decl(enumIteratorBasics1.ts, 2, 6))
+}
+
+const AlsoTest = Test;
+>AlsoTest : Symbol(AlsoTest, Decl(enumIteratorBasics1.ts, 6, 5))
+>Test : Symbol(Test, Decl(enumIteratorBasics1.ts, 0, 0))
+
+for (const member of AlsoTest) {
+>member : Symbol(member, Decl(enumIteratorBasics1.ts, 8, 10))
+>AlsoTest : Symbol(AlsoTest, Decl(enumIteratorBasics1.ts, 6, 5))
+
+  const x: string = member[0];
+>x : Symbol(x, Decl(enumIteratorBasics1.ts, 9, 7))
+>member : Symbol(member, Decl(enumIteratorBasics1.ts, 8, 10))
+>0 : Symbol(0)
+
+  const y: Test = member[1];
+>y : Symbol(y, Decl(enumIteratorBasics1.ts, 10, 7))
+>Test : Symbol(Test, Decl(enumIteratorBasics1.ts, 0, 0))
+>member : Symbol(member, Decl(enumIteratorBasics1.ts, 8, 10))
+>1 : Symbol(1)
+
+  console.log(x, y);
+>console.log : Symbol(Console.log, Decl(lib.dom.d.ts, --, --))
+>console : Symbol(console, Decl(lib.dom.d.ts, --, --))
+>log : Symbol(Console.log, Decl(lib.dom.d.ts, --, --))
+>x : Symbol(x, Decl(enumIteratorBasics1.ts, 9, 7))
+>y : Symbol(y, Decl(enumIteratorBasics1.ts, 10, 7))
+}
+
+const members: [string, Test][] = [...AlsoTest];
+>members : Symbol(members, Decl(enumIteratorBasics1.ts, 14, 5))
+>Test : Symbol(Test, Decl(enumIteratorBasics1.ts, 0, 0))
+>AlsoTest : Symbol(AlsoTest, Decl(enumIteratorBasics1.ts, 6, 5))
+
+console.log(members);
+>console.log : Symbol(Console.log, Decl(lib.dom.d.ts, --, --))
+>console : Symbol(console, Decl(lib.dom.d.ts, --, --))
+>log : Symbol(Console.log, Decl(lib.dom.d.ts, --, --))
+>members : Symbol(members, Decl(enumIteratorBasics1.ts, 14, 5))
+
diff --git a/tests/baselines/reference/enumIteratorBasics1.types b/tests/baselines/reference/enumIteratorBasics1.types
@@ -0,0 +1,56 @@
+=== tests/cases/compiler/enumIteratorBasics1.ts ===
+enum Test {
+>Test : Test
+
+  Foo,
+>Foo : Test.Foo
+
+  Bar,
+>Bar : Test.Bar
+
+  Baz,
+>Baz : Test.Baz
+}
+
+const AlsoTest = Test;
+>AlsoTest : typeof Test
+>Test : typeof Test
+
+for (const member of AlsoTest) {
+>member : [string, Test]
+>AlsoTest : typeof Test
+
+  const x: string = member[0];
+>x : string
+>member[0] : string
+>member : [string, Test]
+>0 : 0
+
+  const y: Test = member[1];
+>y : Test
+>member[1] : Test
+>member : [string, Test]
+>1 : 1
+
+  console.log(x, y);
+>console.log(x, y) : void
+>console.log : (...data: any[]) => void
+>console : Console
+>log : (...data: any[]) => void
+>x : string
+>y : Test
+}
+
+const members: [string, Test][] = [...AlsoTest];
+>members : [string, Test][]
+>[...AlsoTest] : [string, Test][]
+>...AlsoTest : [string, Test]
+>AlsoTest : typeof Test
+
+console.log(members);
+>console.log(members) : void
+>console.log : (...data: any[]) => void
+>console : Console
+>log : (...data: any[]) => void
+>members : [string, Test][]
+
diff --git a/tests/baselines/reference/enumIteratorConstEnum1.errors.txt b/tests/baselines/reference/enumIteratorConstEnum1.errors.txt
@@ -0,0 +1,19 @@
+tests/cases/compiler/enumIteratorConstEnum1.ts(8,22): error TS2475: 'const' enums can only be used in property or index access expressions or the right hand side of an import declaration or export assignment or type query.
+tests/cases/compiler/enumIteratorConstEnum1.ts(8,22): error TS2488: Type 'typeof ConstTest' must have a '[Symbol.iterator]()' method that returns an iterator.
+
+
+==== tests/cases/compiler/enumIteratorConstEnum1.ts (2 errors) ====
+    const enum ConstTest {
+      Foo,
+      Bar,
+      Baz,
+    }
+    
+    // Should be an error.
+    for (const member of ConstTest) {
+                         ~~~~~~~~~
+!!! error TS2475: 'const' enums can only be used in property or index access expressions or the right hand side of an import declaration or export assignment or type query.
+                         ~~~~~~~~~
+!!! error TS2488: Type 'typeof ConstTest' must have a '[Symbol.iterator]()' method that returns an iterator.
+    }
+    
diff --git a/tests/baselines/reference/enumIteratorConstEnum1.js b/tests/baselines/reference/enumIteratorConstEnum1.js
@@ -0,0 +1,16 @@
+//// [enumIteratorConstEnum1.ts]
+const enum ConstTest {
+  Foo,
+  Bar,
+  Baz,
+}
+
+// Should be an error.
+for (const member of ConstTest) {
+}
+
+
+//// [enumIteratorConstEnum1.js]
+// Should be an error.
+for (const member of ConstTest) {
+}
diff --git a/tests/baselines/reference/enumIteratorConstEnum1.symbols b/tests/baselines/reference/enumIteratorConstEnum1.symbols
@@ -0,0 +1,20 @@
+=== tests/cases/compiler/enumIteratorConstEnum1.ts ===
+const enum ConstTest {
+>ConstTest : Symbol(ConstTest, Decl(enumIteratorConstEnum1.ts, 0, 0))
+
+  Foo,
+>Foo : Symbol(ConstTest.Foo, Decl(enumIteratorConstEnum1.ts, 0, 22))
+
+  Bar,
+>Bar : Symbol(ConstTest.Bar, Decl(enumIteratorConstEnum1.ts, 1, 6))
+
+  Baz,
+>Baz : Symbol(ConstTest.Baz, Decl(enumIteratorConstEnum1.ts, 2, 6))
+}
+
+// Should be an error.
+for (const member of ConstTest) {
+>member : Symbol(member, Decl(enumIteratorConstEnum1.ts, 7, 10))
+>ConstTest : Symbol(ConstTest, Decl(enumIteratorConstEnum1.ts, 0, 0))
+}
+
diff --git a/tests/baselines/reference/enumIteratorConstEnum1.types b/tests/baselines/reference/enumIteratorConstEnum1.types
@@ -0,0 +1,20 @@
+=== tests/cases/compiler/enumIteratorConstEnum1.ts ===
+const enum ConstTest {
+>ConstTest : ConstTest
+
+  Foo,
+>Foo : ConstTest.Foo
+
+  Bar,
+>Bar : ConstTest.Bar
+
+  Baz,
+>Baz : ConstTest.Baz
+}
+
+// Should be an error.
+for (const member of ConstTest) {
+>member : any
+>ConstTest : typeof ConstTest
+}
+
diff --git a/tests/cases/compiler/enumIteratorBasics1.ts b/tests/cases/compiler/enumIteratorBasics1.ts
@@ -0,0 +1,17 @@
+// @target: ES6
+enum Test {
+  Foo,
+  Bar,
+  Baz,
+}
+
+const AlsoTest = Test;
+
+for (const member of AlsoTest) {
+  const x: string = member[0];
+  const y: Test = member[1];
+  console.log(x, y);
+}
+
+const members: [string, Test][] = [...AlsoTest];
+console.log(members);
diff --git a/tests/cases/compiler/enumIteratorConstEnum1.ts b/tests/cases/compiler/enumIteratorConstEnum1.ts
