Add support for enhanced for loops

Author: bryanlohxz

src/types/ast/astExtractor/statement-extractor.ts

@@ -30,10 +30,12 @@ import {
   ExpressionStatementCtx,
   LocalVariableTypeCtx,
   VariableDeclaratorListCtx,
-  VariableDeclaratorCtx
+  VariableDeclaratorCtx,
+  EnhancedForStatementCtx
 } from 'java-parser'
 import {
   BasicForStatement,
+  EnhancedForStatement,
   ExpressionStatement,
   IfStatement,
   MethodInvocation,
@@ -46,6 +48,7 @@ import { Location } from '../types'
 import { ExpressionExtractor } from './expression-extractor'
 import { BlockStatementExtractor } from './block-statement-extractor'
 import { TypeExtractor } from './type-extractor'
+import { getLocation } from './utils'
 
 export class StatementExtractor extends BaseJavaCstVisitorWithDefaults {
   constructor() {
@@ -395,4 +398,18 @@ export class StatementExtractor extends BaseJavaCstVisitorWithDefaults {
     })
     return declarations
   }
+
+  enhancedForStatement(ctx: EnhancedForStatementCtx): EnhancedForStatement {
+    const blockStatementExtractor = new BlockStatementExtractor()
+    const expressionExtractor = new ExpressionExtractor()
+    const statementExtractor = new StatementExtractor()
+    return {
+      kind: 'EnhancedForStatement',
+      localVariableType: blockStatementExtractor.visit(ctx.localVariableType),
+      variableDeclaratorId: blockStatementExtractor.visit(ctx.variableDeclaratorId),
+      expression: expressionExtractor.extract(ctx.expression[0]),
+      statement: statementExtractor.extract(ctx.statement[0]),
+      location: getLocation(ctx.For[0])
+    }
+  }
 }

src/types/ast/types/blocks-and-statements.ts

@@ -65,6 +65,11 @@ export interface EmptyStatement extends BaseNode {
 
 export interface EnhancedForStatement extends BaseNode {
   kind: 'EnhancedForStatement'
+  localVariableType: LocalVariableType
+  variableDeclaratorId: VariableDeclaratorId
+  expression: Expression
+  statement: Statement
+  location?: Location
 }
 
 export interface ExplicitConstructorInvocation extends BaseNode {

src/types/checker/__tests__/forStatements.test.ts

@@ -71,6 +71,35 @@ const testcases: {
       for (int i = 0; i < arr.length; i++) {}
     `,
     result: { type: null, errors: [] }
+  },
+  {
+    input: `
+      int[] arr = {1, 2, 3, 4, 5};
+      for (int num : arr) {}
+    `,
+    result: { type: null, errors: [] }
+  },
+  {
+    input: `
+      int[] arr = {1, 2, 3, 4, 5};
+      for (double num : arr) {}
+    `,
+    result: { type: null, errors: [] }
+  },
+  {
+    input: `
+      int[] arr = {1, 2, 3, 4, 5};
+      for (String num : arr) {}
+    `,
+    result: { type: null, errors: [new IncompatibleTypesError()] }
+  },
+  {
+    input: `
+      int[] arr = {1, 2, 3, 4, 5};
+      for (int num : arr) {}
+      String test = num;
+    `,
+    result: { type: null, errors: [new CannotFindSymbolError()] }
   }
 ]
 

src/types/checker/index.ts

@@ -7,6 +7,7 @@ import {
   ArrayRequiredError,
   BadOperandTypesError,
   IncompatibleTypesError,
+  NotApplicableToExpressionTypeError,
   VariableAlreadyDefinedError
 } from '../errors'
 import {
@@ -213,6 +214,24 @@ export const check = (node: Node, frame: Frame = Frame.globalFrame()): Result =>
     case 'EmptyStatement': {
       return OK_RESULT
     }
+    case 'EnhancedForStatement': {
+      const variableType = frame.getType(node.localVariableType)
+      if (variableType instanceof Error) return newResult(null, [variableType])
+      const expressionCheck = check(node.expression, frame)
+      if (expressionCheck.errors.length > 0) return newResult(null, expressionCheck.errors)
+      const expressionType = expressionCheck.currentType
+      if (!(expressionType instanceof ArrayType))
+        return newResult(null, [new NotApplicableToExpressionTypeError()])
+      const arrayContentType = expressionType.getContentType()
+      if (!variableType.canBeAssigned(arrayContentType))
+        return newResult(null, [new IncompatibleTypesError()])
+      const forExpressionFrame = frame.newChildFrame()
+      const error = forExpressionFrame.setVariable(node.variableDeclaratorId, variableType)
+      if (error) return newResult(null, [error])
+      const statementCheck = check(node.statement, forExpressionFrame)
+      if (statementCheck.errors.length > 0) return newResult(null, statementCheck.errors)
+      return OK_RESULT
+    }
     case 'ExpressionName': {
       const type = frame.getVariable(node.name)
       if (type instanceof Error) return newResult(null, [type])

src/types/errors.ts

@@ -34,6 +34,12 @@ export class FloatTooSmallError extends Error {
   }
 }
 
+export class NotApplicableToExpressionTypeError extends Error {
+  constructor() {
+    super('not applicable to expression type')
+  }
+}
+
 export class IllegalUnderscoreError extends Error {
   constructor() {
     super('illegal underscore')

src/types/inputs/Main.class

@@ -1,5 +1,7 @@
 public class Main {
   public static void main(String[] args) {
-    int[] numbers = new int[-1];
+    int[] arr = {1, 2, 3};
+    for (int num : arr) {}
+    System.out.println(num);
   }
 }