x86-64 assembler in C

asm_x64.c

#include <assert.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <sys/mman.h>

typedef int64_t word;
typedef uint64_t uword;
typedef unsigned char byte;
typedef uword (*Function)();

static const int kBitsPerByte = 8;

#define INLINE

// x64 encoding

enum Reg {
  RAX,
  RCX,
  RDX,
  RBX,
  RSP,
  RBP,
  RSI,
  RDI,
  R8,
  R9,
  R10,
  R11,
  R12,
  R13,
  R14,
  R15,
};

enum XmmReg {
  XMM0,
  XMM1,
  XMM2,
  XMM3,
  XMM4,
  XMM5,
  XMM6,
  XMM7,
  XMM8,
  XMM9,
  XMM10,
  XMM11,
  XMM12,
  XMM13,
  XMM14,
  XMM15,
};

enum Scale { X1, X2, X4, X8 };

enum Cond {
  O,
  NO,
  B,
  NB,
  E,
  NE,
  NA,
  A,
  S,
  NS,
  P,
  NP,
  L,
  NL,
  NG,
  G,

  LE = NG,
  GE = NL,
  BE = NA,
  AE = NB,
};

enum Mode { INDIRECT, INDIRECT_DISP8, INDIRECT_DISP32, DIRECT };

INLINE uword rexw(uword rx, uword base, uword index) {
  assert(rx < 16);
  assert(base < 16);
  assert(index < 16);
  return 0x48 | (base >> 3) | ((index >> 3) << 1) | ((rx >> 3) << 2); // 1
}

INLINE uword mod_rx_rm(uword mod, uword rx, uword rm) {
  assert(mod < 4);
  assert(rx < 16);
  assert(rm < 16);
  return (rm & 7) | ((rx & 7) << 3) | (mod << 6); // 1
}

INLINE uword direct(uword rx, uword reg) {
  return mod_rx_rm(DIRECT, rx, reg); // 1
}

INLINE uword indirect(uword rx, uword base) {
  assert((base & 7) != RSP);
  assert((base & 7) != RBP);
  return mod_rx_rm(INDIRECT, rx, base); // 1
}

INLINE uword indirect_rip_disp32(uword rx, uword disp) {
  return mod_rx_rm(INDIRECT, rx, RBP) | (disp << 8); // 5
}

INLINE uword indirect_disp8(uword rx, uword base, uword disp) {
  assert((base & 7) != RSP);
  return mod_rx_rm(INDIRECT_DISP8, rx, base) | (disp << 8); // 2
}

INLINE uword indirect_disp32(uword rx, uword base, uword disp) {
  assert((base & 7) != RSP);
  return mod_rx_rm(INDIRECT_DISP32, rx, base) | (disp << 8); // 5
}

INLINE uword indirect_index(uword rx, uword base, uword index, uword scale) {
  assert((base & 7) != RBP);
  return mod_rx_rm(INDIRECT, rx, RSP) |
         (mod_rx_rm(scale, index, base) << 8); // 2
}

INLINE uword indirect_index_disp8(uword rx, uword base, uword index,
                                  uword scale, uword disp) {
  return mod_rx_rm(INDIRECT_DISP8, rx, RSP) |
         (mod_rx_rm(scale, index, base) << 8) | (disp << 16); // 3
}

INLINE uword indirect_index_disp32(uword rx, uword base, uword index,
                                   uword scale, uword disp) {
  return mod_rx_rm(INDIRECT_DISP32, rx, RSP) |
         (mod_rx_rm(scale, index, base) << 8) | (disp << 16); // 6
}

INLINE int isimm8(uword imm) { return imm + 128 < 256; }

INLINE int isdisp8(uword disp) { return disp + 128 < 256; }

INLINE int isrel8(uint32_t rel) { return rel + 128 < 256; }

#define here (buf->address + Buffer_len(buf))

// Buffer

typedef unsigned char byte;

typedef enum {
  kWritable,
  kExecutable,
} BufferState;

typedef struct {
  byte *address;
  BufferState state;
  word len;
  word capacity;
  word entrypoint;
} Buffer;

byte *Buffer_alloc_writable(word capacity) {
  byte *result = mmap(/*addr=*/NULL, capacity, PROT_READ | PROT_WRITE,
                      MAP_ANONYMOUS | MAP_PRIVATE,
                      /*filedes=*/-1, /*off=*/0);
  assert(result != MAP_FAILED);
  return result;
}

