Enhance support for switch statements

src/ast/__tests__/switch-statement-extractor.test.ts

@@ -429,7 +429,6 @@ describe("extract SwitchStatement correctly", () => {
     };
 
     const ast = parse(programStr);
-    console.log(JSON.stringify(ast, null, 2));
     expect(ast).toEqual(expectedAst);
   });
 });

src/ast/astExtractor/ast-extractor.ts

@@ -1,11 +1,11 @@
 import { BaseJavaCstVisitorWithDefaults, CstNode, TypeDeclarationCtx } from "java-parser";
 
-import { NormalClassDeclaration } from "../types/classes";
+import { ClassDeclaration } from "../types/classes";
 import { AST } from "../types/packages-and-modules";
 import { ClassExtractor } from "./class-extractor";
 
 export class ASTExtractor extends BaseJavaCstVisitorWithDefaults {
-  private topLevelClassOrInterfaceDeclarations: NormalClassDeclaration[] = [];
+  private topLevelClassOrInterfaceDeclarations: ClassDeclaration[] = [];
 
   extract(cst: CstNode): AST {
     this.visit(cst);

src/ast/types/ast.ts

@@ -12,6 +12,7 @@ import {
 } from "./blocks-and-statements";
 import {
   ConstructorDeclaration,
+  EnumDeclaration,
   FieldDeclaration,
   MethodDeclaration,
   NormalClassDeclaration,
@@ -29,6 +30,7 @@ interface NodeMap {
   MethodInvocation: MethodInvocation;
   ReturnStatement: ReturnStatement;
   NormalClassDeclaration: NormalClassDeclaration;
+  EnumDeclaration: EnumDeclaration;
   ClassInstanceCreationExpression: ClassInstanceCreationExpression;
   ConstructorDeclaration: ConstructorDeclaration;
   ExplicitConstructorInvocation: ExplicitConstructorInvocation;

src/ast/types/classes.ts

@@ -1,7 +1,7 @@
 import { BaseNode } from "./ast";
 import { Block, VariableDeclarator } from "./blocks-and-statements";
 
-export type ClassDeclaration = NormalClassDeclaration;
+export type ClassDeclaration = NormalClassDeclaration | EnumDeclaration;
 
 export interface NormalClassDeclaration extends BaseNode {
   kind: "NormalClassDeclaration";
@@ -11,6 +11,26 @@ export interface NormalClassDeclaration extends BaseNode {
   classBody: Array<ClassBodyDeclaration>;
 }
 
+export interface EnumDeclaration extends BaseNode {
+  kind: "EnumDeclaration";
+  classModifier: Array<ClassModifier>;
+  typeIdentifier: Identifier;
+  enumBody: EnumBody;
+}
+
+export interface EnumBody extends BaseNode {
+  kind: "EnumBody";
+  constants: Array<EnumConstant>;
+  bodyMembers?: Array<ClassBodyDeclaration>;
+}
+
+export interface EnumConstant extends BaseNode {
+  kind: "EnumConstant";
+  name: Identifier;
+  arguments?: Array<any>;
+  classBody?: Array<ClassBodyDeclaration>;
+}
+
 export type ClassModifier =
   | "public"
   | "protected"
@@ -20,9 +40,10 @@ export type ClassModifier =
   | "final"
   | "sealed"
   | "non-sealed"
-  | "strictfp";
+  | "strictfp"
+  | "enum";
 
-export type ClassBodyDeclaration = ClassMemberDeclaration | ConstructorDeclaration;
+export type ClassBodyDeclaration = ClassMemberDeclaration | ConstructorDeclaration | EnumDeclaration;
 export type ClassMemberDeclaration = MethodDeclaration | FieldDeclaration;
 
 export interface ConstructorDeclaration extends BaseNode {

src/compiler/__tests__/index.ts

@@ -9,6 +9,7 @@ import { methodInvocationTest } from "./tests/methodInvocation.test";
 import { importTest } from "./tests/import.test";
 import { arrayTest } from "./tests/array.test";
 import { classTest } from "./tests/class.test";
+import { enumTest } from "./tests/enum.test";
 import { typeConversionTest } from "./tests/typeConversion.test";
 
 describe("compiler tests", () => {
@@ -23,5 +24,6 @@ describe("compiler tests", () => {
   importTest();
   arrayTest();
   classTest();
+  enumTest();
   typeConversionTest();
 })

src/compiler/__tests__/tests/enum.test.ts

@@ -0,0 +1,212 @@
+import {
+  runTest,
+  testCase,
+} from "../__utils__/test-utils";
+import { compileFromSource } from "../../index";
+
+const testCases: testCase[] = [
+  {
+    comment: "member enum constant access",
+    program: `
+      public class Main {
+        public enum Day {
+          SUNDAY,
+          MONDAY,
+          TUESDAY,
+          WEDNESDAY,
+          THURSDAY,
+          FRIDAY,
+          SATURDAY
+        }
+
+        public static void main(String[] args) {
+          Day day = Day.SUNDAY;
+        }
+      }
+    `,
+    expectedLines: [],
+  },
+  {
+    comment: "member enum switch selects the matching constant",
+    program: `
+      public class Main {
+        public enum Light { RED, YELLOW, GREEN }
+
+        public static void main(String[] args) {
+          Light light = Light.GREEN;
+          switch (light) {
+            case RED:
+              System.out.println("stop");
+              break;
+            case GREEN:
+              System.out.println("go");
+              break;
+            default:
+              System.out.println("wait");
+          }
+        }
+      }
+    `,
+    expectedLines: ["go"],
+  },
+  {
+    comment: "member enum switch matches the first of seven constants",
+    program: `
+      class Main {
+        public enum Day {
+          SUNDAY, MONDAY, TUESDAY, WEDNESDAY,
+          THURSDAY, FRIDAY, SATURDAY
+        }
+
+        public static void main(String[] args) {
+          Day day = Day.SUNDAY;
+          System.out.println(10);
+          switch (day) {
+            case SUNDAY:
+              System.out.println(0);
+              break;
+            case MONDAY:
+              System.out.println(1);
+              break;
+            case TUESDAY:
+              System.out.println(2);
+              break;
+            case WEDNESDAY:
+              System.out.println(3);
+              break;
+            case THURSDAY:
+              System.out.println(4);
+              break;
+            case FRIDAY:
+              System.out.println(5);
+              break;
+            case SATURDAY:
+              System.out.println(6);
+              break;
+            default:
+              break;
+          }
+        }
+      }
+    `,
+    expectedLines: ["10", "0"],
+  },
+  {
+    comment: "enum switch and synthetic methods",
+    program: `
+      public enum Color {
+        RED,
+        BLUE
+      }
+
+      public class Main {
+        public static void main(String[] args) {
+          Color red = Color.valueOf("RED");
+          System.out.println(Color.RED.ordinal());
+          System.out.println(Color.BLUE.name());
+          System.out.println(red.toString());
+
+          Color selector = Color.BLUE;
+          switch (selector) {
+            case RED:
+              System.out.println("bad");
+              break;
+            case BLUE:
+              System.out.println("ok");
+              break;
+            default:
+              System.out.println("default");
+          }
+        }
+      }
+    `,
+    expectedLines: ["0", "BLUE", "RED", "ok"],
+  },
+  {
+    comment: "enum values returns cloned array",
+    program: `
+      public enum Direction {
+        NORTH,
+        SOUTH
+      }
+
+      public class Main {
+        public static void main(String[] args) {
+          Direction[] copy = Direction.values();
+          copy[0] = Direction.SOUTH;
+          Direction[] fresh = Direction.values();
+
+          switch (fresh[0]) {
+            case NORTH:
+              System.out.println("fresh");
+              break;
+            default:
+              System.out.println("bad");
+          }
+
+          switch (copy[0]) {
+            case SOUTH:
+              System.out.println("mutated");
+              break;
+            default:
+              System.out.println("bad");
+          }
+        }
+      }
+    `,
+    expectedLines: ["fresh", "mutated"],
+  },
+  {
+    comment: "enum constructors and instance fields",
+    program: `
+      public enum Planet {
+        EARTH(1),
+        MARS(2);
+
+        private int moons;
+
+        private Planet(int moons) {
+          this.moons = moons;
+        }
+
+        public int moons() {
+          return this.moons;
+        }
+      }
+
+      public class Main {
+        public static void main(String[] args) {
+          Planet mars = Planet.valueOf("MARS");
+          System.out.println(Planet.EARTH.moons());
+          System.out.println(mars.moons());
+        }
+      }
+    `,
+    expectedLines: ["1", "2"],
+  },
+];
+
+export const enumTest = () => describe("enums", () => {
+  for (let testCase of testCases) {
+    const { comment: comment, program: program, expectedLines: expectedLines } = testCase;
+    it(comment, () => runTest(program, expectedLines));
+  }
+
+  it("rejects qualified enum switch labels", () => {
+    expect(() =>
+      compileFromSource(`
+        class Main {
+          enum Day { SUNDAY }
+
+          public static void main(String[] args) {
+            Day day = Day.SUNDAY;
+            switch (day) {
+              case Day.SUNDAY:
+                break;
+            }
+          }
+        }
+      `)
+    ).toThrow(SyntaxError);
+  });
+});

src/compiler/code-generator.ts

@@ -1595,14 +1595,37 @@ const codeGenerators: { [type: string]: (node: Node, cg: CodeGenerator) => Compi
     const { stackSize: exprStackSize, resultType } = compile(expression, cg)
     let maxStack = exprStackSize
 
+    // If the expression is an enum type, invoke ordinal() to convert to int and then continue
+    let _resultType = resultType
+    let enumTypeName: string | null = null
+    if (_resultType && _resultType.startsWith('L') && _resultType !== 'Ljava/lang/String;') {
+      const clean = _resultType.replace(/^L|;$/g, '')
+      try {
+        const classInfo = cg.symbolTable.queryClass(clean)
+        if (classInfo.isEnum) {
+          // Generated enums provide their own ordinal() method so they run without java.lang.Enum.
+          cg.code.push(
+            OPCODE.INVOKEVIRTUAL,
+            0,
+            cg.constantPoolManager.indexMethodrefInfo(clean, 'ordinal', '()I')
+          )
+          _resultType = 'I'
+          enumTypeName = clean
+          maxStack = Math.max(maxStack, exprStackSize + 1)
+        }
+      } catch (e) {
+        // ignore: not a known class
+      }
+    }
+
     const caseLabels: Label[] = cases.map(() => cg.generateNewLabel())
     const defaultLabel = cg.generateNewLabel()
     const endLabel = cg.generateNewLabel()
 
     // Track the switch statement's end label
     cg.switchLabels.push(endLabel)
 
-    if (['I', 'B', 'S', 'C'].includes(resultType)) {
+    if (['I', 'B', 'S', 'C'].includes(_resultType)) {
       const caseValues: number[] = []
       const caseLabelMap: Map<number, Label> = new Map()
       let hasDefault = false
@@ -1611,7 +1634,18 @@ const codeGenerators: { [type: string]: (node: Node, cg: CodeGenerator) => Compi
       cases.forEach((caseGroup, index) => {
         caseGroup.labels.forEach(label => {
           if (label.kind === 'CaseLabel') {
-            const value = parseInt((label.expression as Literal).literalType.value)
+            const value =
+              label.expression.kind === 'ExpressionName' && enumTypeName
+                ? (() => {
+                    const fields = cg.symbolTable.queryField(
+                      `${enumTypeName}.${label.expression.name}`
+                    )
+                    const field = fields[fields.length - 1] as FieldInfo
+                    if (field.parentClassName !== enumTypeName || field.ordinal === undefined)
+                      throw new Error(`Invalid enum switch label: ${label.expression.name}`)
+                    return field.ordinal
+                  })()
+                : parseInt((label.expression as Literal).literalType.value)
             caseValues.push(value)
             caseLabelMap.set(value, caseLabels[index])
           } else if (label.kind === 'DefaultLabel') {
@@ -1781,7 +1815,7 @@ const codeGenerators: { [type: string]: (node: Node, cg: CodeGenerator) => Compi
       }
 
       endLabel.offset = cg.code.length
-    } else if (resultType === 'Ljava/lang/String;') {
+    } else if (_resultType === 'Ljava/lang/String;') {
       // **String Switch Handling**
       const hashCaseMap: Map<number, Label> = new Map()
 
@@ -1938,7 +1972,7 @@ const codeGenerators: { [type: string]: (node: Node, cg: CodeGenerator) => Compi
       endLabel.offset = cg.code.length
     } else {
       throw new Error(
-        `Switch statements only support byte, short, int, char, or String types. Found: ${resultType}`
+        `Switch statements only support byte, short, int, char, String, or enum types. Found: ${_resultType}`
       )
     }