Implement memory load/store in patio-to-cranelift

7 files changed, 387 insertions, 15 deletions

diff --git a/patio-to-cranelift/src/context.rs b/patio-to-cranelift/src/context.rs
index ecdecdc9..a0785a77 100644
--- a/patio-to-cranelift/src/context.rs
+++ b/patio-to-cranelift/src/context.rs
@@ -112,6 +112,43 @@ impl CompilerContext {
             .define_function(func_id, &mut self.codegen_context)
             .unwrap();
     }
+
+    pub(crate) fn declare_data(
+        &mut self,
+        name: &str,
+        linkage: cranelift_module::Linkage,
+        mutability: Mutability,
+    ) {
+        self.module
+            .declare_data(name, linkage, mutability.is_mutable(), false)
+            .unwrap();
+    }
+
+    pub(crate) fn define_data(&mut self, name: &str, bytes: Box<[u8]>) {
+        // TODO: deduplicate finding data_id from name
+        let data_id = match self.module.declarations().get_name(name) {
+            Some(cranelift_module::FuncOrDataId::Data(data_id)) => data_id,
+            Some(cranelift_module::FuncOrDataId::Func(_)) | None => {
+                todo!("no data with name: {name}")
+            }
+        };
+        self.module
+            .define_data(
+                data_id,
+                &cranelift_module::DataDescription {
+                    // TODO: need to handle relocations?
+                    init: cranelift_module::Init::Bytes { contents: bytes },
+                    function_decls: Default::default(),
+                    data_decls: Default::default(),
+                    function_relocs: Default::default(),
+                    data_relocs: Default::default(),
+                    custom_segment_section: None,
+                    align: None,
+                    used: false,
+                },
+            )
+            .unwrap();
+    }
 }
 
 pub(crate) struct CompilerFunctionContext<'a> {
@@ -169,4 +206,29 @@ impl<'a> CompilerFunctionContext<'a> {
     pub(crate) fn stack_slot_type(&mut self, name: &str) -> cranelift_codegen::ir::Type {
         self.stack_slot_name_to_type[name]
     }
+
+    pub(crate) fn data_ptr(&mut self, name: &str) -> cranelift_codegen::ir::GlobalValue {
+        let data_id = match self.module.declarations().get_name(name) {
+            Some(cranelift_module::FuncOrDataId::Data(data_id)) => data_id,
+            Some(cranelift_module::FuncOrDataId::Func(_)) | None => {
+                todo!("no data with name: {name}")
+            }
+        };
+        self.module
+            .declare_data_in_func(data_id, &mut self.function_builder.func)
+    }
+}
+
+pub(crate) enum Mutability {
+    Immutable,
+    Mutable,
+}
+
+impl Mutability {
+    fn is_mutable(&self) -> bool {
+        match self {
+            Mutability::Immutable => false,
+            Mutability::Mutable => true,
+        }
+    }
 }
