[REWRITE] added back prolog/epilog

build.rs

@@ -8,7 +8,7 @@ fn visit_dirs(dir: &Path) -> io::Result<Vec<PathBuf>> {
             let entry = entry?;
             let path = entry.path();
             if path.is_dir() {
-                paths.extend_from_slice(&visit_dirs(&path)?);
+                paths.extend(visit_dirs(&path)?);
             } else {
                 if entry.path().extension().map_or(false, |ext| ext == "td") {
                     paths.push(entry.path());

src/Optimizations/Analysis/CfgAnalysis.rs

@@ -73,7 +73,7 @@ impl CFGAnalysis {
         for (block, mut preds) in self.predecessors.clone() {
             let new_preds = self.preds_rec(&block, &mut preds);
 
-            self.predecessors.get_mut(&block).unwrap().extend_from_slice(&new_preds); 
+            self.predecessors.get_mut(&block).unwrap().extend(new_preds); 
         }
     }
 

src/Target/instr.rs

@@ -21,4 +21,10 @@ impl Display for Box<dyn McInstr> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(f, "{}", self.asm())
     }
+}
+
+impl PartialEq for Box<dyn McInstr> {
+    fn eq(&self, other: &Self) -> bool {
+        self.asm() == other.asm()
+    }
 }

src/Target/registry.rs

@@ -125,13 +125,15 @@ impl TargetRegistry {
 
             let mut encoded = Vec::new();
             for instr in instrs {
-                encoded.extend_from_slice(&instr.encode());
+                encoded.extend(instr.encode());
 
                 if let Some(mut reloc) = instr.relocation() {
-                    reloc.at += encoded.len();
+                    reloc.at += encoded.len() + machine_code_len - 1;
 
                     reloc.from = func.name.to_owned();
 
+                    println!("{reloc:?}");
+
                     links.push(reloc);
                 }
 
@@ -210,15 +212,18 @@ impl TargetRegistry {
             let mut first = true;
 
             for (block, insts) in block_instrs {
-                printer.begin_block(&block.name, !first);
+                if !block.name.as_str().contains("the epilog") && !block.name.as_str().contains("the prolog") {
+                    printer.begin_block(&block.name, !first);
+                    first = false;
+                } else {
+                    printer.begin_block(&block.name, false);
+                }
 
                 for inst in insts {
                     printer.print_inst(&inst);
                 }
 
                 printer.end_block();
-
-                first = false;
             }
 
             printer.end_func();

src/Target/x86/asm/encode.rs

@@ -61,6 +61,9 @@ impl X86Instr {
             X86Mnemonic::Cdq => instr.encode_cdq(),
             X86Mnemonic::Cqo => instr.encode_cqo(),
             X86Mnemonic::Call => instr.encode_call(),
+            X86Mnemonic::Movq => instr.encode_movq(),
+            X86Mnemonic::Push => instr.encode_push(),
+            X86Mnemonic::Pop => instr.encode_pop(),
         }
     }
 
@@ -1225,6 +1228,63 @@ impl X86Instr {
             BlockEncoderOptions::DONT_FIX_BRANCHES
         ).expect("encoding error").code_buffer
     }
+
+    fn encode_movq(&self) -> Vec<u8> {
+        todo!()
+    }
+
+    fn encode_push(&self) -> Vec<u8> {
+        let operand = self.op1.expect("expected op1");
+
+        let instr = match operand {
+            X86Operand::Reg(reg) => match reg.size {
+                X86RegSize::Word => Instruction::with1::<iced_x86::Register>(Code::Push_r16, reg.into()),
+                X86RegSize::Dword => Instruction::with1::<iced_x86::Register>(Code::Push_r32, reg.into()),
+                X86RegSize::Qword => Instruction::with1::<iced_x86::Register>(Code::Push_r64, reg.into()),
+                X86RegSize::SimdVec | X86RegSize::Byte => panic!("invalid size for push"),
+            },
+            X86Operand::Const(constant) => Instruction::with1(Code::Push_imm16, constant as i32),
+            X86Operand::MemDispl(mem) => match mem.size {
+                X86RegSize::Word => Instruction::with1::<iced_x86::MemoryOperand>(Code::Push_rm16, mem.into()),
+                X86RegSize::Dword => Instruction::with1::<iced_x86::MemoryOperand>(Code::Push_rm32, mem.into()),
+                X86RegSize::Qword => Instruction::with1::<iced_x86::MemoryOperand>(Code::Push_rm64, mem.into()),
+                X86RegSize::SimdVec | X86RegSize::Byte => panic!("invalid size for push"),
+            },
+            _ => panic!("invalid variant: {self}"),
+        }.expect("invalid instruction");
+
+        BlockEncoder::encode(
+            64, 
+            InstructionBlock::new(&[instr], 0), 
+            BlockEncoderOptions::DONT_FIX_BRANCHES
+        ).expect("encoding error").code_buffer
+    }
+
+    fn encode_pop(&self) -> Vec<u8> {
+        let operand = self.op1.expect("expected op1");
+
+        let instr = match operand {
+            X86Operand::Reg(reg) => match reg.size {
+                X86RegSize::Word => Instruction::with1::<iced_x86::Register>(Code::Pop_r16, reg.into()),
+                X86RegSize::Dword => Instruction::with1::<iced_x86::Register>(Code::Pop_r32, reg.into()),
+                X86RegSize::Qword => Instruction::with1::<iced_x86::Register>(Code::Pop_r64, reg.into()),
+                X86RegSize::SimdVec | X86RegSize::Byte => panic!("invalid size for pop"),
+            },
+            X86Operand::MemDispl(mem) => match mem.size {
+                X86RegSize::Word => Instruction::with1::<iced_x86::MemoryOperand>(Code::Pop_rm16, mem.into()),
+                X86RegSize::Dword => Instruction::with1::<iced_x86::MemoryOperand>(Code::Pop_rm32, mem.into()),
+                X86RegSize::Qword => Instruction::with1::<iced_x86::MemoryOperand>(Code::Pop_rm64, mem.into()),
+                X86RegSize::SimdVec | X86RegSize::Byte => panic!("invalid size for pop"),
+            },
+            _ => panic!("invalid variant: {self}"),
+        }.expect("invalid instruction");
+
+        BlockEncoder::encode(
+            64, 
+            InstructionBlock::new(&[instr], 0), 
+            BlockEncoderOptions::DONT_FIX_BRANCHES
+        ).expect("encoding error").code_buffer
+    }
 }
 
 impl Into<iced_x86::Register> for X86Reg {

src/Target/x86/asm/mod.rs

@@ -26,54 +26,42 @@ pub enum X86Mnemonic {
     Movss,
     Movsd,
     Movdqa,
-
     Ret,
-
     Add,
     Addss,
     Paddq,
     Paddd,
-
     Sub,
     Psubq,
     Psubd,
     Psubw,
     Psubb,
-
     Lea,
-
     Jmp,
     Je,
-
     Sete,
     Setne,
     Setl,
     Setle,
     Setg,
     Setge,
     Cmp,
-
     Pinsrb,
     Pinsrw,
     Pinsrd,
     Pinsrq,
     Insertps,
-
     Imul,
-
     And,
     Or,
     Xor,
     Sar,
     Shr,
     Sal,
     Shl,
-
     Neg,
-
     Movsx,
     Movsxd,
-
     Cvtsi2sd,
     Cvtsi2ss,
     Cvtss2si,
@@ -87,6 +75,9 @@ pub enum X86Mnemonic {
     Cdq,
     Cqo,
     Call,
+    Movq,
+    Push,
+    Pop,
 }
 
 /// A X86 assembly operand
@@ -236,6 +227,9 @@ impl std::fmt::Display for X86Instr {
             X86Mnemonic::Cdq => "cdq",
             X86Mnemonic::Cqo => "cqo",
             X86Mnemonic::Call => "call",
+            X86Mnemonic::Movq => "movq",
+            X86Mnemonic::Push => "push",
+            X86Mnemonic::Pop => "pop",
         })?;
 
         if let Some(op) = &self.op1 {
@@ -300,8 +294,6 @@ impl McInstr for X86Instr {
 
             let Some(rel) = mem.rip_rel else { unreachable!() };
 
-
-
             return Some(crate::Obj::Link {
                 from: String::new(),
                 to: crate::Target::x86::get_rel(rel),
@@ -314,4 +306,10 @@ impl McInstr for X86Instr {
 
         None
     }
+}
+
+impl Into<Box<dyn McInstr>> for X86Instr {
+    fn into(self) -> Box<dyn McInstr> {
+        Box::new( self.to_owned() )
+    }
 }

src/Target/x86/lower.rs

@@ -1,14 +1,15 @@
 //use opt::X86BasicOpt;
 
-use std::collections::HashMap;
+use std::collections::{HashMap, VecDeque};
 
 use crate::ydbg;
 use crate::CodeGen::dag::DagNode;
 use crate::CodeGen::reg::TargetReg;
 use crate::CodeGen::regalloc_iterated_col::ItRegCoalAlloc;
-use crate::IR::BlockId;
+use crate::IR::{BlockId, TypeMetadata};
 use crate::{CodeGen::dag, Target::instr::McInstr};
 use super::asm::*;
+use super::reg::X86Reg;
 
 #[allow(warnings)]
 mod auto_gen {
@@ -133,7 +134,7 @@ mod auto_gen {
 }
 
 pub(super) fn x86_lower(func: &mut dag::DagFunction, alloc: &mut ItRegCoalAlloc, module: &mut crate::IR::Module) -> Vec<(BlockId, Vec<Box<dyn McInstr>>)> {
-    let mut blocks = Vec::new();
+    let mut blocks = VecDeque::new();
 
     for (name, nodes) in &mut func.blocks {
         let mut asm: Vec<X86Instr> = Vec::new();
@@ -142,18 +143,37 @@ pub(super) fn x86_lower(func: &mut dag::DagFunction, alloc: &mut ItRegCoalAlloc,
 
             let mut overwrittes = HashMap::new();
 
-            for overwrite in auto_gen::overwrittes(&node) {
+            for overwrite in x86_overwrittesr(&node) {
                 // 1. Check if the value which is overwritten is currently in use
-                if !alloc.regs.contains(&crate::CodeGen::reg::Reg::new_x86(overwrite)) {
-                    continue;
+                let mut is_used = true;
+
+                for reg in &alloc.regs {
+                    if let crate::CodeGen::reg::TargetReg::X86(x86) = &reg.reg {
+                        if x86.variant == overwrite.variant {
+                            ydbg!("alloc.regs -> {:?}", alloc.regs);
+                            is_used = false;
+                            break;
+                        }
+                    }
                 }
 
+                if !is_used { continue; }
+
+
+                ydbg!("[X86] saving overwritten register: {overwrite}");
+
                 // 2. Allocate new spill location
-                // we feed random values into the mem processor cuz they will be ignored
-                let stack = super::alloc::mem_proc(alloc, &DagNode::new(dag::DagOpCode::Copy, crate::IR::TypeMetadata::Void), crate::IR::TypeMetadata::i64);
+                let ty = if overwrite.is_gr() { TypeMetadata::i64 } else if overwrite.is_fp() { TypeMetadata::f64 } else { todo!() };
+                let stack = super::alloc::mem_proc(alloc, &DagNode::new(dag::DagOpCode::Copy, crate::IR::TypeMetadata::Void), ty);
 
                 // We now create our instruction
-                let instr = super::asm::X86Instr::with2(X86Mnemonic::Mov, stack.into(), X86Operand::Reg(overwrite));
+
+                // at first we find the mnemonic
+                let mnemonic = if overwrite.is_gr() { X86Mnemonic::Mov } 
+                else if overwrite.is_fp() { X86Mnemonic::Movq } 
+                else { todo!() }; 
+
+                let instr = super::asm::X86Instr::with2(mnemonic, stack.into(), X86Operand::Reg(overwrite));
 
                 asm.push(instr);
 
@@ -197,10 +217,45 @@ pub(super) fn x86_lower(func: &mut dag::DagFunction, alloc: &mut ItRegCoalAlloc,
                 mc_instrs.push(Box::new(instr));
         }
 
-        blocks.push((name.to_owned(), mc_instrs));
+        blocks.push_back((name.to_owned(), mc_instrs));
     };
 
-    blocks
+
+    // Here we add the prolog
+
+    let mut prolog = Vec::new();
+
+    if alloc.stack.abs() > 16 {
+        prolog.push(X86Instr::with1(X86Mnemonic::Push, X86Operand::Reg(X86Reg::Rbp())).into());
+        prolog.push(X86Instr::with2(X86Mnemonic::Sub, X86Operand::Reg(X86Reg::Rsp()), X86Operand::Const(alloc.stack.abs() as i64)).into());
+        prolog.push(X86Instr::with2(X86Mnemonic::Mov, X86Operand::Reg(X86Reg::Rbp()), X86Operand::Reg(X86Reg::Rsp())).into());
+    }
+
+    blocks.push_front((BlockId("the prolog (we can use this name cuz who really is going to type this out)?".to_string()), prolog));
+
+    // and the epilog
+
+    for (_, block) in &mut blocks {
+        if block.contains(&X86Instr::with0(X86Mnemonic::Ret).into()) {
+            // We now know that the block contains a return instruction
+            // we now need to add the epilog to it
+            // BEFORE the ret instruction
+            block.pop(); // remove return instruction
+
+            let mut epilog = Vec::new();
+
+            if alloc.stack.abs() > 16 {
+                epilog.push(X86Instr::with2(X86Mnemonic::Sub, X86Operand::Reg(X86Reg::Rsp()), X86Operand::Const(alloc.stack.abs() as i64)).into());
+                epilog.push(X86Instr::with1(X86Mnemonic::Pop, X86Operand::Reg(X86Reg::Rbp())).into());
+            }
+
+            block.extend(epilog);
+
+            block.push(X86Instr::with0(X86Mnemonic::Ret).into()); // and add the return instruction back
+        }
+    }
+
+    blocks.into()
 }
 
 pub(super) fn x86_tmps(node: &DagNode) -> Vec<dag::DagTmpInfo> {
@@ -221,6 +276,17 @@ fn call_overwrittes() -> Vec<super::reg::X86Reg> {
     overwrittes
 }
 
+fn x86_overwrittesr(node: &DagNode) -> Vec<X86Reg> {
+    // This should handle most of the overwrittes, expect the one of the call node
+    let mut overwrittes = auto_gen::overwrittes(node);
+    // so if the node is a call node, we add the callee saved registers here
+    if let dag::DagOpCode::Call(_) = node.get_opcode() {
+        overwrittes.extend(call_overwrittes()); 
+    }
+
+    overwrittes
+}
+
 pub(super) fn ov_proc(node: &DagNode) -> Vec<crate::CodeGen::reg::Reg> {
     // This should handle most of the overwrittes, expect the one of the call node
     let mut overwrittes = auto_gen::overwrittes(node);
@@ -235,5 +301,7 @@ pub(super) fn ov_proc(node: &DagNode) -> Vec<crate::CodeGen::reg::Reg> {
         reg_ov.push(crate::CodeGen::reg::Reg::new_x86(ov));
     }
 
+    ydbg!("[X86] overwrittes for node {node}: {reg_ov:?}");
+
     reg_ov
 }