enable returning basic symbols with storage manager

compiler/gen_dev/src/generic64/aarch64.rs

@@ -1,4 +1,4 @@
-use crate::generic64::{Assembler, CallConv, RegTrait, SymbolStorage};
+use crate::generic64::{storage::StorageManager, Assembler, CallConv, RegTrait, SymbolStorage};
 use crate::Relocation;
 use bumpalo::collections::Vec;
 use packed_struct::prelude::*;
@@ -75,7 +75,7 @@ pub struct AArch64Call {}
 
 const STACK_ALIGNMENT: u8 = 16;
 
-impl CallConv<AArch64GeneralReg, AArch64FloatReg> for AArch64Call {
+impl CallConv<AArch64GeneralReg, AArch64FloatReg, AArch64Assembler> for AArch64Call {
     const BASE_PTR_REG: AArch64GeneralReg = AArch64GeneralReg::FP;
     const STACK_PTR_REG: AArch64GeneralReg = AArch64GeneralReg::ZRSP;
 
@@ -268,6 +268,21 @@ impl CallConv<AArch64GeneralReg, AArch64FloatReg> for AArch64Call {
         todo!("Storing args for AArch64");
     }
 
+    fn return_complex_symbol<'a>(
+        buf: &mut Vec<'a, u8>,
+        storage_manager: &mut StorageManager<
+            'a,
+            AArch64GeneralReg,
+            AArch64FloatReg,
+            AArch64Assembler,
+            AArch64Call,
+        >,
+        sym: &Symbol,
+        layout: &Layout<'a>,
+    ) {
+        todo!("Returning complex symbols for AArch64");
+    }
+
     fn return_struct<'a>(
         _buf: &mut Vec<'a, u8>,
         _struct_offset: i32,

compiler/gen_dev/src/generic64/mod.rs

@@ -1,4 +1,7 @@
-use crate::{Backend, Env, Relocation};
+use crate::{
+    single_register_floats, single_register_integers, single_register_layouts, Backend, Env,
+    Relocation,
+};
 use bumpalo::collections::Vec;
 use roc_builtins::bitcode::{FloatWidth, IntWidth};
 use roc_collections::all::{MutMap, MutSet};
@@ -10,9 +13,9 @@ use roc_mono::layout::{Builtin, Layout};
 use roc_target::TargetInfo;
 use std::marker::PhantomData;
 
-pub mod aarch64;
+pub(crate) mod aarch64;
-mod storage;
+pub(crate) mod storage;
-pub mod x86_64;
+pub(crate) mod x86_64;
 
 // TODO: StorageManager is still not fully integrated.
 // General pieces needed:
@@ -26,7 +29,9 @@ pub mod x86_64;
 // - look into fixing join to no longer use multiple backends???
 use storage::StorageManager;
 
-pub trait CallConv<GeneralReg: RegTrait, FloatReg: RegTrait> {
+pub trait CallConv<GeneralReg: RegTrait, FloatReg: RegTrait, ASM: Assembler<GeneralReg, FloatReg>>:
+    Sized
+{
     const BASE_PTR_REG: GeneralReg;
     const STACK_PTR_REG: GeneralReg;
 
@@ -53,12 +58,14 @@ pub trait CallConv<GeneralReg: RegTrait, FloatReg: RegTrait> {
 
     fn setup_stack<'a>(
         buf: &mut Vec<'a, u8>,
+        // TODO: This should deal with float regs as well.
         general_saved_regs: &[GeneralReg],
         requested_stack_size: i32,
         fn_call_stack_size: i32,
     ) -> i32;
     fn cleanup_stack<'a>(
         buf: &mut Vec<'a, u8>,
+        // TODO: This should deal with float regs as well.
         general_saved_regs: &[GeneralReg],
         aligned_stack_size: i32,
         fn_call_stack_size: i32,
@@ -86,6 +93,15 @@ pub trait CallConv<GeneralReg: RegTrait, FloatReg: RegTrait> {
         ret_layout: &Layout<'a>,
     ) -> u32;
 
+    /// return_complex_symbol returns the specified complex/non-primative symbol.
+    /// It uses the layout to determine how the data should be returned.
+    fn return_complex_symbol<'a>(
+        buf: &mut Vec<'a, u8>,
+        storage_manager: &mut StorageManager<'a, GeneralReg, FloatReg, ASM, Self>,
+        sym: &Symbol,
+        layout: &Layout<'a>,
+    );
+
     // return_struct returns a struct currently on the stack at `struct_offset`.
     // It does so using registers and stack as necessary.
     fn return_struct<'a>(
@@ -106,7 +122,7 @@ pub trait CallConv<GeneralReg: RegTrait, FloatReg: RegTrait> {
 /// Thus, some backends will need to use mulitiple instructions to preform a single one of this calls.
 /// Generally, I prefer explicit sources, as opposed to dst being one of the sources. Ex: `x = x + y` would be `add x, x, y` instead of `add x, y`.
 /// dst should always come before sources.
-pub trait Assembler<GeneralReg: RegTrait, FloatReg: RegTrait> {
+pub trait Assembler<GeneralReg: RegTrait, FloatReg: RegTrait>: Sized {
     fn abs_reg64_reg64(buf: &mut Vec<'_, u8>, dst: GeneralReg, src: GeneralReg);
     fn abs_freg64_freg64(
         buf: &mut Vec<'_, u8>,
@@ -262,7 +278,7 @@ pub struct Backend64Bit<
     GeneralReg: RegTrait,
     FloatReg: RegTrait,
     ASM: Assembler<GeneralReg, FloatReg>,
-    CC: CallConv<GeneralReg, FloatReg>,
+    CC: CallConv<GeneralReg, FloatReg, ASM>,
 > {
     phantom_asm: PhantomData<ASM>,
     phantom_cc: PhantomData<CC>,
@@ -313,7 +329,7 @@ pub fn new_backend_64bit<
     GeneralReg: RegTrait,
     FloatReg: RegTrait,
     ASM: Assembler<GeneralReg, FloatReg>,
-    CC: CallConv<GeneralReg, FloatReg>,
+    CC: CallConv<GeneralReg, FloatReg, ASM>,
 >(
     env: &'a Env,
     target_info: TargetInfo,
@@ -355,7 +371,7 @@ impl<
         GeneralReg: RegTrait,
         FloatReg: RegTrait,
         ASM: Assembler<GeneralReg, FloatReg>,
-        CC: CallConv<GeneralReg, FloatReg>,
+        CC: CallConv<GeneralReg, FloatReg, ASM>,
     > Backend<'a> for Backend64Bit<'a, GeneralReg, FloatReg, ASM, CC>
 {
     fn env(&self) -> &Env<'a> {
@@ -1104,71 +1120,30 @@ impl<
     }
 
     fn return_symbol(&mut self, sym: &Symbol, layout: &Layout<'a>) {
-        let val = self.symbol_storage_map.get(sym);
+        if self.storage_manager.is_stored_primitive(sym) {
-        match val {
+            // Just load it to the correct type of reg as a stand alone value.
-            Some(SymbolStorage::GeneralReg(reg)) if *reg == CC::GENERAL_RETURN_REGS[0] => {}
+            match layout {
-            Some(SymbolStorage::GeneralReg(reg)) => {
+                single_register_integers!() => {
-                // If it fits in a general purpose register, just copy it over to.
+                    self.storage_manager.load_to_specified_general_reg(
-                // Technically this can be optimized to produce shorter instructions if less than 64bits.
+                        &mut self.buf,
-                ASM::mov_reg64_reg64(&mut self.buf, CC::GENERAL_RETURN_REGS[0], *reg);
+                        sym,
-            }
+                        CC::GENERAL_RETURN_REGS[0],
-            Some(SymbolStorage::FloatReg(reg)) if *reg == CC::FLOAT_RETURN_REGS[0] => {}
+                    );
-            Some(SymbolStorage::FloatReg(reg)) => {
-                ASM::mov_freg64_freg64(&mut self.buf, CC::FLOAT_RETURN_REGS[0], *reg);
-            }
-            Some(SymbolStorage::Base { offset, size, .. }) => match layout {
-                Layout::Builtin(Builtin::Int(IntWidth::I64 | IntWidth::U64)) => {
-                    ASM::mov_reg64_base32(&mut self.buf, CC::GENERAL_RETURN_REGS[0], *offset);
                 }
-                Layout::Builtin(Builtin::Float(FloatWidth::F64)) => {
+                single_register_floats!() => {
-                    ASM::mov_freg64_base32(&mut self.buf, CC::FLOAT_RETURN_REGS[0], *offset);
+                    self.storage_manager.load_to_specified_float_reg(
+                        &mut self.buf,
+                        sym,
+                        CC::FLOAT_RETURN_REGS[0],
+                    );
                 }
-                Layout::Builtin(Builtin::Str) => {
+                _ => {
-                    if self.symbol_storage_map.contains_key(&Symbol::RET_POINTER) {
+                    internal_error!("All primitive valuse should fit in a single register");
-                        // This will happen on windows, return via pointer here.
-                        todo!("Returning strings via pointer");
-                    } else {
-                        ASM::mov_reg64_base32(&mut self.buf, CC::GENERAL_RETURN_REGS[0], *offset);
-                        ASM::mov_reg64_base32(
-                            &mut self.buf,
-                            CC::GENERAL_RETURN_REGS[1],
-                            *offset + 8,
-                        );
-                    }
                 }
-                Layout::Struct(field_layouts) => {
-                    let (offset, size) = (*offset, *size);
-                    // Nothing to do for empty struct
-                    if size > 0 {
-                        let ret_reg = if self.symbol_storage_map.contains_key(&Symbol::RET_POINTER)
-                        {
-                            Some(
-                                self.storage_manager
-                                    .load_to_general_reg(&mut self.buf, &Symbol::RET_POINTER),
-                            )
-                        } else {
-                            None
-                        };
-                        CC::return_struct(&mut self.buf, offset, size, field_layouts, ret_reg);
-                    }
-                }
-                x => todo!("returning symbol with layout, {:?}", x),
-            },
-            Some(x) => todo!("returning symbol storage, {:?}", x),
-            None if layout == &Layout::Struct(&[]) => {
-                // Empty struct is not defined and does nothing.
-            }
-            None => {
-                internal_error!("Unknown return symbol: {:?}", sym);
             }
+            return;
         }
-        let inst_loc = self.buf.len() as u64;
+        CC::return_complex_symbol(&mut self.buf, &mut self.storage_manager, sym, layout)
-        let offset = ASM::jmp_imm32(&mut self.buf, 0x1234_5678) as u64;
-        self.relocs.push(Relocation::JmpToReturn {
-            inst_loc,
-            inst_size: self.buf.len() as u64 - inst_loc,
-            offset,
-        });
     }
 }
 
@@ -1179,7 +1154,7 @@ impl<
         FloatReg: RegTrait,
         GeneralReg: RegTrait,
         ASM: Assembler<GeneralReg, FloatReg>,
-        CC: CallConv<GeneralReg, FloatReg>,
+        CC: CallConv<GeneralReg, FloatReg, ASM>,
     > Backend64Bit<'a, GeneralReg, FloatReg, ASM, CC>
 {
     // Updates a jump instruction to a new offset and returns the number of bytes written.

compiler/gen_dev/src/generic64/storage.rs

@@ -72,7 +72,7 @@ pub struct StorageManager<
     GeneralReg: RegTrait,
     FloatReg: RegTrait,
     ASM: Assembler<GeneralReg, FloatReg>,
-    CC: CallConv<GeneralReg, FloatReg>,
+    CC: CallConv<GeneralReg, FloatReg, ASM>,
 > {
     phantom_cc: PhantomData<CC>,
     phantom_asm: PhantomData<ASM>,
@@ -112,7 +112,7 @@ pub fn new_storage_manager<
     GeneralReg: RegTrait,
     FloatReg: RegTrait,
     ASM: Assembler<GeneralReg, FloatReg>,
-    CC: CallConv<GeneralReg, FloatReg>,
+    CC: CallConv<GeneralReg, FloatReg, ASM>,
 >(
     env: &'a Env,
     target_info: TargetInfo,
@@ -140,7 +140,7 @@ impl<
         FloatReg: RegTrait,
         GeneralReg: RegTrait,
         ASM: Assembler<GeneralReg, FloatReg>,
-        CC: CallConv<GeneralReg, FloatReg>,
+        CC: CallConv<GeneralReg, FloatReg, ASM>,
     > StorageManager<'a, GeneralReg, FloatReg, ASM, CC>
 {
     pub fn reset(&mut self) {
@@ -160,6 +160,19 @@ impl<
         self.stack_size = 0;
     }
 
+    // Returns true if the symbol is storing a primitive value.
+    pub fn is_stored_primitive(&self, sym: &Symbol) -> bool {
+        let storage = if let Some(storage) = self.symbol_storage_map.get(sym) {
+            storage
+        } else {
+            internal_error!("Unknown symbol: {}", sym);
+        };
+        matches!(
+            storage,
+            Reg(_) | Stack(Primitive { .. } | ReferencedPrimitive { .. })
+        )
+    }
+
     // Get a general register from the free list.
     // Will free data to the stack if necessary to get the register.
     fn get_general_reg(&mut self, buf: &mut Vec<'a, u8>) -> GeneralReg {
@@ -348,6 +361,103 @@ impl<
         }
     }
 
+    // Loads the symbol to the specified register.
+    // It will fail if the symbol is stored in a float register.
+    // This is only made to be used in special cases where exact regs are needed (function args and returns).
+    // It will not try to free the register first.
+    // This will not track the symbol change (it makes no assumptions about the new reg).
+    pub fn load_to_specified_general_reg(
+        &self,
+        buf: &mut Vec<'a, u8>,
+        sym: &Symbol,
+        reg: GeneralReg,
+    ) {
+        let storage = if let Some(storage) = self.symbol_storage_map.get(sym) {
+            storage
+        } else {
+            internal_error!("Unknown symbol: {}", sym);
+        };
+        match storage {
+            Reg(General(old_reg))
+            | Stack(Primitive {
+                reg: Some(General(old_reg)),
+                ..
+            }) => {
+                debug_assert_ne!(*old_reg, reg);
+                ASM::mov_reg64_reg64(buf, reg, *old_reg);
+            }
+            Reg(Float(_))
+            | Stack(Primitive {
+                reg: Some(Float(_)),
+                ..
+            }) => {
+                internal_error!("Cannot load floating point symbol into GeneralReg: {}", sym)
+            }
+            Stack(Primitive {
+                reg: None,
+                base_offset,
+            }) => {
+                debug_assert_eq!(base_offset % 8, 0);
+                ASM::mov_reg64_base32(buf, reg, *base_offset);
+            }
+            Stack(ReferencedPrimitive { .. }) => {
+                todo!("loading referenced primitives")
+            }
+            Stack(Complex { .. }) => {
+                internal_error!("Cannot load large values into general registers: {}", sym)
+            }
+            NoData => {
+                internal_error!("Cannot load no data into general registers: {}", sym)
+            }
+        }
+    }
+
+    // Loads the symbol to the specified register.
+    // It will fail if the symbol is stored in a general register.
+    // This is only made to be used in special cases where exact regs are needed (function args and returns).
+    // It will not try to free the register first.
+    // This will not track the symbol change (it makes no assumptions about the new reg).
+    pub fn load_to_specified_float_reg(&self, buf: &mut Vec<'a, u8>, sym: &Symbol, reg: FloatReg) {
+        let storage = if let Some(storage) = self.symbol_storage_map.get(sym) {
+            storage
+        } else {
+            internal_error!("Unknown symbol: {}", sym);
+        };
+        match storage {
+            Reg(Float(old_reg))
+            | Stack(Primitive {
+                reg: Some(Float(old_reg)),
+                ..
+            }) => {
+                debug_assert_ne!(*old_reg, reg);
+                ASM::mov_freg64_freg64(buf, reg, *old_reg);
+            }
+            Reg(General(_))
+            | Stack(Primitive {
+                reg: Some(General(_)),
+                ..
+            }) => {
+                internal_error!("Cannot load general symbol into FloatReg: {}", sym)
+            }
+            Stack(Primitive {
+                reg: None,
+                base_offset,
+            }) => {
+                debug_assert_eq!(base_offset % 8, 0);
+                ASM::mov_freg64_base32(buf, reg, *base_offset);
+            }
+            Stack(ReferencedPrimitive { .. }) => {
+                todo!("loading referenced primitives")
+            }
+            Stack(Complex { .. }) => {
+                internal_error!("Cannot load large values into float registers: {}", sym)
+            }
+            NoData => {
+                internal_error!("Cannot load no data into general registers: {}", sym)
+            }
+        }
+    }
+
     // Loads a field from a struct or tag union.
     // This is lazy by default. It will not copy anything around.
     pub fn load_field_at_index(

compiler/gen_dev/src/generic64/x86_64.rs

@@ -1,4 +1,4 @@
-use crate::generic64::{Assembler, CallConv, RegTrait, SymbolStorage};
+use crate::generic64::{storage::StorageManager, Assembler, CallConv, RegTrait, SymbolStorage};
 use crate::{
     single_register_floats, single_register_integers, single_register_layouts, Relocation,
 };
@@ -70,7 +70,7 @@ pub struct X86_64SystemV {}
 
 const STACK_ALIGNMENT: u8 = 16;
 
-impl CallConv<X86_64GeneralReg, X86_64FloatReg> for X86_64SystemV {
+impl CallConv<X86_64GeneralReg, X86_64FloatReg, X86_64Assembler> for X86_64SystemV {
     const BASE_PTR_REG: X86_64GeneralReg = X86_64GeneralReg::RBP;
     const STACK_PTR_REG: X86_64GeneralReg = X86_64GeneralReg::RSP;
 
@@ -441,6 +441,88 @@ impl CallConv<X86_64GeneralReg, X86_64FloatReg> for X86_64SystemV {
         stack_offset as u32
     }
 
+    fn return_complex_symbol<'a>(
+        buf: &mut Vec<'a, u8>,
+        storage_manager: &mut StorageManager<
+            'a,
+            X86_64GeneralReg,
+            X86_64FloatReg,
+            X86_64Assembler,
+            X86_64SystemV,
+        >,
+        sym: &Symbol,
+        layout: &Layout<'a>,
+    ) {
+        // Complex types.
+        //     let val = self.symbol_storage_map.get(sym);
+        //     match val {
+        //         Some(SymbolStorage::GeneralReg(reg)) if *reg == CC::GENERAL_RETURN_REGS[0] => {}
+        //         Some(SymbolStorage::GeneralReg(reg)) => {
+        //             // If it fits in a general purpose register, just copy it over to.
+        //             // Technically this can be optimized to produce shorter instructions if less than 64bits.
+        //             ASM::mov_reg64_reg64(&mut self.buf, CC::GENERAL_RETURN_REGS[0], *reg);
+        //         }
+        //         Some(SymbolStorage::FloatReg(reg)) if *reg == CC::FLOAT_RETURN_REGS[0] => {}
+        //         Some(SymbolStorage::FloatReg(reg)) => {
+        //             ASM::mov_freg64_freg64(&mut self.buf, CC::FLOAT_RETURN_REGS[0], *reg);
+        //         }
+        //         Some(SymbolStorage::Base { offset, size, .. }) => match layout {
+        //             Layout::Builtin(Builtin::Int(IntWidth::I64 | IntWidth::U64)) => {
+        //                 ASM::mov_reg64_base32(&mut self.buf, CC::GENERAL_RETURN_REGS[0], *offset);
+        //             }
+        //             Layout::Builtin(Builtin::Float(FloatWidth::F64)) => {
+        //                 ASM::mov_freg64_base32(&mut self.buf, CC::FLOAT_RETURN_REGS[0], *offset);
+        //             }
+        //             Layout::Builtin(Builtin::Str) => {
+        //                 if self.symbol_storage_map.contains_key(&Symbol::RET_POINTER) {
+        //                     // This will happen on windows, return via pointer here.
+        //                     todo!("Returning strings via pointer");
+        //                 } else {
+        //                     ASM::mov_reg64_base32(&mut self.buf, CC::GENERAL_RETURN_REGS[0], *offset);
+        //                     ASM::mov_reg64_base32(
+        //                         &mut self.buf,
+        //                         CC::GENERAL_RETURN_REGS[1],
+        //                         *offset + 8,
+        //                     );
+        //                 }
+        //             }
+        //             Layout::Struct(field_layouts) => {
+        //                 let (offset, size) = (*offset, *size);
+        //                 // Nothing to do for empty struct
+        //                 if size > 0 {
+        //                     let ret_reg = if self.symbol_storage_map.contains_key(&Symbol::RET_POINTER)
+        //                     {
+        //                         Some(
+        //                             self.storage_manager
+        //                                 .load_to_general_reg(&mut self.buf, &Symbol::RET_POINTER),
+        //                         )
+        //                     } else {
+        //                         None
+        //                     };
+        //                     CC::return_struct(&mut self.buf, offset, size, field_layouts, ret_reg);
+        //                 }
+        //             }
+        //             x => todo!("returning symbol with layout, {:?}", x),
+        //         },
+        //         Some(x) => todo!("returning symbol storage, {:?}", x),
+        //         None if layout == &Layout::Struct(&[]) => {
+        //             // Empty struct is not defined and does nothing.
+        //         }
+        //         None => {
+        //             internal_error!("Unknown return symbol: {:?}", sym);
+        //         }
+        //     }
+        //     let inst_loc = self.buf.len() as u64;
+        //     let offset = ASM::jmp_imm32(&mut self.buf, 0x1234_5678) as u64;
+        //     self.relocs.push(Relocation::JmpToReturn {
+        //         inst_loc,
+        //         inst_size: self.buf.len() as u64 - inst_loc,
+        //         offset,
+        //     });
+        // }
+        todo!("Returning complex symbols for X86_64");
+    }
+
     fn return_struct<'a>(
         _buf: &mut Vec<'a, u8>,
         _struct_offset: i32,
@@ -458,7 +540,7 @@ impl CallConv<X86_64GeneralReg, X86_64FloatReg> for X86_64SystemV {
     }
 }
 
-impl CallConv<X86_64GeneralReg, X86_64FloatReg> for X86_64WindowsFastcall {
+impl CallConv<X86_64GeneralReg, X86_64FloatReg, X86_64Assembler> for X86_64WindowsFastcall {
     const BASE_PTR_REG: X86_64GeneralReg = X86_64GeneralReg::RBP;
     const STACK_PTR_REG: X86_64GeneralReg = X86_64GeneralReg::RSP;
 
@@ -765,6 +847,21 @@ impl CallConv<X86_64GeneralReg, X86_64FloatReg> for X86_64WindowsFastcall {
         stack_offset as u32
     }
 
+    fn return_complex_symbol<'a>(
+        buf: &mut Vec<'a, u8>,
+        storage_manager: &mut StorageManager<
+            'a,
+            X86_64GeneralReg,
+            X86_64FloatReg,
+            X86_64Assembler,
+            X86_64WindowsFastcall,
+        >,
+        sym: &Symbol,
+        layout: &Layout<'a>,
+    ) {
+        todo!("Returning symbols for X86_64");
+    }
+
     fn return_struct<'a>(
         _buf: &mut Vec<'a, u8>,
         _struct_offset: i32,