rt: extend runtime to allow adding multiple code blocks

4 files changed, 54 insertions, 25 deletions

diff --git a/examples/add.rs b/examples/add.rs
index ef010b6..8cd1984 100644
--- a/examples/add.rs
+++ b/examples/add.rs
@@ -22,8 +22,8 @@ fn main() {
     let code = asm.into_code();
     std::fs::write("jit.asm", &code).unwrap();
 
-    let rt = Runtime::new(&code);
-    let add42 = unsafe { rt.as_fn::<extern "C" fn(u32) -> u32>() };
+    let mut rt = Runtime::new();
+    let add42 = unsafe { rt.add_code::<extern "C" fn(u32) -> u32>(code) };
 
     let res = add42(5);
     assert_eq!(res, 47);
diff --git a/examples/fib.rs b/examples/fib.rs
index 7acbb50..534dc13 100644
--- a/examples/fib.rs
+++ b/examples/fib.rs
@@ -62,8 +62,8 @@ fn main() {
     std::fs::write("jit.asm", &code).unwrap();
 
     // Move code into executable page and get function pointer to it.
-    let rt = Runtime::new(&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>(code) };
 
     for n in 0..15 {
         let fib_jit = fib(n);
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");
         }
     }
 }