void Buffer_init(Buffer *result, word capacity) {
  result->address = Buffer_alloc_writable(capacity);
  assert(result->address != MAP_FAILED);
  result->state = kWritable;
  result->len = 0;
  result->capacity = capacity;
  result->entrypoint = 0;
}

word Buffer_len(Buffer *buf) { return buf->len; }

void Buffer_deinit(Buffer *buf) {
  munmap(buf->address, buf->capacity);
  buf->address = NULL;
  buf->len = 0;
  buf->capacity = 0;
  buf->entrypoint = 0;
}

int Buffer_make_executable(Buffer *buf) {
  int result = mprotect(buf->address, buf->len, PROT_EXEC);
  buf->state = kExecutable;
  return result;
}

byte Buffer_at8(Buffer *buf, word pos) { return buf->address[pos]; }

void Buffer_at_put8(Buffer *buf, word pos, byte b) { buf->address[pos] = b; }

word max(word left, word right) { return left > right ? left : right; }

void Buffer_ensure_capacity(Buffer *buf, word additional_capacity) {
  if (buf->len + additional_capacity <= buf->capacity) {
    return;
  }
  word new_capacity =
      max(buf->capacity * 2, buf->capacity + additional_capacity);
  byte *address = Buffer_alloc_writable(new_capacity);
  memcpy(address, buf->address, buf->len);
  int result = munmap(buf->address, buf->capacity);
  assert(result == 0 && "munmap failed");
  buf->address = address;
  buf->capacity = new_capacity;
}

void Buffer_write8(Buffer *buf, byte b) {
  Buffer_ensure_capacity(buf, sizeof b);
  Buffer_at_put8(buf, buf->len++, b);
}

void Buffer_write32(Buffer *buf, int32_t value) {
  for (uword i = 0; i < sizeof(value); i++) {
    Buffer_write8(buf, (value >> (i * kBitsPerByte)) & 0xff);
  }
}

void Buffer_write64(Buffer *buf, uint64_t value) {
  for (uword i = 0; i < sizeof(value); i++) {
    Buffer_write8(buf, (value >> (i * kBitsPerByte)) & 0xff);
  }
}

void Buffer_at_put32(Buffer *buf, word offset, int32_t value) {
  for (uword i = 0; i < sizeof(value); i++) {
    Buffer_at_put8(buf, offset + i, (value >> (i * kBitsPerByte)) & 0xff);
  }
}

void Buffer_at_put64(Buffer *buf, word offset, uint64_t value) {
  for (uword i = 0; i < sizeof(value); i++) {
    Buffer_at_put8(buf, offset + i, (value >> (i * kBitsPerByte)) & 0xff);
  }
}

void Buffer_write_arr(Buffer *buf, const byte *arr, word arr_size) {
  Buffer_ensure_capacity(buf, arr_size);
  for (word i = 0; i < arr_size; i++) {
    Buffer_write8(buf, arr[i]);
  }
}

void Buffer_dump(Buffer *buf, FILE *fp) {
  for (word i = 0; i < Buffer_len(buf); i++) {
    fprintf(fp, "%.2x ", buf->address[i]);
  }
  fprintf(fp, "\n");
}

// End Buffer

INLINE void emit(Buffer *buf, uword data, int len) {
  assert(len <= 8);
  byte arr[sizeof data];
  memcpy(arr, &data, len);
  Buffer_write_arr(buf, arr, len);
}

INLINE void emit3(Buffer *buf, uword data1, int len1, uword data2, int len2,
                  uword data3, int len3) {
  emit(buf, data1 | (data2 << (8 * len1)) | (data3 << (8 * (len1 + len2))),
       len1 + len2 + len3);
}

INLINE void emit_instr(Buffer *buf, int64_t op, int oplen, uword rx, uword base,
                       uword index, uword ext, int extlen) {
  emit3(buf, rexw(rx, base, index), 1, op, oplen, ext, extlen);
}

typedef struct {
  bool isripdisp;
  bool isindex;
  uword base;
  uword index;
  uword scale;
  uword disp;
} Mem;

