Tidy up processing of statements in switch

2 files changed, 39 insertions, 10 deletions

diff --git a/src/main/java/org/zwobble/couscous/frontends/java/JavaStatementReader.java b/src/main/java/org/zwobble/couscous/frontends/java/JavaStatementReader.java
index d03c7c0..99e8d7e 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/JavaStatementReader.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/JavaStatementReader.java
@@ -11,7 +11,8 @@ import java.util.Map;
 import java.util.Optional;
 import java.util.stream.Collectors;
 
-import static com.google.common.collect.Iterables.*;
+import static com.google.common.collect.Iterables.filter;
+import static com.google.common.collect.Iterables.tryFind;
 import static com.google.common.collect.Maps.immutableEntry;
 import static org.zwobble.couscous.ast.ExpressionStatementNode.expressionStatement;
 import static org.zwobble.couscous.ast.IfStatementNode.ifStatement;
@@ -21,7 +22,8 @@ import static org.zwobble.couscous.ast.WhileNode.whileLoop;
 import static org.zwobble.couscous.frontends.java.JavaExpressionReader.coerceExpression;
 import static org.zwobble.couscous.frontends.java.JavaTypes.typeOf;
 import static org.zwobble.couscous.util.ExtraLists.*;
-import static org.zwobble.couscous.util.ExtraLists.concat;
+import static org.zwobble.couscous.util.Tails.tail;
+import static org.zwobble.couscous.util.Tails.tails;
 import static org.zwobble.couscous.util.UpToAndIncludingIterable.upToAndIncluding;
 
 class JavaStatementReader {
@@ -100,10 +102,10 @@ class JavaStatementReader {
         List<Statement> statements = switchStatement.statements();
         List<Map.Entry<Optional<ExpressionNode>, List<StatementNode>>> cases = new ArrayList<>();
 
-        for (int statementIndex = 0; statementIndex < statements.size(); statementIndex++) {
-            Statement statement = statements.get(statementIndex);
-            if (statement.getNodeType() == ASTNode.SWITCH_CASE) {
-                cases.add(readSwitchCase(statements, statementIndex));
+        for (List<Statement> tail : tails(statements)) {
+            Statement first = tail.get(0);
+            if (first.getNodeType() == ASTNode.SWITCH_CASE) {
+                cases.add(readSwitchCase((SwitchCase) first, tail(tail)));
             }
         }
 
@@ -129,12 +131,13 @@ class JavaStatementReader {
             handle);
     }
 
-    private Map.Entry<Optional<ExpressionNode>,List<StatementNode>> readSwitchCase(List<Statement> statements, int statementIndex) {
-        SwitchCase caseStatement = (SwitchCase)statements.get(statementIndex);
-
+    private Map.Entry<Optional<ExpressionNode>,List<StatementNode>> readSwitchCase(
+        SwitchCase caseStatement,
+        List<Statement> statements)
+    {
         Iterable<Statement> statementsForCase = upToAndIncluding(
             filter(
-                skip(statements, statementIndex),
+                statements,
                 statement -> statement.getNodeType() != ASTNode.SWITCH_CASE),
             JavaStatementReader::isEndOfCase);
         return immutableEntry(
diff --git a/src/main/java/org/zwobble/couscous/util/Tails.java b/src/main/java/org/zwobble/couscous/util/Tails.java
new file mode 100644
index 0000000..f9ec724
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/util/Tails.java
@@ -0,0 +1,26 @@
+package org.zwobble.couscous.util;
+
+import java.util.Iterator;
+import java.util.List;
+
+public class Tails {
+    public static <T> Iterable<List<T>> tails(List<T> list) {
+        return () -> new Iterator<List<T>>() {
+            private int startIndex = 0;
+
+            @Override
+            public boolean hasNext() {
+                return startIndex < list.size();
+            }
+
+            @Override
+            public List<T> next() {
+                return list.subList(startIndex++, list.size());
+            }
+        };
+    }
+
+    public static <T> List<T> tail(List<T> list) {
+        return list.subList(1, list.size());
+    }
+}