Add methods to VariableCollection to retrieve types.

VariableCollection can now retrieve all the types used, and report them for
use with the TypeDeclarationGenerator interface.

src/main/java/ca/mcscert/jtet/expression/VariableCollection.java

@@ -163,6 +163,47 @@ final public class VariableCollection {
         return getAllVariables().values();
     }
 
+    /**
+     * Helper function to make a set of all variable types from a given collection.
+     * @param variables Variables to extract types from.
+     * @return The collected type set.
+     */
+    private Set<VariableType> getTypesFromVariableCollection(Collection<Variable> variables) {
+        final Set<VariableType> ret = new LinkedHashSet<VariableType>();
+
+        for (Variable var  : variables) {
+            ret.add(var.type());
+        }
+
+        return ret;
+    }
+
+    /**
+     * Returns a set containing all the variable types used in the input variables.
+     * @return A set of used variable types.
+     */
+    public Set<VariableType> getInputVariableTypes() {
+        return getTypesFromVariableCollection(getInputVariablesList());
+    }
+
+    /**
+     * Returns a set containing all the variable types used in the output variables.
+     * @return A set of used variable types.
+     */
+    public Set<VariableType> getOutputVariableTypes() {
+        return getTypesFromVariableCollection(getOutputVariablesList());
+    }
+
+    /**
+     * Returns a set containing all the variable types used.
+     * <p>
+     * This is equivalant to the union of {@link #getInputVariableTypes()} and {@link #getOutputVariableTypes()}.
+     * @return A set of used variable types.
+     */
+    public Set<VariableType> getAllVariableTypes() {
+        return getTypesFromVariableCollection(getAllVariablesList());
+    }
+
     private final Map<String,Variable> m_inputVariables = new LinkedHashMap<String, Variable>();
     private final Map<String,Variable> m_outputVariables = new LinkedHashMap<String, Variable>();
     private final Map<String,Variable> m_variablesAndEnumeratedValues = new HashMap<String, Variable>();

src/test/java/ca/mcscert/jtet/expression/test/VariableCollectionTest.java

@@ -137,4 +137,41 @@ public class VariableCollectionTest {
 
         new VariableCollection(new PartialVariableCollection(var, enums)); // Should throw as e1 the variable conflicts with the enumerated name e1.
     }
+
+    @Test
+    public void VariableTypeFetchingWorks() {
+        final EnumerationVariableType type1 = new EnumerationVariableType("myenum");
+        type1.enumerationValues().add("e1_1");
+
+        final EnumerationVariableType type2 = new EnumerationVariableType("myenum");
+        type2.enumerationValues().add("e2_2");
+
+        final Variable var1e = new Variable("v1e", type1);
+        final Variable var1r = new Variable("v1r", new RealVariableType());
+        final Variable var2e = new Variable("v2e", type2);
+
+        final Set<VariableType> inputTypeSet = new HashSet<VariableType>(2);
+        inputTypeSet.add(type1);
+        inputTypeSet.add(new RealVariableType());
+
+        final Set<VariableType> outputTypeSet = new HashSet<VariableType>(1);
+        outputTypeSet.add(type2);
+
+        final Set<VariableType> allTypeSet = new HashSet<VariableType>(2);
+        allTypeSet.addAll(inputTypeSet);
+        allTypeSet.addAll(outputTypeSet);
+
+        final Map<String, Variable> vars1 = new HashMap<String, Variable>();
+        vars1.put(var1e.name(), var1e);
+        vars1.put(var1r.name(), var1r);
+
+        final Map<String, Variable> vars2 = new HashMap<String, Variable>();
+        vars2.put(var2e.name(), var2e);
+
+        final VariableCollection vc = new VariableCollection(new PartialVariableCollection(vars1), new PartialVariableCollection(vars2));
+
+        assertThat(vc.getInputVariableTypes(), equalTo(inputTypeSet));
+        assertThat(vc.getOutputVariableTypes(), equalTo(outputTypeSet));
+        assertThat(vc.getAllVariableTypes(), equalTo(allTypeSet));
+    }
 }