Enhance support for switch statements

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,107 @@
+import {
+  runTest,
+  testCase,
+} from "../__utils__/test-utils";
+
+const testCases: testCase[] = [
+  {
+    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 Color.RED:
+              System.out.println("bad");
+              break;
+            case Color.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 Direction.NORTH:
+              System.out.println("fresh");
+              break;
+            default:
+              System.out.println("bad");
+          }
+
+          switch (copy[0]) {
+            case Direction.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));
+  }
+});

src/compiler/binary-writer.ts

@@ -49,7 +49,9 @@ export class BinaryWriter {
     fs.writeFileSync(filename, binary)
   }
 
-  private normalizeClassFile(classFile: ClassFile | Class | Array<ClassFile> | Array<Class>): ClassFile {
+  private normalizeClassFile(
+    classFile: ClassFile | Class | Array<ClassFile> | Array<Class>
+  ): ClassFile {
     if (Array.isArray(classFile)) {
       if (classFile.length === 0) {
         throw new Error('BinaryWriter expected a non-empty array of classes')

src/compiler/code-generator.ts

@@ -1,4 +1,5 @@
 import { OPCODE } from '../ClassFile/constants/instructions'
+import { ACCESS_FLAGS } from '../ClassFile/types'
 import { ExceptionHandler, AttributeInfo } from '../ClassFile/types/attributes'
 import { FIELD_FLAGS } from '../ClassFile/types/fields'
 import { METHOD_FLAGS } from '../ClassFile/types/methods'
@@ -190,19 +191,19 @@ const EMPTY_TYPE: string = ''
 function areClassTypesCompatible(fromType: string, toType: string, cg: CodeGenerator): boolean {
   const cleanFrom = fromType.replace(/^L|;$/g, '')
   const cleanTo = toType.replace(/^L|;$/g, '')
-  if (cleanFrom === cleanTo) return true;
+  if (cleanFrom === cleanTo) return true
 
   try {
-    let current = cg.symbolTable.queryClass(cleanFrom);
+    let current = cg.symbolTable.queryClass(cleanFrom)
     while (current.parentClassName) {
-      const parentClean = current.parentClassName;
-      if (parentClean === cleanTo) return true;
-      current = cg.symbolTable.queryClass(parentClean);
+      const parentClean = current.parentClassName
+      if (parentClean === cleanTo) return true
+      current = cg.symbolTable.queryClass(parentClean)
     }
   } catch (e) {
-    return false;
+    return false
   }
-  return false;
+  return false
 }
 
 function handleImplicitTypeConversion(fromType: string, toType: string, cg: CodeGenerator): number {
@@ -838,30 +839,30 @@ const codeGenerators: { [type: string]: (node: Node, cg: CodeGenerator) => Compi
     // --- Handle super. calls ---
     if (n.identifier.startsWith('super.')) {
       candidateMethods = cg.symbolTable.queryMethod(n.identifier.slice(6)) as MethodInfos
-      candidateMethods = candidateMethods.filter(method =>
-        method.className == cg.symbolTable.queryClass(cg.currentClass).parentClassName)
-      cg.code.push(OPCODE.ALOAD, 0);
+      candidateMethods = candidateMethods.filter(
+        method => method.className == cg.symbolTable.queryClass(cg.currentClass).parentClassName
+      )
+      cg.code.push(OPCODE.ALOAD, 0)
     }
     // --- Handle qualified calls (e.g. System.out.println or p.show) ---
     else if (n.identifier.includes('.')) {
-      const lastDot = n.identifier.lastIndexOf('.');
-      const receiverStr = n.identifier.slice(0, lastDot);
+      const lastDot = n.identifier.lastIndexOf('.')
+      const receiverStr = n.identifier.slice(0, lastDot)
 
       if (receiverStr === 'this') {
         candidateMethods = cg.symbolTable.queryMethod(n.identifier.slice(5)) as MethodInfos
-        candidateMethods = candidateMethods.filter(method =>
-          method.className == cg.currentClass)
-        cg.code.push(OPCODE.ALOAD, 0);
+        candidateMethods = candidateMethods.filter(method => method.className == cg.currentClass)
+        cg.code.push(OPCODE.ALOAD, 0)
       } else {
-        const recvRes = compile({ kind: 'ExpressionName', name: receiverStr }, cg);
-        maxStack = Math.max(maxStack, recvRes.stackSize);
+        const recvRes = compile({ kind: 'ExpressionName', name: receiverStr }, cg)
+        maxStack = Math.max(maxStack, recvRes.stackSize)
         candidateMethods = cg.symbolTable.queryMethod(n.identifier).pop() as MethodInfos
       }
     }
     // --- Handle unqualified calls ---
     else {
       candidateMethods = cg.symbolTable.queryMethod(n.identifier) as MethodInfos
-      unqualifiedCall = true;
+      unqualifiedCall = true
     }
 
     // Filter candidate methods by matching the argument list.
@@ -903,11 +904,15 @@ const codeGenerators: { [type: string]: (node: Node, cg: CodeGenerator) => Compi
             .slice(1, methodMatches[i].typeDescriptor.indexOf(')'))
             .match(/(\[+[BCDFIJSZ])|(\[+L[^;]+;)|[BCDFIJSZ]|L[^;]+;/g) || []
         if (
-          candParams.map((p, idx) => isSubtype(p, currParams[idx], cg)).reduce((a, b) => a && b, true)
+          candParams
+            .map((p, idx) => isSubtype(p, currParams[idx], cg))
+            .reduce((a, b) => a && b, true)
         ) {
           selectedMethod = methodMatches[i]
         } else if (
-          !currParams.map((p, idx) => isSubtype(p, candParams[idx], cg)).reduce((a, b) => a && b, true)
+          !currParams
+            .map((p, idx) => isSubtype(p, candParams[idx], cg))
+            .reduce((a, b) => a && b, true)
         ) {
           throw new AmbiguousMethodCallError(n.identifier + argDescs.join(','))
         }
@@ -1266,7 +1271,7 @@ const codeGenerators: { [type: string]: (node: Node, cg: CodeGenerator) => Compi
     try {
       info = cg.symbolTable.queryVariable(name)
     } catch (e) {
-      return { stackSize: 1, resultType: 'Ljava/lang/Class;' };
+      return { stackSize: 1, resultType: 'Ljava/lang/Class;' }
     }
     if (Array.isArray(info)) {
       const fieldInfos = info
@@ -1375,14 +1380,35 @@ 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
+    if (_resultType && _resultType.startsWith('L') && _resultType !== 'Ljava/lang/String;') {
+      const clean = _resultType.replace(/^L|;$/g, '')
+      try {
+        const classInfo = cg.symbolTable.queryClass(clean)
+        if (classInfo.accessFlags & ACCESS_FLAGS.ACC_ENUM) {
+          // call java.lang.Enum.ordinal() (returns int)
+          cg.code.push(
+            OPCODE.INVOKEVIRTUAL,
+            0,
+            cg.constantPoolManager.indexMethodrefInfo('java/lang/Enum', 'ordinal', '()I')
+          )
+          _resultType = 'I'
+          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
@@ -1470,7 +1496,12 @@ const codeGenerators: { [type: string]: (node: Node, cg: CodeGenerator) => Compi
         // Generate lookup table (pairs of case values and corresponding labels)
         caseValues.forEach((value, index) => {
           // push 4-byte key
-          cg.code.push((value >> 24) & 0xff, (value >> 16) & 0xff, (value >> 8) & 0xff, value & 0xff)
+          cg.code.push(
+            (value >> 24) & 0xff,
+            (value >> 16) & 0xff,
+            (value >> 8) & 0xff,
+            value & 0xff
+          )
           // reserve 4 bytes for the branch target
           cg.code.push(0, 0, 0, 0)
           // label offset starts after the 4-byte key
@@ -1556,7 +1587,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()
 
@@ -1613,7 +1644,12 @@ const codeGenerators: { [type: string]: (node: Node, cg: CodeGenerator) => Compi
       // Populate LOOKUPSWITCH
       const hashLabels: Label[] = []
       hashCaseMap.forEach((label, hashCode) => {
-        cg.code.push((hashCode >> 24) & 0xff, (hashCode >> 16) & 0xff, (hashCode >> 8) & 0xff, hashCode & 0xff)
+        cg.code.push(
+          (hashCode >> 24) & 0xff,
+          (hashCode >> 16) & 0xff,
+          (hashCode >> 8) & 0xff,
+          hashCode & 0xff
+        )
         // reserve 4 bytes for the branch target
         cg.code.push(0, 0, 0, 0)
         // label offset starts after the 4-byte key
@@ -1708,7 +1744,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}`
       )
     }