feat: implement native compilation for goto

src/jit.rs

@@ -1,13 +1,16 @@
 //! JIT compiler for coldrew targeting x86_64.
 use std::collections::{HashMap, VecDeque};
+use std::io::Read;
 
 use crate::bytecode::OPCode;
 use crate::runtime::{Frame, ProgramCounter, Value};
 use crate::trace::Trace;
 
 use dynasmrt::x64::Assembler;
-use dynasmrt::{dynasm, DynamicLabel};
-use dynasmrt::{AssemblyOffset, DynasmApi, ExecutableBuffer};
+use dynasmrt::{
+    dynasm, AssemblyOffset, DynamicLabel, DynasmApi, DynasmLabelApi,
+    ExecutableBuffer,
+};
 
 /// Intel x86-64 registers, ordered by their syntactic order in the Intel
 /// manuals. The usage of the registers follows the System ADM64 ABI.
@@ -281,7 +284,7 @@ impl JitCache {
             // Record the instruction program counter to a new label.
             let inst_label = ops.new_dynamic_label();
             let _ = self.labels.insert(entry.pc(), inst_label);
-            println!("Record label {} @ {}", inst_label.get_id(), entry.pc());
+            println!("Recorded label {} @ {}", inst_label.get_id(), entry.pc());
             match entry.instruction().get_mnemonic() {
                 // Load operation loads a constant from the locals array at
                 // the position given by the opcode's operand.
@@ -301,6 +304,11 @@ impl JitCache {
                         _ => unreachable!("Operand to iload (index in locals) must be int in current implementation")
                     };
                     let dst = self.first_available_register();
+
+                    #[cfg(target_arch = "x86_64")]
+                    dynasm!(ops
+                        ; =>inst_label
+                    );
                     Self::emit_mov(
                         &mut ops,
                         &dst,
@@ -320,6 +328,9 @@ impl JitCache {
                             _ => unreachable!("Operand to istore (index in locals) must be int in current implementation")
                     };
                     if let Some(src) = self.free_register() {
+                        dynasm!(ops
+                            ; =>inst_label
+                        );
                         Self::emit_mov(
                             &mut ops,
                             &Operand::Memory(Register::Rdi, 8 * value),
@@ -336,22 +347,42 @@ impl JitCache {
                 }
                 OPCode::IAdd => {
                     println!("Compiling an iadd");
+                    #[cfg(target_arch = "x86_64")]
+                    dynasm!(ops
+                        ; =>inst_label
+                    );
                     self.emit_arithmetic(&mut ops, Op::Add);
                 }
                 OPCode::ISub => {
                     println!("Compiling an isub");
+                    #[cfg(target_arch = "x86_64")]
+                    dynasm!(ops
+                        ; =>inst_label
+                    );
                     self.emit_arithmetic(&mut ops, Op::Sub);
                 }
                 OPCode::IMul => {
                     println!("Compiling an imul");
+                    #[cfg(target_arch = "x86_64")]
+                    dynasm!(ops
+                        ; =>inst_label
+                    );
                     self.emit_arithmetic(&mut ops, Op::IMul);
                 }
                 OPCode::IDiv => {
                     println!("Compiling an idiv");
+                    #[cfg(target_arch = "x86_64")]
+                    dynasm!(ops
+                        ; =>inst_label
+                    );
                     self.emit_arithmetic(&mut ops, Op::IDiv);
                 }
                 OPCode::IRem => {
                     println!("Compiling an iadd");
+                    #[cfg(target_arch = "x86_64")]
+                    dynasm!(ops
+                        ; =>inst_label
+                    );
                     self.emit_arithmetic(&mut ops, Op::IRem);
                 }
                 OPCode::IInc => {
@@ -365,6 +396,9 @@ impl JitCache {
                         _ => unreachable!("Second operand to iinc (constant for increment) must be int in current implementation")
                     };
                     #[cfg(target_arch = "x86_64")]
+                    dynasm!(ops
+                        ; =>inst_label
+                    );
                     dynasm!(ops
                         ; add [Rq(Register::Rdi as u8) + 8 * index], constant as _
                     );
@@ -385,6 +419,22 @@ impl JitCache {
                     // }
                     // If the Goto target is outside then abondon this trace.
                     //
+                    let target = match entry.instruction().nth(0) {
+                        Some(Value::Int(x)) => x,
+                            _ => unreachable!("First operand to goto (relative offset) must be int")
+                    };
+                    println!("Compiling a goto");
+                    if let Some(label) = self.labels.get(&ProgramCounter::new(
+                        entry.pc().get_method_index(),
+                        (entry.pc().get_instruction_index() as isize
+                            + target as isize) as usize,
+                    )) {
+                        println!("Found target label {}", label.get_id());
+                        #[cfg(target_arch = "x86_64")]
+                        dynasm!(ops
+                            ; jmp =>*label
+                        );
+                    }
                 }
                 _ => println!(
                     "Found opcode : {:}",
@@ -402,6 +452,19 @@ impl JitCache {
         println!("Exit PC {}", exit_pc);
         // println!("Compiled trace @ {pc}");
         let buf = ops.finalize().unwrap();
+        // Because you can't have randomness without YAFC
+        // HAXOR START
+        println!("Dumping jit result to disk");
+        use std::fs::File;
+        use std::io::Write;
+        let dummy = vec![1, 2, 3, 4, 5];
+        let addr = &dummy as *const _ as usize;
+        let filename = format!("test-{addr}.bin");
+        let mut testfile = File::create(&filename).unwrap();
+        let _ = testfile.write_all(&buf.to_vec());
+        println!("Dumped jitted code to {filename}");
+        // HAXOR DONE
+
         let native_trace = NativeTrace(offset, buf);
         self.traces.insert(pc, native_trace);
         // println!("Added trace to native traces");

src/profiler.rs

@@ -19,7 +19,7 @@ impl Profiler {
     pub fn new() -> Profiler {
         Profiler {
             threshold: 1,
-            last_pc: ProgramCounter::new(),
+            last_pc: ProgramCounter::default(),
             records: HashMap::new(),
         }
     }

src/runtime.rs

@@ -244,8 +244,8 @@ impl Instruction {
 /// and method we're executing.
 #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
 pub struct ProgramCounter {
-    instruction_index: usize,
     method_index: usize,
+    instruction_index: usize,
 }
 
 impl fmt::Display for ProgramCounter {
@@ -259,10 +259,11 @@ impl fmt::Display for ProgramCounter {
 }
 
 impl ProgramCounter {
-    pub fn new() -> Self {
+    /// Create a new program counter object at the given method and instruction index.
+    pub fn new(method_index: usize, instruction_index: usize) -> Self {
         Self {
-            instruction_index: 0,
-            method_index: 0,
+            method_index,
+            instruction_index,
         }
     }
 
@@ -282,7 +283,7 @@ impl ProgramCounter {
 
 impl Default for ProgramCounter {
     fn default() -> Self {
-        Self::new()
+        Self::new(0, 0)
     }
 }
 
@@ -391,13 +392,13 @@ impl Runtime {
                 // Cache the trace.
                 self.traces.insert(pc, recorded_trace.clone());
                 // Dump trace to stdout.
-                for entry in recorded_trace.trace {
+                for entry in &recorded_trace.trace {
                     println!("{entry}");
                 }
                 // Compile recorded trace.
-                // self.jit_cache.compile(&recorded_trace);
+                self.jit_cache.compile(&recorded_trace);
             }
-            if self.jit_cache.has_native_trace(pc) {
+            if false {
                 println!("Entering the Jit @ {pc}");
                 // If we have a native trace at this pc run it
                 // and capture the return value which is the next

src/trace.rs

@@ -62,8 +62,8 @@ impl Default for Recorder {
 impl Recorder {
     pub fn new() -> Self {
         Self {
-            trace_start: ProgramCounter::new(),
-            loop_header: ProgramCounter::new(),
+            trace_start: ProgramCounter::default(),
+            loop_header: ProgramCounter::default(),
             is_recording: false,
             last_instruction_was_branch: false,
             trace: Vec::new(),