Implement loops into patio-to-cranelift

4 files changed, 117 insertions, 5 deletions

diff --git a/patio-to-cranelift/src/codegen/context.rs b/patio-to-cranelift/src/codegen/context.rs
index 33e18e07..a9333e26 100644
--- a/patio-to-cranelift/src/codegen/context.rs
+++ b/patio-to-cranelift/src/codegen/context.rs
@@ -102,6 +102,7 @@ impl CompilerContext {
             local_name_to_value: Default::default(),
             local_name_to_type: Default::default(),
             var_name_to_variable: Default::default(),
+            loop_blocks: Default::default(),
         };
         generate_function_body(&mut context);
 
@@ -159,6 +160,7 @@ pub(crate) struct CompilerFunctionContext<'a> {
     local_name_to_value: HashMap<String, cranelift_codegen::ir::Value>,
     local_name_to_type: HashMap<String, cranelift_codegen::ir::Type>,
     var_name_to_variable: HashMap<String, cranelift::prelude::Variable>,
+    loop_blocks: Vec<Loop>,
 }
 
 impl<'a> CompilerFunctionContext<'a> {
@@ -224,6 +226,25 @@ impl<'a> CompilerFunctionContext<'a> {
     pub(crate) fn current_block(&self) -> cranelift_codegen::ir::Block {
         self.function_builder.current_block().unwrap()
     }
+
+    pub(crate) fn enter_loop(
+        &mut self,
+        start_block: cranelift_codegen::ir::Block,
+        end_block: cranelift_codegen::ir::Block,
+    ) {
+        self.loop_blocks.push(Loop {
+            start_block,
+            end_block,
+        });
+    }
+
+    pub(crate) fn current_loop(&self) -> &Loop {
+        self.loop_blocks.last().unwrap()
+    }
+
+    pub(crate) fn exit_loop(&mut self) {
+        self.loop_blocks.pop();
+    }
 }
 
 pub(crate) enum Mutability {
@@ -239,3 +260,8 @@ impl Mutability {
         }
     }
 }
+
+pub(crate) struct Loop {
+    pub(crate) start_block: cranelift_codegen::ir::Block,
+    pub(crate) end_block: cranelift_codegen::ir::Block,
+}
diff --git a/patio-to-cranelift/src/codegen/mod.rs b/patio-to-cranelift/src/codegen/mod.rs
index 316c9e2a..438a0893 100644
--- a/patio-to-cranelift/src/codegen/mod.rs
+++ b/patio-to-cranelift/src/codegen/mod.rs
@@ -18,8 +18,9 @@ use crate::{
         instructions::{
             PatioBlock, PatioBoolToInt, PatioCall, PatioCallPointer, PatioDataPtrToInt, PatioIf,
             PatioInstruction, PatioIntAdd, PatioIntEquals, PatioIntNotEqual, PatioIntSub,
-            PatioIntToBool, PatioIntToDataPtr, PatioLocalStore, PatioMemoryLoad, PatioMemoryStore,
-            PatioResume, PatioReturn, PatioReturnVoid, PatioSwitch, PatioVarLoad, PatioVarStore,
+            PatioIntToBool, PatioIntToDataPtr, PatioLocalStore, PatioLoop, PatioLoopBreak,
+            PatioLoopContinue, PatioMemoryLoad, PatioMemoryStore, PatioResume, PatioReturn,
+            PatioReturnVoid, PatioSwitch, PatioVarLoad, PatioVarStore,
         },
         operands::{PatioLocalRef, PatioOperand},
         types::PatioType,
@@ -335,6 +336,18 @@ fn compile_instruction(instruction: &PatioInstruction, context: &mut CompilerFun
             compile_local_store(local_store, context);
         }
 
+        PatioInstruction::Loop(loop_) => {
+            compile_loop(loop_, context);
+        }
+
+        PatioInstruction::LoopBreak(loop_break) => {
+            compile_loop_break(loop_break, context);
+        }
+
+        PatioInstruction::LoopContinue(loop_continue) => {
+            compile_loop_continue(loop_continue, context);
+        }
+
         PatioInstruction::MemoryLoad(memory_load) => {
             compile_memory_load(memory_load, context);
         }
@@ -574,6 +587,52 @@ fn compile_local_store(local_store: &PatioLocalStore, context: &mut CompilerFunc
     def_destination(&local_store.destination, cranelift_value, context);
 }
 
+fn compile_loop(loop_: &PatioLoop, context: &mut CompilerFunctionContext<'_>) {
+    let cranelift_loop_block = context.create_block();
+    let cranelift_end_block = context.create_block();
+
+    context
+        .function_builder
+        .ins()
+        .jump(cranelift_loop_block, []);
+
+    context.switch_to_block(cranelift_loop_block);
+    context.enter_loop(cranelift_loop_block, cranelift_end_block);
+    for instruction in &loop_.body {
+        compile_instruction(instruction, context);
+    }
+    context.exit_loop();
+    context.seal_block(cranelift_loop_block);
+    context.function_builder.ins().jump(cranelift_end_block, []);
+
+    context.seal_block(cranelift_end_block);
+    context.switch_to_block(cranelift_end_block);
+}
+
+fn compile_loop_break(_: &PatioLoopBreak, context: &mut CompilerFunctionContext<'_>) {
+    let cranelift_loop_end_block = context.current_loop().end_block;
+    context
+        .function_builder
+        .ins()
+        .jump(cranelift_loop_end_block, []);
+
+    // TODO: avoid creating orphaned blocks
+    let cranelift_unreachable_block = context.create_block();
+    context.switch_to_block(cranelift_unreachable_block);
+}
+
+fn compile_loop_continue(_: &PatioLoopContinue, context: &mut CompilerFunctionContext<'_>) {
+    let cranelift_loop_start_block = context.current_loop().start_block;
+    context
+        .function_builder
+        .ins()
+        .jump(cranelift_loop_start_block, []);
+
+    // TODO: avoid creating orphaned blocks
+    let cranelift_unreachable_block = context.create_block();
+    context.switch_to_block(cranelift_unreachable_block);
+}
+
 fn compile_memory_load(memory_load: &PatioMemoryLoad, context: &mut CompilerFunctionContext<'_>) {
     let value_layout =
         types::data_layout_element_at(&memory_load.layout, &memory_load.element_offsets);
diff --git a/patio-to-cranelift/src/patio/flatbuffers.rs b/patio-to-cranelift/src/patio/flatbuffers.rs
index 900a2c4c..34e0a1a1 100644
--- a/patio-to-cranelift/src/patio/flatbuffers.rs
+++ b/patio-to-cranelift/src/patio/flatbuffers.rs
@@ -9,9 +9,9 @@ use crate::{
         instructions::{
             PatioBlock, PatioBoolToInt, PatioCall, PatioCallPointer, PatioDataPtrToInt, PatioIf,
             PatioInstruction, PatioIntAdd, PatioIntEquals, PatioIntNotEqual, PatioIntSub,
-            PatioIntToBool, PatioIntToDataPtr, PatioLocalStore, PatioMemoryLoad, PatioMemoryStore,
-            PatioResume, PatioReturn, PatioReturnVoid, PatioSwitch, PatioSwitchCase, PatioVarLoad,
-            PatioVarStore,
+            PatioIntToBool, PatioIntToDataPtr, PatioLocalStore, PatioLoop, PatioLoopBreak,
+            PatioLoopContinue, PatioMemoryLoad, PatioMemoryStore, PatioResume, PatioReturn,
+            PatioReturnVoid, PatioSwitch, PatioSwitchCase, PatioVarLoad, PatioVarStore,
         },
         operands::{PatioLocalRef, PatioOperand},
         types::{PatioDataLayout, PatioType},
@@ -272,6 +272,22 @@ impl From<patio_generated::patio::PatioInstruction<'_>> for PatioInstruction {
                 })
             }
 