INLINE void emit_op_rx_mem(Buffer *buf, uword op, int oplen, uword rx,
                           Mem mem) {
  uword addr;
  int addrlen;
  if (mem.isripdisp) {
    addr = indirect_rip_disp32(rx, mem.disp);
    addrlen = 5;
  } else if (mem.isindex) {
    if (mem.disp || (mem.base & 7) == RBP) {
      if (isdisp8(mem.disp)) {
        addr =
            indirect_index_disp8(rx, mem.base, mem.index, mem.scale, mem.disp);
        addrlen = 3;
      } else {
        addr =
            indirect_index_disp32(rx, mem.base, mem.index, mem.scale, mem.disp);
        addrlen = 6;
      }
    } else {
      addr = indirect_index(rx, mem.base, mem.index, mem.scale);
      addrlen = 2;
    }
  } else {
    if (mem.disp || (mem.base & 7) == RBP) {
      if (isdisp8(mem.disp)) {
        addr = indirect_disp8(rx, mem.base, mem.disp);
        addrlen = 2;
      } else {
        addr = indirect_disp32(rx, mem.base, mem.disp);
        addrlen = 5;
      }
    } else {
      addr = indirect(rx, mem.base);
      addrlen = 1;
    }
  }
  emit_instr(buf, op, oplen, rx, mem.base, mem.index, addr, addrlen);
}

INLINE void emit_op_reg_reg(Buffer *buf, uword op, int oplen, uword dest_reg,
                            uword src_reg) {
  emit_instr(buf, op, oplen, dest_reg, src_reg, 0, direct(dest_reg, src_reg),
             1);
}

INLINE void emit_op_reg_mem(Buffer *buf, uword op, int oplen, uword dest_reg,
                            Mem src_mem) {
  emit_op_rx_mem(buf, op, oplen, dest_reg, src_mem);
}

INLINE void emit_op_mem_reg(Buffer *buf, uword op, int oplen, Mem dest_mem,
                            uword src_reg) {
  emit_op_rx_mem(buf, op, oplen, src_reg, dest_mem);
}

INLINE void emit_op_reg(Buffer *buf, uword op, int oplen, uword rx, uword reg) {
  emit_op_reg_reg(buf, op, oplen, rx, reg);
}

INLINE void emit_op_mem(Buffer *buf, uword op, int oplen, uword rx, Mem mem) {
  emit_op_rx_mem(buf, op, oplen, rx, mem);
}

INLINE void emit_sse_op_reg_reg(Buffer *buf, uword op, int oplen, uword prefix,
                                uword dest_reg, uword src_reg) {
  emit(buf, prefix, 1);
  emit_op_reg_reg(buf, op, oplen, dest_reg, src_reg);
}

INLINE void emit_sse_op_reg_mem(Buffer *buf, uword op, int oplen, uword prefix,
                                uword dest_reg, Mem src_mem) {
  emit(buf, prefix, 1);
  emit_op_rx_mem(buf, op, oplen, dest_reg, src_mem);
}

INLINE void emit_sse_op_mem_reg(Buffer *buf, uword op, int oplen, uword prefix,
                                Mem dest_mem, uword src_reg) {
  emit(buf, prefix, 1);
  emit_op_rx_mem(buf, op, oplen, src_reg, dest_mem);
}

INLINE void emit_op_reg_imm(Buffer *buf, uword op8, uword op32, int oplen,
                            uword rx8, uword rx32, uword dest_reg,
                            uword src_imm) {
  uword op, rx;
  int immlen;
  if ((isimm8(src_imm) && op8) || !op32) {
    assert(op8);
    op = op8;
    rx = rx8;
    immlen = 1;
  } else {
    op = op32;
    rx = rx32;
    immlen = 4;
  }
  emit_instr(buf, op, oplen, rx, dest_reg, 0,
             direct(rx, dest_reg) | (src_imm << 8), 1 + immlen);
}

INLINE void emit_op_mem_imm(Buffer *buf, uword op8, uword op32, int oplen,
                            uword rx8, uword rx32, Mem dest_mem,
                            uword src_imm) {
  if ((isimm8(src_imm) && op8) || !op32) {
    assert(op8);
    emit_op_rx_mem(buf, op8, oplen, rx8, dest_mem);
    emit(buf, src_imm, 1);
  } else {
    emit_op_rx_mem(buf, op32, oplen, rx32, dest_mem);
    emit(buf, src_imm, 4);
  }
}

