2 files changed, 39 insertions, 46 deletions

diff --git a/modelica-parser-lalrpop/src/ast.rs b/modelica-parser-lalrpop/src/ast.rs
index 7eb9dda..f31681e 100644
--- a/modelica-parser-lalrpop/src/ast.rs
+++ b/modelica-parser-lalrpop/src/ast.rs
@@ -1,7 +1,8 @@
 
 use std::clone::Clone;
 use std::fmt::{Debug, Formatter, Error};
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
+use std::iter::FromIterator;
 
 
 #[derive(Clone, PartialEq)]
@@ -146,7 +147,7 @@ impl ModelicaModel {
     // the sole element on the LHS of an equation.
     // Bugs:
     //   if a var is on LHS and RHS of same equation
-    pub fn get_free_vars(&self) -> Vec<String> {
+    pub fn get_free_vars(&self) -> HashSet<String> {
         // Start with components, and remove constants and parameters
         let vars = self.components.iter().filter(|v| match v.prefix {
                 Some(ComponentPrefix::Constant) | Some(ComponentPrefix::Parameter) => false,
@@ -164,39 +165,34 @@ impl ModelicaModel {
         }
         let vars = vars.filter(|v| !outputs.contains(&v.name));
 
-        vars.map(|c| c.name.clone()).collect()
+        HashSet::from_iter(vars.map(|c| c.name.clone()))
     }
 
 }
 
-fn union_strings(a: &Vec<String>, b: &Vec<String>) -> Vec<String> {
-    let mut u = a.clone();
-    for e in b {
-        if !(u.contains(&e)) {
-            u.push(e.clone());
-        }
-    }
-    u
-}
-
 impl Expr {
 
     // Order is undefined
-    // TODO: should return a HashSet, not a Vec
-    pub fn identifiers(&self) -> Vec<String> {
+    pub fn identifiers(&self) -> HashSet<String> {
         use self::Expr::*;
         match *self {
-            Integer(_) | Float(_) | Boolean(_) | StringLiteral(_) => vec![],
-            Ident(ref s) => vec![s.clone()],
+            Integer(_) | Float(_) | Boolean(_) | StringLiteral(_) => HashSet::new(),
+            Ident(ref s) => {
+                let mut hs = HashSet::new();
+                hs.insert(s.to_string());
+                hs
+            },
             Der(ref e) | Sign(ref e) => e.identifiers(),
             MathUnaryExpr(_, ref e) => e.identifiers(),
             BinExpr(_, ref e1, ref e2) => {
-                union_strings(&e1.identifiers(), &e2.identifiers())
+                let mut hs = e1.identifiers();
+                hs.extend(e2.identifiers());
+                hs
             },
             Array(ref el) => {
-                let mut all: Vec<String> = vec![];
+                let mut all: HashSet<String> = HashSet::new();
                 for e in el {
-                    all.append(&mut e.identifiers());
+                    all.extend(e.identifiers());
                 }
                 all
             }
@@ -204,15 +200,17 @@ impl Expr {
     }
 
     pub fn contains(&self, ident: &str) -> bool{
-        self.identifiers().contains(&ident.to_string())
+        self.identifiers().contains(ident)
     }
 }
 
 impl SimpleEquation {
 
     // Order is undefined
-    pub fn identifiers(&self) -> Vec<String> {
-        union_strings(&self.lhs.identifiers(), &self.rhs.identifiers())
+    pub fn identifiers(&self) -> HashSet<String> {
+        let mut hs = self.lhs.identifiers();
+        hs.extend(self.rhs.identifiers());
+        hs
     }
 
     pub fn contains(&self, ident: &str) -> bool{
diff --git a/modelica-parser-lalrpop/tests/ast.rs b/modelica-parser-lalrpop/tests/ast.rs
index 51c1014..d551142 100644
--- a/modelica-parser-lalrpop/tests/ast.rs
+++ b/modelica-parser-lalrpop/tests/ast.rs
@@ -5,52 +5,47 @@ use std::collections::HashSet;
 use std::iter::FromIterator;
 use modelica_parser::*;
 
-fn set_eq(a: Vec<String>, b: Vec<String>) -> bool {
-    let set_a: HashSet<String> = HashSet::from_iter(a);
-    let set_b: HashSet<String> = HashSet::from_iter(b);
-    return set_a == set_b;
-}
 
 #[test]
 fn test_expr_identifiers() {
     use modelica_parser::Expr::*;
 
-    assert!(set_eq(
-        vec![],
-        Integer(0).identifiers()));
-    assert!(set_eq(
-        vec!["x".to_string()],
-        Ident("x".to_string()).identifiers()));
-    assert!(set_eq(
-        vec!["x".to_string(), "y".to_string(), "z".to_string()],
+    assert_eq!(
+        HashSet::new(),
+        Integer(0).identifiers());
+    assert_eq!(
+        HashSet::from_iter(vec!["x".to_string()]),
+        Ident("x".to_string()).identifiers());
+    assert_eq!(
+        HashSet::from_iter(vec!["x".to_string(), "y".to_string(), "z".to_string()]),
         BinExpr(BinOperator::Add,
             Box::new(Sign(Box::new(Ident("z".to_string())))),
             Box::new(BinExpr(BinOperator::Add,
                 Box::new(Sign(Box::new(Ident("x".to_string())))),
-                Box::new(Sign(Box::new(Ident("y".to_string()))))))).identifiers()));
-    assert!(set_eq(
-        vec!["z".to_string()],
+                Box::new(Sign(Box::new(Ident("y".to_string()))))))).identifiers());
+    assert_eq!(
+        HashSet::from_iter(vec!["z".to_string()]),
         BinExpr(BinOperator::Add,
             Box::new(Ident("z".to_string())),
-            Box::new(Ident("z".to_string()))).identifiers()));
+            Box::new(Ident("z".to_string()))).identifiers());
 }
 
 #[test]
 fn test_eqn_identifiers() {
     use modelica_parser::Expr::*;
 
-    assert!(set_eq(
-        vec![],
+    assert_eq!(
+        HashSet::new(),
         SimpleEquation{
             lhs: Integer(0),
             rhs: Integer(0),
-        }.identifiers()));
-    assert!(set_eq(
-        vec!["z".to_string()],
+        }.identifiers());
+    assert_eq!(
+        HashSet::from_iter(vec!["z".to_string()]),
         SimpleEquation{
             lhs: Ident("z".to_string()),
             rhs: BinExpr(BinOperator::Add,
                 Box::new(Ident("z".to_string())),
                 Box::new(Ident("z".to_string()))),
-        }.identifiers()));
+        }.identifiers());
 }