basic AST substituting (finishing first stab at solve_for)

3 files changed, 50 insertions, 4 deletions

diff --git a/src/bin/mt-tool.rs b/src/bin/mt-tool.rs
index 317ac20..37b1705 100644
--- a/src/bin/mt-tool.rs
+++ b/src/bin/mt-tool.rs
@@ -56,9 +56,9 @@ fn main() {
             .arg_from_usage("<DIR> 'model to load'"))
         .subcommand(SubCommand::with_name("solve_for")
             .about("")
-            .arg_from_usage("<DIR> 'model to load'
-                             --dep <VAR>... 'dependent variable'
-                             --indep <VAR>... 'independent variable'"))
+            .arg_from_usage("<DIR> 'model to load'")
+            .arg_from_usage("--dep <VAR>... 'dependent variable'")
+            .arg_from_usage("--indep <VAR>... 'independent variable'"))
         .get_matches();
 
     
diff --git a/src/modelica_model.rs b/src/modelica_model.rs
index 81680d6..7b77123 100644
--- a/src/modelica_model.rs
+++ b/src/modelica_model.rs
@@ -134,6 +134,7 @@ impl ModelicaModelExt for ModelicaModel {
         let mut solved: Vec<SimpleEquation> = vec![];
         let mut unsolved_vars = dep_vars.clone();
         while unsolved_eqns.len() > 0 {
+            println!("vars: {:?}    eqns: {:?}", &unsolved_vars, &unsolved_eqns);
             let next_i = unsolved_eqns.iter()
                 .position(|ref x| intersect_strings(&unsolved_vars, &x.identifiers()).len() == 1);
             let eqn = match next_i {
@@ -143,13 +144,20 @@ impl ModelicaModelExt for ModelicaModel {
                 Some(i) => unsolved_eqns.remove(i),
             };
             let ref var = intersect_strings(&unsolved_vars, &eqn.identifiers())[0];
+            println!("solving: {}; {:?}", var, eqn);
             let eqn = eqn.rebalance_for(var.to_string()).expect("rebalance for success");
+            println!("got: {:?}", eqn);
+            // Replace all other references to var with RHS of solved equation
+            unsolved_eqns = unsolved_eqns.iter().map(|ref e|
+                SimpleEquation{
+                    lhs: substitute_with(&e.lhs, &Expr::Ident(var.to_string()), &eqn.rhs),
+                    rhs: substitute_with(&e.rhs, &Expr::Ident(var.to_string()), &eqn.rhs),
+                }).collect();
             solved.push(eqn);
             let var_i = unsolved_vars.iter().position(|ref x| x == &var).unwrap();
             unsolved_vars.remove(var_i);
         }
 
-        // TODO: sort output equations by LHS
         Ok(ModelicaModel {
             name: self.name.clone(),
             description: self.description.clone(),
@@ -160,6 +168,27 @@ impl ModelicaModelExt for ModelicaModel {
     }
 }
 
+// Recurses through 'original', replacing all instances of 'a' with 'b'
+fn substitute_with(original: &Expr, a: &Expr, b: &Expr) -> Expr {
+    use modelica_parser::ast::Expr::*;
+    println!("original: {:?}  replacing: {:?}   with: {:?}", original, a, b);
+    if *original == *a {
+        return b.clone();
+    }
+    match *original {
+        Integer(_) | Float(_) | Boolean(_) | StringLiteral(_) | Ident(_) => original.clone(),
+        Der(ref e) => Der(Box::new(substitute_with(e, a, b))),
+        Sign(ref e) => Sign(Box::new(substitute_with(e, a, b))),
+        MathUnaryExpr(muf, ref e) =>
+            MathUnaryExpr(muf, Box::new(substitute_with(e, a, b))),
+        BinExpr(bo, ref e1, ref e2) =>
+            BinExpr(bo,
+                    Box::new(substitute_with(e1, a, b)),
+                    Box::new(substitute_with(e2, a, b))),
+        Array(ref l) => Array(l.iter().map(|ref e| substitute_with(e, a, b)).collect()),
+    }
+}
+
 fn intersect_strings(a: &Vec<String>, b: &Vec<String>) -> Vec<String> {
     let mut both = vec![];
     for e in a {
diff --git a/tests/lib.rs b/tests/lib.rs
index bf28476..77359de 100644
--- a/tests/lib.rs
+++ b/tests/lib.rs
@@ -32,3 +32,20 @@ fn test_load_model_entry() {
 fn test_search_models() {
     assert_eq!(search_models(Path::new("./examples/")).len() > 1, true);
 }
+
+#[test]
+fn test substitute_with() {
+    use modelica_parser::ast::Expr::*;
+
+    let y = BinExpr(BinOperator::Add,
+                Box::new(Ident("y".to_string())),
+                Box::new(Ident("y".to_string())));
+    let z = BinExpr(BinOperator::Add,
+                Box::new(Ident("z".to_string())),
+                Box::new(Ident("z".to_string())));
+
+    assert_eq!(z,
+               substitute_with(&y,
+                               Ident("y".to_string()),
+                               Ident("z".to_string())));
+}