jit.c

#include <stdio.h>
#include <stdlib.h>
#include <sys/mman.h>
#include <string.h>

#define JIT_ALLOC_FUNC(size) \
  (u8 *)mmap(NULL, size, PROT_READ|PROT_WRITE|PROT_EXEC, \
    (MAP_PRIVATE|MAP_ANON), -1, 0)

typedef unsigned char u8;
typedef int OBJ;
typedef int (JitFunc)();

typedef struct {
  u8 *start, *ptr;
  long size;
} JitContext;

static void jit_x86_put(JitContext *jit, OBJ val, size_t len) {
  if (jit->size - (jit->ptr - jit->start) < len) {
    size_t dist = jit->ptr - jit->start;
    jit->size += 4096;
    jit->start = (u8*)realloc((u8*)jit->start, sizeof(u8)*jit->size);
    jit->ptr = jit->start + dist;
  }

  if (len == sizeof(u8))
    *jit->ptr = (u8)val;
  else if (len == sizeof(int))
    *((int *)jit->ptr) = (int)val;
  else if (len == sizeof(OBJ))
    *((OBJ *)jit->ptr) = val;
  jit->ptr += len;
}

#define RBP(x)   (0x100 - ((x + 1) * sizeof(OBJ)))
#define RBPI(x)  (0x100 - ((x + 1) * sizeof(int)))
#define X86(i)   jit_x86_put(jit, (OBJ)i, sizeof(u8))
#define X86I(i)  jit_x86_put(jit, (OBJ)(i), sizeof(int))

enum {
  LOADA,
  LOADB,
  ADD,
  RETURN
};

JitFunc *jit_compile(u8 *bytecode) {
  JitContext _jit;
  JitContext *jit = &_jit;
  jit->size = 4096;
  jit->start = jit->ptr = malloc(4096);
  u8 *ip = bytecode;
  
  X86(0x55);              // push   %ebp
  
  while (1) {
    switch (*ip) {
      case LOADA:
        X86(0xB8); X86I(*++ip);     // mov    $0x7,%eax
        break;
      
      case LOADB:
        X86(0xBA); X86I(*++ip);     // mov    $0x2,%edx
        break;
        
      case ADD:
        X86(0x01); X86(0xD0);   // add    %edx,%eax
        break;
        
      case RETURN:
        X86(0x89); X86(0xE5);   // mov    %esp,%ebp
        X86(0xC9);              // leave
        X86(0xC3);              // ret
        goto assemble;
    }
    ip++;
  }
  
assemble:
  jit->size = jit->ptr - jit->start;
  JitFunc *func = (JitFunc *)JIT_ALLOC_FUNC(jit->size);
  memcpy(func, jit->start, sizeof(u8)*jit->size);
  
  return func;
}

int main(int argc, char const *argv[]) {
  u8 bytecode[] = {
    LOADA, 1,
    LOADB, 2,
    ADD,
    LOADB, 4,
    ADD,
    RETURN,
  };
  int ret = jit_compile(bytecode)();
  printf("> %d\n", ret);
  return 0;
}