// x64 instruction emitters

#define X64_OP_REG(name, op, oplen, rx)                                        \
  INLINE void name##_reg(Buffer *buf, uword reg) {                             \
    emit_op_reg(buf, op, oplen, rx, reg);                                      \
  }

#define X64_OP_MEM(name, op, oplen, rx)                                        \
  INLINE void name##_mem(Buffer *buf, Mem mem) {                               \
    emit_op_mem(buf, op, oplen, rx, mem);                                      \
  }

#define X64_OP_REG_REG(name, op, oplen)                                        \
  INLINE void name##_reg_reg(Buffer *buf, uword dest_reg, uword src_reg) {     \
    emit_op_reg_reg(buf, op, oplen, dest_reg, src_reg);                        \
  }

#define X64_OP_REG_MEM(name, op, oplen)                                        \
  INLINE void name##_reg_mem(Buffer *buf, uword dest_reg, Mem src_mem) {       \
    emit_op_reg_mem(buf, op, oplen, dest_reg, src_mem);                        \
  }

#define X64_OP_MEM_REG(name, op, oplen)                                        \
  INLINE void name##_mem_reg(Buffer *buf, Mem dest_mem, uword src_reg) {       \
    emit_op_mem_reg(buf, op, oplen, dest_mem, src_reg);                        \
  }

#define X64_OP_REG_IMM(name, op8, op32, oplen, rx8, rx32)                      \
  INLINE void name##_reg_imm(Buffer *buf, uword dest_reg, uword src_imm) {     \
    emit_op_reg_imm(buf, op8, op32, oplen, rx8, rx32, dest_reg, src_imm);      \
  }

#define X64_OP_MEM_IMM(name, op8, op32, oplen, rx8, rx32)                      \
  INLINE void name##_mem_imm(Buffer *buf, Mem dest_mem, uword src_imm) {       \
    emit_op_mem_imm(buf, op8, op32, oplen, rx8, rx32, dest_mem, src_imm);      \
  }

#define SSE_OP_REG_REG(name, op, oplen, prefix)                                \
  INLINE void name##_reg_reg(Buffer *buf, uword dest_reg, uword src_reg) {     \
    emit_sse_op_reg_reg(buf, op, oplen, prefix, dest_reg, src_reg);            \
  }

#define SSE_OP_MEM_REG(name, op, oplen, prefix)                                \
  INLINE void name##_mem_reg(Buffer *buf, Mem dest_mem, uword src_reg) {       \
    emit_sse_op_mem_reg(buf, op, oplen, prefix, dest_mem, src_reg);            \
  }

#define SSE_OP_REG_MEM(name, op, oplen, prefix)                                \
  INLINE void name##_reg_mem(Buffer *buf, uword dest_reg, Mem src_mem) {       \
    emit_sse_op_reg_mem(buf, op, oplen, prefix, dest_reg, src_mem);            \
  }

#define X64_OP_RM(name, op, oplen, rx)                                         \
  X64_OP_REG(name, op, oplen, rx)                                              \
  X64_OP_MEM(name, op, oplen, rx)

#define X64_OP_REG_RM(name, op, oplen)                                         \
  X64_OP_REG_REG(name, op, oplen)                                              \
  X64_OP_REG_MEM(name, op, oplen)

#define X64_OP_RM_IMM(name, op8, op32, oplen, rx8, rx32)                       \
  X64_OP_REG_IMM(name, op8, op32, oplen, rx8, rx32)                            \
  X64_OP_MEM_IMM(name, op8, op32, oplen, rx8, rx32)

#define SSE_OP_REG_RM(name, op, oplen, prefix)                                 \
  SSE_OP_REG_REG(name, op, oplen, prefix)                                      \
  SSE_OP_REG_MEM(name, op, oplen, prefix)

// x64 instruction definitions

#define X64_UNARY_TABLE(_)                                                     \
  _(neg, 0xF7, 0x03)                                                           \
  _(idiv, 0xF7, 0x07)                                                          \
  //  _(name,   rm,      rx)

