equation rebalancing

2 files changed, 105 insertions, 3 deletions

diff --git a/src/modelica_ast.rs b/src/modelica_ast.rs
index 6fd29f8..33af4ac 100644
--- a/src/modelica_ast.rs
+++ b/src/modelica_ast.rs
@@ -1,4 +1,5 @@
 
+use std::clone::Clone;
 use std::fmt::{Debug, Formatter, Error};
 use std::collections::HashMap;
 
@@ -23,7 +24,7 @@ pub enum ComponentPrefix {
     Constant,
 }
 
-#[derive(PartialEq)]
+#[derive(Clone, PartialEq)]
 pub struct Component {
     pub prefix: Option<ComponentPrefix>,
     pub specifier: String,
@@ -39,13 +40,13 @@ pub struct Connection {
     pub b: String,
 }
 
-#[derive(PartialEq)]
+#[derive(Clone, PartialEq)]
 pub struct SimpleEquation {
     pub lhs: Expr,
     pub rhs: Expr,
 }
 
-#[derive(PartialEq)]
+#[derive(Clone, PartialEq)]
 pub enum Expr {
     Integer(i64),
     Float(f64),
@@ -127,6 +128,10 @@ impl Expr {
             },
         }
     }
+
+    pub fn contains(&self, ident: &str) -> bool{
+        self.identifiers().contains(&ident.to_string())
+    }
 }
 
 impl SimpleEquation {
@@ -135,6 +140,81 @@ impl SimpleEquation {
     pub fn identifiers(&self) -> Vec<String> {
         union_vecs(&self.lhs.identifiers(), &self.rhs.identifiers())
     }
+
+    pub fn contains(&self, ident: &str) -> bool{
+        let s = &ident.to_string();
+        self.lhs.identifiers().contains(s) || self.rhs.identifiers().contains(s)
+    }
+
+    pub fn rebalance_for(&self, ident: String) -> Result<SimpleEquation,String> {
+        let lvars = self.lhs.identifiers();
+        let rvars = self.rhs.identifiers();
+
+        let ret = match (lvars.contains(&ident), rvars.contains(&ident)) {
+            (true, true) => Err("SymbolicError: NaiveImplementation".to_string()),
+            (false, false) => Err("SymbolicError: VariableNotFound".to_string()),
+            (true, false) => self.simplify_lhs(&ident),
+            (false, true) =>
+                SimpleEquation{lhs: self.rhs.clone(),
+                               rhs: self.lhs.clone()}.simplify_lhs(&ident),
+        };
+        match ret {
+            Ok(eqn) => {
+                if eqn.rhs.contains(&ident) {
+                    Err("SymbolicError: NaiveImplementation".to_string())
+                } else {
+                    Ok(eqn)
+                }},
+            Err(_) => ret,
+        }
+    }
+
+    pub fn simplify_lhs(&self, ident: &str) -> Result<SimpleEquation,String> {
+        use self::Expr::*;
+        use self::BinOperator::*;
+        match self.lhs {
+            Ident(ref s) if s == ident => Ok((*self).clone()),
+            Ident(_) | Integer(_) | Float(_) =>
+                Err("SymbolicError: InternalError: expected var on LHS".to_string()),
+            Der(_) | Abs(_) =>
+                Err("SymbolicError: NaiveImplementation: can't simplify der() or abs()".to_string()),
+            // TODO: create a macro for the below...
+            BinExpr(Multiply, ref a, ref b) if a.contains(ident) => {
+                SimpleEquation{
+                    lhs: *a.clone(),
+                    rhs: BinExpr(Divide, Box::new(self.rhs.clone()), b.clone())}.simplify_lhs(&ident) },
+            BinExpr(Multiply, ref a, ref b) if b.contains(ident) => {
+                SimpleEquation{
+                    lhs: *b.clone(),
+                    rhs: BinExpr(Divide, Box::new(self.rhs.clone()), a.clone())}.simplify_lhs(&ident) },
+            BinExpr(Divide, ref a, ref b) if a.contains(ident) => {
+                SimpleEquation{
+                    lhs: *a.clone(),
+                    rhs: BinExpr(Multiply, Box::new(self.rhs.clone()), b.clone())}.simplify_lhs(&ident) },
+            BinExpr(Divide, ref a, ref b) if b.contains(ident) => {
+                SimpleEquation{
+                    lhs: *b.clone(),
+                    rhs: BinExpr(Divide, a.clone(), Box::new(self.rhs.clone()))}.simplify_lhs(&ident) },
+            BinExpr(Add, ref a, ref b) if a.contains(ident) => {
+                SimpleEquation{
+                    lhs: *a.clone(),
+                    rhs: BinExpr(Subtract, Box::new(self.rhs.clone()), b.clone())}.simplify_lhs(&ident) },
+            BinExpr(Add, ref a, ref b) if b.contains(ident) => {
+                SimpleEquation{
+                    lhs: *b.clone(),
+                    rhs: BinExpr(Subtract, Box::new(self.rhs.clone()), a.clone())}.simplify_lhs(&ident) },
+            BinExpr(Subtract, ref a, ref b) if a.contains(ident) => {
+                SimpleEquation{
+                    lhs: *a.clone(),
+                    rhs: BinExpr(Add, Box::new(self.rhs.clone()), b.clone())}.simplify_lhs(&ident) },
+            BinExpr(Subtract, ref a, ref b) if b.contains(ident) => {
+                SimpleEquation{
+                    lhs: *b.clone(),
+                    rhs: BinExpr(Subtract, a.clone(), Box::new(self.rhs.clone()))}.simplify_lhs(&ident) },
+            BinExpr(_, _, _) => Err("SymbolicError: NotImplemented BinOperator (or else couldn't find var...)".to_string()),
+            // in case we add opers: _ => Err("NotImplemented".to_string()),
+        }
+    }
 }
 
 
diff --git a/tests/modelica_ast.rs b/tests/modelica_ast.rs
index b6233f3..d39615e 100644
--- a/tests/modelica_ast.rs
+++ b/tests/modelica_ast.rs
@@ -54,3 +54,25 @@ fn test_eqn_identifiers() {
                 Box::new(Ident("z".to_string()))),
         }.identifiers()));
 }
+
+#[test]
+fn test_rebalance_for() {
+    use modelthing::modelica_ast::Expr::*;
+
+    // z = a - 1.2345
+    // a = z + 1.2345
+    let done = SimpleEquation{
+        lhs: Ident("z".to_string()),
+        rhs: BinExpr(BinOperator::Subtract,
+            Box::new(Ident("a".to_string())),
+            Box::new(Float(1.2345)))};
+    assert_eq!(done,
+        done.rebalance_for("z".to_string()).unwrap());
+    assert_eq!(SimpleEquation{
+        lhs: Ident("a".to_string()),
+        rhs: BinExpr(BinOperator::Add,
+            Box::new(Ident("z".to_string())),
+            Box::new(Float(1.2345)))},
+        done.rebalance_for("a".to_string()).unwrap());
+
+}