Use only where appropriate

2 files changed, 12 insertions, 4 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 d44bb39..b75106a 100644
--- a/src/main/java/org/zwobble/couscous/ast/visitors/NodeTransformer.java
+++ b/src/main/java/org/zwobble/couscous/ast/visitors/NodeTransformer.java
@@ -15,6 +15,7 @@ import java.util.Set;
 import java.util.function.Function;
 
 import static org.zwobble.couscous.util.ExtraIterables.lazyMap;
+import static org.zwobble.couscous.util.ExtraIterables.only;
 import static org.zwobble.couscous.util.ExtraLists.*;
 
 public class NodeTransformer {
@@ -103,10 +104,7 @@ public class NodeTransformer {
             return transformExpression((ExpressionNode) node);
         } else if (node instanceof StatementNode) {
             List<StatementNode> statements = ExtraLists.copyOf(transformStatement((StatementNode) node));
-            if (statements.size() != 1) {
-                throw new UnsupportedOperationException("TODO: add support for compound statements");
-            }
-            return statements.get(0);
+            return only(statements, () -> new UnsupportedOperationException("TODO: add support for compound statements"));
         } else if (node instanceof TypeNode) {
             return transformTypeDeclaration((TypeNode) node);
         } else {
diff --git a/src/main/java/org/zwobble/couscous/util/ExtraIterables.java b/src/main/java/org/zwobble/couscous/util/ExtraIterables.java
index fcc13d2..3cdb031 100644
--- a/src/main/java/org/zwobble/couscous/util/ExtraIterables.java
+++ b/src/main/java/org/zwobble/couscous/util/ExtraIterables.java
@@ -211,4 +211,14 @@ public class ExtraIterables {
             return value;
         }
     }
+
+    public static <T> T only(Iterable<T> iterable, Supplier<RuntimeException> orElse) {
+        Iterator<T> iterator = iterable.iterator();
+        T value = iterator.next();
+        if (iterator.hasNext()) {
+            throw orElse.get();
+        } else {
+            return value;
+        }
+    }
 }