identifiers() for simpleequations

1 files changed, 42 insertions, 14 deletions

diff --git a/src/modelica_ast.rs b/src/modelica_ast.rs
index 2828f81..d93d81d 100644
--- a/src/modelica_ast.rs
+++ b/src/modelica_ast.rs
@@ -1,6 +1,6 @@
 
 use std::fmt::{Debug, Formatter, Error};
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 #[derive(PartialEq)]
 pub struct ModelicaModel {
@@ -102,6 +102,15 @@ impl ModelicaModel {
 
         vars.map(|c| c.name.clone()).collect()
     }
+
+fn union_vecs(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 {
@@ -114,28 +123,23 @@ impl Expr {
             Ident(ref s) => vec![s.clone()],
             Der(ref e) | Abs(ref e) => e.identifiers(),
             BinExpr(_, ref e1, ref e2) => {
-                let mut all = e1.identifiers();
-                for i in e2.identifiers() {
-                    if !(all.contains(&i)) {
-                        all.push(i.clone());
-                    }
-                }
-                all
+                union_vecs(&e1.identifiers(), &e2.identifiers())
             },
         }
     }
 }
 
+#[cfg(test)]
+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 self::Expr::*;
 
-    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;
-    }
-
     assert!(set_eq(
         vec![],
         Integer(0).identifiers()));
@@ -157,8 +161,32 @@ fn test_expr_identifiers() {
 }
 
 impl SimpleEquation {
+
+    // Order is undefined
+    pub fn identifiers(&self) -> Vec<String> {
+        union_vecs(&self.lhs.identifiers(), &self.rhs.identifiers())
+    }
 }
 
+#[test]
+fn test_eqn_identifiers() {
+    use self::Expr::*;
+
+    assert!(set_eq(
+        vec![],
+        SimpleEquation{
+            lhs: Integer(0),
+            rhs: Integer(0),
+        }.identifiers()));
+    assert!(set_eq(
+        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()));
+}
 
 //// Debug Implementations