#define X64_BINARY_TABLE(_)                                                    \
  _(add, 0x03, 0x01, 0x83, 0x00, 0x81, 0x00)                                   \
  _(and, 0x23, 0x21, 0x83, 0x04, 0x81, 0x04)                                   \
  _(mov, 0x8B, 0x89, 0x00, 0x00, 0xC7, 0x00)                                   \
  //  _(name,   reg_rm,  rm_reg, rm_imm8, rm_imm8x, rm_imm32, rm_imm32x)

#define SSE_BINARY_TABLE(_)                                                    \
  _(mulss, 0x580F, 0xF3)                                                       \
  _(andss, 0x590F, 0xF3)                                                       \
  _(movss, 0x100F, 0xF3)                                                       \
  //  _(name,   reg_rm,  prefix)

#define X64_UNARY_OPS(name, rm, rx) X64_OP_RM(name, rm, 1, rx)

#define X64_BINARY_OPS(name, reg_rm, rm_reg, rm_imm8, rm_imm8x, rm_imm32,      \
                       rm_imm32x)                                              \
  X64_OP_REG_RM(name, reg_rm, 1)                                               \
  X64_OP_MEM_REG(name, rm_reg, 1)                                              \
  X64_OP_RM_IMM(name, rm_imm8, rm_imm32, 1, rm_imm8x, rm_imm32x)

#define SSE_BINARY_OPS(name, reg_rm, prefix)                                   \
  SSE_OP_REG_RM(name, reg_rm, 2, prefix)

X64_UNARY_TABLE(X64_UNARY_OPS)
X64_BINARY_TABLE(X64_BINARY_OPS)
SSE_BINARY_TABLE(SSE_BINARY_OPS)
SSE_OP_MEM_REG(movss, 0x110F, 2, 0xF3)
X64_OP_REG_REG(imul, 0xAF0F, 2)
X64_OP_REG_IMM(imul, 0x6B, 0x69, 1, 0, 0)
X64_OP_RM(shl, 0xD3, 1, 0x04)
X64_OP_REG_IMM(shl, 0xC1, 0, 1, 0x04, 0)
X64_OP_MEM_REG(mov8, 0x88, 1)
X64_OP_MEM_REG(mov32, 0x89, 1)
X64_OP_REG_RM(movsx8, 0xBE0F, 2)
X64_OP_REG_RM(movsx16, 0xBF0F, 2)
X64_OP_REG_RM(movsx32, 0xBF0F, 2)
X64_OP_REG_RM(movzx8, 0xB60F, 2)
X64_OP_REG_RM(movzx16, 0xB70F, 2)

INLINE void mov16_mem_reg(Buffer *buf, Mem dest_mem, uword src_reg) {
  emit(buf, 0x66, 1);
  byte *start = here;
  emit_op_mem_reg(buf, 0x89, 1, dest_mem, src_reg);
  *start &= ~8;
}

INLINE void movzx32_reg_reg(Buffer *buf, uword dest_reg, uword src_reg) {
  byte *start = here;
  emit_op_reg_reg(buf, 0x8B, 1, dest_reg, src_reg);
  *start &= ~8;
}

INLINE void movzx32_reg_mem(Buffer *buf, uword dest_reg, Mem src_mem) {
  byte *start = here;
  emit_op_rx_mem(buf, 0x8B, 1, dest_reg, src_mem);
  *start &= ~8;
}

INLINE void cmov_reg_reg_if(Buffer *buf, uword cond, uword dest_reg,
                            uword src_reg) {
  emit_op_reg_reg(buf, 0x400F | (cond << 8), 2, dest_reg, src_reg);
}

INLINE void set8_reg_if(Buffer *buf, uword cond, uword dest_reg) {
  assert(cond < 16);
  emit_op_reg(buf, 0x900F | (cond << 8), 2, 0, dest_reg);
}

INLINE uint32_t *jmp(Buffer *buf, byte *target) {
  uint32_t rel = (uint32_t)(target - (here + 2));
  if (isrel8(rel)) {
    emit(buf, 0xEB | (rel << 8), 2);
    return 0;
  } else {
    emit(buf, 0xE9 | ((rel - 3) << 8), 5);
    return (uint32_t *)(here - 4);
  }
}

