Merge pull request #335 from CohenArthur/187-mutate-var-in-block

Fix variable mutation in block

Author: CohenArthur

src/context.rs

@@ -198,10 +198,12 @@ impl Context {
 
     /// Replace a variable or create it if it does not exist
     pub fn replace_variable(&mut self, var: Var) -> Result<(), Error> {
-        // Remove the variable if it exists
-        let _ = self.remove_variable(&var);
+        match self.scope_map.get_variable_mut(var.name()) {
+            None => self.add_variable(var)?,
+            Some(var_ref) => var_ref.set_instance(var.instance()),
+        }
 
-        self.add_variable(var)
+        Ok(())
     }
 
     /// Get a mutable reference on an existing function

src/context/scope_map.rs

@@ -22,6 +22,11 @@ impl<V, F, T> Scope<V, F, T> {
         self.variables.get(name)
     }
 
+    /// Get a mutable reference on a variable from the scope map if it has been inserted already
+    pub fn get_variable_mut(&mut self, name: &str) -> Option<&mut V> {
+        self.variables.get_mut(name)
+    }
+
     /// Get a reference on a function from the scope map if is has been inserted already
     pub fn get_function(&self, name: &str) -> Option<&F> {
         self.functions.get(name)
@@ -144,6 +149,17 @@ impl<V, F, T> ScopeMap<V, F, T> {
         None
     }
 
+    pub fn get_variable_mut(&mut self, name: &str) -> Option<&mut V> {
+        for scope in self.scopes.iter_mut() {
+            match scope.get_variable_mut(name) {
+                Some(v) => return Some(v),
+                None => continue,
+            };
+        }
+
+        None
+    }
+
     /// Maybe get a function in any available scopes
     pub fn get_function(&self, name: &str) -> Option<&F> {
         // FIXME: Use find for code quality?

src/instruction/var_assignment.rs

@@ -157,10 +157,10 @@ impl TypeCheck for VarAssign {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::jinko_fail;
     use crate::parser::{Construct, Token};
     use crate::value::{JkInt, JkString};
     use crate::ToObjectInstance;
+    use crate::{jinko, jinko_fail};
 
     #[test]
     fn non_mutable() {
@@ -223,4 +223,36 @@ mod tests {
             mut a0 = 14;
         };
     }
+
+    #[test]
+    fn assign_mutable_in_block_187() {
+        let ctx = jinko! {
+            mut x = 1;
+            {
+                x = 0;
+            }
+        };
+
+        assert_eq!(
+            ctx.get_variable("x").unwrap().instance(),
+            JkInt::from(0).to_instance()
+        );
+    }
+
+    #[test]
+    fn assign_mutable_in_function_187() {
+        let ctx = jinko! {
+            mut x = 1;
+            func change_global() {
+                x = 0;
+            }
+
+            change_global()
+        };
+
+        assert_eq!(
+            ctx.get_variable("x").unwrap().instance(),
+            JkInt::from(0).to_instance()
+        );
+    }
 }

tests/ft/regression/187.jk

@@ -0,0 +1,6 @@
+mut x = 1;
+{
+    x = 0;
+}
+
+x

tests/ft/regression/187_fn.jk

@@ -0,0 +1,13 @@
+mut global_x = 1;
+
+func mutate_global() {
+    func inner() {
+        global_x = 0;
+    }
+
+    inner()
+}
+
+mutate_global();
+
+x

tests/ft/regression/regression.yml

@@ -4,3 +4,13 @@ tests:
     args:
       - "tests/ft/regression/string_double_free.jk"
     exit_code: 0
+  - name: "Variable mutation in block #187"
+    binary: "target/debug/jinko"
+    args:
+      - "tests/ft/regression/187.jk"
+    exit_code: 0
+  - name: "Variable mutation in function #187"
+    binary: "target/debug/jinko"
+    args:
+      - "tests/ft/regression/187_fn.jk"
+    exit_code: 0