Transform subnodes in NodeTransformer

2 files changed, 9 insertions, 4 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/NodeTypes.java b/src/main/java/org/zwobble/couscous/ast/NodeTypes.java
index 41ba8af..9a9a244 100644
--- a/src/main/java/org/zwobble/couscous/ast/NodeTypes.java
+++ b/src/main/java/org/zwobble/couscous/ast/NodeTypes.java
@@ -3,6 +3,9 @@ package org.zwobble.couscous.ast;
 import com.google.common.collect.ImmutableMap;
 import org.zwobble.couscous.ast.sugar.SwitchCaseNode;
 import org.zwobble.couscous.ast.sugar.SwitchNode;
+import org.zwobble.couscous.backends.csharp.CsharpNodeTypes;
+import org.zwobble.couscous.backends.csharp.primitives.PrimitiveInstanceMethodCall;
+import org.zwobble.couscous.backends.csharp.primitives.PrimitiveStaticMethodCall;
 
 import java.util.Map;
 
@@ -80,6 +83,8 @@ public class NodeTypes {
         .put(TypeCoercionNode.class, TYPE_COERCION)
         .put(VariableReferenceNode.class, VARIABLE_REFERENCE)
         .put(WhileNode.class, WHILE)
+        .put(PrimitiveInstanceMethodCall.class, CsharpNodeTypes.PRIMITIVE_INSTANCE_METHOD_CALL)
+        .put(PrimitiveStaticMethodCall.class, CsharpNodeTypes.PRIMITIVE_STATIC_METHOD_CALL)
         .build();
 
     public static int forClass(Class<?> nodeClass) {
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 78bd6ef..296e5c7 100644
--- a/src/main/java/org/zwobble/couscous/ast/visitors/NodeTransformer.java
+++ b/src/main/java/org/zwobble/couscous/ast/visitors/NodeTransformer.java
@@ -126,8 +126,8 @@ public class NodeTransformer {
     }
 
     public ExpressionNode transformExpression(ExpressionNode expression) {
-        return transformExpression.apply(expression)
-            .orElseGet(() -> expression.transform(this));
+        ExpressionNode transformed = expression.transform(this);
+        return transformExpression.apply(transformed).orElse(transformed);
     }
 
     public Receiver transformReceiver(Receiver receiver) {
@@ -139,8 +139,8 @@ public class NodeTransformer {
     }
 
     public final Iterable<StatementNode> transformStatement(StatementNode statement) {
-        return transformStatement.apply(statement)
-            .orElseGet(() -> list(statement.transform(this)));
+        StatementNode transformed = statement.transform(this);
+        return transformStatement.apply(transformed).orElseGet(() -> list(transformed));
     }
 
     public FormalTypeParameterNode transformFormalTypeParameter(FormalTypeParameterNode typeParameter) {