INLINE uint32_t *jmp_if(Buffer *buf, uword cond, byte *target) {
  assert(cond < 16);
  uint32_t rel = (uint32_t)(target - (here + 2));
  if (isrel8(rel)) {
    emit(buf, 0x70 | cond | (rel << 8), 2);
    return 0;
  } else {
    emit(buf, 0x800F | (cond << 8) | ((rel - 4) << 16), 6);
    return (uint32_t *)(here - 4);
  }
}

INLINE void patch_rel(uint32_t *rel_ptr, byte *target) {
  *rel_ptr = (uint32_t)(target - ((byte *)rel_ptr + 4));
}

INLINE void ret(Buffer *buf) { emit(buf, 0xc3, 1); }

// End x64 encoding

// Addressing helpers

INLINE Mem base(uword base) { return (Mem){.base = base}; }

INLINE Mem base_disp(uword base, uword disp) {
  return (Mem){.base = base, .disp = disp};
}

INLINE Mem base_index(uword base, uword index) {
  return (Mem){.base = base, .index = index, .isindex = true};
}

INLINE Mem base_index_disp(uword base, uword index, uword disp) {
  return (Mem){.base = base, .index = index, .disp = disp, .isindex = true};
}

INLINE Mem base_index_scale(uword base, uword index, uword scale) {
  return (Mem){.base = base, .index = index, .scale = scale, .isindex = true};
}

INLINE Mem base_index_scale_disp(uword base, uword index, uword scale,
                                 uword disp) {
  return (Mem){.base = base,
               .index = index,
               .scale = scale,
               .disp = disp,
               .isindex = true};
}

INLINE Mem rip_disp(uword disp) {
  return (Mem){.disp = disp, .isripdisp = true};
}

// End Addressing helpers

void large_example(Buffer *buf) {
  mov_reg_imm(buf, RAX, 4);
  mov_reg_reg(buf, RAX, R9);
  mov8_mem_reg(buf, base(RAX), R9);
  mov16_mem_reg(buf, base(RAX), R9);
  mov32_mem_reg(buf, base(RAX), R9);
  movss_mem_reg(buf, base(RAX), XMM10);
  add_reg_mem(buf, RAX, rip_disp(0x1234));
  neg_reg(buf, R9);
  idiv_reg(buf, RAX);
  imul_reg_reg(buf, RDX, R9);
  mulss_reg_reg(buf, XMM9, XMM10);
  mov_reg_mem(buf, R9, base(R10));
  mov_reg_imm(buf, RAX, 0x12345678);
  mov_reg_imm(buf, RAX, -128);
  mov_mem_imm(buf, base(RAX), 0x12345678);
  mulss_reg_reg(buf, XMM0, XMM9);
  mulss_reg_reg(buf, XMM0, XMM9);
  mulss_reg_mem(buf, XMM3, base_index_scale(RBX, RCX, X8));
  shl_reg(buf, RAX);
  shl_reg_imm(buf, R9, 0xA);
  movzx8_reg_reg(buf, RAX, R9);
  movzx16_reg_reg(buf, RAX, R9);
  movzx32_reg_reg(buf, RAX, R9);
  movzx8_reg_mem(buf, RAX, base_index_scale(RBX, RCX, X8));
  movzx16_reg_mem(buf, RAX, base_index_scale(RBX, RCX, X8));
  movzx32_reg_mem(buf, RAX, base_index_scale(RBX, RCX, X8));
  byte *start = here;
  uint32_t *rel_ptr = jmp(buf, 0);
  set8_reg_if(buf, NE, R9);
  cmov_reg_reg_if(buf, LE, RAX, R9);
  jmp_if(buf, E, start);
  patch_rel(rel_ptr, here);
  and_reg_reg(buf, RAX, R9);
}

void small_example(Buffer *buf) {
  mov_reg_imm(buf, RAX, 42);
  ret(buf);
}

uword execute(Buffer *buf) {
  Function function = *(Function *)(&buf->address);
  return function();
}

int main() {
  Buffer buf;
  Buffer_init(&buf, 1);
  small_example(&buf);
  fprintf(stderr, "Code: ");
  Buffer_dump(&buf, stderr);
  Buffer_make_executable(&buf);
  uword result = execute(&buf);
  fprintf(stderr, "Result: %ld\n", result);
  return 0;
}

Makefile

all:
	gcc -Wall -Wextra -pedantic -g asm_x64.c -o asm_x64