From 07d4fe416a9fe4d9b9f83878b04f97375ac71df0 Mon Sep 17 00:00:00 2001
From: Jakub Doka <jakub.doka2@gmail.com>
Date: Mon, 16 Dec 2024 13:27:15 +0100
Subject: [PATCH] forgot to add defer handling to unrolled loops

Signed-off-by: Jakub Doka <jakub.doka2@gmail.com>
---
 lang/src/son.rs | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/lang/src/son.rs b/lang/src/son.rs
index 8960c3300..09b4c6771 100644
--- a/lang/src/son.rs
+++ b/lang/src/son.rs
@@ -2328,6 +2328,7 @@ pub enum CtLoopState {
 enum Loop {
     Comptime {
         state: Option<(CtLoopState, Pos)>,
+        defer_base: usize,
     },
     Runtime {
         node: Nid,
@@ -3951,7 +3952,9 @@ impl<'a> Codegen<'a> {
             Expr::Loop { unrolled: true, body, pos } => {
                 let mut loop_fuel = 100;
 
-                self.ci.loops.push(Loop::Comptime { state: None });
+                self.ci
+                    .loops
+                    .push(Loop::Comptime { state: None, defer_base: self.ci.defers.len() });
 
                 loop {
                     if loop_fuel == 0 {
@@ -3964,7 +3967,7 @@ impl<'a> Codegen<'a> {
 
                     let terminated = self.expr(body).is_none();
 
-                    let Some(&Loop::Comptime { state }) = self.ci.loops.last() else {
+                    let Some(&Loop::Comptime { state, .. }) = self.ci.loops.last() else {
                         unreachable!()
                     };
 
@@ -5119,8 +5122,8 @@ impl<'a> Codegen<'a> {
             return None;
         };
 
-        match loob {
-            &mut Loop::Comptime { state: Some((_, prev)) } => {
+        match *loob {
+            Loop::Comptime { state: Some((_, prev)), .. } => {
                 self.error(
                     pos,
                     "reached multiple control flow keywords inside an unrolled loop, \
@@ -5128,10 +5131,11 @@ impl<'a> Codegen<'a> {
                 );
                 self.error(prev, "previous reachable control flow found here");
             }
-            Loop::Comptime { state: state @ None } => {
-                *state = Some(([CtLoopState::Continued, CtLoopState::Terminated][id], pos))
+            Loop::Comptime { state: ref mut state @ None, defer_base } => {
+                *state = Some(([CtLoopState::Continued, CtLoopState::Terminated][id], pos));
+                self.gen_defers(defer_base)?;
             }
-            &mut Loop::Runtime { defer_base, .. } => {
+            Loop::Runtime { defer_base, .. } => {
                 self.gen_defers(defer_base)?;
 
                 let Loop::Runtime { ctrl: lctrl, ctrl_scope, scope, .. } =