2 files changed, 33 insertions, 18 deletions

diff --git a/src/lisp/eval.rs b/src/lisp/eval.rs
index 784cba0..c5e4ed0 100644
--- a/src/lisp/eval.rs
+++ b/src/lisp/eval.rs
@@ -153,14 +153,9 @@ where
         match args {
             [LispExpr::Ident(id), expr] => {
                 let value = self.eval(&expr)?;
-                let local_env = self.app.lisp_env.last_mut();
-                if let Some(env) = local_env {
-                    env.insert(id.into(), value)
-                        .ok_or_else(|| EvalError::UnboundVariable(id.into()).into())
-                } else {
-                    error!("Unable to set in global env!");
-                    Err(EvalError::BadForm.into())
-                }
+                let loc = lookup_mut(&mut self.app.lisp_env, id)?;
+                *loc = value;
+                Ok(LispExpr::Unit)
             }
             _ => {
                 error!("Invalid usage of `set!`");
@@ -334,16 +329,22 @@ pub fn create_lambda(cdr: &[LispExpr]) -> Result<LispExpr, LispError> {
 }
 
 pub fn lookup(env_list: &[Environment], key: &str) -> Result<LispExpr, LispError> {
-    if env_list.is_empty() {
-        Err(EvalError::UnboundVariable(key.into()).into())
-    } else {
-        let local_env = env_list.last().unwrap();
-        if let Some(val) = local_env.get(key) {
-            Ok(val.clone())
-        } else {
-            lookup(&env_list[..env_list.len() - 1], key)
-        }
-    }
+    Ok(env_list
+        .iter()
+        .rev()
+        .find_map(|local_env| local_env.get(key).cloned())
+        .ok_or(EvalError::UnboundVariable(key.into()))?)
+}
+
+pub fn lookup_mut<'a, 'b>(
+    env_list: &'a mut [Environment],
+    key: &'b str,
+) -> Result<&'a mut LispExpr, LispError> {
+    Ok(env_list
+        .iter_mut()
+        .rev()
+        .find_map(|local_env| local_env.get_mut(key))
+        .ok_or(EvalError::UnboundVariable(key.into()))?)
 }
 
 pub fn completions<'a>(env_list: &'a [Environment], prefix: &str) -> Vec<&'a str> {
diff --git a/src/lisp/test.lisp b/src/lisp/test.lisp
index a167f28..5eb1aaa 100644
--- a/src/lisp/test.lisp
+++ b/src/lisp/test.lisp
@@ -17,3 +17,17 @@
 (assert-eq
   ((lambda (x) x) 2)
   2)
+
+(define x 2)
+(define y 5)
+(set! x 3)
+(assert-eq x 3)
+
+(define (square-x x)
+  (let [(y x)]
+    (begin
+      (set! y (* x x))
+      y)))
+
+(assert-eq (square-x 2) 4)
+(assert-eq y 5)