Merge branch 'remove-continue' into demo

CavefulGames · Sep 4, 2024 · d973bfd · d973bfd
2 parents e8b6ab2 + 7a8cf3d
commit d973bfd
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 29 deletions.
diff --git a/src/rules/remove_continue.rs b/src/rules/remove_continue.rs
@@ -104,37 +104,45 @@ fn continues_with_breaks_to_breaks(block: &mut Block, break_variable_name: &str)
     }
 }
 
+impl Processor {
+    fn process(&self, block: &mut Block) {
+        let (has_continue, has_break) = continue_break_exists(block);
+
+        if has_continue {
+            let (mut stmts, break_variable_handler) = if has_break {
+                let var = TypedIdentifier::new(self.break_variable_name.as_str());
+                let value = Expression::False(None);
+                let local_assign_stmt = LocalAssignStatement::new(vec![var], vec![value]);
+
+                let break_block = Block::new(vec![], Some(LastStatement::new_break()));
+                let break_variable_handler = IfStatement::create(
+                    Identifier::new(self.break_variable_name.as_str()),
+                    break_block,
+                );
+                continues_with_breaks_to_breaks(block, self.break_variable_name.as_str());
+                (vec![local_assign_stmt.into()], Some(break_variable_handler))
+            } else {
+                continues_to_breaks(block);
+                (Vec::new(), None)
+            };
+            let repeat_stmt = RepeatStatement::new(block.clone(), true);
+            stmts.push(repeat_stmt.into());
+            if let Some(break_variable_handler) = break_variable_handler {
+                stmts.push(break_variable_handler.into());
+            }
+            *block = Block::new(stmts, None);
+        }
+    }
+}
+
 impl NodeProcessor for Processor {
     fn process_statement(&mut self, statement: &mut Statement) {
-        if let Statement::NumericFor(numeric_for) = statement {
-            let block = numeric_for.mutate_block();
-
-            let (has_continue, has_break) = continue_break_exists(block);
-
-            if has_continue {
-                let (mut stmts, break_variable_handler) = if has_break {
-                    let var = TypedIdentifier::new(self.break_variable_name.as_str());
-                    let value = Expression::False(None);
-                    let local_assign_stmt = LocalAssignStatement::new(vec![var], vec![value]);
-
-                    let break_block = Block::new(vec![], Some(LastStatement::new_break()));
-                    let break_variable_handler = IfStatement::create(
-                        Identifier::new(self.break_variable_name.as_str()),
-                        break_block,
-                    );
-                    continues_with_breaks_to_breaks(block, self.break_variable_name.as_str());
-                    (vec![local_assign_stmt.into()], Some(break_variable_handler))
-                } else {
-                    continues_to_breaks(block);
-                    (Vec::new(), None)
-                };
-                let repeat_stmt = RepeatStatement::new(block.clone(), true);
-                stmts.push(repeat_stmt.into());
-                if let Some(break_variable_handler) = break_variable_handler {
-                    stmts.push(break_variable_handler.into());
-                }
-                *block = Block::new(stmts, None);
-            }
+        match statement {
+            Statement::NumericFor(numeric_for) => self.process(numeric_for.mutate_block()),
+            Statement::GenericFor(generic_for) => self.process(generic_for.mutate_block()),
+            Statement::Repeat(repeat_stmt) => self.process(repeat_stmt.mutate_block()),
+            Statement::While(while_stmt) => self.process(while_stmt.mutate_block()),
+            _ => (),
         }
     }
 }

diff --git a/tests/rule_tests/remove_continue.rs b/tests/rule_tests/remove_continue.rs
@@ -4,6 +4,9 @@ test_rule!(
     remove_continue,
     RemoveContinue::default(),
     continue_inside_numeric_for("for i = 1, 10 do continue end") => "for i = 1, 10 do repeat break until true end",
+    continue_inside_generic_for("for i,v in {a,b,c} do continue end") => "for i,v in {a,b,c} do repeat break until true end",
+    continue_inside_repeat("repeat continue until true") => "repeat repeat break until true until true",
+    continue_inside_while("while true do continue end") => "while true do repeat break until true end"
 );
 
 #[test]