fully integrate storage manager removing old code

compiler/gen_dev/src/generic64/aarch64.rs

@@ -159,13 +159,14 @@ impl CallConv<AArch64GeneralReg, AArch64FloatReg, AArch64Assembler> for AArch64C
     #[inline(always)]
     fn setup_stack(
         buf: &mut Vec<'_, u8>,
-        saved_regs: &[AArch64GeneralReg],
+        saved_general_regs: &[AArch64GeneralReg],
+        saved_float_regs: &[AArch64FloatReg],
         requested_stack_size: i32,
         fn_call_stack_size: i32,
     ) -> i32 {
         // Full size is upcast to i64 to make sure we don't overflow here.
         let full_stack_size = match requested_stack_size
-            .checked_add(8 * saved_regs.len() as i32 + 8) // The extra 8 is space to store the frame pointer.
+            .checked_add(8 * (saved_general_regs.len() + saved_float_regs.len()) as i32 + 8) // The extra 8 is space to store the frame pointer.
             .and_then(|size| size.checked_add(fn_call_stack_size))
         {
             Some(size) => size,
@@ -203,10 +204,14 @@ impl CallConv<AArch64GeneralReg, AArch64FloatReg, AArch64Assembler> for AArch64C
                 AArch64Assembler::mov_stack32_reg64(buf, offset, AArch64GeneralReg::FP);
 
                 offset = aligned_stack_size - fn_call_stack_size;
-                for reg in saved_regs {
+                for reg in saved_general_regs {
                     offset -= 8;
                     AArch64Assembler::mov_base32_reg64(buf, offset, *reg);
                 }
+                for reg in saved_float_regs {
+                    offset -= 8;
+                    AArch64Assembler::mov_base32_freg64(buf, offset, *reg);
+                }
                 aligned_stack_size
             } else {
                 0
@@ -219,7 +224,8 @@ impl CallConv<AArch64GeneralReg, AArch64FloatReg, AArch64Assembler> for AArch64C
     #[inline(always)]
     fn cleanup_stack(
         buf: &mut Vec<'_, u8>,
-        saved_regs: &[AArch64GeneralReg],
+        saved_general_regs: &[AArch64GeneralReg],
+        saved_float_regs: &[AArch64FloatReg],
         aligned_stack_size: i32,
         fn_call_stack_size: i32,
     ) {
@@ -232,10 +238,14 @@ impl CallConv<AArch64GeneralReg, AArch64FloatReg, AArch64Assembler> for AArch64C
             AArch64Assembler::mov_reg64_stack32(buf, AArch64GeneralReg::FP, offset);
 
             offset = aligned_stack_size - fn_call_stack_size;
-            for reg in saved_regs {
+            for reg in saved_general_regs {
                 offset -= 8;
                 AArch64Assembler::mov_reg64_base32(buf, *reg, offset);
             }
+            for reg in saved_float_regs {
+                offset -= 8;
+                AArch64Assembler::mov_freg64_base32(buf, *reg, offset);
+            }
             AArch64Assembler::add_reg64_reg64_imm32(
                 buf,
                 AArch64GeneralReg::ZRSP,

compiler/gen_dev/src/generic64/mod.rs

@@ -1,7 +1,7 @@
 use crate::{single_register_floats, single_register_integers, Backend, Env, Relocation};
 use bumpalo::collections::Vec;
 use roc_builtins::bitcode::{FloatWidth, IntWidth};
-use roc_collections::all::{MutMap, MutSet};
+use roc_collections::all::MutMap;
 use roc_error_macros::internal_error;
 use roc_module::symbol::{Interns, Symbol};
 use roc_mono::code_gen_help::CodeGenHelp;
@@ -53,15 +53,15 @@ pub trait CallConv<GeneralReg: RegTrait, FloatReg: RegTrait, ASM: Assembler<Gene
 
     fn setup_stack<'a>(
         buf: &mut Vec<'a, u8>,
-        // TODO: This should deal with float regs as well.
         general_saved_regs: &[GeneralReg],
+        float_saved_regs: &[FloatReg],
         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],
+        float_saved_regs: &[FloatReg],
         aligned_stack_size: i32,
         fn_call_stack_size: i32,
     );
@@ -236,29 +236,6 @@ pub trait Assembler<GeneralReg: RegTrait, FloatReg: RegTrait>: Sized {
     fn ret(buf: &mut Vec<'_, u8>);
 }
 
-#[derive(Clone, Debug, PartialEq)]
-pub enum SymbolStorage<GeneralReg: RegTrait, FloatReg: RegTrait> {
-    GeneralReg(GeneralReg),
-    FloatReg(FloatReg),
-    Base {
-        offset: i32,
-        size: u32,
-        owned: bool,
-    },
-    BaseAndGeneralReg {
-        reg: GeneralReg,
-        offset: i32,
-        size: u32,
-        owned: bool,
-    },
-    BaseAndFloatReg {
-        reg: FloatReg,
-        offset: i32,
-        size: u32,
-        owned: bool,
-    },
-}
-
 pub trait RegTrait: Copy + PartialEq + Eq + std::hash::Hash + std::fmt::Debug + 'static {
     fn value(&self) -> u8;
 }
@@ -275,7 +252,6 @@ pub struct Backend64Bit<
     phantom_asm: PhantomData<ASM>,
     phantom_cc: PhantomData<CC>,
     env: &'a Env<'a>,
-    target_info: TargetInfo,
     interns: &'a mut Interns,
     helper_proc_gen: CodeGenHelp<'a>,
     helper_proc_symbols: Vec<'a, (Symbol, ProcLayout<'a>)>,
@@ -292,27 +268,6 @@ pub struct Backend64Bit<
     join_map: MutMap<JoinPointId, u64>,
 
     storage_manager: StorageManager<'a, GeneralReg, FloatReg, ASM, CC>,
-    symbol_storage_map: MutMap<Symbol, SymbolStorage<GeneralReg, FloatReg>>,
-
-    // This should probably be smarter than a vec.
-    // There are certain registers we should always use first. With pushing and popping, this could get mixed.
-    general_free_regs: Vec<'a, GeneralReg>,
-    float_free_regs: Vec<'a, FloatReg>,
-
-    // The last major thing we need is a way to decide what reg to free when all of them are full.
-    // Theoretically we want a basic lru cache for the currently loaded symbols.
-    // For now just a vec of used registers and the symbols they contain.
-    general_used_regs: Vec<'a, (GeneralReg, Symbol)>,
-    float_used_regs: Vec<'a, (FloatReg, Symbol)>,
-
-    // used callee saved regs must be tracked for pushing and popping at the beginning/end of the function.
-    general_used_callee_saved_regs: MutSet<GeneralReg>,
-    float_used_callee_saved_regs: MutSet<FloatReg>,
-
-    free_stack_chunks: Vec<'a, (i32, u32)>,
-    stack_size: u32,
-    // The amount of stack space needed to pass args for function calling.
-    fn_call_stack_size: u32,
 }
 
 /// new creates a new backend that will output to the specific Object.
@@ -331,7 +286,6 @@ pub fn new_backend_64bit<
         phantom_asm: PhantomData,
         phantom_cc: PhantomData,
         env,
-        target_info,
         interns,
         helper_proc_gen: CodeGenHelp::new(env.arena, target_info, env.module_id),
         helper_proc_symbols: bumpalo::vec![in env.arena],
@@ -342,18 +296,8 @@ pub fn new_backend_64bit<
         last_seen_map: MutMap::default(),
         layout_map: MutMap::default(),
         free_map: MutMap::default(),
-        symbol_storage_map: MutMap::default(),
         literal_map: MutMap::default(),
         join_map: MutMap::default(),
-        general_free_regs: bumpalo::vec![in env.arena],
-        general_used_regs: bumpalo::vec![in env.arena],
-        general_used_callee_saved_regs: MutSet::default(),
-        float_free_regs: bumpalo::vec![in env.arena],
-        float_used_regs: bumpalo::vec![in env.arena],
-        float_used_callee_saved_regs: MutSet::default(),
-        free_stack_chunks: bumpalo::vec![in env.arena],
-        stack_size: 0,
-        fn_call_stack_size: 0,
         storage_manager: storage::new_storage_manager(env, target_info),
     }
 }
