package yokohama.unit.ast_junit;

import java.util.Set;
import java.util.TreeSet;
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) {
        TreeSet<Pair<ClassType, String>> varDecls = new TreeSet<>();
        varDecls.addAll(Stream.concat(
                testMethod.getBefore().stream().flatMap(this::visitStatement),
                testMethod.getStatements().stream().flatMap(this::visitStatement))
                .collect(Collectors.toSet()));
        return varDecls;

    }

    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());
    }
}

@@ -50,37 +50,11 @@
 til.
-Set;%0D%0Aimport java.util.TreeSe
+Lis
 t;%0D%0A
@@ -200,18 +200,19 @@
  public 
-Se
+Lis
 t%3CPair%3CC
@@ -284,94 +284,15 @@
     
-T
 re
-eSet%3CPair%3CClassType, String%3E%3E varDecls = new TreeSet%3C%3E();%0D%0A        varDecls.addAll(
+turn 
 Stre
@@ -515,38 +515,162 @@
 t())
-);%0D%0A        return varDecls;%0D%0A
+%0D%0A                .stream()%0D%0A                .sorted((o1, o2) -%3E o1.getSecond().compareTo(o2.getSecond()))%0D%0A                .collect(Collectors.toList());
 %0D%0A  

Use Stream.sorted to sort var decls instead of TreeSet (since Pair is not Comparable)