package yokohama.unit.ast_junit;

import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import yokohama.unit.util.Pair;

public class VarDeclVisitor {
    public Set<Pair<ClassType, String>> visitTestMethod(TestMethod testMethod) {
        return Stream.concat(
                testMethod.getBefore().stream().flatMap(this::visitStatement),
                testMethod.getStatements().stream().flatMap(this::visitStatement))
                .collect(Collectors.toSet());
    }

    public Stream<Pair<ClassType, String>> visitStatement(Statement statement) { 
        return statement.accept(
                isStatement -> Stream.<Pair<ClassType, String>>empty(),
                isNotStatement -> Stream.<Pair<ClassType, String>>empty(),
                actionStatement -> Stream.<Pair<ClassType, String>>empty(),
                varInitStatement ->
                        Stream.of(new Pair<ClassType, String>(varInitStatement.getType(), varInitStatement.getName())),
                returnIsStatement -> Stream.<Pair<ClassType, String>>empty(),
                returnIsNotStatement -> Stream.<Pair<ClassType, String>>empty(),
                invokeVoidStatement -> Stream.<Pair<ClassType, String>>empty(),
                tryStatement ->
                        Stream.concat(
                                tryStatement.getTryStatements().stream().flatMap(this::visitStatement),
                                tryStatement.getFinallyStatements().stream().flatMap(this::visitStatement)),
                varDeclStatement ->
                        Stream.of(new Pair<ClassType, String>(varDeclStatement.getClassType(), varDeclStatement.getName())),
                varAssignStatement -> Stream.<Pair<ClassType, String>>empty());
    }
}

@@ -52,16 +52,43 @@
 l.Set;%0D%0A
+import java.util.TreeSet;%0D%0A
 import j
@@ -116,16 +116,16 @@
 ctors;%0D%0A
-
 import j
@@ -305,23 +305,102 @@
         
+T
 re
-turn 
+eSet%3CPair%3CClassType, String%3E%3E varDecls = new TreeSet%3C%3E();%0D%0A        varDecls.addAll(
 Stream.c
@@ -615,17 +615,46 @@
 toSet())
-;
+);%0D%0A        return varDecls;%0D%0A
 %0D%0A    %7D%0D

Move variable declaration to the beginning of method