+            patio_generated::patio::PatioInstructionValue::PatioLoop => {
+                let loop_ = value.value_as_patio_loop().unwrap();
+
+                Self::Loop(PatioLoop {
+                    body: loop_.body().iter().map(Into::into).collect(),
+                })
+            }
+
+            patio_generated::patio::PatioInstructionValue::PatioLoopBreak => {
+                Self::LoopBreak(PatioLoopBreak)
+            }
+
+            patio_generated::patio::PatioInstructionValue::PatioLoopContinue => {
+                Self::LoopContinue(PatioLoopContinue)
+            }
+
             patio_generated::patio::PatioInstructionValue::PatioMemoryLoad => {
                 let memory_load = value.value_as_patio_memory_load().unwrap();
 
diff --git a/patio-to-cranelift/src/patio/instructions.rs b/patio-to-cranelift/src/patio/instructions.rs
index fa8d098b..e9a3824e 100644
--- a/patio-to-cranelift/src/patio/instructions.rs
+++ b/patio-to-cranelift/src/patio/instructions.rs
@@ -17,6 +17,9 @@ pub(crate) enum PatioInstruction {
     IntToBool(PatioIntToBool),
     IntToDataPtr(PatioIntToDataPtr),
     LocalStore(PatioLocalStore),
+    Loop(PatioLoop),
+    LoopBreak(PatioLoopBreak),
+    LoopContinue(PatioLoopContinue),
     MemoryLoad(PatioMemoryLoad),
     MemoryStore(PatioMemoryStore),
     Resume(PatioResume),
@@ -100,6 +103,14 @@ pub(crate) struct PatioLocalStore {
     pub(crate) value: PatioOperand,
 }
 
+pub(crate) struct PatioLoop {
+    pub(crate) body: Vec<PatioInstruction>,
+}
+
+pub(crate) struct PatioLoopBreak;
+
+pub(crate) struct PatioLoopContinue;
+
 pub(crate) struct PatioMemoryLoad {
     pub(crate) destination: PatioLocalRef,
     pub(crate) pointer: PatioOperand,