Allow single compilation unit to generate multiple classes

5 files changed, 26 insertions, 18 deletions

diff --git a/pom.xml b/pom.xml
index c2cff19..b7c9f4f 100644
--- a/pom.xml
+++ b/pom.xml
@@ -22,6 +22,12 @@
       <scope>test</scope>
     </dependency>
     <dependency>
+      <groupId>org.hamcrest</groupId>
+      <artifactId>hamcrest-all</artifactId>
+      <version>1.3</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
       <groupId>com.google.guava</groupId>
       <artifactId>guava</artifactId>
       <version>18.0</version>
diff --git a/src/main/java/org/zwobble/couscous/frontends/java/JavaFrontend.java b/src/main/java/org/zwobble/couscous/frontends/java/JavaFrontend.java
index 2a842c9..6e9945e 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/JavaFrontend.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/JavaFrontend.java
@@ -14,7 +14,7 @@ public class JavaFrontend implements Frontend {
     @Override
     public List<ClassNode> readSourceDirectory(Path directoryPath) throws IOException {
         JavaReader reader = new JavaReader();
-        return ExtraLists.map(
+        return ExtraLists.flatMap(
             findJavaFiles(directoryPath),
             javaFile -> reader.readClassFromFile(directoryPath, javaFile));
     }
diff --git a/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java b/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java
index 4f07149..446bc31 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java
@@ -8,25 +8,27 @@ import org.zwobble.couscous.ast.TypeName;
 
 import java.io.IOException;
 import java.nio.file.Path;
+import java.util.List;
 import java.util.Optional;
 
+import static java.util.Arrays.asList;
 import static org.zwobble.couscous.frontends.java.JavaTypes.typeOf;
 
 public class JavaReader {
     private final JavaParser parser = new JavaParser();
     
-    public ClassNode readClassFromFile(Path root, Path sourcePath) throws IOException {
+    public List<ClassNode> readClassFromFile(Path root, Path sourcePath) throws IOException {
         CompilationUnit ast = parser.parseCompilationUnit(root, sourcePath);
         return readCompilationUnit(ast);
     }
     
-    private static ClassNode readCompilationUnit(CompilationUnit ast) {
+    private static List<ClassNode> readCompilationUnit(CompilationUnit ast) {
         String name = generateClassName(ast);
         ClassNodeBuilder classBuilder = new ClassNodeBuilder(name);
         TypeDeclaration type = (TypeDeclaration)ast.types().get(0);
         readFields(type, classBuilder);
         readMethods(type, classBuilder);
-        return classBuilder.build();
+        return asList(classBuilder.build());
     }
     
     private static void readFields(TypeDeclaration type, ClassNodeBuilder classBuilder) {
diff --git a/src/main/java/org/zwobble/couscous/util/ExtraLists.java b/src/main/java/org/zwobble/couscous/util/ExtraLists.java
index 8c282e8..1951868 100644
--- a/src/main/java/org/zwobble/couscous/util/ExtraLists.java
+++ b/src/main/java/org/zwobble/couscous/util/ExtraLists.java
@@ -20,14 +20,14 @@ public class ExtraLists {
     }
 
     @SuppressWarnings("unchecked")
-    public static <T, R, E extends Exception> List<R> map(
+    public static <T, R, E extends Exception> List<R> flatMap(
             Stream<T> stream,
-            CheckedFunction<T, R, E> function) throws E {
+            CheckedFunction<T, Iterable<R>, E> function) throws E {
         try {
             return stream
-                .map(value -> {
+                .flatMap(value -> {
                     try {
-                        return function.apply(value);
+                        return StreamSupport.stream(function.apply(value).spliterator(), false);
                     } catch (Exception exception) {
                         throw new StashedException(exception);
                     }
diff --git a/src/test/java/org/zwobble/couscous/tests/frontends/java/JavaReaderTests.java b/src/test/java/org/zwobble/couscous/tests/frontends/java/JavaReaderTests.java
index 9d21ed2..96c7b3b 100644
--- a/src/test/java/org/zwobble/couscous/tests/frontends/java/JavaReaderTests.java
+++ b/src/test/java/org/zwobble/couscous/tests/frontends/java/JavaReaderTests.java
@@ -1,9 +1,5 @@
 package org.zwobble.couscous.tests.frontends.java;
 
-import java.nio.file.Files;
-import java.nio.file.Path;
-import java.util.List;
-
 import org.junit.Test;
 import org.zwobble.couscous.ast.*;
 import org.zwobble.couscous.frontends.java.JavaReader;
@@ -12,9 +8,15 @@ import org.zwobble.couscous.values.IntegerValue;
 import org.zwobble.couscous.values.ObjectValues;
 import org.zwobble.couscous.values.StringValue;
 
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+
 import static java.util.Arrays.asList;
 import static java.util.Collections.emptyList;
+import static org.hamcrest.Matchers.hasSize;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThat;
 import static org.zwobble.couscous.ast.AnnotationNode.annotation;
 import static org.zwobble.couscous.ast.AssignmentNode.assign;
 import static org.zwobble.couscous.ast.AssignmentNode.assignStatement;
@@ -27,11 +29,7 @@ import static org.zwobble.couscous.ast.LiteralNode.literal;
 import static org.zwobble.couscous.ast.LocalVariableDeclarationNode.localVariableDeclaration;
 import static org.zwobble.couscous.ast.MethodCallNode.*;
 import static org.zwobble.couscous.ast.ReturnNode.returns;
-import static org.zwobble.couscous.ast.StaticMethodCallNode.boxBoolean;
-import static org.zwobble.couscous.ast.StaticMethodCallNode.boxInt;
-import static org.zwobble.couscous.ast.StaticMethodCallNode.same;
-import static org.zwobble.couscous.ast.StaticMethodCallNode.staticMethodCall;
-import static org.zwobble.couscous.ast.StaticMethodCallNode.unboxInt;
+import static org.zwobble.couscous.ast.StaticMethodCallNode.*;
 import static org.zwobble.couscous.ast.TernaryConditionalNode.ternaryConditional;
 import static org.zwobble.couscous.ast.ThisReferenceNode.thisReference;
 import static org.zwobble.couscous.ast.VariableReferenceNode.reference;
@@ -630,7 +628,9 @@ public class JavaReaderTests {
                 Files.write(sourcePath, asList(javaClass));
                 
                 JavaReader reader = new JavaReader();
-                return reader.readClassFromFile(directoryPath, sourcePath);
+                List<ClassNode> classes = reader.readClassFromFile(directoryPath, sourcePath);
+                assertThat(classes, hasSize(1));
+                return classes.get(0);
             } finally {
                 deleteRecursively(directoryPath.toFile());
             }