Allow transformation of type declarations to return multiple types

1 files changed, 3 insertions, 3 deletions

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 2db12c6..b2255d2 100644
--- a/src/main/java/org/zwobble/couscous/ast/visitors/NodeTransformer.java
+++ b/src/main/java/org/zwobble/couscous/ast/visitors/NodeTransformer.java
@@ -25,7 +25,7 @@ public class NodeTransformer {
     }
 
     public static List<TypeNode> apply(NodeTransformer transformer, List<TypeNode> types) {
-        return eagerMap(types, transformer::transformTypeDeclaration);
+        return eagerFlatMap(types, transformer::transformTypeDeclaration);
     }
 
     public static NodeTransformer replaceExpressions(Map<ExpressionNode, ExpressionNode> replacements) {
@@ -178,8 +178,8 @@ public class NodeTransformer {
         return field.transformSubtree(this);
     }
 
-    public TypeNode transformTypeDeclaration(TypeNode typeNode) {
-        return typeNode.transformSubtree(this);
+    public List<TypeNode> transformTypeDeclaration(TypeNode typeNode) {
+        return list(typeNode.transformSubtree(this));
     }
 
     public ClassNode transformClass(ClassNode classNode) {