From c130627b63e8563c7064aadc13669fca27b305b5 Mon Sep 17 00:00:00 2001
From: Nilstrieb <48135649+Nilstrieb@users.noreply.github.com>
Date: Fri, 15 Apr 2022 21:04:17 +0200
Subject: [PATCH] fix pass_group to use recursively optimize loop if loop is
 the first statement

---
 rust2/src/lib.rs  |  3 +++
 rust2/src/opts.rs | 12 ++++++++++--
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/rust2/src/lib.rs b/rust2/src/lib.rs
index fee2c28..c60ac54 100644
--- a/rust2/src/lib.rs
+++ b/rust2/src/lib.rs
@@ -40,6 +40,9 @@ where
 
     let code = codegen::generate(&cg_alloc, &optimized_ir);
 
+    println!("{:#?}", optimized_ir);
+    return Ok(());
+
     drop(optimized_ir);
     drop(ir_alloc);
 
diff --git a/rust2/src/opts.rs b/rust2/src/opts.rs
index 710801c..51dc2d4 100644
--- a/rust2/src/opts.rs
+++ b/rust2/src/opts.rs
@@ -53,7 +53,7 @@ fn ast_to_ir<'ir>(alloc: &'ir Bump, ast: &[(Instr<'_>, Span)]) -> Ir<'ir> {
             Instr::Out => StmtKind::Out,
             Instr::In => StmtKind::In,
             Instr::Loop(body) => {
-                let ir_body = ast_to_ir(alloc, &body);
+                let ir_body = ast_to_ir(alloc, body);
                 StmtKind::Loop(ir_body)
             }
         };
@@ -72,7 +72,15 @@ fn pass_group<'ir>(alloc: &'ir Bump, ir: Ir<'ir>) -> Ir<'ir> {
         .into_iter()
         .fold(new_stmts, |mut stmts: BumpVec<'ir, Stmt<'ir>>, next| {
             let Some(old) = stmts.last_mut() else {
-                stmts.push(next);
+                if let StmtKind::Loop(body) = next.kind {
+                    let new_body = pass_group(alloc, body);
+                    stmts.push(Stmt {
+                        span: next.span,
+                        kind: StmtKind::Loop(new_body),
+                    });
+                } else {
+                    stmts.push(next);
+                }
                 return stmts;
             };