From 4a80838151a9945438739ab937c415939e2ccf5b Mon Sep 17 00:00:00 2001
From: Johannes Stoelp <johannes.stoelp@gmail.com>
Date: Tue, 5 Dec 2023 00:55:50 +0100
Subject: rt: extend runtime to allow adding multiple code blocks

---
 src/lib.rs |  4 ++--
 src/rt.rs  | 67 ++++++++++++++++++++++++++++++++++++++++++++------------------
 2 files changed, 50 insertions(+), 21 deletions(-)

(limited to 'src')

diff --git a/src/lib.rs b/src/lib.rs
index de12c57..0cc86fb 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -61,8 +61,8 @@
 //!     asm.ret();
 //!
 //!     // Move code into executable page and get function pointer to it.
-//!     let rt = Runtime::new(&asm.into_code());
-//!     let fib = unsafe { rt.as_fn::<extern "C" fn(u64) -> u64>() };
+//!     let mut rt = Runtime::new();
+//!     let fib = unsafe { rt.add_code::<extern "C" fn(u64) -> u64>(&asm.into_code()) };
 //!
 //!     for n in 0..15 {
 //!         let fib_jit = fib(n);
diff --git a/src/rt.rs b/src/rt.rs
index 20d1de7..7def0ea 100644
--- a/src/rt.rs
+++ b/src/rt.rs
@@ -1,59 +1,88 @@
 //! A simple runtime which can be used to execute emitted instructions.
 
-use core::ffi::c_void;
+use core::slice;
 use nix::sys::mman::{mmap, mprotect, munmap, MapFlags, ProtFlags};
 
 /// A simple `mmap`ed runtime with executable pages.
 pub struct Runtime {
-    buf: *mut c_void,
+    buf: *mut u8,
     len: usize,
+    idx: usize,
 }
 
 impl Runtime {
     /// Create a new [Runtime].
-    pub fn new(code: impl AsRef<[u8]>) -> Runtime {
+    pub fn new() -> Runtime {
         // Allocate a single page.
         let len = core::num::NonZeroUsize::new(4096).unwrap();
         let buf = unsafe {
             mmap(
                 None,
                 len,
-                ProtFlags::PROT_WRITE,
+                ProtFlags::PROT_WRITE | ProtFlags::PROT_READ | ProtFlags::PROT_EXEC,
                 MapFlags::MAP_PRIVATE | MapFlags::MAP_ANONYMOUS,
                 0, /* fd */
                 0, /* off */
             )
-            .unwrap()
+            .unwrap() as *mut u8
         };
-        {
-            // Copy over code.
-            let code = code.as_ref();
-            assert!(code.len() < len.get());
-            unsafe { std::ptr::copy_nonoverlapping(code.as_ptr(), buf.cast(), len.get()) };
-        }
-        unsafe {
-            // Remove write permissions from code buffer and allow to read-execute from it.
-            mprotect(buf, len.get(), ProtFlags::PROT_READ | ProtFlags::PROT_EXEC)
-                .expect("Failed to RX mprotect Runtime code buffer");
-        }
 
         Runtime {
             buf,
             len: len.get(),
+            idx: 0,
         }
     }
 
+    /// Write protect the underlying code page(s).
+    pub fn protect(&mut self) {
+        unsafe {
+            // Remove write permissions from code buffer and allow to read-execute from it.
+            mprotect(
+                self.buf.cast(),
+                self.len,
+                ProtFlags::PROT_READ | ProtFlags::PROT_EXEC,
+            )
+            .expect("Failed to RX mprotect Runtime code buffer");
+        }
+    }
+
+    /// Add block of code to the runtime and get function pointer back.
+    pub unsafe fn add_code<F>(&mut self, code: impl AsRef<[u8]>) -> F {
+        // Get pointer to start of next free byte.
+        assert!(self.idx < self.len);
+        let fn_start = self.buf.add(self.idx);
+
+        // Copy over code.
+        let code = code.as_ref();
+        assert!(code.len() < (self.len - self.idx));
+        unsafe { std::ptr::copy_nonoverlapping(code.as_ptr(), fn_start, code.len()) };
+
+        // Increment index to next free byte.
+        self.idx += code.len();
+
+        // Return function to newly added code.
+        Self::as_fn::<F>(fn_start)
+    }
+
     /// Reinterpret the block of code as `F`.
     #[inline]
-    pub unsafe fn as_fn<F>(&self) -> F {
-        unsafe { std::mem::transmute_copy(&self.buf) }
+    unsafe fn as_fn<F>(fn_start: *mut u8) -> F {
+        unsafe { std::mem::transmute_copy(&fn_start) }
+    }
+
+    /// Dump the currently added code to a file called `jit.asm`. The disassembly can be inspected
+    /// as `ndisasm -b 64 jit.asm`.
+    pub fn dump(&self) {
+        let code = unsafe { slice::from_raw_parts(self.buf, self.idx) };
+        std::fs::write("jit.asm", code).unwrap();
     }
 }
 
 impl Drop for Runtime {
     fn drop(&mut self) {
         unsafe {
-            munmap(self.buf, self.len).expect("Failed to munmap Runtime");
+            munmap(self.buf.cast(), self.len).expect("Failed to munmap Runtime");
         }
     }
 }
-- 
cgit v1.2.3