fix double pull on tmp reg

src/riscv/codegen.mbt

@@ -195,7 +195,11 @@ fn CodegenBlock::assign_f(
   self.insert_asms(todo(target))
 }
 
-fn CodegenBlock::pull_val_f(self : CodegenBlock, val : Value) -> FReg {
+fn CodegenBlock::pull_val_f(
+  self : CodegenBlock,
+  val : Value,
+  ~freg_swap : FReg = freg_swap
+) -> FReg {
   match val {
     Var(var) => {
       guard let F(reg) = self.allocation[var].unwrap() else {
@@ -217,7 +221,11 @@ fn CodegenBlock::pull_val_f(self : CodegenBlock, val : Value) -> FReg {
   }
 }
 
-fn CodegenBlock::pull_val_i(self : CodegenBlock, val : Value) -> Reg {
+fn CodegenBlock::pull_val_i(
+  self : CodegenBlock,
+  val : Value,
+  ~reg_swap : Reg = reg_swap
+) -> Reg {
   match val {
     Var(var) => {
       if var.ty.val == Unit {
@@ -453,9 +461,8 @@ fn CodegenBlock::resolve_arg_regs(
 }
 
 fn CodegenBlock::new_label(self : CodegenBlock, prefix : String) -> String {
-  let result = "\{prefix}_\{self.cfg.counter}"
-  self.cfg.counter += 1
-  result
+  let var = self.cfg.new_named(prefix)
+  var.to_string()
 }
 
 // PERF: proritize the use of saved regs in stead of putting onto stack
@@ -585,6 +592,17 @@ fn CodegenBlock::codegen(self : CodegenBlock) -> Unit {
   for inst in block.insts {
     self.insert_asm(Comment(inst.to_string()))
     match inst {
+      Copy(bind, copied) =>
+        match get_reg_ty(copied) {
+          F64 => {
+            let freg_copied = self.pull_val_f(copied)
+            self.assign_f(bind, fn(freg) { [FmvD(freg, freg_copied)] })
+          }
+          _ => {
+            let reg_copied = self.pull_val_i(copied)
+            self.assign_i(bind, fn(reg) { [Mv(reg, reg_copied)] })
+          }
+        }
       Load(var) => {
         let offset = self.spilled_offset[var].unwrap()
         match get_reg_ty(Var(var)) {
@@ -749,12 +767,14 @@ fn CodegenBlock::codegen(self : CodegenBlock) -> Unit {
         let reg_ptr = self.pull_val_i(ptr)
         self.store_val(rhs, { base: reg_ptr, offset: 0 })
       }
-      Prim(bind, Math(op, Int), [lhs, rhs]) => {
-        let reg_lhs = self.pull_val_i(lhs)
-        let reg_rhs = self.pull_val_i(rhs)
+      Prim(bind, Math(op, Int), [lhs, rhs]) =>
         self.assign_i(
           bind,
           fn(reg) {
+            let reg_lhs = self.pull_val_i(lhs)
+            // NOTE: whenever we pull 2 regs and both of them maybe tmp, 
+            // we replace the second tmp with target
+            let reg_rhs = self.pull_val_i(rhs, reg_swap=reg)
             match op {
               Add => [Add(reg, reg_lhs, reg_rhs)]
               Sub => [Sub(reg, reg_lhs, reg_rhs)]
@@ -763,14 +783,12 @@ fn CodegenBlock::codegen(self : CodegenBlock) -> Unit {
             }
           },
         )
-      }
-      // BUG: any binop is prune to using same reg
-      Prim(bind, Math(op, Double), [lhs, rhs]) => {
-        let reg_lhs = self.pull_val_f(lhs)
-        let reg_rhs = self.pull_val_f(rhs)
+      Prim(bind, Math(op, Double), [lhs, rhs]) =>
         self.assign_f(
           bind,
           fn(reg) {
+            let reg_lhs = self.pull_val_f(lhs)
+            let reg_rhs = self.pull_val_f(rhs, freg_swap=reg)
             match op {
               Add => [FaddD(reg, reg_lhs, reg_rhs)]
               Sub => [FsubD(reg, reg_lhs, reg_rhs)]
@@ -779,20 +797,25 @@ fn CodegenBlock::codegen(self : CodegenBlock) -> Unit {
             }
           },
         )
-      }
       Prim(bind, Eq, [lhs, rhs]) =>
         match lhs.get_type() {
           Double => {
             let reg_lhs = self.pull_val_f(lhs)
             let reg_rhs = self.pull_val_f(rhs)
+            // NOTE: since we've reserve freg first, this won't happen
+            if reg_lhs == reg_rhs {
+              @util.die("pulling same reg for comparing floats")
+            }
             self.assign_i(bind, fn(reg) { [FeqD(reg, reg_lhs, reg_rhs)] })
           }
           _ => {
             let reg_lhs = self.pull_val_i(lhs)
-            let reg_rhs = self.pull_val_i(rhs)
             self.assign_i(
               bind,
-              fn(reg) { [Xor(reg_swap, reg_lhs, reg_rhs), Seqz(reg, reg_swap)] },
+              fn(reg) {
+                let reg_rhs = self.pull_val_i(rhs, reg_swap=reg)
+                [Xor(reg_swap, reg_lhs, reg_rhs), Seqz(reg, reg_swap)]
+              },
             )
           }
         }
@@ -801,15 +824,19 @@ fn CodegenBlock::codegen(self : CodegenBlock) -> Unit {
           Double => {
             let reg_lhs = self.pull_val_f(lhs)
             let reg_rhs = self.pull_val_f(rhs)
+            // NOTE: since we've reserve freg first, this won't happen
+            if reg_lhs == reg_rhs {
+              @util.die("pulling same reg for comparing floats")
+            }
             self.assign_i(bind, fn(reg) { [FleD(reg, reg_lhs, reg_rhs)] })
           }
           _ => {
             let reg_lhs = self.pull_val_i(lhs)
-            let reg_rhs = self.pull_val_i(rhs)
             self.assign_i(
               bind,
               // HACK: we're using 32bits int so this won't overflow
               fn(reg) {
+                let reg_rhs = self.pull_val_i(rhs, reg_swap=reg)
                 [
                   Sub(reg, reg_lhs, reg_rhs),
                   Addi(reg, reg, -1),

src/riscv/collect_labels.mbt

@@ -33,7 +33,7 @@ fn collect_externals(cfg : @ssacfg.SsaCfg) -> ExternalLabels {
           collect_label_var(bind)
           vals.each(collect_label_val)
         }
-        KthTuple(bind, val, _) => {
+        KthTuple(bind, val, _) | Copy(bind, val) => {
           collect_label_var(bind)
           collect_label_val(val)
         }

src/riscv/emit.mbt

@@ -3,15 +3,13 @@ let heap_ptr_label : Ref[String] = { val: "" }
 let stub_label : Ref[String] = { val: "" }
 
 pub fn emit(cfg : @ssacfg.SsaCfg) -> Array[AssemblyFunction] {
+  let cfg = reserve_freg(cfg)
   let output = []
-  heap_ptr_label.val = "heap_\{cfg.counter}"
+  heap_ptr_label.val = cfg.new_named("heap").to_string()
 
   // generate stub for CPS / C-Calling convention interop
-  cfg.counter += 1
-  stub_label.val = "c_stub_\{cfg.counter}"
-  cfg.counter += 1
-  let stub_resolved_i = "c_stub_resolved_i_\{cfg.counter}"
-  cfg.counter += 1
+  stub_label.val = cfg.new_named("c_stub").to_string()
+  let stub_resolved_i = cfg.new_named("c_stub_resolved_i").to_string()
   let body : Array[RvAsm] = []
   let externals = collect_externals(cfg)
   // generate stubs for function returning ints
@@ -58,8 +56,7 @@ pub fn emit(cfg : @ssacfg.SsaCfg) -> Array[AssemblyFunction] {
     ],
   )
   // generates stub for function returning floats
-  let stub_resolved_f = "c_stub_resolved_f_\{cfg.counter}"
-  cfg.counter += 1
+  let stub_resolved_f = cfg.new_named("c_stub_resolved_f").to_string()
   for external in externals.returns_f {
     guard let @typing.Type::Fun(args, _) = external.ty.val else {
       _ => @util.die("external non function \{external}")

src/riscv/interference_graph_build.mbt

@@ -105,6 +105,10 @@ fn LiveVarAnalysis::collect_inst(
       self.var_set.remove(bind)
       self.collect_val(tup)
     }
+    Copy(bind, copied) => {
+      self.var_set.remove(bind)
+      self.collect_val(copied)
+    }
     Prim(bind, _, args) => {
       self.var_set.remove(bind)
       args.each(fn(v) { self.collect_val(v) })

src/riscv/reg_spill.mbt

@@ -54,6 +54,15 @@ fn reg_spill_block(blk : @ssacfg.Block, spilled_var : Var) -> @ssacfg.Block {
           insts_new.push(Store(spilled_var))
         }
       }
+      Copy(bind, copied) => {
+        if val_spilled(copied) {
+          insts_new.push(Load(spilled_var))
+        }
+        insts_new.push(inst)
+        if bind == spilled_var {
+          insts_new.push(Store(spilled_var))
+        }
+      }
       Prim(bind, _, args) => {
         if vals_spilled(args) {
           insts_new.push(Load(spilled_var))

src/riscv/reserve_freg.mbt

@@ -0,0 +1,35 @@
+fn reserve_freg(cfg : @ssacfg.SsaCfg) -> @ssacfg.SsaCfg {
+  for item in cfg.blocks {
+    let (_, block) = item
+    let insts_backup = block.insts
+    let insts = []
+    for inst in insts_backup {
+      match inst {
+        Prim(bind, Eq | Le as op, [lhs, rhs]) as inst =>
+          match lhs.get_type() {
+            Double => {
+              let lhs_is_imm = match lhs {
+                Double(_) => true
+                _ => false
+              }
+              let rhs_is_imm = match rhs {
+                Double(_) => true
+                _ => false
+              }
+              if lhs_is_imm && rhs_is_imm {
+                let tmp = cfg.new_named("tmp", ty=Double)
+                insts.push(@ssacfg.Inst::Copy(tmp, lhs))
+                insts.push(Prim(bind, op, [Var(tmp), rhs]))
+              } else {
+                insts.push(Prim(bind, op, [lhs, rhs]))
+              }
+            }
+            _ => insts.push(inst)
+          }
+        inst => insts.push(inst)
+      }
+    }
+    block.insts = insts
+  }
+  cfg
+}

src/ssacfg/ssa_ir.mbt

@@ -10,14 +10,16 @@ pub enum PrimOp {
 } derive(Show)
 
 pub enum Inst {
-  // cps ir
+  // Cps IR
   MakeTuple(Var, Array[Value])
   KthTuple(Var, Value, Int)
-  // primitives
+  // Primitives
   Prim(Var, PrimOp, Array[Value])
-  // for reg allocation spilling
+  // For reg allocation spilling
   Store(Var)
   Load(Var)
+  // So we can deal with the case where 1 tmp reg is not enough
+  Copy(Var, Value)
 } derive(Show)
 
 pub enum PCInst {
@@ -32,7 +34,7 @@ pub enum PCInst {
 pub struct Block {
   from : Var?
   fn_name : Var // for TCO
-  insts : Array[Inst]
+  mut insts : Array[Inst] //
   last_inst : Ref[PCInst]
 } derive(Show)
 
@@ -63,7 +65,7 @@ fn SsaCfg::new_fn_block(
   { fn_name, from: None, insts: [], last_inst: { val: Exit } }
 }
 
-fn SsaCfg::new_named(
+pub fn SsaCfg::new_named(
   self : SsaCfg,
   name : String,
   ~ty : @typing.Type = Unit