Minimal Interpreter for SSA Bytecode — Explicitly Track CFG Paths

ssa-interp2.c

#include <stdlib.h>
#include <stddef.h>
#include <stdbool.h>
#include <stdio.h>
#include <string.h>
#include <inttypes.h>
#include <assert.h>

typedef int64_t value_t;
typedef uint64_t instr_t;

enum opcode {
  NOP, YIELD, PHI, CALL, RET, MOV, CMP, JZ, LOADI, ADD,
};

#define MAX_CALL_DEPTH 256

struct vm {
  value_t r[1024];
  uint64_t rsp[MAX_CALL_DEPTH];  // return stack pointer
  uint64_t r_assigned[16];  // register assignment bitmap
  uint64_t pc,  // program counter
    ces; // CFG edge stack--LSB is 0 if the next phi instruction should use its
         // first operand, 1 for its second
  instr_t *instrs;
  value_t *sp;  // stack pointer
};

void vm_init(struct vm *vm) {
  memset(vm->r, 0, sizeof vm->r);
  memset(vm->r_assigned, 0, sizeof vm->r_assigned);
  memset(vm->rsp, 0, sizeof vm->rsp);
  vm->pc = 0;
  vm->ces = 0;
  vm->sp = calloc(MAX_CALL_DEPTH*1024, sizeof(value_t));
}

void vm_destroy(struct vm *vm) {
  free(vm->sp);
}

inline __attribute__((always_inline))
void vm_set(struct vm *vm, uint32_t r, value_t x) {
  vm->r[r] = x;
  vm->r_assigned[r/64] |= (1 << r%64);
}

inline __attribute__((always_inline))
value_t vm_get(struct vm *vm, uint32_t r) {
  return vm->r[r];
}

#define RSP_PUSH(addr) (*((*(uint64_t **)&vm->rsp)++) = (addr))
#define RSP_POP() (*(--(*(uint64_t **)&vm->rsp)))

#define SP_PUSH(val) (*(vm->sp++) = (val))
#define SP_POP() (*(--vm->sp))

#define vm_jump(addr) do{                       \
    instr_t instr = vm->instrs[(addr)];         \
    enum opcode opc = instr & 0xff;             \
    imm1 = (instr >>  8) & 0xffff;              \
    imm2 = (instr >> 24) & 0xffff;              \
    imm3 =  instr >> 40;                        \
    goto *dispatch[opc];                        \
  }while(0)

#define jump(addr) vm_jump(vm->pc = (addr))

#define next vm_jump(++vm->pc)

value_t vm_exec(struct vm *vm) {
  void *dispatch[] = {
    [NOP] = &&I_NOP, [YIELD] = &&I_YIELD, [PHI] = &&I_PHI,
    [CALL] = &&I_CALL, [RET] = &&I_RET, [MOV] = &&I_MOV,
    [CMP] = &&I_CMP, [JZ] = &&I_JZ, [LOADI] = &&I_LOADI,
    [ADD] = &&I_ADD,
  };

  uint32_t imm1, imm2, imm3;

 I_NOP:
  next;

 I_YIELD:
  return vm_get(vm, imm1);

 I_PHI:
  vm_set(vm, imm1, (vm->ces & 1) ? vm_get(vm, imm3) : vm_get(vm, imm2));
  vm->ces >>= 1;
  next;

 I_CALL:
  RSP_PUSH(vm->pc);
  for (uint32_t i = 0; i < sizeof vm->r / sizeof vm->r[0]; i++) {
    if (vm->r_assigned[i/64] & (1 << i%64))
      SP_PUSH(vm->r[i]);
  }
  for (uint32_t i = 0; i < 16; i++) {
    SP_PUSH((value_t)vm->r_assigned[i]);
    vm->r_assigned[i] = 0;
  }
  jump(vm_get(vm, imm1));

 I_RET:
  for (uint32_t i = 0; i < 16; i++)
    vm->r_assigned[i] = (uint64_t)SP_POP();
  for (uint32_t i = 0; i < sizeof vm->r / sizeof vm->r[0]; i++) {
    if (vm->r_assigned[i/64] & (1 << i%64))
      vm->r[i] = SP_POP();
  }
  jump(RSP_POP());
  next;

 I_MOV:
  vm_set(vm, imm1, vm_get(vm, imm2));
  vm->r_assigned[imm2/64] &= ~(1 << imm2%64);
  next;

 I_CMP: {
  value_t x = vm_get(vm, imm2), y = vm_get(vm, imm3);
  vm_set(vm, imm1, (x > y) - (x < y));
  next;
 }

 I_JZ:
  if (imm1 == 0)  // don't bother to push the phi-resolution stack if reg is const 0
    jump(vm->r[imm2]);
  else if (vm->r[imm1] == 0) {
    vm->ces <<= 1;
    jump(vm->r[imm2]);
  }
  else {
    vm->ces = ((vm->ces << 1) | 1);
    jump(vm->r[imm3]);
  }
  next;

 I_LOADI:
  vm_set(vm, imm1, (value_t)imm2 | (value_t)imm3);
  next;

 I_ADD:
  vm_set(vm, imm1, vm_get(vm, imm2) + vm_get(vm, imm3));
  next;
}

#define I(opc, imm1, imm2, imm3) ((UINT64_C(imm3)<<40) | (UINT64_C(imm2)<<24) | (UINT64_C(imm1)<<8) | (opc))

int main(int argc, char **argv) {
  static instr_t prog[] = {
    I(NOP, 0, 0, 0),    //  0
    I(LOADI, 1, 0, 10), //  1
    I(YIELD, 1, 0, 0),  //  2
    I(LOADI, 2, 0, 20), //  3
    I(YIELD, 2, 0, 0),  //  4
    I(ADD, 3, 1, 2),    //  5
    I(YIELD, 3, 0, 0),  //  6
    I(LOADI, 4, 0, 30), //  7
    I(CMP, 5, 3, 4),    //  8
    I(LOADI, 6, 0, 16), //  9
    I(LOADI, 7, 0, 18), // 10
    I(LOADI, 12, 0, 13),// 11
    I(JZ, 5, 6, 7),     // 12
    I(PHI, 10, 8, 9),   // 13
    I(YIELD, 10, 0, 0), // 14
    I(YIELD, 0, 0, 0),  // 15
    I(LOADI, 8, 0, 42), // 16
    I(JZ, 0, 12, 0),    // 17
    I(LOADI, 9, 0, 50), // 18
    I(JZ, 0, 12, 0)     // 19
  };

  struct vm vm;
  vm_init(&vm);
  vm.instrs = prog;
  uint64_t out = 0;
  do {
    out = vm_exec(&vm);
    printf("out: %" PRIu64 "\n", out);
  } while (out != 0);
  vm_destroy(&vm);

  return 0;
}