Simplify generation of anonymous classes

7 files changed, 107 insertions, 49 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/CallableNode.java b/src/main/java/org/zwobble/couscous/ast/CallableNode.java
index 79d2e84..6d6fc7b 100644
--- a/src/main/java/org/zwobble/couscous/ast/CallableNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/CallableNode.java
@@ -5,5 +5,4 @@ import java.util.List;
 public interface CallableNode extends Node {
     List<FormalArgumentNode> getArguments();
     List<StatementNode> getBody();
-    
 }
diff --git a/src/main/java/org/zwobble/couscous/ast/ClassNode.java b/src/main/java/org/zwobble/couscous/ast/ClassNode.java
index c252bb3..55453f4 100644
--- a/src/main/java/org/zwobble/couscous/ast/ClassNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/ClassNode.java
@@ -8,7 +8,6 @@ public class ClassNode implements Node {
     public static ClassNodeBuilder builder(String name) {
         return new ClassNodeBuilder(name);
     }
-    
     public static ClassNode declareClass(
             TypeName name,
             List<FieldDeclarationNode> fields,
diff --git a/src/main/java/org/zwobble/couscous/ast/ClassNodeBuilder.java b/src/main/java/org/zwobble/couscous/ast/ClassNodeBuilder.java
index 75aa5a3..5981e34 100644
--- a/src/main/java/org/zwobble/couscous/ast/ClassNodeBuilder.java
+++ b/src/main/java/org/zwobble/couscous/ast/ClassNodeBuilder.java
@@ -13,12 +13,17 @@ public class ClassNodeBuilder {
     private final ImmutableList.Builder<FieldDeclarationNode> fields;
     private Optional<ConstructorNode> constructor;
     private final ImmutableList.Builder<MethodNode> methods;
-    
-    public ClassNodeBuilder(String name) {
-        this.name = TypeName.of(name);
+
+    public ClassNodeBuilder(TypeName name) {
+        this.name = name;
         this.fields = ImmutableList.builder();
         this.constructor = Optional.empty();
         this.methods = ImmutableList.builder();
+
+    }
+
+    public ClassNodeBuilder(String name) {
+        this(TypeName.of(name));
     }
 
     public ClassNodeBuilder field(String name, TypeName type) {
@@ -30,14 +35,6 @@ public class ClassNodeBuilder {
         return this;
     }
 
-    public ClassNodeBuilder callable(CallableNode callable) {
-        if (callable instanceof ConstructorNode) {
-            return constructor((ConstructorNode) callable);
-        } else {
-            return method((MethodNode) callable);
-        }
-    }
-
     public ClassNodeBuilder constructor(ConstructorNode constructor) {
         this.constructor = Optional.of(constructor);
         return this;
@@ -85,7 +82,7 @@ public class ClassNodeBuilder {
         return ClassNode.declareClass(
             name,
             fields.build(),
-            constructor.orElse(constructorBuilder().build()),
+            constructor.orElse(ConstructorNode.DEFAULT),
             methods.build());
     }
 
diff --git a/src/main/java/org/zwobble/couscous/ast/ConstructorNode.java b/src/main/java/org/zwobble/couscous/ast/ConstructorNode.java
index 2e8b85a..0d4bad8 100644
--- a/src/main/java/org/zwobble/couscous/ast/ConstructorNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/ConstructorNode.java
@@ -1,11 +1,13 @@
 package org.zwobble.couscous.ast;
 
-import java.util.List;
-
 import org.zwobble.couscous.ast.visitors.NodeMapper;
-import org.zwobble.couscous.ast.visitors.NodeVisitor;
+
+import java.util.Collections;
+import java.util.List;
 
 public class ConstructorNode implements CallableNode {
+    public static final ConstructorNode DEFAULT = constructor(Collections.emptyList(), Collections.emptyList());
+
     public static ConstructorNode constructor(
             List<FormalArgumentNode> arguments,
             List<StatementNode> body) {
diff --git a/src/main/java/org/zwobble/couscous/ast/MethodNode.java b/src/main/java/org/zwobble/couscous/ast/MethodNode.java
index 3e97a36..5763b11 100644
--- a/src/main/java/org/zwobble/couscous/ast/MethodNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/MethodNode.java
@@ -4,6 +4,7 @@ import com.google.common.collect.ImmutableList;
 import org.zwobble.couscous.ast.visitors.NodeMapper;
 
 import java.util.List;
+import java.util.function.Function;
 
 public class MethodNode implements CallableNode {
     public static Builder staticMethod(String name) {
@@ -115,6 +116,10 @@ public class MethodNode implements CallableNode {
         return body;
     }
 
+    public MethodNode mapBody(Function<List<StatementNode>, List<StatementNode>> function) {
+        return new MethodNode(annotations, isStatic, name, arguments, function.apply(body));
+    }
+
     @Override
     public <T> T accept(NodeMapper<T> visitor) {
         return visitor.visit(this);
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 c8dc854..04ee748 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java
@@ -43,7 +43,12 @@ public class JavaReader {
     private ClassNode readCompilationUnit(CompilationUnit ast) {
         String name = generateClassName(ast);
         TypeDeclaration type = (TypeDeclaration)ast.types().get(0);
-        return readTypeDeclarationBody(name, type.bodyDeclarations());
+        TypeDeclarationBody body = readTypeDeclarationBody(type.bodyDeclarations());
+        return ClassNode.declareClass(
+            TypeName.of(name),
+            body.getFields(),
+            body.getConstructor(),
+            body.getMethods());
     }
 
     GeneratedClosure readLambda(LambdaExpression expression) {
@@ -105,23 +110,19 @@ public class JavaReader {
 
     GeneratedClosure readAnonymousClass(AnonymousClassDeclaration declaration) {
         String className = generateClassName(declaration);
-        ClassNode classNode = readTypeDeclarationBody(className, declaration.bodyDeclarations());
-        List<VariableDeclaration> freeVariables = findFreeVariables(asList(classNode));
+        TypeDeclarationBody bodyDeclarations = readTypeDeclarationBody(declaration.bodyDeclarations());
+        List<VariableDeclaration> freeVariables = findFreeVariables(bodyDeclarations.getNodes());
         List<StatementNode> restoreCaptures = generateCaptureRestoration(className, freeVariables);
 
-        // TODO: map over methods rather than rebuilding entire class
-        ClassNodeBuilder classBuilder = new ClassNodeBuilder(className);
-        classBuilder.constructor(buildConstructor(TypeName.of(className), freeVariables));
-        classNode.getFields().forEach(classBuilder::field);
-        classNode.getMethods().forEach(method ->
-            classBuilder.method(MethodNode.method(
-                method.getAnnotations(),
-                method.isStatic(),
-                method.getName(),
-                method.getArguments(),
-                concat(restoreCaptures, method.getBody()))));
-
-        classes.add(classBuilder.build());
+        List<MethodNode> methods = eagerMap(bodyDeclarations.getMethods(), method ->
+            method.mapBody(body -> concat(restoreCaptures, body)));
+
+        ClassNode classNode = ClassNode.declareClass(
+            TypeName.of(className),
+            bodyDeclarations.getFields(),
+            buildConstructor(TypeName.of(className), freeVariables),
+            methods);
+        classes.add(classNode);
         return new GeneratedClosure(classNode.getName(), freeVariables);
     }
 
@@ -137,25 +138,61 @@ public class JavaReader {
         return type.getQualifiedName() + "_Anonymous_" + (anonymousClassCount++);
     }
 
-    private ClassNode readTypeDeclarationBody(String name, List<Object> bodyDeclarations) {
-        ClassNodeBuilder classBuilder = new ClassNodeBuilder(name);
-        readFields(ofType(bodyDeclarations, FieldDeclaration.class), classBuilder);
-        readMethods(ofType(bodyDeclarations, MethodDeclaration.class))
-            .forEach(classBuilder::callable);
-        return classBuilder.build();
-    }
+    private static class TypeDeclarationBody {
+        private final List<FieldDeclarationNode> fields;
+        private final ConstructorNode constructor;
+        private final List<MethodNode> methods;
+
+        public TypeDeclarationBody(List<FieldDeclarationNode> fields, ConstructorNode constructor, List<MethodNode> methods) {
+            this.fields = fields;
+            this.constructor = constructor;
+            this.methods = methods;
+        }
+
+        public List<FieldDeclarationNode> getFields() {
+            return fields;
+        }
+
+        public ConstructorNode getConstructor() {
+            return constructor;
+        }
 
-    private void readFields(List<FieldDeclaration> fields, ClassNodeBuilder classBuilder) {
-        for (FieldDeclaration field : fields) {
-            readField(field, classBuilder);
+        public List<MethodNode> getMethods() {
+            return methods;
+        }
+
+        public List<Node> getNodes() {
+            return concat(fields, asList(constructor), methods);
         }
     }
 
-    private void readField(FieldDeclaration field, ClassNodeBuilder classBuilder) {
-        for (Object fragment : field.fragments()) {
-            String name = ((VariableDeclarationFragment)fragment).getName().getIdentifier();
-            classBuilder.field(name, typeOf(field.getType()));
+    private TypeDeclarationBody readTypeDeclarationBody(List<Object> bodyDeclarations) {
+        ImmutableList.Builder<MethodNode> methods = ImmutableList.builder();
+        ConstructorNode constructor = ConstructorNode.DEFAULT;
+
+        for (CallableNode callable : readMethods(ofType(bodyDeclarations, MethodDeclaration.class))) {
+            if (callable instanceof ConstructorNode) {
+                constructor = (ConstructorNode) callable;
+            } else {
+                methods.add((MethodNode) callable);
+            }
         }
+        return new TypeDeclarationBody(
+            readFields(ofType(bodyDeclarations, FieldDeclaration.class)),
+            constructor,
+            methods.build());
+    }
+
+    private List<FieldDeclarationNode> readFields(List<FieldDeclaration> fields) {
+        return eagerFlatMap(fields, this::readField);
+    }
+
+    private List<FieldDeclarationNode> readField(FieldDeclaration field) {
+        @SuppressWarnings("unchecked")
+        List<VariableDeclarationFragment> fragments = field.fragments();
+        TypeName type = typeOf(field.getType());
+        return eagerMap(fragments, fragment ->
+            FieldDeclarationNode.field(fragment.getName().getIdentifier(), type));
     }
 
     private List<CallableNode> readMethods(List<MethodDeclaration> methods) {
diff --git a/src/main/java/org/zwobble/couscous/util/ExtraLists.java b/src/main/java/org/zwobble/couscous/util/ExtraLists.java
index de28f21..e0c183a 100644
--- a/src/main/java/org/zwobble/couscous/util/ExtraLists.java
+++ b/src/main/java/org/zwobble/couscous/util/ExtraLists.java
@@ -2,6 +2,7 @@ package org.zwobble.couscous.util;
 
 import java.util.List;
 import java.util.function.Function;
+import java.util.function.Predicate;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
@@ -15,10 +16,14 @@ public class ExtraLists {
         return Stream.concat(list.stream(), Stream.of(value)).collect(Collectors.toList());
     }
 
-    public static <T> List<T> concat(List<T> first, List<T> second) {
+    public static <T> List<T> concat(List<? extends T> first, List<? extends T> second) {
         return Stream.concat(first.stream(), second.stream()).collect(Collectors.toList());
     }
 
+    public static <T> List<T> concat(List<? extends T> first, List<? extends T> second, List<? extends T> third) {
+        return Stream.concat(Stream.concat(first.stream(), second.stream()), third.stream()).collect(Collectors.toList());
+    }
+
     @SuppressWarnings("unchecked")
     public static <T, R extends T> List<R> ofType(List<T> list, Class<R> clazz) {
         return (List<R>)list.stream()
@@ -44,7 +49,13 @@ public class ExtraLists {
             throw (E)exception.exception;
         }
     }
-    
+
+    public static <T> List<T> eagerFilter(Iterable<T> iterable, Predicate<T> predicate) {
+        return StreamSupport.stream(iterable.spliterator(), false)
+            .filter(predicate)
+            .collect(Collectors.toList());
+    }
+
     public static <T, R> List<R> eagerMap(
             Iterable<T> iterable,
             Function<T, R> function) {
@@ -52,6 +63,14 @@ public class ExtraLists {
             .map(function)
             .collect(Collectors.toList());
     }
+
+    public static <T, R> List<R> eagerFlatMap(
+        Iterable<T> iterable,
+        Function<T, Iterable<R>> function) {
+        return StreamSupport.stream(iterable.spliterator(), false)
+            .flatMap(element -> StreamSupport.stream(function.apply(element).spliterator(), false))
+            .collect(Collectors.toList());
+    }
     
     @FunctionalInterface
     public interface CheckedFunction<T, R, E extends Exception> {