@@ -388,25 +332,11 @@ impl<
     fn reset(&mut self, name: String, is_self_recursive: SelfRecursive) {
         self.proc_name = Some(name);
         self.is_self_recursive = Some(is_self_recursive);
-        self.stack_size = 0;
-        self.free_stack_chunks.clear();
-        self.fn_call_stack_size = 0;
         self.last_seen_map.clear();
         self.layout_map.clear();
         self.join_map.clear();
         self.free_map.clear();
-        self.symbol_storage_map.clear();
         self.buf.clear();
-        self.general_used_callee_saved_regs.clear();
-        self.general_free_regs.clear();
-        self.general_used_regs.clear();
-        self.general_free_regs
-            .extend_from_slice(CC::GENERAL_DEFAULT_FREE_REGS);
-        self.float_used_callee_saved_regs.clear();
-        self.float_free_regs.clear();
-        self.float_used_regs.clear();
-        self.float_free_regs
-            .extend_from_slice(CC::FLOAT_DEFAULT_FREE_REGS);
         self.helper_proc_symbols.clear();
         self.storage_manager.reset();
     }
@@ -435,13 +365,14 @@ impl<
         let mut out = bumpalo::vec![in self.env.arena];
 
         // Setup stack.
-        let mut used_regs = bumpalo::vec![in self.env.arena];
+        let used_general_regs = self.storage_manager.general_used_callee_saved_regs();
-        used_regs.extend(&self.general_used_callee_saved_regs);
+        let used_float_regs = self.storage_manager.float_used_callee_saved_regs();
         let aligned_stack_size = CC::setup_stack(
             &mut out,
-            &used_regs,
+            &used_general_regs,
-            self.stack_size as i32,
+            &used_float_regs,
-            self.fn_call_stack_size as i32,
+            self.storage_manager.stack_size() as i32,
+            self.storage_manager.fn_call_stack_size() as i32,
         );
         let setup_offset = out.len();
 
