From 5cb207d8b37d171718e0bf76d201d5cfbb728a3a Mon Sep 17 00:00:00 2001
From: Mark <markjordanthom@gmail.com>
Date: Tue, 28 Nov 2023 14:38:48 -0700
Subject: [PATCH] allocate cut variables properly (#2183)

---
 src/allocator.rs        |  7 +------
 src/codegen.rs          |  8 ++------
 src/debray_allocator.rs | 25 +++++++++++++++----------
 3 files changed, 18 insertions(+), 22 deletions(-)

diff --git a/src/allocator.rs b/src/allocator.rs
index c2300f90..e20bf2cf 100644
--- a/src/allocator.rs
+++ b/src/allocator.rs
@@ -36,12 +36,7 @@ pub(crate) trait Allocator {
         is_new_var: bool,
     );
 
-    fn mark_cut_var<'a, Target: CompilationTarget<'a>>(
-        &mut self,
-        var_num: usize,
-        term_loc: GenContext,
-        code: &mut CodeDeque,
-    ) -> RegType;
+    fn mark_cut_var(&mut self, var_num: usize, chunk_num: usize) -> RegType;
 
     fn mark_var<'a, Target: CompilationTarget<'a>>(
         &mut self,
diff --git a/src/codegen.rs b/src/codegen.rs
index 993a122e..09722fbd 100644
--- a/src/codegen.rs
+++ b/src/codegen.rs
@@ -926,16 +926,12 @@ impl<'b> CodeGenerator<'b> {
                         match term {
                             &QueryTerm::GetLevel(var_num) => {
                                 let code = branch_code_stack.code(code);
-                                let r = self
-                                    .marker
-                                    .mark_cut_var::<QueryInstruction>(var_num, term_loc, code);
+                                let r = self.marker.mark_cut_var(var_num, chunk_num);
                                 code.push_back(instr!("get_level", r));
                             }
                             &QueryTerm::GetCutPoint { var_num, prev_b } => {
                                 let code = branch_code_stack.code(code);
-                                let r = self
-                                    .marker
-                                    .mark_cut_var::<FactInstruction>(var_num, term_loc, code);
+                                let r = self.marker.mark_cut_var(var_num, chunk_num);
 
                                 code.push_back(if prev_b {
                                     instr!("get_prev_level", r)
diff --git a/src/debray_allocator.rs b/src/debray_allocator.rs
index 7154cea7..8bf858c3 100644
--- a/src/debray_allocator.rs
+++ b/src/debray_allocator.rs
@@ -840,18 +840,23 @@ impl Allocator for DebrayAllocator {
         self.in_use.insert(o);
     }
 
-    fn mark_cut_var<'a, Target: CompilationTarget<'a>>(
-        &mut self,
-        var_num: usize,
-        term_loc: GenContext,
-        code: &mut CodeDeque,
-    ) -> RegType {
+    fn mark_cut_var(&mut self, var_num: usize, chunk_num: usize) -> RegType {
         match self.get_binding(var_num) {
-            RegType::Perm(0) => RegType::Perm(self.alloc_perm_var(var_num, term_loc.chunk_num())),
+            RegType::Perm(0) => RegType::Perm(self.alloc_perm_var(var_num, chunk_num)),
             RegType::Temp(0) => {
-                let cell = Cell::default();
-                self.mark_var::<Target>(var_num, Level::Shallow, &cell, term_loc, code);
-                cell.get().norm()
+                let t = self.alloc_reg_to_non_var();
+
+                match &mut self.var_data.records[var_num].allocation {
+                    VarAlloc::Temp {
+                        temp_reg, safety, ..
+                    } => {
+                        *temp_reg = t;
+                        *safety = VarSafetyStatus::GloballyUnneeded;
+                    }
+                    _ => unreachable!(),
+                };
+
+                RegType::Temp(t)
             }
             r => r,
         }
-- 
2.54.0