Skip to content

Commit 532d09d

Browse files
ark120202Perryvw
authored andcommitted
Simplify and improve enum transform (#685)
* Simplify and improve enum transform * Tests * Fix enum toString test not working * Don't generate reverse mapping for nil member values
1 parent 688517d commit 532d09d

File tree

5 files changed

+110
-163
lines changed

5 files changed

+110
-163
lines changed

src/LuaTransformer.ts

Lines changed: 53 additions & 91 deletions
Original file line numberDiff line numberDiff line change
@@ -1727,115 +1727,77 @@ export class LuaTransformer {
17271727
return result;
17281728
}
17291729

1730-
public transformEnumDeclaration(enumDeclaration: ts.EnumDeclaration): StatementVisitResult {
1731-
const type = this.checker.getTypeAtLocation(enumDeclaration);
1732-
1733-
// Const enums should never appear in the resulting code
1734-
if (type.symbol.getFlags() & ts.SymbolFlags.ConstEnum) {
1730+
public transformEnumDeclaration(node: ts.EnumDeclaration): StatementVisitResult {
1731+
if (ts.getCombinedModifierFlags(node) & ts.ModifierFlags.Const && !this.options.preserveConstEnums) {
17351732
return undefined;
17361733
}
17371734

1735+
const type = this.checker.getTypeAtLocation(node);
17381736
const membersOnly = tsHelper.getCustomDecorators(type, this.checker).has(DecoratorKind.CompileMembersOnly);
1739-
17401737
const result: tstl.Statement[] = [];
17411738

17421739
if (!membersOnly) {
1743-
const name = this.transformIdentifier(enumDeclaration.name);
1740+
const name = this.transformIdentifier(node.name);
17441741
const table = tstl.createTableExpression();
1745-
result.push(...this.createLocalOrExportedOrGlobalDeclaration(name, table, enumDeclaration));
1742+
result.push(...this.createLocalOrExportedOrGlobalDeclaration(name, table, node));
17461743
}
17471744

1748-
for (const enumMember of this.computeEnumMembers(enumDeclaration)) {
1749-
const memberName = this.transformPropertyName(enumMember.name);
1750-
if (membersOnly) {
1751-
const enumSymbol = this.checker.getSymbolAtLocation(enumDeclaration.name);
1752-
const exportScope = enumSymbol ? this.getSymbolExportScope(enumSymbol) : undefined;
1745+
const enumReference = this.transformExpression(node.name);
1746+
for (const member of node.members) {
1747+
const memberName = this.transformPropertyName(member.name);
17531748

1754-
if (tstl.isIdentifier(memberName)) {
1755-
result.push(
1756-
...this.createLocalOrExportedOrGlobalDeclaration(
1757-
memberName,
1758-
enumMember.value,
1759-
enumDeclaration,
1760-
undefined,
1761-
exportScope
1762-
)
1763-
);
1764-
} else {
1765-
result.push(
1766-
...this.createLocalOrExportedOrGlobalDeclaration(
1767-
tstl.createIdentifier(enumMember.name.getText(), enumMember.name),
1768-
enumMember.value,
1769-
enumDeclaration,
1770-
undefined,
1771-
exportScope
1772-
)
1773-
);
1749+
let valueExpression: tstl.Expression | undefined;
1750+
const constEnumValue = this.tryGetConstEnumValue(member);
1751+
if (constEnumValue) {
1752+
valueExpression = constEnumValue;
1753+
} else if (member.initializer) {
1754+
if (ts.isIdentifier(member.initializer)) {
1755+
const symbol = this.checker.getSymbolAtLocation(member.initializer);
1756+
if (
1757+
symbol &&
1758+
symbol.valueDeclaration &&
1759+
ts.isEnumMember(symbol.valueDeclaration) &&
1760+
symbol.valueDeclaration.parent === node
1761+
) {
1762+
const otherMemberName = this.transformPropertyName(symbol.valueDeclaration.name);
1763+
valueExpression = tstl.createTableIndexExpression(enumReference, otherMemberName);
1764+
}
17741765
}
1775-
} else {
1776-
const enumTable = this.transformIdentifierExpression(enumDeclaration.name);
1777-
const property = tstl.createTableIndexExpression(enumTable, memberName);
1778-
result.push(tstl.createAssignmentStatement(property, enumMember.value, enumMember.original));
17791766

1780-
const valueIndex = tstl.createTableIndexExpression(enumTable, enumMember.value);
1781-
result.push(tstl.createAssignmentStatement(valueIndex, memberName, enumMember.original));
1782-
}
1783-
}
1784-
1785-
return result;
1786-
}
1787-
1788-
protected computeEnumMembers(
1789-
node: ts.EnumDeclaration
1790-
): Array<{ name: ts.PropertyName; value: tstl.Expression; original: ts.Node }> {
1791-
let numericValue = 0;
1792-
let hasStringInitializers = false;
1793-
1794-
const valueMap = new Map<ts.PropertyName, ExpressionVisitResult>();
1795-
1796-
return node.members.map(member => {
1797-
let valueExpression: ExpressionVisitResult;
1798-
if (member.initializer) {
1799-
if (ts.isNumericLiteral(member.initializer)) {
1800-
numericValue = Number(member.initializer.text);
1801-
valueExpression = this.transformNumericLiteral(member.initializer);
1802-
numericValue++;
1803-
} else if (ts.isStringLiteral(member.initializer)) {
1804-
hasStringInitializers = true;
1805-
valueExpression = this.transformStringLiteral(member.initializer);
1806-
} else {
1807-
if (ts.isIdentifier(member.initializer)) {
1808-
const [isEnumMember, originalName] = tsHelper.isEnumMember(node, member.initializer);
1809-
if (isEnumMember === true && originalName !== undefined) {
1810-
if (valueMap.has(originalName)) {
1811-
valueExpression = valueMap.get(originalName)!;
1812-
} else {
1813-
throw new Error(`Expected valueMap to contain ${originalName}`);
1814-
}
1815-
} else {
1816-
valueExpression = this.transformExpression(member.initializer);
1817-
}
1818-
} else {
1819-
valueExpression = this.transformExpression(member.initializer);
1820-
}
1767+
if (!valueExpression) {
1768+
valueExpression = this.transformExpression(member.initializer);
18211769
}
1822-
} else if (hasStringInitializers) {
1823-
throw TSTLErrors.HeterogeneousEnum(node);
18241770
} else {
1825-
valueExpression = tstl.createNumericLiteral(numericValue);
1826-
numericValue++;
1771+
valueExpression = tstl.createNilLiteral();
18271772
}
18281773

1829-
valueMap.set(member.name, valueExpression);
1774+
if (membersOnly) {
1775+
const enumSymbol = this.checker.getSymbolAtLocation(node.name);
1776+
const exportScope = enumSymbol ? this.getSymbolExportScope(enumSymbol) : undefined;
18301777

1831-
const enumMember = {
1832-
name: member.name,
1833-
original: member,
1834-
value: valueExpression,
1835-
};
1778+
result.push(
1779+
...this.createLocalOrExportedOrGlobalDeclaration(
1780+
tstl.isIdentifier(memberName)
1781+
? memberName
1782+
: tstl.createIdentifier(member.name.getText(), member.name),
1783+
valueExpression,
1784+
node,
1785+
undefined,
1786+
exportScope
1787+
)
1788+
);
1789+
} else {
1790+
const memberAccessor = tstl.createTableIndexExpression(enumReference, memberName);
1791+
result.push(tstl.createAssignmentStatement(memberAccessor, valueExpression, member));
18361792

1837-
return enumMember;
1838-
});
1793+
if (!tstl.isStringLiteral(valueExpression) && !tstl.isNilLiteral(valueExpression)) {
1794+
const reverseMemberAccessor = tstl.createTableIndexExpression(enumReference, memberAccessor);
1795+
result.push(tstl.createAssignmentStatement(reverseMemberAccessor, memberName, member));
1796+
}
1797+
}
1798+
}
1799+
1800+
return result;
18391801
}
18401802

18411803
protected transformGeneratorFunction(
@@ -4707,7 +4669,7 @@ export class LuaTransformer {
47074669
}
47084670

47094671
private tryGetConstEnumValue(
4710-
node: ts.PropertyAccessExpression | ts.ElementAccessExpression
4672+
node: ts.EnumMember | ts.PropertyAccessExpression | ts.ElementAccessExpression
47114673
): tstl.Expression | undefined {
47124674
const value = this.checker.getConstantValue(node);
47134675
if (typeof value === "string") {

src/TSHelper.ts

Lines changed: 0 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -871,26 +871,6 @@ export function isStandardLibraryType(type: ts.Type, name: string | undefined, p
871871
return isStandardLibraryDeclaration(declaration, program);
872872
}
873873

874-
export function isEnumMember(
875-
enumDeclaration: ts.EnumDeclaration,
876-
value: ts.Expression
877-
): [true, ts.PropertyName] | [false, undefined] {
878-
if (ts.isIdentifier(value)) {
879-
const enumMember = enumDeclaration.members.find(m => ts.isIdentifier(m.name) && m.name.text === value.text);
880-
if (enumMember !== undefined) {
881-
if (enumMember.initializer && ts.isIdentifier(enumMember.initializer)) {
882-
return isEnumMember(enumDeclaration, enumMember.initializer);
883-
} else {
884-
return [true, enumMember.name];
885-
}
886-
} else {
887-
return [false, undefined];
888-
}
889-
} else {
890-
return [false, undefined];
891-
}
892-
}
893-
894874
export function isWithinLiteralAssignmentStatement(node: ts.Node): boolean {
895875
if (!node.parent) {
896876
return false;

src/TSTLErrors.ts

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -31,13 +31,6 @@ export const InvalidInstanceOfLuaTable = (node: ts.Node) =>
3131
export const ForbiddenLuaTableUseException = (description: string, node: ts.Node) =>
3232
new TranspileError(`Invalid @luaTable usage: ${description}`, node);
3333

34-
export const HeterogeneousEnum = (node: ts.Node) =>
35-
new TranspileError(
36-
`Invalid heterogeneous enum. Enums should either specify no member values, ` +
37-
`or specify values (of the same type) for all members.`,
38-
node
39-
);
40-
4134
export const InvalidDecoratorArgumentNumber = (name: string, got: number, expected: number, node: ts.Node) =>
4235
new TranspileError(`${name} expects ${expected} argument(s) but got ${got}.`, node);
4336

test/unit/enum.spec.ts

Lines changed: 45 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
import * as TSTLErrors from "../../src/TSTLErrors";
21
import * as util from "../util";
32

43
// TODO: string.toString()
@@ -10,21 +9,18 @@ const serializeEnum = (identifier: string) => `(() => {
109
return mappedTestEnum;
1110
})()`;
1211

13-
// TODO: Move to namespace tests?
14-
test("in a namespace", () => {
15-
util.testModule`
16-
namespace Test {
17-
export enum TestEnum {
18-
A,
19-
B,
12+
describe("initializers", () => {
13+
test("string", () => {
14+
util.testFunction`
15+
enum TestEnum {
16+
A = "A",
17+
B = "B",
2018
}
21-
}
2219
23-
export const result = ${serializeEnum("Test.TestEnum")}
24-
`.expectToMatchJsResult();
25-
});
20+
return ${serializeEnum("TestEnum")}
21+
`.expectToMatchJsResult();
22+
});
2623

27-
describe("initializers", () => {
2824
test("expression", () => {
2925
util.testFunction`
3026
const value = 6;
@@ -37,6 +33,18 @@ describe("initializers", () => {
3733
`.expectToMatchJsResult();
3834
});
3935

36+
test("expression with side effect", () => {
37+
util.testFunction`
38+
let value = 0;
39+
enum TestEnum {
40+
A = value++,
41+
B = A,
42+
}
43+
44+
return ${serializeEnum("TestEnum")}
45+
`.expectToMatchJsResult();
46+
});
47+
4048
test("inference", () => {
4149
util.testFunction`
4250
enum TestEnum {
@@ -61,7 +69,7 @@ describe("initializers", () => {
6169
`.expectToMatchJsResult();
6270
});
6371

64-
test("other member reference", () => {
72+
test("member reference", () => {
6573
util.testFunction`
6674
enum TestEnum {
6775
A,
@@ -72,38 +80,38 @@ describe("initializers", () => {
7280
return ${serializeEnum("TestEnum")}
7381
`.expectToMatchJsResult();
7482
});
75-
});
7683

77-
test("invalid heterogeneous enum", () => {
78-
util.testFunction`
79-
enum TestEnum {
80-
A,
81-
B = "B",
82-
C,
83-
}
84-
`
85-
.disableSemanticCheck()
86-
.expectToHaveDiagnosticOfError(TSTLErrors.HeterogeneousEnum(util.nodeStub));
84+
test("string literal member reference", () => {
85+
util.testFunction`
86+
enum TestEnum {
87+
["A"],
88+
"B" = A,
89+
C = B,
90+
}
91+
92+
return ${serializeEnum("TestEnum")}
93+
`.expectToMatchJsResult();
94+
});
8795
});
8896

8997
describe("const enum", () => {
9098
const expectToBeConst: util.TapCallback = builder =>
9199
expect(builder.getMainLuaCodeChunk()).not.toContain("TestEnum");
92100

93-
test.each(["", "declare"])("%s without initializer", () => {
94-
util.testFunction`
95-
const enum TestEnum {
101+
test.each(["", "declare "])("%swithout initializer", modifier => {
102+
util.testModule`
103+
${modifier} const enum TestEnum {
96104
A,
97105
B,
98106
}
99107
100-
return TestEnum.A;
108+
export const A = TestEnum.A;
101109
`
102110
.tap(expectToBeConst)
103111
.expectToMatchJsResult();
104112
});
105113

106-
test("with initializer", () => {
114+
test("with string initializer", () => {
107115
util.testFunction`
108116
const enum TestEnum {
109117
A = "ONE",
@@ -131,26 +139,18 @@ describe("const enum", () => {
131139
});
132140
});
133141

134-
test("enum toString", () => {
135-
const code = `
142+
test("toString", () => {
143+
util.testFunction`
136144
enum TestEnum {
137145
A,
138146
B,
139147
C,
140148
}
141-
let test = TestEnum.A;
142-
return test.toString();`;
143-
expect(util.transpileAndExecute(code)).toBe(0);
144-
});
145149
146-
test("enum concat", () => {
147-
const code = `
148-
enum TestEnum {
149-
A,
150-
B,
151-
C,
150+
function foo(value: TestEnum) {
151+
return value.toString();
152152
}
153-
let test = TestEnum.A;
154-
return test + "_foobar";`;
155-
expect(util.transpileAndExecute(code)).toBe("0_foobar");
153+
154+
return foo(TestEnum.A);
155+
`.expectToMatchJsResult();
156156
});

test/unit/namespaces.spec.ts

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,3 +119,15 @@ test("`import =` on a namespace", () => {
119119
export const result = importedFunc();
120120
`.expectToMatchJsResult();
121121
});
122+
123+
test("enum in a namespace", () => {
124+
util.testModule`
125+
namespace Test {
126+
export enum TestEnum {
127+
A,
128+
}
129+
}
130+
131+
export const result = Test.TestEnum.A;
132+
`.expectToMatchJsResult();
133+
});

0 commit comments

Comments
 (0)