Add enhanced for statement to Java AST

5 files changed, 194 insertions, 1 deletions

diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/java/JavaWriter.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/java/JavaWriter.java
index 38dde1f..4a21d88 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/java/JavaWriter.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/java/JavaWriter.java
@@ -395,6 +395,10 @@ public class JavaWriter implements AutoCloseable {
 
     void writeBlockStatement(JavaBlockStatement statement) throws IOException {
         switch (statement) {
+            case JavaEnhancedForStatement enhancedForStatement -> {
+                writeEnhancedForStatement(enhancedForStatement);
+            }
+
             case JavaExpressionStatement expressionStatement -> {
                 writeExpressionStatement(expressionStatement);
             }
@@ -409,6 +413,32 @@ public class JavaWriter implements AutoCloseable {
         }
     }
 
+    private void writeEnhancedForStatement(JavaEnhancedForStatement enhancedForStatement) throws IOException {
+        this.writer.write("for (");
+
+        if (enhancedForStatement.elementType().isPresent()) {
+            this.writeTypeRef(enhancedForStatement.elementType().get());
+        } else {
+            this.writer.write("var");
+        }
+
+        this.writer.write(" ");
+        this.writeIdentifier(enhancedForStatement.elementName());
+        this.writer.write(" : ");
+
+        this.writeTopLevelExpression(enhancedForStatement.elements());
+
+        this.writer.write(") {");
+        this.writer.indent();
+        this.writer.newLine();
+
+        this.writeBlock(enhancedForStatement.body());
+
+        this.writer.dedent();
+        this.writer.newLine();
+        this.writer.write("}");
+    }
+
     private void writeExpressionStatement(JavaExpressionStatement expressionStatement) throws IOException {
         writeTopLevelExpression(expressionStatement.expression());
         this.writer.write(";");
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/java/ast/JavaBlockStatement.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/java/ast/JavaBlockStatement.java
index 53e74f6..76098cb 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/java/ast/JavaBlockStatement.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/java/ast/JavaBlockStatement.java
@@ -5,7 +5,7 @@ package org.zwobble.hobgoblin.compiler.output.lang.java.ast;
 // Custom area start: org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaBlockStatement imports
 // Custom area end: org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaBlockStatement imports
 
-public sealed interface JavaBlockStatement permits org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaExpressionStatement, org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaLocalVariableDeclaration, org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaReturn {
+public sealed interface JavaBlockStatement permits org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement, org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaExpressionStatement, org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaLocalVariableDeclaration, org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaReturn {
     public interface Builder {
         public org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaBlockStatement build();
 
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/java/ast/JavaEnhancedForStatement.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/java/ast/JavaEnhancedForStatement.java
new file mode 100644
index 0000000..4d99c2b
--- /dev/null
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/java/ast/JavaEnhancedForStatement.java
@@ -0,0 +1,56 @@
+// Generated by hobgoblin.
+
+package org.zwobble.hobgoblin.compiler.output.lang.java.ast;
+
+// Custom area start: org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement imports
+// Custom area end: org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement imports
+
+public record JavaEnhancedForStatement(java.util.Optional<org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaTypeRef> elementType, org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaIdentifier elementName, org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaExpression elements, org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaBlock body) implements org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaBlockStatement {
+    public static org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement.Builder arbitrary() {
+        return new org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement.Builder(java.util.Optional.empty(), org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaIdentifier.arbitrary().build(), org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaBooleanLiteral.arbitrary().build(), org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaBlock.arbitrary().build());
+    }
+
+    public record Builder(java.util.Optional<org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaTypeRef> elementType, org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaIdentifier elementName, org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaExpression elements, org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaBlock body) implements org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaBlockStatement.Builder {
+        public org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement build() {
+            return new org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement(elementType, elementName, elements, body);
+        }
+
+        public org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement.Builder withElementType(java.util.Optional<org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaTypeRef> elementType) {
+            return new org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement.Builder(elementType, elementName, elements, body);
+        }
+
+        public org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement.Builder withElementType(org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaTypeRef elementType) {
+            return new org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement.Builder(java.util.Optional.of(elementType), elementName, elements, body);
+        }
+
+        public org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement.Builder withElementName(org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaIdentifier elementName) {
+            return new org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement.Builder(elementType, elementName, elements, body);
+        }
+
+        public org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement.Builder withElementName(org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaIdentifier.Builder elementName) {
+            return new org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement.Builder(elementType, elementName.build(), elements, body);
+        }
+
+        public org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement.Builder withElements(org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaExpression elements) {
+            return new org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement.Builder(elementType, elementName, elements, body);
+        }
+
+        public org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement.Builder withElements(org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaExpression.Builder elements) {
+            return new org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement.Builder(elementType, elementName, elements.build(), body);
+        }
+
+        public org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement.Builder withBody(org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaBlock body) {
+            return new org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement.Builder(elementType, elementName, elements, body);
+        }
+
+        public org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement.Builder withBody(org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaBlock.Builder body) {
+            return new org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement.Builder(elementType, elementName, elements, body.build());
+        }
+
+        // Custom area start: org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement.Builder body
+        // Custom area end: org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement.Builder body
+    }
+
+    // Custom area start: org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement body
+    // Custom area end: org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement body
+}
diff --git a/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/java/JavaWriterTests.java b/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/java/JavaWriterTests.java
index 5464116..96dfe16 100644
--- a/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/java/JavaWriterTests.java
+++ b/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/java/JavaWriterTests.java
@@ -836,6 +836,56 @@ public class JavaWriterTests {
         assertThat(string, equalTo("return 42;\nreturn 47;"));
     }
 
+    // === Enhanced for statements ===
+
+    @Test
+    public void enhancedForStatementUsingExplicitType() throws IOException {
+        var java = JavaEnhancedForStatement.arbitrary()
+            .withElementType(JavaTypeRef.topLevel(JavaPackageName.EMPTY, JavaIdentifier.of("Fruit")))
+            .withElementName(JavaIdentifier.of("fruit"))
+            .withElements(JavaRef.arbitrary().withName(JavaIdentifier.of("fruits")))
+            .withBody(JavaBlock.arbitrary()
+                .addStatement(JavaExpressionStatement.arbitrary().withExpression(
+                    JavaStaticMethodCall.arbitrary()
+                        .withType(JavaTypeRef.topLevel(JavaPackageName.EMPTY, JavaIdentifier.of("System")))
+                        .withMethodName(JavaIdentifier.of("println"))
+                        .addArg(JavaRef.arbitrary().withName(JavaIdentifier.of("fruit")))
+                ))
+            )
+            .build();
+
+        var string = write(writer -> writer.writeBlockStatement(java));
+
+        assertThat(string, equalTo("""
+            for (Fruit fruit : fruits) {
+                System.println(fruit);
+            }"""));
+    }
+
+    @Test
+    public void enhancedForStatementUsingVar() throws IOException {
+        var java = JavaEnhancedForStatement.arbitrary()
+            .withElementType(Optional.empty())
+            .withElementName(JavaIdentifier.of("fruit"))
+            .withElements(JavaRef.arbitrary().withName(JavaIdentifier.of("fruits")))
+            .withBody(JavaBlock.arbitrary()
+                .addStatement(JavaExpressionStatement.arbitrary().withExpression(
+                    JavaStaticMethodCall.arbitrary()
+                        .withType(JavaTypeRef.topLevel(JavaPackageName.EMPTY, JavaIdentifier.of("System")))
+                        .withMethodName(JavaIdentifier.of("println"))
+                        .addArg(JavaRef.arbitrary().withName(JavaIdentifier.of("fruit")))
+                ))
+            )
+            .build();
+
+        var string = write(writer -> writer.writeBlockStatement(java));
+
+        assertThat(string, equalTo("""
+            for (var fruit : fruits) {
+                System.println(fruit);
+            }"""));
+    }
+
     // === Expression statements ===
 
     @Test
diff --git a/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/java/ast/JavaEnhancedForStatementMatcher.java b/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/java/ast/JavaEnhancedForStatementMatcher.java
new file mode 100644
index 0000000..ab98443
--- /dev/null
+++ b/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/java/ast/JavaEnhancedForStatementMatcher.java
@@ -0,0 +1,57 @@
+// Generated by hobgoblin.
+
+package org.zwobble.hobgoblin.compiler.output.lang.java.ast;
+
+// Custom area start: org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatementMatcher imports
+// Custom area end: org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatementMatcher imports
+
+public class JavaEnhancedForStatementMatcher implements org.zwobble.precisely.Matcher<java.lang.Object> {
+    public static org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatementMatcher isJavaEnhancedForStatement() {
+        return new org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatementMatcher(java.util.List.of());
+    }
+
+    private final java.util.List<org.zwobble.precisely.Matcher<? super org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement>> submatchers;
+
+    private JavaEnhancedForStatementMatcher(java.util.List<org.zwobble.precisely.Matcher<? super org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement>> submatchers) {
+        this.submatchers = submatchers;
+    }
+
+    public org.zwobble.precisely.MatchResult match(java.lang.Object actual) {
+        return this.toMatcher().match(actual);
+    }
+
+    public org.zwobble.precisely.TextTree describe() {
+        return this.toMatcher().describe();
+    }
+
+    private org.zwobble.precisely.Matcher<java.lang.Object> toMatcher() {
+        return org.zwobble.precisely.Matchers.instanceOf(org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement.class, this.submatchers);
+    }
+
+    public org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatementMatcher withElementType(org.zwobble.precisely.Matcher<? super java.util.Optional<org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaTypeRef>> elementType) {
+        var submatchers = new java.util.ArrayList<org.zwobble.precisely.Matcher<? super org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement>>(this.submatchers);
+        submatchers.add(org.zwobble.precisely.Matchers.has("elementType", org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement::elementType, elementType));
+        return new org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatementMatcher(submatchers);
+    }
+
+    public org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatementMatcher withElementName(org.zwobble.precisely.Matcher<? super org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaIdentifier> elementName) {
+        var submatchers = new java.util.ArrayList<org.zwobble.precisely.Matcher<? super org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement>>(this.submatchers);
+        submatchers.add(org.zwobble.precisely.Matchers.has("elementName", org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement::elementName, elementName));
+        return new org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatementMatcher(submatchers);
+    }
+
+    public org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatementMatcher withElements(org.zwobble.precisely.Matcher<? super org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaExpression> elements) {
+        var submatchers = new java.util.ArrayList<org.zwobble.precisely.Matcher<? super org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement>>(this.submatchers);
+        submatchers.add(org.zwobble.precisely.Matchers.has("elements", org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement::elements, elements));
+        return new org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatementMatcher(submatchers);
+    }
+
+    public org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatementMatcher withBody(org.zwobble.precisely.Matcher<? super org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaBlock> body) {
+        var submatchers = new java.util.ArrayList<org.zwobble.precisely.Matcher<? super org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement>>(this.submatchers);
+        submatchers.add(org.zwobble.precisely.Matchers.has("body", org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatement::body, body));
+        return new org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatementMatcher(submatchers);
+    }
+
+    // Custom area start: org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatementMatcher body
+    // Custom area end: org.zwobble.hobgoblin.compiler.output.lang.java.ast.JavaEnhancedForStatementMatcher body
+}