@@ -492,9 +423,10 @@ impl<
         // Cleanup stack.
         CC::cleanup_stack(
             &mut out,
-            &used_regs,
+            &used_general_regs,
+            &used_float_regs,
             aligned_stack_size,
-            self.fn_call_stack_size as i32,
+            self.storage_manager.fn_call_stack_size() as i32,
         );
         ASM::ret(&mut out);
 

compiler/gen_dev/src/generic64/storage.rs

@@ -176,6 +176,26 @@ impl<
         self.fn_call_stack_size = 0;
     }
 
+    pub fn stack_size(&self) -> u32 {
+        self.stack_size
+    }
+
+    pub fn fn_call_stack_size(&self) -> u32 {
+        self.fn_call_stack_size
+    }
+
+    pub fn general_used_callee_saved_regs(&self) -> Vec<'a, GeneralReg> {
+        let mut used_regs = bumpalo::vec![in self.env.arena];
+        used_regs.extend(&self.general_used_callee_saved_regs);
+        used_regs
+    }
+
+    pub fn float_used_callee_saved_regs(&self) -> Vec<'a, FloatReg> {
+        let mut used_regs = bumpalo::vec![in self.env.arena];
+        used_regs.extend(&self.float_used_callee_saved_regs);
+        used_regs
+    }
+
     // Returns true if the symbol is storing a primitive value.
     pub fn is_stored_primitive(&self, sym: &Symbol) -> bool {
         matches!(
@@ -250,6 +270,7 @@ impl<
         self.general_free_regs.push(reg);
     }
 
+    #[allow(dead_code)]
     // This claims a temporary float register and enables is used in the passed in function.
     // Temporary registers are not safe across call instructions.
     pub fn with_tmp_float_reg<F: FnOnce(&mut Self, &mut Vec<'a, u8>, FloatReg)>(
@@ -634,7 +655,7 @@ impl<
                 match self
                     .general_used_regs
                     .iter()
-                    .position(|(used_reg, sym)| reg == *used_reg)
+                    .position(|(used_reg, _sym)| reg == *used_reg)
                 {
                     Some(position) => {
                         let (_, sym) = self.general_used_regs.remove(position);
@@ -652,7 +673,7 @@ impl<
                 match self
                     .float_used_regs
                     .iter()
-                    .position(|(used_reg, sym)| reg == *used_reg)
+                    .position(|(used_reg, _sym)| reg == *used_reg)
                 {
                     Some(position) => {
                         let (_, sym) = self.float_used_regs.remove(position);
@@ -790,7 +811,7 @@ impl<
             // Claim a location for every join point parameter to be loaded at.
             match layout {
                 single_register_integers!() => {
-                    let reg = self.claim_general_reg(buf, symbol);
+                    self.claim_general_reg(buf, symbol);
                 }
                 single_register_floats!() => {
                     self.claim_float_reg(buf, symbol);

compiler/gen_dev/src/generic64/x86_64.rs

@@ -163,13 +163,15 @@ impl CallConv<X86_64GeneralReg, X86_64FloatReg, X86_64Assembler> for X86_64Syste
     #[inline(always)]
     fn setup_stack<'a>(
         buf: &mut Vec<'a, u8>,
-        general_saved_regs: &[X86_64GeneralReg],
+        saved_general_regs: &[X86_64GeneralReg],
+        saved_float_regs: &[X86_64FloatReg],
         requested_stack_size: i32,
         fn_call_stack_size: i32,
     ) -> i32 {
         x86_64_generic_setup_stack(
             buf,
-            general_saved_regs,
+            saved_general_regs,
+            saved_float_regs,
             requested_stack_size,
             fn_call_stack_size,
         )
@@ -178,13 +180,15 @@ impl CallConv<X86_64GeneralReg, X86_64FloatReg, X86_64Assembler> for X86_64Syste
     #[inline(always)]
     fn cleanup_stack<'a>(
         buf: &mut Vec<'a, u8>,
-        general_saved_regs: &[X86_64GeneralReg],
+        saved_general_regs: &[X86_64GeneralReg],
+        saved_float_regs: &[X86_64FloatReg],
         aligned_stack_size: i32,
         fn_call_stack_size: i32,
     ) {
         x86_64_generic_cleanup_stack(
             buf,
-            general_saved_regs,
+            saved_general_regs,
+            saved_float_regs,
             aligned_stack_size,
             fn_call_stack_size,
         )
@@ -514,21 +518,35 @@ impl CallConv<X86_64GeneralReg, X86_64FloatReg, X86_64Assembler> for X86_64Windo
     #[inline(always)]
     fn setup_stack<'a>(
         buf: &mut Vec<'a, u8>,
-        saved_regs: &[X86_64GeneralReg],
+        saved_general_regs: &[X86_64GeneralReg],
+        saved_float_regs: &[X86_64FloatReg],
         requested_stack_size: i32,
         fn_call_stack_size: i32,
     ) -> i32 {
-        x86_64_generic_setup_stack(buf, saved_regs, requested_stack_size, fn_call_stack_size)
+        x86_64_generic_setup_stack(
+            buf,
+            saved_general_regs,
+            saved_float_regs,
+            requested_stack_size,
+            fn_call_stack_size,
+        )
     }
 
     #[inline(always)]
     fn cleanup_stack<'a>(
         buf: &mut Vec<'a, u8>,
-        saved_regs: &[X86_64GeneralReg],
+        saved_general_regs: &[X86_64GeneralReg],
+        saved_float_regs: &[X86_64FloatReg],
         aligned_stack_size: i32,
         fn_call_stack_size: i32,
     ) {
-        x86_64_generic_cleanup_stack(buf, saved_regs, aligned_stack_size, fn_call_stack_size)
+        x86_64_generic_cleanup_stack(
+            buf,
+            saved_general_regs,
+            saved_float_regs,
+            aligned_stack_size,
+            fn_call_stack_size,
+        )
     }
 
     #[inline(always)]
@@ -706,7 +724,8 @@ impl X86_64WindowsFastcall {
 #[inline(always)]
 fn x86_64_generic_setup_stack<'a>(
     buf: &mut Vec<'a, u8>,
-    saved_regs: &[X86_64GeneralReg],
+    saved_general_regs: &[X86_64GeneralReg],
+    saved_float_regs: &[X86_64FloatReg],
     requested_stack_size: i32,
     fn_call_stack_size: i32,
 ) -> i32 {
@@ -714,7 +733,7 @@ fn x86_64_generic_setup_stack<'a>(
     X86_64Assembler::mov_reg64_reg64(buf, X86_64GeneralReg::RBP, X86_64GeneralReg::RSP);
 
     let full_stack_size = match requested_stack_size
-        .checked_add(8 * saved_regs.len() as i32)
+        .checked_add(8 * (saved_general_regs.len() + saved_float_regs.len()) as i32)
         .and_then(|size| size.checked_add(fn_call_stack_size))
     {
         Some(size) => size,
@@ -741,10 +760,14 @@ fn x86_64_generic_setup_stack<'a>(
 
             // Put values at the top of the stack to avoid conflicts with previously saved variables.
             let mut offset = aligned_stack_size - fn_call_stack_size;
-            for reg in saved_regs {
+            for reg in saved_general_regs {
                 X86_64Assembler::mov_base32_reg64(buf, -offset, *reg);
                 offset -= 8;
             }
+            for reg in saved_float_regs {
+                X86_64Assembler::mov_base32_freg64(buf, -offset, *reg);
+                offset -= 8;
+            }
             aligned_stack_size
         } else {
             0
@@ -758,16 +781,21 @@ fn x86_64_generic_setup_stack<'a>(
 #[allow(clippy::unnecessary_wraps)]
 fn x86_64_generic_cleanup_stack<'a>(
     buf: &mut Vec<'a, u8>,
-    saved_regs: &[X86_64GeneralReg],
+    saved_general_regs: &[X86_64GeneralReg],
+    saved_float_regs: &[X86_64FloatReg],
     aligned_stack_size: i32,
     fn_call_stack_size: i32,
 ) {
     if aligned_stack_size > 0 {
         let mut offset = aligned_stack_size - fn_call_stack_size;
-        for reg in saved_regs {
+        for reg in saved_general_regs {
             X86_64Assembler::mov_reg64_base32(buf, *reg, -offset);
             offset -= 8;
         }
+        for reg in saved_float_regs {
+            X86_64Assembler::mov_freg64_base32(buf, *reg, -offset);
+            offset -= 8;
+        }
         X86_64Assembler::add_reg64_reg64_imm32(
             buf,
             X86_64GeneralReg::RSP,