From 924b388f2c6969007faf63ced8dfbfa59ec3e5bd Mon Sep 17 00:00:00 2001
From: Jeff <dev@jeffa.io>
Date: Mon, 12 Feb 2024 14:19:07 -0500
Subject: [PATCH] Fix bug with loop contexts

---
 src/abstract_tree/block.rs | 8 +++-----
 src/abstract_tree/for.rs   | 2 +-
 2 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/src/abstract_tree/block.rs b/src/abstract_tree/block.rs
index e0fe1e8..ff73a67 100644
--- a/src/abstract_tree/block.rs
+++ b/src/abstract_tree/block.rs
@@ -67,12 +67,10 @@ impl AbstractTree for Block {
     }
 
     fn validate(&self, _source: &str, _context: &Context) -> Result<(), ValidationError> {
+        self.context.inherit_from(_context)?;
+
         for statement in &self.statements {
-            if let Statement::Return(inner_statement) = statement {
-                return inner_statement.validate(_source, &self.context);
-            } else {
-                statement.validate(_source, &self.context)?;
-            }
+            statement.validate(_source, &self.context)?;
         }
 
         Ok(())
diff --git a/src/abstract_tree/for.rs b/src/abstract_tree/for.rs
index a768150..63aacf0 100644
--- a/src/abstract_tree/for.rs
+++ b/src/abstract_tree/for.rs
@@ -72,7 +72,7 @@ impl AbstractTree for For {
         };
         let key = self.item_id.inner().clone();
 
-        context.set_type(key, item_type)?;
+        self.block.context().set_type(key, item_type)?;
         self.block.validate(_source, context)
     }