diff --git a/patio-to-cranelift/src/lib.rs b/patio-to-cranelift/src/lib.rs
index 687687f0..36b4b13f 100644
--- a/patio-to-cranelift/src/lib.rs
+++ b/patio-to-cranelift/src/lib.rs
@@ -1,9 +1,9 @@
-use std::{path::Path, str::FromStr};
+use std::{io::Write, path::Path, str::FromStr};
 
 use cranelift::prelude::InstBuilder;
 
 use crate::{
-    context::{CompilerContext, CompilerFunctionContext},
+    context::{CompilerContext, CompilerFunctionContext, Mutability},
     errors::CompileError,
     patio::{
         compilation_unit::{
@@ -14,11 +14,12 @@ use crate::{
         instructions::{
             PatioBlock, PatioBoolToInt, PatioCall, PatioIf, PatioInstruction, PatioIntAdd,
             PatioIntEquals, PatioIntNotEqual, PatioIntSub, PatioIntToBool, PatioLocalStore,
-            PatioResume, PatioReturn, PatioReturnVoid, PatioVarLoad, PatioVarStore,
+            PatioMemoryLoad, PatioMemoryStore, PatioResume, PatioReturn, PatioReturnVoid,
+            PatioVarLoad, PatioVarStore,
         },
         operands::{PatioLocalRef, PatioOperand},
         types::{PatioDataLayout, PatioType},
-        values::PatioIntValue,
+        values::{PatioDataPtrValue, PatioDataValue, PatioIntValue},
     },
 };
 
@@ -72,7 +73,7 @@ fn declare_compilation_unit_element(
         PatioCompilationUnitElement::DataElementImmutableDefinition(
             data_element_immutable_definition,
         ) => {
-            define_data_element_immutable_definition(data_element_immutable_definition, context);
+            declare_data_element_immutable_definition(data_element_immutable_definition, context);
         }
 
         PatioCompilationUnitElement::FunctionDeclaration(function_declaration) => {
@@ -90,7 +91,11 @@ fn define_compilation_unit_element(
     context: &mut CompilerContext,
 ) {
     match element {
-        PatioCompilationUnitElement::DataElementImmutableDefinition(_) => {}
+        PatioCompilationUnitElement::DataElementImmutableDefinition(
+            data_element_immutable_definition,
+        ) => {
+            define_data_element_immutable_definition(data_element_immutable_definition, context);
+        }
 
         PatioCompilationUnitElement::FunctionDeclaration(_) => {}
 
@@ -100,10 +105,26 @@ fn define_compilation_unit_element(
     }
 }
 
+fn declare_data_element_immutable_definition(
+    data_element_immutable_definition: &PatioDataElementImmutableDefinition,
+    context: &mut CompilerContext,
+) {
+    context.declare_data(
+        &data_element_immutable_definition.name,
+        patio_visibility_to_cranelift_linkage(data_element_immutable_definition.visibility),
+        Mutability::Immutable,
+    );
+}
+
 fn define_data_element_immutable_definition(
     data_element_immutable_definition: &PatioDataElementImmutableDefinition,
     context: &mut CompilerContext,
 ) {
+    let bytes = data_value_to_bytes(
+        &data_element_immutable_definition.value,
+        &context.target_triple,
+    );
+    context.define_data(&data_element_immutable_definition.name, bytes);
 }
 
 fn declare_function_declaration(
@@ -227,6 +248,14 @@ fn compile_instruction(instruction: &PatioInstruction, context: &mut CompilerFun
             compile_local_store(local_store, context);
         }
 
+        PatioInstruction::MemoryLoad(memory_load) => {
+            compile_memory_load(memory_load, context);
+        }
+
+        PatioInstruction::MemoryStore(memory_store) => {
+            compile_memory_store(memory_store, context);
+        }
+
         PatioInstruction::Resume(resume) => {
             compile_resume(resume, context);
         }
@@ -406,6 +435,53 @@ fn compile_local_store(local_store: &PatioLocalStore, context: &mut CompilerFunc
     def_destination(&local_store.destination, cranelift_value, context);
 }
 
+fn compile_memory_load(memory_load: &PatioMemoryLoad, context: &mut CompilerFunctionContext<'_>) {
+    let value_layout = data_layout_element_at(&memory_load.layout, &memory_load.element_offsets);
+    let value_type = data_layout_to_type(value_layout);
+    let cranelift_type = patio_type_to_cranelift_type(value_type, &context.target_triple);
+
+    let cranelift_pointer = compile_operand(&memory_load.pointer, context);
+    let cranelift_flags = cranelift_codegen::ir::MemFlags::new();
+    let byte_offset = data_layout_byte_offset(
+        &memory_load.layout,
+        &memory_load.element_offsets,
+        &context.target_triple,
+    );
+    let cranelift_offset =
+        cranelift_codegen::ir::immediates::Offset32::new(byte_offset.try_into().unwrap());
+    let cranelift_value = context.function_builder.ins().load(
+        cranelift_type,
+        cranelift_flags,
+        cranelift_pointer,
+        cranelift_offset,
+    );
+
+    def_destination(&memory_load.destination, cranelift_value, context);
+}
+
+fn compile_memory_store(
+    memory_store: &PatioMemoryStore,
+    context: &mut CompilerFunctionContext<'_>,
+) {
+    let cranelift_value = compile_operand(&memory_store.value, context);
+    let cranelift_pointer = compile_operand(&memory_store.pointer, context);
+    let cranelift_flags = cranelift_codegen::ir::MemFlags::new();
+    let byte_offset = data_layout_byte_offset(
+        &memory_store.layout,
+        &memory_store.element_offsets,
+        &context.target_triple,
+    );
+    let cranelift_offset =
+        cranelift_codegen::ir::immediates::Offset32::new(byte_offset.try_into().unwrap());
+
+    context.function_builder.ins().store(
+        cranelift_flags,
+        cranelift_value,
+        cranelift_pointer,
+        cranelift_offset,
+    );
+}
+
 fn compile_resume(resume: &PatioResume, context: &mut CompilerFunctionContext<'_>) {
     let cranelift_value = compile_operand(&resume.value, context);
 
@@ -466,6 +542,8 @@ fn compile_operand(
                 .iconst(CRANELIFT_BOOL_TYPE, integer_value)
         }
 
+        PatioOperand::DataPtrConst(value) => compile_data_ptr_value(&value, context),
+
         PatioOperand::IntConst(value) => {
             let cranelift_int_type = cranelift_int_type(&context.target_triple);
             compile_int_value(&value, cranelift_int_type, context)
@@ -486,6 +564,24 @@ fn compile_operand(
     }
 }
 
+fn compile_data_ptr_value(
+    data_ptr_value: &PatioDataPtrValue,
+    context: &mut CompilerFunctionContext<'_>,
+) -> cranelift_codegen::ir::Value {
+    match data_ptr_value {
+        PatioDataPtrValue::DataElement { name } => {
+            let global_value = context.data_ptr(name);
+            let data_ptr_type =
+                cranelift::prelude::Type::triple_pointer_type(&context.target_triple);
+            context
+                .function_builder
+                .ins()
+                .global_value(data_ptr_type, global_value)
+        }
+        PatioDataPtrValue::Null => todo!(),
+    }
+}
+
 fn compile_int_value(
     int_value: &PatioIntValue,
     cranelift_type: cranelift::prelude::Type,
@@ -918,3 +1014,71 @@ mod data_layout_element_at_test {
         assert_eq!(element, &PatioDataLayout::Byte);
     }
 }
+
+fn data_layout_to_type(layout: &PatioDataLayout) -> PatioType {
+    match layout {
+        PatioDataLayout::Array { .. } => panic!("array cannot be represented as type"),
+        PatioDataLayout::Bool => PatioType::Bool,
+        PatioDataLayout::Byte => PatioType::Byte,
+        PatioDataLayout::DataPtr => PatioType::DataPtr,
+        PatioDataLayout::FunctionPtr => PatioType::FunctionPtr,
+        PatioDataLayout::Int => PatioType::Int,
+        PatioDataLayout::Int32 => PatioType::Int32,
+        PatioDataLayout::Struct { .. } => panic!("struct cannot be represented as type"),
+    }
+}
+
+fn data_value_to_bytes(
+    value: &PatioDataValue,
+    target_triple: &target_lexicon::Triple,
+) -> Box<[u8]> {
+    let mut bytes = Vec::with_capacity(
+        data_layout_byte_size(&value.layout(), target_triple)
+            .try_into()
+            .unwrap(),
+    );
+
+    write_data_value(value, &mut bytes, target_triple);
+
+    bytes.into_boxed_slice()
+}
+
+fn write_data_value(
+    value: &PatioDataValue,
+    bytes: &mut Vec<u8>,
+    target_triple: &target_lexicon::Triple,
+) {
+    match value {
+        PatioDataValue::Array {
+            element_layout: _,
+            elements: _,
+        } => todo!(),
+        PatioDataValue::Bool(_) => todo!(),
+        PatioDataValue::Byte(_) => todo!(),
+        PatioDataValue::DataPtr(_) => todo!(),
+        PatioDataValue::Int(int_value) => {
+            // TODO: handle 32-bit int
+            let value: i64 = eval_int_value(int_value, target_triple).into();
+            bytes.write_all(&value.to_le_bytes()).unwrap();
+        }
+        PatioDataValue::Int32(_) => todo!(),
+        PatioDataValue::Struct { elements: _ } => todo!(),
+    }
+}
+
+fn eval_int_value(value: &PatioIntValue, target_triple: &target_lexicon::Triple) -> i32 {
+    match value {
+        PatioIntValue::ByteOffsetOf {
+            layout,
+            element_offsets,
+        } => data_layout_byte_offset(layout, element_offsets, target_triple)
+            .try_into()
+            .unwrap(),
+
+        PatioIntValue::ByteSizeOf { layout } => data_layout_byte_size(layout, target_triple)
+            .try_into()
+            .unwrap(),
+
+        PatioIntValue::Literal { value } => *value,
+    }
+}
diff --git a/patio-to-cranelift/src/patio/flatbuffers.rs b/patio-to-cranelift/src/patio/flatbuffers.rs
index e6cdbd39..337b9960 100644
--- a/patio-to-cranelift/src/patio/flatbuffers.rs
+++ b/patio-to-cranelift/src/patio/flatbuffers.rs
@@ -8,15 +8,16 @@ use crate::{
         instructions::{
             PatioBlock, PatioBoolToInt, PatioCall, PatioIf, PatioInstruction, PatioIntAdd,
             PatioIntEquals, PatioIntNotEqual, PatioIntSub, PatioIntToBool, PatioLocalStore,
-            PatioResume, PatioReturn, PatioReturnVoid, PatioVarLoad, PatioVarStore,
+            PatioMemoryLoad, PatioMemoryStore, PatioResume, PatioReturn, PatioReturnVoid,
+            PatioVarLoad, PatioVarStore,
         },
         operands::{PatioLocalRef, PatioOperand},
         types::{PatioDataLayout, PatioType},
-        values::{PatioDataValue, PatioIntValue},
+        values::{PatioDataPtrValue, PatioDataValue, PatioIntValue},
     },
     patio_generated::{
         self,
-        patio::{PatioDataLayoutValue, PatioIntValueValue, PatioTypeValue},
+        patio::{PatioDataLayoutValue, PatioDataValueValue, PatioIntValueValue, PatioTypeValue},
     },
 };
 
@@ -222,6 +223,28 @@ impl From<patio_generated::patio::PatioInstruction<'_>> for PatioInstruction {
                 })
             }
 
+            patio_generated::patio::PatioInstructionValue::PatioMemoryLoad => {
+                let memory_load = value.value_as_patio_memory_load().unwrap();
+
+                Self::MemoryLoad(PatioMemoryLoad {
+                    destination: memory_load.destination().into(),
+                    pointer: memory_load.pointer().into(),
+                    layout: memory_load.layout().into(),
+                    element_offsets: memory_load.element_offsets().iter().collect(),
+                })
+            }
+
+            patio_generated::patio::PatioInstructionValue::PatioMemoryStore => {
+                let memory_store = value.value_as_patio_memory_store().unwrap();
+
+                Self::MemoryStore(PatioMemoryStore {
+                    pointer: memory_store.pointer().into(),
+                    layout: memory_store.layout().into(),
+                    element_offsets: memory_store.element_offsets().iter().collect(),
+                    value: memory_store.value().into(),
+                })
+            }
+
             patio_generated::patio::PatioInstructionValue::PatioResume => {
                 let resume = value.value_as_patio_resume().unwrap();
 
@@ -293,7 +316,10 @@ impl From<patio_generated::patio::PatioOperand<'_>> for PatioOperand {
                 Self::ByteConst(byte_const.value().try_into().unwrap())
             }
 
-            patio_generated::patio::PatioOperandValue::PatioDataPtrConst => Self::DataPtrConst,
+            patio_generated::patio::PatioOperandValue::PatioDataPtrConst => {
+                let data_ptr_const = value.value_as_patio_data_ptr_const().unwrap();
+                Self::DataPtrConst(data_ptr_const.value().into())
+            }
 
             patio_generated::patio::PatioOperandValue::PatioFunctionPtrConst => {
                 Self::FunctionPtrConst
@@ -331,6 +357,29 @@ impl From<patio_generated::patio::PatioLocalRef<'_>> for PatioLocalRef {
     }
 }
 
+// == Data ptr values ==
+
+impl From<patio_generated::patio::PatioDataPtrValue<'_>> for PatioDataPtrValue {
+    fn from(value: patio_generated::patio::PatioDataPtrValue<'_>) -> Self {
+        match value.value_type() {
+            patio_generated::patio::PatioDataPtrValueValue::PatioDataPtrValueDataElement => {
+                let data_element = value.value_as_patio_data_ptr_value_data_element().unwrap();
+                Self::DataElement {
+                    name: data_element.name().to_string(),
+                }
+            }
+
+            patio_generated::patio::PatioDataPtrValueValue::PatioDataPtrValueNull => Self::Null,
+
+            data_ptr_value_type => todo!(
+                "Unhandled data ptr value type: {:?} {}",
+                data_ptr_value_type.variant_name(),
+                data_ptr_value_type.0
+            ),
+        }
+    }
+}
+
 // == Int values ==
 
 impl From<patio_generated::patio::PatioIntValue<'_>> for PatioIntValue {
@@ -441,6 +490,46 @@ impl From<patio_generated::patio::PatioDataLayout<'_>> for PatioDataLayout {
 impl From<patio_generated::patio::PatioDataValue<'_>> for PatioDataValue {
     fn from(value: patio_generated::patio::PatioDataValue<'_>) -> Self {
         match value.value_type() {
+            PatioDataValueValue::PatioDataValueArray => {
+                let array = value.value_as_patio_data_value_array().unwrap();
+                Self::Array {
+                    element_layout: array.element_layout().into(),
+                    elements: array.elements().iter().map(Into::into).collect(),
+                }
+            }
+
+            PatioDataValueValue::PatioDataValueBool => {
+                let bool = value.value_as_patio_data_value_bool().unwrap();
+                Self::Bool(bool.value())
+            }
+
+            PatioDataValueValue::PatioDataValueByte => {
+                let byte = value.value_as_patio_data_value_byte().unwrap();
+                Self::Byte(byte.value().try_into().unwrap())
+            }
+
+            PatioDataValueValue::PatioDataValueDataPtr => {
+                let data_ptr = value.value_as_patio_data_value_data_ptr().unwrap();
+                Self::DataPtr(data_ptr.value().into())
+            }
+
+            PatioDataValueValue::PatioDataValueInt => {
+                let int = value.value_as_patio_data_value_int().unwrap();
+                Self::Int(int.value().into())
+            }
+
+            PatioDataValueValue::PatioDataValueInt32 => {
+                let int32 = value.value_as_patio_data_value_int_32().unwrap();
+                Self::Int32(int32.value().into())
+            }
+
+            PatioDataValueValue::PatioDataValueStruct => {
+                let struct_ = value.value_as_patio_data_value_struct().unwrap();
+                Self::Struct {
+                    elements: struct_.elements().iter().map(Into::into).collect(),
+                }
+            }
+
             value_type => {
                 todo!(
                     "Unhandled data value type: {:?} {}",
diff --git a/patio-to-cranelift/src/patio/instructions.rs b/patio-to-cranelift/src/patio/instructions.rs
index 23268a22..39c39964 100644
--- a/patio-to-cranelift/src/patio/instructions.rs
+++ b/patio-to-cranelift/src/patio/instructions.rs
@@ -1,4 +1,7 @@
-use crate::patio::operands::{PatioLocalRef, PatioOperand};
+use crate::patio::{
+    operands::{PatioLocalRef, PatioOperand},
+    types::PatioDataLayout,
+};
 
 pub(crate) enum PatioInstruction {
     Block(PatioBlock),
@@ -11,6 +14,8 @@ pub(crate) enum PatioInstruction {
     IntSub(PatioIntSub),
     IntToBool(PatioIntToBool),
     LocalStore(PatioLocalStore),
+    MemoryLoad(PatioMemoryLoad),
+    MemoryStore(PatioMemoryStore),
     Resume(PatioResume),
     Return(PatioReturn),
     ReturnVoid(PatioReturnVoid),
@@ -75,6 +80,20 @@ pub(crate) struct PatioLocalStore {
     pub(crate) value: PatioOperand,
 }
 
+pub(crate) struct PatioMemoryLoad {
+    pub(crate) destination: PatioLocalRef,
+    pub(crate) pointer: PatioOperand,
+    pub(crate) layout: PatioDataLayout,
+    pub(crate) element_offsets: Vec<u32>,
+}
+
+pub(crate) struct PatioMemoryStore {
+    pub(crate) pointer: PatioOperand,
+    pub(crate) layout: PatioDataLayout,
+    pub(crate) element_offsets: Vec<u32>,
+    pub(crate) value: PatioOperand,
+}
+
 pub(crate) struct PatioResume {
     pub(crate) value: PatioOperand,
 }
diff --git a/patio-to-cranelift/src/patio/operands.rs b/patio-to-cranelift/src/patio/operands.rs
index 65dd6a72..05f9eec3 100644
--- a/patio-to-cranelift/src/patio/operands.rs
+++ b/patio-to-cranelift/src/patio/operands.rs
@@ -1,10 +1,10 @@
-use crate::patio::values::PatioIntValue;
+use crate::patio::values::{PatioDataPtrValue, PatioIntValue};
 
 #[derive(Debug)]
 pub(crate) enum PatioOperand {
     BoolConst(bool),
     ByteConst(u8),
-    DataPtrConst,
+    DataPtrConst(PatioDataPtrValue),
     FunctionPtrConst,
     IntConst(PatioIntValue),
     Int32Const(PatioIntValue),
diff --git a/patio-to-cranelift/src/patio/types.rs b/patio-to-cranelift/src/patio/types.rs
index 524edd52..c513ed4e 100644
--- a/patio-to-cranelift/src/patio/types.rs
+++ b/patio-to-cranelift/src/patio/types.rs
@@ -9,7 +9,7 @@ pub(crate) enum PatioType {
     Void,
 }
 
-#[derive(Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq)]
 pub(crate) enum PatioDataLayout {
     Array {
         element: Box<PatioDataLayout>,
diff --git a/patio-to-cranelift/src/patio/values.rs b/patio-to-cranelift/src/patio/values.rs
index 79dd548c..fe511437 100644
--- a/patio-to-cranelift/src/patio/values.rs
+++ b/patio-to-cranelift/src/patio/values.rs
@@ -1,6 +1,44 @@
 use crate::patio::types::PatioDataLayout;
 
-pub(crate) enum PatioDataValue {}
+pub(crate) enum PatioDataValue {
+    Array {
+        element_layout: PatioDataLayout,
+        elements: Vec<PatioDataValue>,
+    },
+    Bool(bool),
+    Byte(u8),
+    DataPtr(PatioDataPtrValue),
+    Int(PatioIntValue),
+    Int32(PatioIntValue),
+    Struct {
+        elements: Vec<PatioDataValue>,
+    },
+}
+
+impl PatioDataValue {
+    pub(crate) fn layout(&self) -> PatioDataLayout {
+        match self {
+            PatioDataValue::Array {
+                element_layout,
+                elements,
+            } => PatioDataLayout::array(element_layout.clone(), elements.len().try_into().unwrap()),
+            PatioDataValue::Bool(_) => PatioDataLayout::Bool,
+            PatioDataValue::Byte(_) => PatioDataLayout::Byte,
+            PatioDataValue::DataPtr(_) => PatioDataLayout::DataPtr,
+            PatioDataValue::Int(_) => PatioDataLayout::Int,
+            PatioDataValue::Int32(_) => PatioDataLayout::Int32,
+            PatioDataValue::Struct { elements } => {
+                PatioDataLayout::struct_(elements.iter().map(|element| element.layout()).collect())
+            }
+        }
+    }
+}
+
+#[derive(Debug)]
+pub(crate) enum PatioDataPtrValue {
+    DataElement { name: String },
+    Null,
+}
 
 #[derive(Debug)]
 pub(crate) enum PatioIntValue {