Implement switch in patio-to-cranelift

4 files changed, 128 insertions, 17 deletions

diff --git a/patio-to-cranelift/src/codegen/context.rs b/patio-to-cranelift/src/codegen/context.rs
index 09b08173..213e1a6f 100644
--- a/patio-to-cranelift/src/codegen/context.rs
+++ b/patio-to-cranelift/src/codegen/context.rs
@@ -221,6 +221,26 @@ impl<'a> CompilerFunctionContext<'a> {
     pub(crate) fn pointer_type(&self) -> cranelift_codegen::ir::Type {
         cranelift_codegen::ir::Type::triple_pointer_type(&self.target_triple)
     }
+
+    pub(crate) fn create_block(&mut self) -> cranelift_codegen::ir::Block {
+        self.function_builder.create_block()
+    }
+
+    pub(crate) fn switch_to_block(&mut self, block: cranelift_codegen::ir::Block) {
+        self.function_builder.switch_to_block(block);
+    }
+
+    pub(crate) fn seal_current_block(&mut self) {
+        self.function_builder
+            .seal_block(self.function_builder.current_block().unwrap());
+    }
+
+    pub(crate) fn block_call(
+        &mut self,
+        block: cranelift_codegen::ir::Block,
+    ) -> cranelift_codegen::ir::BlockCall {
+        self.function_builder.func.dfg.block_call(block, [])
+    }
 }
 
 pub(crate) enum Mutability {
diff --git a/patio-to-cranelift/src/codegen/mod.rs b/patio-to-cranelift/src/codegen/mod.rs
index 20c20874..83422307 100644
--- a/patio-to-cranelift/src/codegen/mod.rs
+++ b/patio-to-cranelift/src/codegen/mod.rs
@@ -18,7 +18,7 @@ use crate::{
             PatioBlock, PatioBoolToInt, PatioCall, PatioCallPointer, PatioDataPtrToInt, PatioIf,
             PatioInstruction, PatioIntAdd, PatioIntEquals, PatioIntNotEqual, PatioIntSub,
             PatioIntToBool, PatioIntToDataPtr, PatioLocalStore, PatioMemoryLoad, PatioMemoryStore,
-            PatioResume, PatioReturn, PatioReturnVoid, PatioVarLoad, PatioVarStore,
+            PatioResume, PatioReturn, PatioReturnVoid, PatioSwitch, PatioVarLoad, PatioVarStore,
         },
         operands::{PatioLocalRef, PatioOperand},
         types::PatioType,
@@ -298,6 +298,10 @@ fn compile_instruction(instruction: &PatioInstruction, context: &mut CompilerFun
             compile_return_void(return_void, context);
         }
 
+        PatioInstruction::Switch(switch) => {
+            compile_switch(switch, context);
+        }
+
         PatioInstruction::VarLoad(var_load) => {
             compile_var_load(var_load, context);
         }
@@ -394,9 +398,9 @@ fn compile_data_ptr_to_int(
 
 fn compile_if(if_: &PatioIf, context: &mut CompilerFunctionContext<'_>) {
     let cranelift_condition = compile_operand(&if_.condition, context);
-    let cranelift_then_block = context.function_builder.create_block();
-    let cranelift_else_block = context.function_builder.create_block();
-    let cranelift_end_block = context.function_builder.create_block();
+    let cranelift_then_block = context.create_block();
+    let cranelift_else_block = context.create_block();
+    let cranelift_end_block = context.create_block();
     let cranelift_end_params: Vec<_> = if_
         .destinations
         .iter()
@@ -415,13 +419,9 @@ fn compile_if(if_: &PatioIf, context: &mut CompilerFunctionContext<'_>) {
         cranelift_else_block,
         &[],
     );
-    context
-        .function_builder
-        .seal_block(context.function_builder.current_block().unwrap());
+    context.seal_current_block();
 
-    context
-        .function_builder
-        .switch_to_block(cranelift_then_block);
+    context.switch_to_block(cranelift_then_block);
     let cranelift_then_results = compile_block(&if_.true_body, context);
     context.function_builder.ins().jump(
         cranelift_end_block,
@@ -429,9 +429,7 @@ fn compile_if(if_: &PatioIf, context: &mut CompilerFunctionContext<'_>) {
     );
     context.function_builder.seal_block(cranelift_then_block);
 
-    context
-        .function_builder
-        .switch_to_block(cranelift_else_block);
+    context.switch_to_block(cranelift_else_block);
     let cranelift_else_results = compile_block(&if_.false_body, context);
     context.function_builder.ins().jump(
         cranelift_end_block,
@@ -439,9 +437,7 @@ fn compile_if(if_: &PatioIf, context: &mut CompilerFunctionContext<'_>) {
     );
     context.function_builder.seal_block(cranelift_else_block);
 
-    context
-        .function_builder
-        .switch_to_block(cranelift_end_block);
+    context.switch_to_block(cranelift_end_block);
 
     for (destination, cranelift_end_param) in if_.destinations.iter().zip(cranelift_end_params) {
         def_destination(destination, cranelift_end_param, context);
@@ -585,6 +581,67 @@ fn compile_return_void(_return_void: &PatioReturnVoid, context: &mut CompilerFun
     context.function_builder.ins().return_(&[]);
 }
 
+fn compile_switch(switch: &PatioSwitch, context: &mut CompilerFunctionContext<'_>) {
+    let cranelift_value = compile_operand(&switch.value, context);
+
+    // Create the various blocks.
+    let default_block = context.create_block();
+    let case_blocks: Vec<_> = switch
+        .cases
+        .iter()
+        .map(|_| context.create_block())
+        .collect();
+    let end_block = context.create_block();
+    let cranelift_end_params: Vec<_> = switch
+        .destinations
+        .iter()
+        .map(|destination| {
+            let cranelift_type = context.var_type(&destination.name);
+            context
+                .function_builder
+                .append_block_param(end_block, cranelift_type)
+        })
+        .collect();
+
+    // Create the jump itself.
+    let mut cranelift_switch = cranelift::frontend::Switch::new();
+    for (case, case_block) in switch.cases.iter().zip(&case_blocks) {
+        cranelift_switch.set_entry(case.value.into(), *case_block);
+    }
+    context.seal_current_block();
+    cranelift_switch.emit(
+        &mut context.function_builder,
+        cranelift_value,
+        default_block,
+    );
+
+    // Compile the case and default blocks.
+    for (case, case_block) in switch.cases.iter().zip(&case_blocks) {
+        context.switch_to_block(*case_block);
+        let end_args = compile_block(&case.body, context);
+        context
+            .function_builder
+            .ins()
+            .jump(end_block, &cranelift_block_args(end_args));
+        context.seal_current_block();
+    }
+
+    {
+        context.switch_to_block(default_block);
+        let end_args = compile_block(&switch.default_body, context);
+        context
+            .function_builder
+            .ins()
+            .jump(end_block, &cranelift_block_args(end_args));
+        context.seal_current_block();
+    }
+
+    context.switch_to_block(end_block);
+    for (destination, cranelift_end_param) in switch.destinations.iter().zip(cranelift_end_params) {
+        def_destination(destination, cranelift_end_param, context);
+    }
+}
+
 fn compile_var_load(var_load: &PatioVarLoad, context: &mut CompilerFunctionContext<'_>) {
     let stack_slot = context.stack_slot(&var_load.var.name);
     let type_ = context.stack_slot_type(&var_load.var.name);
diff --git a/patio-to-cranelift/src/patio/flatbuffers.rs b/patio-to-cranelift/src/patio/flatbuffers.rs
index ad667963..69720b19 100644
--- a/patio-to-cranelift/src/patio/flatbuffers.rs
+++ b/patio-to-cranelift/src/patio/flatbuffers.rs
@@ -9,7 +9,8 @@ use crate::{
             PatioBlock, PatioBoolToInt, PatioCall, PatioCallPointer, PatioDataPtrToInt, PatioIf,
             PatioInstruction, PatioIntAdd, PatioIntEquals, PatioIntNotEqual, PatioIntSub,
             PatioIntToBool, PatioIntToDataPtr, PatioLocalStore, PatioMemoryLoad, PatioMemoryStore,
-            PatioResume, PatioReturn, PatioReturnVoid, PatioVarLoad, PatioVarStore,
+            PatioResume, PatioReturn, PatioReturnVoid, PatioSwitch, PatioSwitchCase, PatioVarLoad,
+            PatioVarStore,
         },
         operands::{PatioLocalRef, PatioOperand},
         types::{PatioDataLayout, PatioType},
@@ -303,6 +304,17 @@ impl From<patio_generated::patio::PatioInstruction<'_>> for PatioInstruction {
                 Self::ReturnVoid(PatioReturnVoid)
             }
 
+            patio_generated::patio::PatioInstructionValue::PatioSwitch => {
+                let switch = value.value_as_patio_switch().unwrap();
+
+                Self::Switch(PatioSwitch {
+                    destinations: switch.destinations().iter().map(Into::into).collect(),
+                    value: switch.value().into(),
+                    cases: switch.cases().iter().map(Into::into).collect(),
+                    default_body: switch.default_body().into(),
+                })
+            }
+
             patio_generated::patio::PatioInstructionValue::PatioVarLoad => {
                 let var_load = value.value_as_patio_var_load().unwrap();
 
@@ -339,6 +351,15 @@ impl From<patio_generated::patio::PatioBlock<'_>> for PatioBlock {
     }
 }
 
+impl From<patio_generated::patio::PatioSwitchCase<'_>> for PatioSwitchCase {
+    fn from(value: patio_generated::patio::PatioSwitchCase<'_>) -> Self {
+        Self {
+            value: value.value(),
+            body: value.body().into(),
+        }
+    }
+}
+
 // == Operands ==
 
 impl From<patio_generated::patio::PatioOperand<'_>> for PatioOperand {
diff --git a/patio-to-cranelift/src/patio/instructions.rs b/patio-to-cranelift/src/patio/instructions.rs
index 908a0779..fa8d098b 100644
--- a/patio-to-cranelift/src/patio/instructions.rs
+++ b/patio-to-cranelift/src/patio/instructions.rs
@@ -22,6 +22,7 @@ pub(crate) enum PatioInstruction {
     Resume(PatioResume),
     Return(PatioReturn),
     ReturnVoid(PatioReturnVoid),
+    Switch(PatioSwitch),
     VarLoad(PatioVarLoad),
     VarStore(PatioVarStore),
 }
@@ -123,6 +124,18 @@ pub(crate) struct PatioReturn {
 
 pub(crate) struct PatioReturnVoid;
 
+pub(crate) struct PatioSwitch {
+    pub(crate) destinations: Vec<PatioLocalRef>,
+    pub(crate) value: PatioOperand,
+    pub(crate) cases: Vec<PatioSwitchCase>,
+    pub(crate) default_body: PatioBlock,
+}
+
+pub(crate) struct PatioSwitchCase {
+    pub(crate) value: u32,
+    pub(crate) body: PatioBlock,
+}
+
 pub(crate) struct PatioVarLoad {
     pub(crate) destination: PatioLocalRef,
     pub(crate) var: PatioLocalRef,