stencil c and python

foo.c

#include <stdint.h>
#include <stdio.h>

typedef unsigned char byte;

#define TAILCALL __attribute__((musttail))

#define STENCIL(ret_type, name)                                                \
  extern byte __start_stencil_##name[];                                        \
  extern byte __stop_stencil_##name[];                                         \
  extern byte* section_limits_##name[];                                        \
  byte* section_limits_##name[] = {                                          \
      __start_stencil_##name, __stop_stencil_##name};                          \
  extern __attribute__((noinline, used, section("stencil_" #name))) ret_type name

#define STENCIL_SIZE(name) (__stop_stencil_##name - __start_stencil_##name)

#define HOLE0 0xdeadbeefcafef00d
#define HOLE1 0xcafebeefcafebeef


extern uint64_t get_hole_val(long i) {
    uint64_t HOLES[] = {HOLE0, HOLE1, 0};
    return HOLES[i];
}

STENCIL(uint64_t, ifgt)(uint64_t a, uint64_t b) {
  uint64_t (*iftrue)(uint64_t, uint64_t) =
      (uint64_t(*)(uint64_t, uint64_t))HOLE0;
  uint64_t (*iffalse)(uint64_t, uint64_t) =
      (uint64_t(*)(uint64_t, uint64_t))HOLE1;
  if (a > b)
    TAILCALL return iftrue(a, b);
  else
    TAILCALL return iffalse(a, b);
}

STENCIL(uint64_t, add_const)(uint64_t a) {
    uint64_t (*cont)(uint64_t) =
        (uint64_t(*)(uint64_t))HOLE0;
    TAILCALL return cont(a + HOLE1);
}

STENCIL(uint64_t, ret_arg)(uint64_t a) {
    return a;
}

STENCIL(uint64_t, ret_const)() {
    return HOLE0;
}

jit.py

# clang -O3 foo.c -o stencil.so -shared -fno-pic -fno-pie -mcmodel=large  && python3 jit.py
import os
from cffi import FFI
import struct
import mmap


ffi = FFI()
ffi.cdef("""
  static char* section_limits_ret_const[];
  uint64_t ret_const();
  static char* section_limits_ret_arg[];
  uint64_t ret_arg(uint64_t);
  static char* section_limits_add_const[];
  uint64_t add_const(uint64_t);
  uint64_t get_hole_val(int);
""")

C = ffi.dlopen("./stencil.so")
HOLES = [struct.pack("=Q", C.get_hole_val(i)) for i in range(2)]

class Stencil:
    def __init__(self, name, numholes):
        self.name = name
        self.numholes = numholes
        limits = getattr(C, "section_limits_" + name)
        self.code = ffi.unpack(limits[0], limits[1] - limits[0])
        for i in range(numholes):
            assert self.code.count(HOLES[i]) == 1

    def patch(self, *args):
        assert len(args) == self.numholes
        new = self.code
        for arg, hole in zip(args, HOLES):
            new = new.replace(hole, struct.pack("=Q", arg))
        return new

for name, numholes in [("ret_const", 1),
                       ("ret_arg", 0),
                       ("add_const", 2)]:
    globals()[name] = Stencil(name, numholes)


memory = ffi.from_buffer(mmap.mmap(-1, 1024*1024, prot=mmap.PROT_WRITE | mmap.PROT_EXEC, flags=mmap.MAP_ANON | mmap.MAP_PRIVATE))
start = int(ffi.cast("unsigned long", memory))
# add 42 to argument, jump to next thing
add_const_patched = add_const.patch(start + len(add_const.code), 42)
ffi.memmove(memory, add_const_patched, len(add_const_patched))
# return argument
ret_arg_patched = ret_arg.patch()
ffi.memmove(memory + len(add_const_patched), ret_arg_patched, len(ret_arg_patched))

# dump instructions
with open("code.dat", "wb") as f:
    f.write(ffi.unpack(memory, len(ret_arg_patched) + len(add_const_patched)))
os.system("objdump -b binary  -D code.dat -m i386:x86-64 --adjust-vma=%s" % (start, ))

# call
fptr = ffi.cast(ffi.typeof(C.add_const), memory)
print(fptr(100))