Split switch handling off to desugar step

5 files changed, 114 insertions, 6 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/NodeTypes.java b/src/main/java/org/zwobble/couscous/ast/NodeTypes.java
index a82cd1c..86703ff 100644
--- a/src/main/java/org/zwobble/couscous/ast/NodeTypes.java
+++ b/src/main/java/org/zwobble/couscous/ast/NodeTypes.java
@@ -30,6 +30,8 @@ public class NodeTypes {
     public static final int OPERATION = 10021;
     public static final int RETURN = 10022;
     public static final int STATIC_RECEIVER = 10023;
+    public static final int SWITCH = 10032;
+    public static final int SWITCH_CASE = 10033;
     public static final int TERNARY_CONDITIONAL = 10024;
     public static final int THIS_REFERENCE = 10025;
     public static final int THROW = 10026;
diff --git a/src/main/java/org/zwobble/couscous/ast/Operations.java b/src/main/java/org/zwobble/couscous/ast/Operations.java
index 0ef9b88..308d070 100644
--- a/src/main/java/org/zwobble/couscous/ast/Operations.java
+++ b/src/main/java/org/zwobble/couscous/ast/Operations.java
@@ -4,7 +4,9 @@ import org.zwobble.couscous.types.ScalarType;
 import org.zwobble.couscous.types.Types;
 import org.zwobble.couscous.values.InternalCouscousValue;
 
+import static org.zwobble.couscous.ast.MethodCallNode.methodCall;
 import static org.zwobble.couscous.ast.MethodCallNode.staticMethodCall;
+import static org.zwobble.couscous.ast.TypeCoercionNode.coerce;
 import static org.zwobble.couscous.ast.TypeCoercionNode.typeCoercion;
 import static org.zwobble.couscous.util.ExtraLists.list;
 
@@ -94,4 +96,8 @@ public class Operations {
     public static ExpressionNode instanceOf(ExpressionNode value, ScalarType type) {
         return new InstanceOfNode(value, type);
     }
+
+    public static ExpressionNode methodEquals(ExpressionNode left, ExpressionNode right) {
+        return methodCall(left, "equals", list(coerce(right, Types.OBJECT)), Types.BOOLEAN);
+    }
 }
diff --git a/src/main/java/org/zwobble/couscous/ast/sugar/SwitchCaseNode.java b/src/main/java/org/zwobble/couscous/ast/sugar/SwitchCaseNode.java
index 534f244..3aa9bdc 100644
--- a/src/main/java/org/zwobble/couscous/ast/sugar/SwitchCaseNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/sugar/SwitchCaseNode.java
@@ -1,12 +1,25 @@
 package org.zwobble.couscous.ast.sugar;
 
+import com.google.common.collect.Iterables;
 import org.zwobble.couscous.ast.ExpressionNode;
+import org.zwobble.couscous.ast.Node;
+import org.zwobble.couscous.ast.NodeTypes;
 import org.zwobble.couscous.ast.StatementNode;
+import org.zwobble.couscous.ast.visitors.NodeTransformer;
+import org.zwobble.couscous.util.ExtraIterables;
 
 import java.util.List;
 import java.util.Optional;
 
-public class SwitchCaseNode {
+public class SwitchCaseNode implements Node {
+    public static SwitchCaseNode switchCase(ExpressionNode value, List<StatementNode> statements) {
+        return switchCase(Optional.of(value), statements);
+    }
+
+    public static SwitchCaseNode switchCase(List<StatementNode> statements) {
+        return switchCase(Optional.empty(), statements);
+    }
+
     public static SwitchCaseNode switchCase(Optional<ExpressionNode> value, List<StatementNode> statements) {
         return new SwitchCaseNode(value, statements);
     }
@@ -30,4 +43,25 @@ public class SwitchCaseNode {
     public List<StatementNode> getStatements() {
         return statements;
     }
+
+    @Override
+    public int type() {
+        return NodeTypes.SWITCH_CASE;
+    }
+
+    @Override
+    public Iterable<? extends Node> childNodes() {
+        return Iterables.concat(
+            ExtraIterables.iterable(value),
+            statements
+        );
+    }
+
+    @Override
+    public SwitchCaseNode transform(NodeTransformer transformer) {
+        return new SwitchCaseNode(
+            value.map(transformer::transformExpression),
+            transformer.transformStatements(statements)
+        );
+    }
 }
diff --git a/src/main/java/org/zwobble/couscous/ast/sugar/SwitchNode.java b/src/main/java/org/zwobble/couscous/ast/sugar/SwitchNode.java
new file mode 100644
index 0000000..8e1e226
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/ast/sugar/SwitchNode.java
@@ -0,0 +1,48 @@
+package org.zwobble.couscous.ast.sugar;
+
+import org.zwobble.couscous.ast.ExpressionNode;
+import org.zwobble.couscous.ast.Node;
+import org.zwobble.couscous.ast.NodeTypes;
+import org.zwobble.couscous.ast.StatementNode;
+import org.zwobble.couscous.ast.visitors.NodeTransformer;
+
+import java.util.List;
+
+import static org.zwobble.couscous.util.ExtraIterables.lazyCons;
+import static org.zwobble.couscous.util.ExtraLists.eagerMap;
+
+public class SwitchNode implements StatementNode {
+    private final ExpressionNode value;
+    private final List<SwitchCaseNode> cases;
+
+    public SwitchNode(ExpressionNode value, List<SwitchCaseNode> cases) {
+        this.value = value;
+        this.cases = cases;
+    }
+
+    public ExpressionNode getValue() {
+        return value;
+    }
+
+    public List<SwitchCaseNode> getCases() {
+        return cases;
+    }
+
+    @Override
+    public int type() {
+        return NodeTypes.SWITCH;
+    }
+
+    @Override
+    public Iterable<? extends Node> childNodes() {
+        return lazyCons(value, cases);
+    }
+
+    @Override
+    public StatementNode transform(NodeTransformer transformer) {
+        return new SwitchNode(
+            transformer.transformExpression(value),
+            eagerMap(cases, caseNode -> caseNode.transform(transformer))
+        );
+    }
+}
diff --git a/src/main/java/org/zwobble/couscous/ast/visitors/NodeTransformer.java b/src/main/java/org/zwobble/couscous/ast/visitors/NodeTransformer.java
index 721ed8e..78bd6ef 100644
--- a/src/main/java/org/zwobble/couscous/ast/visitors/NodeTransformer.java
+++ b/src/main/java/org/zwobble/couscous/ast/visitors/NodeTransformer.java
@@ -14,7 +14,7 @@ import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 
-import static org.zwobble.couscous.util.ExtraLists.eagerMap;
+import static org.zwobble.couscous.util.ExtraLists.*;
 
 public class NodeTransformer {
     public static NodeTransformer replaceExpressions(Map<ExpressionNode, ExpressionNode> replacements) {
@@ -29,6 +29,7 @@ public class NodeTransformer {
     }
 
     public static class Builder {
+        private Function<StatementNode, Optional<List<StatementNode>>> transformStatement = statement -> Optional.empty();
         private Function<ExpressionNode, Optional<ExpressionNode>> transformExpression = expression -> Optional.empty();
         private Function<Type, Type> transformType = type -> type;
         private Function<MethodSignature, String> transformMethodName = MethodSignature::getName;
@@ -41,6 +42,13 @@ public class NodeTransformer {
             return this;
         }
 
+        public Builder transformStatement(
+            Function<StatementNode, Optional<List<StatementNode>>> transformStatement)
+        {
+            this.transformStatement = transformStatement;
+            return this;
+        }
+
         public Builder transformType(
             Function<Type, Type> transformType)
         {
@@ -59,22 +67,31 @@ public class NodeTransformer {
         }
 
         public NodeTransformer build() {
-            return new NodeTransformer(transformExpression, transformType, transformMethodName, transformFieldName);
+            return new NodeTransformer(
+                transformExpression,
+                transformStatement,
+                transformType,
+                transformMethodName,
+                transformFieldName
+            );
         }
     }
 
     private final Function<ExpressionNode, Optional<ExpressionNode>> transformExpression;
+    private final Function<StatementNode, Optional<List<StatementNode>>> transformStatement;
     private final Function<Type, Type> transformType;
     private final Function<MethodSignature, String> transformMethodName;
     private final Function<String, String> transformFieldName;
 
     private NodeTransformer(
         Function<ExpressionNode, Optional<ExpressionNode>> transformExpression,
+        Function<StatementNode, Optional<List<StatementNode>>> transformStatement,
         Function<Type, Type> transformType,
         Function<MethodSignature, String> transformMethodName,
         Function<String, String> transformFieldName)
     {
         this.transformExpression = transformExpression;
+        this.transformStatement = transformStatement;
         this.transformType = transformType;
         this.transformMethodName = transformMethodName;
         this.transformFieldName = transformFieldName;
@@ -121,8 +138,9 @@ public class NodeTransformer {
         return VariableDeclaration.var(declaration.getId(), declaration.getName(), transform(declaration.getType()));
     }
 
-    public final StatementNode transformStatement(StatementNode statement) {
-        return statement.transform(this);
+    public final Iterable<StatementNode> transformStatement(StatementNode statement) {
+        return transformStatement.apply(statement)
+            .orElseGet(() -> list(statement.transform(this)));
     }
 
     public FormalTypeParameterNode transformFormalTypeParameter(FormalTypeParameterNode typeParameter) {
@@ -178,7 +196,7 @@ public class NodeTransformer {
     }
 
     public List<StatementNode> transformStatements(List<StatementNode> body) {
-        return transformList(body, this::transformStatement);
+        return eagerFlatMap(body, this::transformStatement);
     }
 
     public List<ExpressionNode> transformExpressions(List<? extends ExpressionNode> expressions) {