Implement StatementNode.replaceExpressions using node transformer

6 files changed, 10 insertions, 40 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/ExpressionStatementNode.java b/src/main/java/org/zwobble/couscous/ast/ExpressionStatementNode.java
index 68acedb..1720fc5 100644
--- a/src/main/java/org/zwobble/couscous/ast/ExpressionStatementNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/ExpressionStatementNode.java
@@ -3,8 +3,6 @@ package org.zwobble.couscous.ast;
 import org.zwobble.couscous.ast.visitors.NodeTransformer;
 import org.zwobble.couscous.ast.visitors.StatementNodeMapper;
 
-import java.util.function.Function;
-
 public class ExpressionStatementNode implements StatementNode {
     public static ExpressionStatementNode expressionStatement(ExpressionNode expression) {
         return new ExpressionStatementNode(expression);
@@ -26,10 +24,6 @@ public class ExpressionStatementNode implements StatementNode {
     }
 
     @Override
-    public StatementNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace) {
-        return new ExpressionStatementNode(replace.apply(expression));
-    }
-
     public StatementNode transform(NodeTransformer transformer) {
         return new ExpressionStatementNode(transformer.visit(expression));
     }
diff --git a/src/main/java/org/zwobble/couscous/ast/IfStatementNode.java b/src/main/java/org/zwobble/couscous/ast/IfStatementNode.java
index 972f0a4..2a34a47 100644
--- a/src/main/java/org/zwobble/couscous/ast/IfStatementNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/IfStatementNode.java
@@ -4,9 +4,6 @@ import org.zwobble.couscous.ast.visitors.NodeTransformer;
 import org.zwobble.couscous.ast.visitors.StatementNodeMapper;
 
 import java.util.List;
-import java.util.function.Function;
-
-import static org.zwobble.couscous.util.ExtraLists.eagerMap;
 
 public class IfStatementNode implements StatementNode {
     public static StatementNode ifStatement(
@@ -47,13 +44,6 @@ public class IfStatementNode implements StatementNode {
     }
 
     @Override
-    public StatementNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace) {
-        return new IfStatementNode(
-            replace.apply(condition),
-            eagerMap(trueBranch, statement -> statement.replaceExpressions(replace)),
-            eagerMap(falseBranch, statement -> statement.replaceExpressions(replace)));
-    }
-
     public StatementNode transform(NodeTransformer transformer) {
         return new IfStatementNode(
             transformer.visit(condition),
diff --git a/src/main/java/org/zwobble/couscous/ast/LocalVariableDeclarationNode.java b/src/main/java/org/zwobble/couscous/ast/LocalVariableDeclarationNode.java
index 4ee0e97..145b37f 100644
--- a/src/main/java/org/zwobble/couscous/ast/LocalVariableDeclarationNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/LocalVariableDeclarationNode.java
@@ -4,8 +4,6 @@ import org.zwobble.couscous.ast.identifiers.Identifier;
 import org.zwobble.couscous.ast.visitors.NodeTransformer;
 import org.zwobble.couscous.ast.visitors.StatementNodeMapper;
 
-import java.util.function.Function;
-
 import static org.zwobble.couscous.ast.VariableDeclaration.var;
 
 public class LocalVariableDeclarationNode implements VariableNode, StatementNode {
@@ -55,10 +53,6 @@ public class LocalVariableDeclarationNode implements VariableNode, StatementNode
     }
 
     @Override
-    public StatementNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace) {
-        return new LocalVariableDeclarationNode(declaration, replace.apply(initialValue));
-    }
-
     public StatementNode transform(NodeTransformer transformer) {
         return new LocalVariableDeclarationNode(
             transformer.transform(declaration),
diff --git a/src/main/java/org/zwobble/couscous/ast/ReturnNode.java b/src/main/java/org/zwobble/couscous/ast/ReturnNode.java
index 5bd5987..c1f0124 100644
--- a/src/main/java/org/zwobble/couscous/ast/ReturnNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/ReturnNode.java
@@ -3,8 +3,6 @@ package org.zwobble.couscous.ast;
 import org.zwobble.couscous.ast.visitors.NodeTransformer;
 import org.zwobble.couscous.ast.visitors.StatementNodeMapper;
 
-import java.util.function.Function;
-
 public class ReturnNode implements StatementNode {
     public static ReturnNode returns(ExpressionNode value) {
         return new ReturnNode(value);
@@ -26,13 +24,7 @@ public class ReturnNode implements StatementNode {
     }
 
     @Override
-    public StatementNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace) {
-        return new ReturnNode(replace.apply(value));
-    }
-
     public StatementNode transform(NodeTransformer transformer) {
-        // TODO: remove duplication with replaceExpressions
-        // (can probably implement replaceExpressions with transform)
         return new ReturnNode(transformer.visit(value));
     }
 
diff --git a/src/main/java/org/zwobble/couscous/ast/StatementNode.java b/src/main/java/org/zwobble/couscous/ast/StatementNode.java
index 4b9c9de..87b1319 100644
--- a/src/main/java/org/zwobble/couscous/ast/StatementNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/StatementNode.java
@@ -1,6 +1,7 @@
 package org.zwobble.couscous.ast;
 
 import org.zwobble.couscous.ast.visitors.NodeMapper;
+import org.zwobble.couscous.ast.visitors.NodeTransformer;
 import org.zwobble.couscous.ast.visitors.StatementNodeMapper;
 
 import java.util.function.Function;
@@ -38,5 +39,13 @@ public interface StatementNode extends Node {
         });
     }
 
-    StatementNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace);
+    StatementNode transform(NodeTransformer transformer);
+    default StatementNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace) {
+        return transform(new NodeTransformer() {
+            @Override
+            public ExpressionNode visit(ExpressionNode value) {
+                return replace.apply(value);
+            }
+        });
+    }
 }
diff --git a/src/main/java/org/zwobble/couscous/ast/WhileNode.java b/src/main/java/org/zwobble/couscous/ast/WhileNode.java
index 34d1883..82e2e87 100644
--- a/src/main/java/org/zwobble/couscous/ast/WhileNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/WhileNode.java
@@ -5,9 +5,6 @@ import org.zwobble.couscous.ast.visitors.NodeTransformer;
 import org.zwobble.couscous.ast.visitors.StatementNodeMapper;
 
 import java.util.List;
-import java.util.function.Function;
-
-import static org.zwobble.couscous.util.ExtraLists.eagerMap;
 
 public class WhileNode implements StatementNode {
     public static WhileNode whileLoop(ExpressionNode condition, List<StatementNode> body) {
@@ -36,12 +33,6 @@ public class WhileNode implements StatementNode {
     }
 
     @Override
-    public StatementNode replaceExpressions(Function<ExpressionNode, ExpressionNode> replace) {
-        return new WhileNode(
-            replace.apply(condition),
-            eagerMap(body, statement -> statement.replaceExpressions(replace)));
-    }
-
     public StatementNode transform(NodeTransformer transformer) {
         return new WhileNode(
             transformer.visit(condition),