Add flattening of variable reassignment (#1)

examples/reassignment.code

@@ -1,4 +1,4 @@
-def reass(x):
+def reassign(x):
   a = x + 5
   b = a + x
   a = 7

src/absy.rs

@@ -37,7 +37,7 @@ impl Prog {
 }
 impl fmt::Display for Prog {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "{}({}):\n{}", self.id, self.arguments.iter().map(|x| format!("{}", x)).collect::<Vec<_>>().join(","), self.statements.iter().map(|x| format!("\t{}", x)).collect::<Vec<_>>().join("\n"))
+        write!(f, "def {}({}):\n{}", self.id, self.arguments.iter().map(|x| format!("{}", x)).collect::<Vec<_>>().join(","), self.statements.iter().map(|x| format!("\t{}", x)).collect::<Vec<_>>().join("\n"))
     }
 }
 impl fmt::Debug for Prog {
@@ -94,6 +94,27 @@ pub enum Expression {
     IfElse(Box<Condition>, Box<Expression>, Box<Expression>),
 }
 impl Expression {
+    pub fn apply_substitution(&self, substitution: &HashMap<String, String>) -> Expression {
+        match *self {
+            ref e @ Expression::NumberLiteral(_) => e.clone(),
+            Expression::VariableReference(ref v) => {
+                let mut new_name = v.to_string();
+                loop {
+                    match substitution.get(&new_name) {
+                        Some(x) => new_name = x.to_string(),
+                        None => return Expression::VariableReference(new_name),
+                    }
+                }
+            },
+            Expression::Add(ref e1, ref e2) => Expression::Add(box e1.apply_substitution(substitution), box e2.apply_substitution(substitution)),
+            Expression::Sub(ref e1, ref e2) => Expression::Sub(box e1.apply_substitution(substitution), box e2.apply_substitution(substitution)),
+            Expression::Mult(ref e1, ref e2) => Expression::Mult(box e1.apply_substitution(substitution), box e2.apply_substitution(substitution)),
+            Expression::Div(ref e1, ref e2) => Expression::Div(box e1.apply_substitution(substitution), box e2.apply_substitution(substitution)),
+            Expression::Pow(ref e1, ref e2) => Expression::Pow(box e1.apply_substitution(substitution), box e2.apply_substitution(substitution)),
+            Expression::IfElse(ref c, ref e1, ref e2) => Expression::IfElse(box c.apply_substitution(substitution), box e1.apply_substitution(substitution), box e2.apply_substitution(substitution)),
+        }
+    }
+
     fn solve(&self, inputs: &mut HashMap<String, i32>) -> i32 {
         match *self {
             Expression::NumberLiteral(x) => x,
@@ -201,6 +222,16 @@ pub enum Condition {
     Gt(Expression, Expression),
 }
 impl Condition {
+    fn apply_substitution(&self, substitution: &HashMap<String, String>) -> Condition {
+        match *self {
+            Condition::Lt(ref lhs, ref rhs) => Condition::Lt(lhs.apply_substitution(substitution), rhs.apply_substitution(substitution)),
+            Condition::Le(ref lhs, ref rhs) => Condition::Le(lhs.apply_substitution(substitution), rhs.apply_substitution(substitution)),
+            Condition::Eq(ref lhs, ref rhs) => Condition::Eq(lhs.apply_substitution(substitution), rhs.apply_substitution(substitution)),
+            Condition::Ge(ref lhs, ref rhs) => Condition::Ge(lhs.apply_substitution(substitution), rhs.apply_substitution(substitution)),
+            Condition::Gt(ref lhs, ref rhs) => Condition::Gt(lhs.apply_substitution(substitution), rhs.apply_substitution(substitution)),
+        }
+    }
+
     fn solve(&self, inputs: &mut HashMap<String, i32>) -> bool {
         match *self {
             Condition::Lt(ref lhs, ref rhs) => lhs.solve(inputs) < rhs.solve(inputs),

src/flatten.rs

@@ -1,4 +1,6 @@
 //
+//
+//
 //@file flatten.rs
 //@author Dennis Kuhnert <dennis.kuhnert@campus.tu-berlin.de>
 //@date 2017
@@ -234,18 +236,45 @@ impl Flattener {
         for def in prog.statements {
             match def {
                 Statement::Return(expr) => {
+                    let expr_subbed = expr.apply_substitution(&self.substitution);
+                    let rhs = self.flatten_expression(&mut statements_flattened, expr_subbed);
                     self.variables.insert("~out".to_string());
-                    let rhs = self.flatten_expression(&mut statements_flattened, expr);
                     statements_flattened.push(Statement::Return(rhs));
                 },
                 Statement::Definition(id, expr) => {
-                    self.variables.insert(id.to_string());
-                    let rhs = self.flatten_expression(&mut statements_flattened, expr);
-                    statements_flattened.push(Statement::Definition(id, rhs));
+                    let expr_subbed = expr.apply_substitution(&self.substitution);
+                    let rhs = self.flatten_expression(&mut statements_flattened, expr_subbed);
+                    statements_flattened.push(Statement::Definition(self.use_variable(id), rhs));
                 },
                 Statement::Condition(..) => unimplemented!(),
             }
         }
+        println!("DEBUG self.variables {:?}", self.variables);
+        println!("DEBUG self.substitution {:?}", self.substitution);
         Prog { id: prog.id, arguments: prog.arguments, statements: statements_flattened }
     }
+
+    /// Proofs if the given name is a not used variable and returns a fresh variable.
+    ///
+    /// # Arguments
+    ///
+    /// * `name` - A String that holds the name of the variable
+    fn use_variable(&mut self, name: String) -> String {
+        let mut i = 0;
+        let mut new_name = name.to_string();
+        loop {
+            if self.variables.contains(&new_name) {
+                new_name = format!("{}_{}", &name, i);
+                i += 1;
+            } else {
+                self.variables.insert(new_name.to_string());
+                if i == 1 {
+                    self.substitution.insert(name, new_name.to_string());
+                } else if i > 1 {
+                    self.substitution.insert(format!("{}_{}", name, i - 2), new_name.to_string());
+                }
+                return new_name;
+            }
+        }
+    }
 }