From 465039cf430729c9179c14301dc4d1b59a404af1 Mon Sep 17 00:00:00 2001
From: Nick Krichevsky <nick@ollien.com>
Date: Fri, 24 May 2024 10:35:39 -0400
Subject: [PATCH] Fix bug where assignment would shadow variables in higher
 scopes

---
 src/eval.rs             | 15 +++++++--------
 src/eval/environment.rs | 41 ++++++++++++++++++++++++++++++++++++++---
 2 files changed, 45 insertions(+), 11 deletions(-)

diff --git a/src/eval.rs b/src/eval.rs
index 52fabae..6ec7c59 100644
--- a/src/eval.rs
+++ b/src/eval.rs
@@ -196,18 +196,17 @@ impl ExprVisitor<Result<EvaluatedValue, ScriptError>> for InterpreterRunner<'_>
     }
 
     fn visit_assign(&mut self, name: &Token, value: &Expr) -> Result<EvaluatedValue, ScriptError> {
-        if self.interpreter.env.get_value(name.lexeme()).is_none() {
-            return Err(ScriptError {
+        let value = self.visit_expr(value)?;
+        let assign_res = self.interpreter.env.update_value(name.lexeme(), value);
+
+        match assign_res {
+            Some(stored) => Ok(stored.clone()),
+            None => Err(ScriptError {
                 line: name.line(),
                 location: String::new(),
                 message: format!("Cannot assign to undefined variable {}", name.lexeme()),
-            });
+            }),
         }
-
-        let value = self.visit_expr(value)?;
-        let inserted_value = self.interpreter.env.set_value(name.lexeme(), value);
-
-        Ok(inserted_value.clone())
     }
 }
 
diff --git a/src/eval/environment.rs b/src/eval/environment.rs
index d61d125..d9091bb 100644
--- a/src/eval/environment.rs
+++ b/src/eval/environment.rs
@@ -1,4 +1,4 @@
-use std::collections::{hash_map::Entry, BTreeMap, HashMap};
+use std::collections::{hash_map, BTreeMap, HashMap};
 
 use super::value::EvaluatedValue;
 
@@ -67,6 +67,25 @@ impl Environment {
         self.current_scope_mut().set(name, value)
     }
 
+    pub fn update_value(&mut self, name: &str, value: EvaluatedValue) -> Option<&EvaluatedValue> {
+        let mut cursor = Some(self.current_scope_key());
+        while let Some(scope_key) = cursor {
+            let scope = self.scopes.get_mut(&scope_key).unwrap();
+            cursor = scope.parent;
+            if scope.get(name).is_some() {
+                let scope_ptr: *mut Scope = scope;
+                // https://github.com/rust-lang/rust/issues/54663
+                // Safety: this pointer isn't aliased anywhere, and is only used to set the field
+                unsafe {
+                    let inserted_value = (*scope_ptr).set(name, value);
+                    return Some(inserted_value);
+                }
+            }
+        }
+
+        None
+    }
+
     fn current_scope(&self) -> &Scope {
         // this can't fail, given we know where these scope ids come from
         self.scopes.get(&self.current_scope_key()).unwrap()
@@ -113,8 +132,8 @@ impl Scope {
 
     fn set<S: Into<String>>(&mut self, name: S, value: EvaluatedValue) -> &EvaluatedValue {
         match self.values.entry(name.into()) {
-            Entry::Vacant(entry) => entry.insert(value),
-            Entry::Occupied(entry) => {
+            hash_map::Entry::Vacant(entry) => entry.insert(value),
+            hash_map::Entry::Occupied(entry) => {
                 let item = entry.into_mut();
                 *item = value;
                 item
@@ -189,4 +208,20 @@ mod tests {
 
         assert_eq!(&EvaluatedValue::Number(42_f64), inserted);
     }
+
+    #[test]
+    fn test_updating_item_in_higher_scope_persists_if_scope_is_left() {
+        let mut global = Environment::new();
+
+        global.enter_scope();
+        global.set_value("foo", EvaluatedValue::String("hello".to_string().into()));
+        global.enter_scope();
+        global.update_value("foo", EvaluatedValue::String("world".to_string().into()));
+        global.exit_scope();
+
+        assert_eq!(
+            Some(&EvaluatedValue::String("world".to_string().into())),
+            global.get_value("foo")
+        );
+    }
 }