From aa60c1f36ae3d6c949e66d087e6484db25a305a1 Mon Sep 17 00:00:00 2001
From: Calvin Rose <calsrose@gmail.com>
Date: Sat, 3 Dec 2022 17:52:23 -0600
Subject: [PATCH] Add ffi jit example.

---
 examples/jitfn/hello.bin   | Bin 0 -> 39 bytes
 examples/jitfn/hello.nasm  |  17 +++++++++++++++++
 examples/jitfn/jitfn.janet |  12 ++++++++++++
 src/core/ffi.c             |  25 +++++++++++++++++--------
 4 files changed, 46 insertions(+), 8 deletions(-)
 create mode 100644 examples/jitfn/hello.bin
 create mode 100644 examples/jitfn/hello.nasm
 create mode 100644 examples/jitfn/jitfn.janet

diff --git a/examples/jitfn/hello.bin b/examples/jitfn/hello.bin
new file mode 100644
index 0000000000000000000000000000000000000000..412a55e7583f33aa412b0571f7418c386be75eeb
GIT binary patch
literal 39
pcmdnN$iTp`A4GWcnsNZCU3@@-pY^auYEDkRjzW2UQBH~?7XX%#2)+OS

literal 0
HcmV?d00001

diff --git a/examples/jitfn/hello.nasm b/examples/jitfn/hello.nasm
new file mode 100644
index 00000000..85cb27a6
--- /dev/null
+++ b/examples/jitfn/hello.nasm
@@ -0,0 +1,17 @@
+BITS 64
+
+;;;
+;;; Code
+;;;
+mov rax, 1          ; write(
+mov rdi, 1          ;   STDOUT_FILENO,
+lea rsi, [rel msg]  ;   msg,
+mov rdx, msglen     ;   sizeof(msg)
+syscall             ; );
+ret                 ; return;
+
+;;;
+;;; Constants
+;;;
+msg: db "Hello, world!", 10
+msglen: equ $ - msg
diff --git a/examples/jitfn/jitfn.janet b/examples/jitfn/jitfn.janet
new file mode 100644
index 00000000..2c668895
--- /dev/null
+++ b/examples/jitfn/jitfn.janet
@@ -0,0 +1,12 @@
+###
+### Relies on NASM being installed to assemble code.
+### Only works on x86-64 Linux.
+###
+### Before running, compile hello.nasm to hello.bin with
+### $ nasm hello.nasm -o hello.bin
+
+(def bin (slurp "hello.bin"))
+(def f (ffi/jitfn bin))
+(def signature (ffi/signature :default :void))
+(ffi/call f signature)
+(print "called a jitted function with FFI!")
diff --git a/src/core/ffi.c b/src/core/ffi.c
index 48e0e569..9794626a 100644
--- a/src/core/ffi.c
+++ b/src/core/ffi.c
@@ -1280,27 +1280,36 @@ static Janet janet_ffi_win64(JanetFFISignature *signature, void *function_pointe
 
 #endif
 
+/* Allocate executable memory chunks in sizes of a page. Ideally we would keep
+ * an allocator around so that multiple JIT allocations would point to the same
+ * region but it isn't really worth it. */
+#define FFI_PAGE_MASK 0xFFF
+
 JANET_CORE_FN(cfun_ffi_jitfn,
               "(ffi/jitfn bytes)",
               "Create an abstract type that can be used as the pointer argument to `ffi/call`. The content "
               "of `bytes` is architecture specific machine code that will be copied into executable memory.") {
     janet_fixarity(argc, 1);
     JanetByteView bytes = janet_getbytes(argv, 0);
+
+    /* Quick hack to align to page boundary, we should query OS. FIXME */
+    size_t alloc_size = ((size_t) bytes.len + FFI_PAGE_MASK) & ~FFI_PAGE_MASK;
+
 #ifdef JANET_FFI_JIT
     JanetFFIJittedFn *fn = janet_abstract_threaded(&janet_type_ffijit, sizeof(JanetFFIJittedFn));
     fn->function_pointer = NULL;
     fn->size = 0;
 #ifdef JANET_WINDOWS
-    void *ptr = VirtualAlloc(NULL, bytes.len, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
+    void *ptr = VirtualAlloc(NULL, alloc_size, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
 #elif defined(MAP_ANONYMOUS)
-    void *ptr = mmap(0, bytes.len, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+    void *ptr = mmap(0, alloc_size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
 #elif defined(MAP_ANON)
     /* macos doesn't have MAP_ANONYMOUS */
-    void *ptr = mmap(0, bytes.len, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANON, -1, 0);
+    void *ptr = mmap(0, alloc_size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANON, -1, 0);
 #else
     /* -std=c99 gets in the way */
     /* #define MAP_ANONYMOUS 0x20 should work, though. */
-    void *ptr = mmap(0, bytes.len, PROT_READ | PROT_WRITE, MAP_PRIVATE, -1, 0);
+    void *ptr = mmap(0, alloc_size, PROT_READ | PROT_WRITE, MAP_PRIVATE, -1, 0);
 #endif
     if (!ptr) {
         janet_panic("failed to memory map writable memory");
@@ -1308,16 +1317,16 @@ JANET_CORE_FN(cfun_ffi_jitfn,
     memcpy(ptr, bytes.bytes, bytes.len);
 #ifdef JANET_WINDOWS
     DWORD old = 0;
-    if (!VirtualProtect(ptr, fn->size, PAGE_EXECUTE_READ, &old)) {
+    if (!VirtualProtect(ptr, alloc_size, PAGE_EXECUTE_READ, &old)) {
         janet_panic("failed to make mapped memory executable");
     }
 #else
-    if (mprotect(ptr, fn->size, PROT_READ | PROT_EXEC) == -1) {
+    if (mprotect(ptr, alloc_size, PROT_READ | PROT_EXEC) == -1) {
         janet_panic("failed to make mapped memory executable");
     }
 #endif
-    fn->size = (size_t) bytes.len;
-    fn->function_pointer = (JanetCFunction) ptr;
+    fn->size = alloc_size;
+    fn->function_pointer = ptr;
     return janet_wrap_abstract(fn);
 #else
     janet_panic("ffi/jitfn not available on this platform");