From d1bcc8d55bce5197a6023d5fa8273576b5ca1fd3 Mon Sep 17 00:00:00 2001
From: Brendan Hansknecht <brendan.hansknecht@gmail.com>
Date: Wed, 16 Feb 2022 19:19:35 -0800
Subject: [PATCH] start storage rewrite

---
 compiler/gen_dev/src/generic64/mod.rs     |  1 +
 compiler/gen_dev/src/generic64/storage.rs | 87 +++++++++++++++++++++++
 2 files changed, 88 insertions(+)
 create mode 100644 compiler/gen_dev/src/generic64/storage.rs

diff --git a/compiler/gen_dev/src/generic64/mod.rs b/compiler/gen_dev/src/generic64/mod.rs
index f30e98edc3..6ebf85c6ee 100644
--- a/compiler/gen_dev/src/generic64/mod.rs
+++ b/compiler/gen_dev/src/generic64/mod.rs
@@ -11,6 +11,7 @@ use roc_target::TargetInfo;
 use std::marker::PhantomData;
 
 pub mod aarch64;
+mod storage;
 pub mod x86_64;
 
 const TARGET_INFO: TargetInfo = TargetInfo::default_x86_64();
diff --git a/compiler/gen_dev/src/generic64/storage.rs b/compiler/gen_dev/src/generic64/storage.rs
new file mode 100644
index 0000000000..148df8d55a
--- /dev/null
+++ b/compiler/gen_dev/src/generic64/storage.rs
@@ -0,0 +1,87 @@
+use crate::generic64::{CallConv, RegTrait};
+use bumpalo::collections::Vec;
+use roc_collections::all::{MutMap, MutSet};
+use roc_error_macros::internal_error;
+use roc_module::symbol::Symbol;
+use std::marker::PhantomData;
+use std::rc::Rc;
+
+#[derive(Clone, Debug, PartialEq)]
+enum RegStorage<GeneralReg: RegTrait, FloatReg: RegTrait> {
+    General(GeneralReg),
+    Float(FloatReg),
+}
+
+#[derive(Clone, Debug, PartialEq)]
+struct StackStorage<'a, GeneralReg: RegTrait, FloatReg: RegTrait> {
+    // The parent of this stack storage.
+    // If the data is owned, there is none.
+    // If this is nested in another stack storage, it will have a parent.
+    parent: Option<Rc<StackStorage<'a, GeneralReg, FloatReg>>>,
+    // Offset from the base pointer in bytes.
+    base_offset: i32,
+    // Size on the stack in bytes.
+    size: u32,
+    // Values of this storage currently loaded into registers.
+    // This is stored in tuples of (offset from start of this storage, register).
+    // This save on constantly reloading a list pointer for example.
+    // TODO: add small vec optimization most of the time this should be 0 or 1 items.
+    refs: Vec<'a, (u32, RegStorage<GeneralReg, FloatReg>)>,
+}
+
+#[derive(Clone, Debug, PartialEq)]
+enum Storage<'a, GeneralReg: RegTrait, FloatReg: RegTrait> {
+    Reg(RegStorage<GeneralReg, FloatReg>),
+    Stack(StackStorage<'a, GeneralReg, FloatReg>),
+}
+
+struct StorageManager<
+    'a,
+    GeneralReg: RegTrait,
+    FloatReg: RegTrait,
+    CC: CallConv<GeneralReg, FloatReg>,
+> {
+    phantom_cc: PhantomData<CC>,
+    symbol_storage_map: MutMap<Symbol, Storage<'a, 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,
+}
+
+impl<'a, FloatReg: RegTrait, GeneralReg: RegTrait, CC: CallConv<GeneralReg, FloatReg>>
+    StorageManager<'a, GeneralReg, FloatReg, CC>
+{
+    // This claims a temporary register and enables is used in the passed in function.
+    // Temporary registers are not safe across call instructions.
+    fn with_tmp_general_reg<F: FnOnce(&mut Self, GeneralReg)>(&mut self, callback: F) {
+        let reg = if let Some(reg) = self.general_free_regs.pop() {
+            if CC::general_callee_saved(&reg) {
+                self.general_used_callee_saved_regs.insert(reg);
+            }
+            reg
+        } else if !self.general_used_regs.is_empty() {
+            let (reg, sym) = self.general_used_regs.remove(0);
+            // self.free_to_stack(&sym);
+            reg
+        } else {
+            internal_error!("completely out of general purpose registers");
+        };
+        callback(self, reg);
+        self.general_free_regs.push(reg);
+    }
+}