Allow single compilation unit to generate multiple classes

3 files changed, 10 insertions, 8 deletions

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);
                     }