From 2dd482c9c2b39d17262d5e47e5db712df8c3f4ca Mon Sep 17 00:00:00 2001
From: sofia <sofia@teascade.net>
Date: Sat, 16 Aug 2025 16:25:08 +0300
Subject: [PATCH] Fix for-loop generation

---
 reid/src/ast/mod.rs     |  2 +-
 reid/src/ast/process.rs | 31 ++++++++++++++++++++++++++++---
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/reid/src/ast/mod.rs b/reid/src/ast/mod.rs
index 817ea41..4f265da 100644
--- a/reid/src/ast/mod.rs
+++ b/reid/src/ast/mod.rs
@@ -208,7 +208,7 @@ pub enum SelfKind {
     None,
 }
 
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, PartialEq)]
 pub enum ReturnType {
     Soft,
     Hard,
diff --git a/reid/src/ast/process.rs b/reid/src/ast/process.rs
index 0cdba3c..cc73355 100644
--- a/reid/src/ast/process.rs
+++ b/reid/src/ast/process.rs
@@ -1,7 +1,7 @@
 use std::path::PathBuf;
 
 use crate::{
-    ast::{self},
+    ast::{self, ReturnType},
     mir::{
         self, CustomTypeKey, FunctionParam, ModuleMap, NamedVariableRef, ReturnKind, SourceModuleId, StmtKind,
         StructField, StructType, WhileStatement,
@@ -265,8 +265,33 @@ impl ast::Block {
                         ),
                         counter_range.as_meta(module_id),
                     );
-                    let mut mir_block = block.into_mir(module_id);
-                    mir_block.statements.push(set_new);
+
+                    let mir_block = if let Some((ret_kind, ret_expr)) = &block.1 {
+                        if *ret_kind == ReturnType::Soft {
+                            if let Some(ret_expr) = ret_expr {
+                                let mir_ret = ret_expr.process(module_id);
+                                let mut clone = block.clone();
+                                clone.1 = None;
+                                let mut mir_block = clone.into_mir(module_id);
+                                mir_block
+                                    .statements
+                                    .push(mir::Statement(StmtKind::Expression(mir_ret.clone()), mir_ret.1));
+                                mir_block.statements.push(set_new);
+                                mir_block
+                            } else {
+                                let mut mir_block = block.into_mir(module_id);
+                                mir_block.statements.push(set_new);
+                                mir_block
+                            }
+                        } else {
+                            block.into_mir(module_id)
+                        }
+                    } else {
+                        let mut mir_block = block.into_mir(module_id);
+                        mir_block.statements.push(set_new);
+                        mir_block
+                    };
+
                     let while_statement = mir::Statement(
                         StmtKind::While(WhileStatement {
                             condition: mir::Expression(