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)
     }