Use only where appropriate

4 files changed, 22 insertions, 15 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;
+        }
+    }
 }
diff --git a/src/test/java/org/zwobble/couscous/tests/frontends/java/ExpressionReadingTests.java b/src/test/java/org/zwobble/couscous/tests/frontends/java/ExpressionReadingTests.java
index e8b391e..e7b2f66 100644
--- a/src/test/java/org/zwobble/couscous/tests/frontends/java/ExpressionReadingTests.java
+++ b/src/test/java/org/zwobble/couscous/tests/frontends/java/ExpressionReadingTests.java
@@ -33,6 +33,7 @@ import static org.zwobble.couscous.tests.frontends.java.JavaReading.*;
 import static org.zwobble.couscous.tests.frontends.java.NodeMatchers.*;
 import static org.zwobble.couscous.types.ParameterizedType.parameterizedType;
 import static org.zwobble.couscous.types.TypeParameter.typeParameter;
+import static org.zwobble.couscous.util.ExtraIterables.only;
 import static org.zwobble.couscous.util.ExtraLists.list;
 
 public class ExpressionReadingTests {
@@ -335,7 +336,7 @@ public class ExpressionReadingTests {
             "java.util.Objects.toString(42)");
         assertEquals(
             typeCoercion(literal(42), Types.OBJECT),
-            expression.getArguments().get(0));
+            only(expression.getArguments()));
     }
 
     @Test
@@ -465,7 +466,7 @@ public class ExpressionReadingTests {
                 "    return name = \"blah\";" +
                 "}");
 
-        ReturnNode returnNode = (ReturnNode) classNode.getMethods().get(0).getBody().get().get(0);
+        ReturnNode returnNode = (ReturnNode) only(only(classNode.getMethods()).getBody().get());
         assertEquals(
             assign(
                 fieldAccess(
@@ -484,7 +485,7 @@ public class ExpressionReadingTests {
                 "    return value = 4;" +
                 "}");
 
-        ReturnNode returnNode = (ReturnNode) classNode.getMethods().get(0).getBody().get().get(0);
+        ReturnNode returnNode = (ReturnNode) only(only(classNode.getMethods()).getBody().get());
         assertEquals(
             assign(
                 fieldAccess(
diff --git a/src/test/java/org/zwobble/couscous/tests/frontends/java/JavaReading.java b/src/test/java/org/zwobble/couscous/tests/frontends/java/JavaReading.java
index eb6e839..e2d3d06 100644
--- a/src/test/java/org/zwobble/couscous/tests/frontends/java/JavaReading.java
+++ b/src/test/java/org/zwobble/couscous/tests/frontends/java/JavaReading.java
@@ -7,9 +7,8 @@ import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.List;
 
-import static org.hamcrest.Matchers.hasSize;
-import static org.junit.Assert.assertThat;
 import static org.zwobble.couscous.tests.util.ExtraFiles.deleteRecursively;
+import static org.zwobble.couscous.util.ExtraIterables.only;
 import static org.zwobble.couscous.util.ExtraLists.list;
 
 public class JavaReading {
@@ -20,7 +19,8 @@ public class JavaReading {
                 "}";
 
         ClassNode classNode = readClass(javaClass);
-        ReturnNode returnStatement = (ReturnNode) classNode.getMethods().get(0).getBody().get().get(0);
+        MethodNode method = only(classNode.getMethods());
+        ReturnNode returnStatement = (ReturnNode) only(method.getBody().get());
         return returnStatement.getValue();
     }
 
@@ -38,20 +38,18 @@ public class JavaReading {
     }
 
     static StatementNode readStatement(String returnType, String statementSource) {
-        // TODO: use only()
-        return readStatements(returnType, statementSource).get(0);
+        return only(readStatements(returnType, statementSource));
     }
 
     static List<StatementNode> readStatements(String returnType, String statementsSource) {
         String javaClass = generateMethodSource(returnType, statementsSource);
         ClassNode classNode = readClass(javaClass);
-        return classNode.getMethods().get(0).getBody().get();
+        return only(classNode.getMethods()).getBody().get();
     }
 
     static ClassNode readClass(String classBody) {
         List<TypeNode> classes = readTypes(classBody);
-        assertThat(classes, hasSize(1));
-        return (ClassNode)classes.get(0);
+        return (ClassNode) only(classes);
     }
 
     static List<TypeNode> readTypes(String classBody) {