Desugar all types at once

3 files changed, 24 insertions, 20 deletions

diff --git a/src/main/java/org/zwobble/couscous/backends/python/PythonBackend.java b/src/main/java/org/zwobble/couscous/backends/python/PythonBackend.java
index 2a3701e..3df2d98 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonBackend.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonBackend.java
@@ -4,6 +4,10 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.io.Resources;
 import org.zwobble.couscous.Backend;
 import org.zwobble.couscous.ast.TypeNode;
+import org.zwobble.couscous.ast.visitors.NodeTransformer;
+import org.zwobble.couscous.transforms.DesugarForEachToFor;
+import org.zwobble.couscous.transforms.DesugarForToWhile;
+import org.zwobble.couscous.transforms.DesugarSwitchToIfElse;
 
 import java.io.File;
 import java.io.IOException;
@@ -34,8 +38,8 @@ public class PythonBackend implements Backend {
     
     @Override
     public void compile(List<TypeNode> classes) throws IOException {
-        for (TypeNode classNode : classes) {
-            compileClass(classNode);
+        for (TypeNode classNode : desugar(classes)) {
+            writeClass(classNode);
         }
         for (String runtimeFile :  RUNTIME_FILES) {
             String path = relativePathForModule(runtimeFile);
@@ -44,7 +48,18 @@ public class PythonBackend implements Backend {
                 Resources.toString(Resources.getResource("org/zwobble/couscous/backends/python/runtime/" + path), StandardCharsets.UTF_8));
         }
     }
-    
+
+    private List<TypeNode> desugar(List<TypeNode> classes) {
+        return NodeTransformer.applyAll(
+            list(
+                DesugarSwitchToIfElse.transformer(),
+                DesugarForEachToFor.transformer(),
+                DesugarForToWhile.transformer()
+            ),
+            classes
+        );
+    }
+
     private Path destinationPathForModule(String moduleName) {
         return root.resolve(packageName).resolve(relativePathForModule(moduleName));
     }
@@ -53,7 +68,7 @@ public class PythonBackend implements Backend {
         return moduleName.replace(".", File.separator) + ".py";
     }
     
-    private void compileClass(TypeNode classNode) throws IOException {
+    private void writeClass(TypeNode classNode) throws IOException {
         writeModule(
             classNode.getName().getQualifiedName(),
             serialize(generateCode(classNode)));
diff --git a/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java b/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
index 2190f66..0ce02e3 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
@@ -7,12 +7,8 @@ import com.google.common.collect.Iterators;
 import org.zwobble.couscous.ast.*;
 import org.zwobble.couscous.ast.structure.NodeStructure;
 import org.zwobble.couscous.ast.visitors.DynamicNodeMapper;
-import org.zwobble.couscous.ast.visitors.NodeTransformer;
 import org.zwobble.couscous.backends.naming.Names;
 import org.zwobble.couscous.backends.python.ast.*;
-import org.zwobble.couscous.transforms.DesugarForEachToFor;
-import org.zwobble.couscous.transforms.DesugarForToWhile;
-import org.zwobble.couscous.transforms.DesugarSwitchToIfElse;
 import org.zwobble.couscous.types.ScalarType;
 import org.zwobble.couscous.types.Type;
 import org.zwobble.couscous.types.Types;
@@ -60,16 +56,7 @@ import static org.zwobble.couscous.util.ExtraIterables.*;
 import static org.zwobble.couscous.util.ExtraLists.*;
 
 public class PythonCodeGenerator {
-    public static PythonModuleNode generateCode(TypeNode sugaredTypeNode) {
-        NodeTransformer switchTransformer = DesugarSwitchToIfElse.transformer();
-        NodeTransformer forTransformer = DesugarForToWhile.transformer();
-        NodeTransformer forEachTransformer = DesugarForEachToFor.transformer();
-        TypeNode typeNode = forTransformer.transformTypeDeclaration(
-            forEachTransformer.transformTypeDeclaration(
-                switchTransformer.transformTypeDeclaration(sugaredTypeNode)
-            )
-        );
-
+    public static PythonModuleNode generateCode(TypeNode typeNode) {
         Iterator<PythonImportNode> imports = generateImports(typeNode).iterator();
 
         PythonImportNode internalsImport = pythonImport(
diff --git a/src/main/java/org/zwobble/couscous/interpreter/JavaProject.java b/src/main/java/org/zwobble/couscous/interpreter/JavaProject.java
index 3e4b182..f813faf 100644
--- a/src/main/java/org/zwobble/couscous/interpreter/JavaProject.java
+++ b/src/main/java/org/zwobble/couscous/interpreter/JavaProject.java
@@ -13,6 +13,8 @@ import org.zwobble.couscous.transforms.DesugarSwitchToIfElse;
 
 import java.util.List;
 
+import static org.zwobble.couscous.util.ExtraLists.list;
+
 public class JavaProject {
     private static final InterpreterType OBJECT_TYPE = IntrinsicInterpreterType.classBuilder("java.lang.Object")
         .build();
@@ -28,8 +30,8 @@ public class JavaProject {
         NodeTransformer switchTransformer = DesugarSwitchToIfElse.transformer();
         NodeTransformer forTransformer = DesugarForToWhile.transformer();
         Iterable<InterpreterType> concreteTypes = Iterables.transform(
-            classNodes,
-            classNode -> new UserDefinedInterpreterType(forTransformer.transformTypeDeclaration(switchTransformer.transformTypeDeclaration(classNode)))
+            NodeTransformer.applyAll(list(switchTransformer, forTransformer), classNodes),
+            UserDefinedInterpreterType::new
         );
         return builder()
                 .addClasses(concreteTypes)