From 9f24797bc8f16ef98e418549f1dd6782c80a4840 Mon Sep 17 00:00:00 2001
From: Tom Schuster <evilpies@gmail.com>
Date: Sun, 11 Oct 2020 19:43:19 +0200
Subject: [PATCH] basic loop support

---
 jit/src/instructions.rs | 33 ++++++++++++++++++---------------
 jit/tests/misc_tests.rs | 17 +++++++++++++++++
 2 files changed, 35 insertions(+), 15 deletions(-)

diff --git a/jit/src/instructions.rs b/jit/src/instructions.rs
index 9f95889c0..c31e3d17a 100644
--- a/jit/src/instructions.rs
+++ b/jit/src/instructions.rs
@@ -97,25 +97,29 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
         }
     }
 
+    fn get_or_create_block(&mut self, label: &Label) -> Block {
+        let builder = &mut self.builder;
+        *self
+            .label_to_block
+            .entry(*label)
+            .or_insert_with(|| builder.create_block())
+    }
+
     pub fn compile(&mut self, bytecode: &CodeObject) -> Result<(), JitCompileError> {
         let offset_to_label: HashMap<&usize, &Label> =
             bytecode.label_map.iter().map(|(k, v)| (v, k)).collect();
 
         for (offset, instruction) in bytecode.instructions.iter().enumerate() {
             if let Some(&label) = offset_to_label.get(&offset) {
-                let builder = &mut self.builder;
-                let block = self
-                    .label_to_block
-                    .entry(*label)
-                    .or_insert_with(|| builder.create_block());
+                let block = self.get_or_create_block(label);
 
                 // If the current block is not terminated/filled just jump
                 // into the new block.
                 if !self.builder.is_filled() {
-                    self.builder.ins().jump(*block, &[]);
+                    self.builder.ins().jump(block, &[]);
                 }
 
-                self.builder.switch_to_block(*block);
+                self.builder.switch_to_block(block);
             }
 
             // Sometimes the bytecode contains instructions after a return
@@ -168,10 +172,8 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
             Instruction::JumpIfFalse { target } => {
                 let cond = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
 
-                let then_block = self.builder.create_block();
-                self.label_to_block.insert(*target, then_block);
-
                 let val = self.boolean_val(cond)?;
+                let then_block = self.get_or_create_block(target);
                 self.builder.ins().brz(val, then_block, &[]);
 
                 let block = self.builder.create_block();
@@ -183,10 +185,8 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
             Instruction::JumpIfTrue { target } => {
                 let cond = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
 
-                let then_block = self.builder.create_block();
-                self.label_to_block.insert(*target, then_block);
-
                 let val = self.boolean_val(cond)?;
+                let then_block = self.get_or_create_block(target);
                 self.builder.ins().brnz(val, then_block, &[]);
 
                 let block = self.builder.create_block();
@@ -196,8 +196,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                 Ok(())
             }
             Instruction::Jump { target } => {
-                let target_block = self.builder.create_block();
-                self.label_to_block.insert(*target, target_block);
+                let target_block = self.get_or_create_block(target);
                 self.builder.ins().jump(target_block, &[]);
 
                 Ok(())
@@ -379,6 +378,10 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                     _ => Err(JitCompileError::NotSupported),
                 }
             }
+            Instruction::SetupLoop { .. } | Instruction::PopBlock => {
+                // TODO: block support
+                Ok(())
+            }
             _ => Err(JitCompileError::NotSupported),
         }
     }
diff --git a/jit/tests/misc_tests.rs b/jit/tests/misc_tests.rs
index 45674c12c..2d944c6dc 100644
--- a/jit/tests/misc_tests.rs
+++ b/jit/tests/misc_tests.rs
@@ -84,3 +84,20 @@ fn test_if_else() {
     assert_eq!(if_else(-1), Ok(42));
     assert_eq!(if_else(100), Ok(42));
 }
+
+#[test]
+fn test_while_loop() {
+    let while_loop = jit_function! { while_loop(a:i64) -> i64 => r##"
+        def while_loop(a: int):
+            b = 0
+            while a > 0:
+                b += 1
+                a -= 1
+            return b
+    "## };
+
+    assert_eq!(while_loop(0), Ok(0));
+    assert_eq!(while_loop(-1), Ok(0));
+    assert_eq!(while_loop(1), Ok(1));
+    assert_eq!(while_loop(10), Ok(10));
+}