codegen.c

// https://flint.cs.yale.edu/cs422/doc/ELF_Format.pdf

#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#include "a64.h"
#include "codegen.h"
#include "elf.h"
#include "parser.h"
#include "utils.h"

typedef struct LocalVar LocalVar;
struct LocalVar {
  char *name;
  size_t stack_offset;
  LocalVar *next;
};

typedef struct {
  StrBuf b;
  StrBuf ro_data;
  StrBuf str_table;
  StrBuf sym_table;
  StrBuf rel_table;
  size_t global_var_len;
  LocalVar *local_vars;
  size_t frame_size;

  Elf64_Ehdr elf_header;
  Elf64_Shdr sh_text;
  Elf64_Shdr sh_ro_data;
  Elf64_Shdr sh_global_vars;
  Elf64_Shdr sh_sym_table;
  Elf64_Shdr sh_rel_table;
  Elf64_Shdr sh_str_table;
  size_t current_section_offset;
  int sh_count;
} State;

static size_t get_str(State *s, const char *str) {
  if (s->str_table.len == 0) {
    strbuf_push(&s->str_table, 0); // Null string at index 0
  }

  for (size_t i = 0; i < s->str_table.len; i++) {
    if (strcmp(s->str_table.ptr + i, str) == 0) {
      return i;
    }
  }

  size_t len = strlen(str) + 1;
  size_t offset = s->str_table.len;
  strbuf_append(&s->str_table, str, len);

  return offset;
}

static size_t start_section(State *s) {
  s->current_section_offset = s->b.len;
  return s->current_section_offset;
}

static size_t end_section(State *s) {
  size_t size = s->b.len - s->current_section_offset;
  s->current_section_offset = 0;

  while (s->b.len % 8 != 0) {
    strbuf_push(&s->b, 0);
  }

  return size;
}

static void start_function(State *s, const char *name) {
  Elf64_Sym sym = {
      .st_name = get_str(s, name),
      .st_info = ELF64_ST_INFO(STB_GLOBAL, STT_FUNC),
      .st_other = 0,
      .st_shndx = 1, // .text section
      .st_value = s->b.len - s->current_section_offset,
      .st_size = 0,
  };
  strbuf_append(&s->sym_table, (char *)&sym, sizeof(sym));
}

static int get_symbol(State *s, const char *name, uint8_t type) {
  // Search existing symbols
  for (size_t i = 0; i < s->sym_table.len; i += sizeof(Elf64_Sym)) {
    Elf64_Sym *sym = (Elf64_Sym *)(s->sym_table.ptr + i);
    const char *sym_name = s->str_table.ptr + sym->st_name;
    if (strcmp(sym_name, name) == 0) {
      return i / sizeof(Elf64_Sym);
    }
  }

  // If not found, create a new extern symbol
  int index = s->sym_table.len / sizeof(Elf64_Sym);
  Elf64_Sym sym = {
      .st_name = get_str(s, name),
      .st_info = ELF64_ST_INFO(STB_GLOBAL, type),
      .st_other = 0,
      .st_shndx = SHN_UNDEF,
      .st_value = 0,
      .st_size = 0,
  };
  strbuf_append(&s->sym_table, (char *)&sym, sizeof(sym));
  return index;
}

static int write_ro_data(State *s, const char *name, const char *data,
                         size_t size) {
  size_t offset = s->ro_data.len;
  strbuf_append(&s->ro_data, data, size);

  Elf64_Sym sym = {
      .st_name = get_str(s, name),
      .st_info = ELF64_ST_INFO(STB_GLOBAL, STT_OBJECT),
      .st_other = 0,
      .st_shndx = 2, // .rodata section
      .st_value = offset,
      .st_size = size,
  };
  strbuf_append(&s->sym_table, (char *)&sym, sizeof(sym));

  return s->sym_table.len / sizeof(Elf64_Sym) - 1; // Return symbol index
}

static void write_global_var(State *s, const char *name, size_t size) {
  size_t offset = s->global_var_len;
  s->global_var_len += size;

  Elf64_Sym sym = {
      .st_name = get_str(s, name),
      .st_info = ELF64_ST_INFO(STB_GLOBAL, STT_OBJECT),
      .st_other = 0,
      .st_shndx = 3,
      .st_value = offset,
      .st_size = size,
  };
  strbuf_append(&s->sym_table, (char *)&sym, sizeof(sym));
}

static void write_rel(State *s, uint32_t type, int symbol_index) {
  Elf64_Rel rel = {
      .r_offset = s->b.len - s->current_section_offset,
      .r_info = ELF64_R_INFO(symbol_index, type),
  };
  strbuf_append(&s->rel_table, (char *)&rel, sizeof(rel));
}

static void write_elf_header(State *s) {
  s->elf_header = (Elf64_Ehdr){
      .e_ident = {ELFMAG0, ELFMAG1, ELFMAG2, ELFMAG3, ELFCLASS64, ELFDATA2LSB,
                  EV_CURRENT},
      .e_type = ET_REL,
      .e_machine = EM_AARCH64,
      .e_version = EV_CURRENT,
      .e_ehsize = sizeof(Elf64_Ehdr),
      .e_shentsize = sizeof(Elf64_Shdr),
  };
  strbuf_append(&s->b, (char *)&s->elf_header, sizeof(s->elf_header));
}

static void update_elf_header(State *s) {
  memcpy(s->b.ptr, &s->elf_header, sizeof(s->elf_header));
}

static char *mangle_fn_call(const Ast *t) {
  StrBuf b = EMPTY_STRBUF;
  strbuf_append(&b, t->value.ident, strlen(t->value.ident));
  strbuf_push(&b, '_');
  for (size_t i = 0; i < t->children_len; i++) {
    const Ast *arg = &t->children[i];
    if (arg->type == AST_STR) {
      strbuf_append(&b, "s", 1);
    } else {
      strbuf_append(&b, "i", 1);
    }
  }
  return strbuf_finish(b);
}

static char *mangle_fn_def(const Ast *t) {
  StrBuf b = EMPTY_STRBUF;
  strbuf_append(&b, t->value.ident, strlen(t->value.ident));
  strbuf_push(&b, '_');
  const Ast *params = &t->children[0];
  for (size_t i = 0; i < params->children_len; i++) {
    strbuf_append(&b, "i", 1);
  }
  return strbuf_finish(b);
}

static void push_local_var(State *s, const char *name, size_t size) {
  LocalVar *var = malloc(sizeof(LocalVar));
  var->name = strdup(name);
  var->stack_offset = s->frame_size;
  var->next = s->local_vars;
  s->local_vars = var;

  s->frame_size += size;
}

static void clear_local_vars(State *s) {
  LocalVar *var = s->local_vars;
  while (var != NULL) {
    LocalVar *next = var->next;
    free(var);
    var = next;
  }
  s->local_vars = NULL;
  s->frame_size = 0;
}

static LocalVar *get_local_var(const State *s, const char *name) {
  LocalVar *var = s->local_vars;
  while (var != NULL) {
    if (strcmp(var->name, name) == 0) {
      return var;
    }
    var = var->next;
  }
  return NULL;
}

static void write_symbol_addr(State *s, int index, uint8_t rd) {
  write_rel(s, R_AARCH64_ADR_PREL_PG_HI21, index);
  a64_adrp(&s->b, rd);
  write_rel(s, R_AARCH64_ADD_ABS_LO12_NC, index);
  a64_add64_imm(&s->b, rd, rd, 0);
}

static void write_code(State *s, const Ast *t, uint8_t rd) {
  switch (t->type) {
  case AST_PROGRAM:
    for (size_t i = 0; i < t->children_len; i++) {
      const Ast *child = &t->children[i];
      if (child->type == AST_BLOCK) {
        start_function(s, "main");
        a64_stp64_pre_imm(&s->b, A64_FP, A64_LR, A64_SP, -16);

        write_code(s, child, A64_X0);

        a64_mov32_imm(&s->b, A64_W0, 0);
        a64_ldp64_post_imm(&s->b, A64_FP, A64_LR, A64_SP, 16);
        a64_ret(&s->b);
      } else {
        write_code(s, child, A64_X0);
      }
    }
    break;

  case AST_VAR:
    if (s->local_vars == NULL) {
      write_global_var(s, t->value.ident, sizeof(int));
    }
    break;

  case AST_FUNCTION: {
    char *name = mangle_fn_def(t);
    start_function(s, name);
    free(name);

    push_local_var(s, t->value.ident, sizeof(int));
    for (size_t i = 0; i < t->children_len; i++) {
      const Ast *child = &t->children[i];
      switch (child->type) {
      case AST_PARAMS:
        for (size_t j = 0; j < child->children_len; j++) {
          const Ast *param = &child->children[j];
          push_local_var(s, param->value.ident, sizeof(int));
        }
        break;
      case AST_VAR:
        push_local_var(s, child->value.ident, sizeof(int));
        break;
      default:
        break;
      }
    }
    while (s->frame_size % 16 != 0) {
      s->frame_size++;
    }

    a64_stp64_pre_imm(&s->b, A64_FP, A64_LR, A64_SP, -16);
    a64_sub64_imm(&s->b, A64_SP, A64_SP, s->frame_size);
    a64_add64_imm(&s->b, A64_FP, A64_SP, s->frame_size);

    for (size_t i = 0; i < t->children_len; i++) {
      write_code(s, &t->children[i], A64_X0);
    }

    a64_ldr32_imm(&s->b, A64_X0, A64_SP, 0);
    a64_add64_imm(&s->b, A64_SP, A64_SP, s->frame_size);
    a64_ldp64_post_imm(&s->b, A64_FP, A64_LR, A64_SP, 16);
    a64_ret(&s->b);

    clear_local_vars(s);
    break;
  }

  case AST_PARAMS: {
    for (size_t i = 0; i < t->children_len; i++) {
      const Ast *param = &t->children[i];
      LocalVar *var = get_local_var(s, param->value.ident);
      a64_str32_imm(&s->b, A64_W0 + i, A64_SP, var->stack_offset);
    }
    break;
  }

  case AST_BLOCK:
    for (size_t i = 0; i < t->children_len; i++) {
      write_code(s, &t->children[i], rd);
    }
    break;

  case AST_FOR: {
    // int i = start;
    int index = get_symbol(s, t->value.ident, STT_OBJECT);
    write_symbol_addr(s, index, rd);
    write_code(s, &t->children[0], rd + 1); // start value
    a64_str32(&s->b, rd + 1, rd);

    // i <= end;
    size_t loop_start = s->b.len;
    write_symbol_addr(s, index, rd);
    a64_ldr32(&s->b, rd + 1, rd);
    write_code(s, &t->children[1], rd + 2); // end
    a64_cmp32(&s->b, rd + 1, rd + 2);
    size_t loop_end_patch = s->b.len;
    a64_b_cond(&s->b, A64_COND_GT);

    // body
    write_code(s, &t->children[2], rd);

    // i++
    write_symbol_addr(s, index, rd);
    a64_ldr32(&s->b, rd + 1, rd);
    a64_add32_imm(&s->b, rd + 1, rd + 1, 1);
    a64_str32(&s->b, rd + 1, rd);

    size_t loop_start_patch = s->b.len;
    a64_b_cond(&s->b, A64_COND_AL);
    size_t loop_end = s->b.len;

    a64_patch_b_cond(&s->b, loop_start_patch, loop_start - loop_start_patch);
    a64_patch_b_cond(&s->b, loop_end_patch, loop_end - loop_end_patch);
    break;
  }

  case AST_IF:
    write_code(s, &t->children[0], rd); // condition
    size_t else_patch = s->b.len;
    a64_cbz32(&s->b, rd);               // if false, jump to else
    write_code(s, &t->children[1], rd); // then branch
    if (t->children_len > 2) {
      size_t end_patch = s->b.len;
      a64_b_cond(&s->b, A64_COND_AL); // jump to end
      size_t else_ = s->b.len;

      write_code(s, &t->children[2], rd); // else branch

      size_t end = s->b.len;
      a64_patch_b_cond(&s->b, else_patch, else_ - else_patch);
      a64_patch_b_cond(&s->b, end_patch, end - end_patch);
    } else {
      size_t end = s->b.len;
      a64_patch_b_cond(&s->b, else_patch, end - else_patch);
    }
    break;

  case AST_CALL: {
    for (size_t i = 0; i < t->children_len; i++) {
      write_code(s, &t->children[i], A64_X0 + i);
    }

    char *mangled = mangle_fn_call(t);
    int index = get_symbol(s, mangled, STT_FUNC);
    free(mangled);
    write_rel(s, R_AARCH64_CALL26, index);
    a64_bl(&s->b);
    break;
  }

  case AST_ASSIGN: {
    write_code(s, &t->children[1], rd);
    LocalVar *var = get_local_var(s, t->children[0].value.ident);
    if (var != NULL) {
      a64_str32_imm(&s->b, rd, A64_SP, var->stack_offset);
    } else {
      int index = get_symbol(s, t->children[0].value.ident, STT_OBJECT);
      write_symbol_addr(s, index, rd + 1);
      a64_str32(&s->b, rd, rd + 1);
    }
    break;
  }

  case AST_STR: {
    int index =
        write_ro_data(s, t->value.str, t->value.str, strlen(t->value.str) + 1);
    write_symbol_addr(s, index, rd);
    break;
  }

  case AST_INT:
    a64_mov32_imm(&s->b, rd, t->value.i);
    break;

  case AST_IDENT: {
    LocalVar *var = get_local_var(s, t->value.ident);
    if (var != NULL) {
      a64_ldr32_imm(&s->b, rd, A64_SP, var->stack_offset);
    } else {
      int index = get_symbol(s, t->value.ident, STT_OBJECT);
      write_symbol_addr(s, index, rd);
      a64_ldr32(&s->b, rd, rd);
    }
    break;
  }

  case AST_OP:
    switch (t->value.op) {
    case OP_EQ:
      write_code(s, &t->children[0], rd);
      write_code(s, &t->children[1], rd + 1);
      a64_cmp32(&s->b, rd, rd + 1);
      a64_cset32(&s->b, rd, A64_COND_EQ);
      break;
    case OP_LT:
      write_code(s, &t->children[0], rd);
      write_code(s, &t->children[1], rd + 1);
      a64_cmp32(&s->b, rd, rd + 1);
      a64_cset32(&s->b, rd, A64_COND_LT);
      break;
    case OP_LE:
      write_code(s, &t->children[0], rd);
      write_code(s, &t->children[1], rd + 1);
      a64_cmp32(&s->b, rd, rd + 1);
      a64_cset32(&s->b, rd, A64_COND_LE);
      break;
    case OP_GT:
      write_code(s, &t->children[0], rd);
      write_code(s, &t->children[1], rd + 1);
      a64_cmp32(&s->b, rd, rd + 1);
      a64_cset32(&s->b, rd, A64_COND_GT);
      break;
    case OP_GE:
      write_code(s, &t->children[0], rd);
      write_code(s, &t->children[1], rd + 1);
      a64_cmp32(&s->b, rd, rd + 1);
      a64_cset32(&s->b, rd, A64_COND_GE);
      break;
    case OP_MOD:
      write_code(s, &t->children[0], rd);
      write_code(s, &t->children[1], rd + 1);
      a64_sdiv32(&s->b, rd + 2, rd, rd + 1);     // rd+2 = lhs / rhs
      a64_msub32(&s->b, rd, rd + 2, rd + 1, rd); // rd = lhs - (lhs / rhs) * rhs
      break;
    }
    break;
  }
}

static void write_text_section(State *s, const Ast *t) {
  size_t offset = start_section(s);

  write_code(s, t, A64_X0);

  size_t size = end_section(s);
  s->sh_text = (Elf64_Shdr){
      .sh_name = get_str(s, ".text"),
      .sh_type = SHT_PROGBITS,
      .sh_flags = SHF_ALLOC | SHF_EXECINSTR,
      .sh_offset = offset,
      .sh_size = size,
      .sh_addralign = 8,
  };
}

static void write_ro_data_section(State *s) {
  size_t offset = start_section(s);

  strbuf_append(&s->b, s->ro_data.ptr, s->ro_data.len);

  size_t size = end_section(s);
  s->sh_ro_data = (Elf64_Shdr){
      .sh_name = get_str(s, ".rodata"),
      .sh_type = SHT_PROGBITS,
      .sh_flags = SHF_ALLOC,
      .sh_offset = offset,
      .sh_size = size,
      .sh_addralign = 8,
  };
}

static void write_global_vars_section(State *s) {
  s->sh_global_vars = (Elf64_Shdr){
      .sh_name = get_str(s, ".bss"),
      .sh_type = SHT_NOBITS,
      .sh_flags = SHF_ALLOC | SHF_WRITE,
      .sh_offset = 0,
      .sh_size = 0,
      .sh_addralign = 8,
  };
}

static void write_sym_table(State *s) {
  size_t offset = start_section(s);

  strbuf_append(&s->b, s->sym_table.ptr, s->sym_table.len);

  size_t size = end_section(s);
  s->sh_sym_table = (Elf64_Shdr){
      .sh_name = get_str(s, ".symtab"),
      .sh_type = SHT_SYMTAB,
      .sh_offset = offset,
      .sh_size = size,
      .sh_link = 6, // .strtab section
      .sh_info = 1, // No local symbols (first symbol is null)
      .sh_entsize = sizeof(Elf64_Sym),
  };
}

static void write_rel_table(State *s) {
  size_t offset = start_section(s);

  strbuf_append(&s->b, s->rel_table.ptr, s->rel_table.len);

  size_t size = end_section(s);
  s->sh_rel_table = (Elf64_Shdr){
      .sh_name = get_str(s, ".rel.text"),
      .sh_type = SHT_REL,
      .sh_offset = offset,
      .sh_size = size,
      .sh_link = 4, // .symtab section
      .sh_info = 1, // .text section
      .sh_entsize = sizeof(Elf64_Rel),
  };
}

static void write_str_table(State *s) {
  size_t offset = start_section(s);

  // Ensure required strings are present
  get_str(s, ".strtab");

  strbuf_append(&s->b, s->str_table.ptr, s->str_table.len);

  size_t size = end_section(s);
  s->sh_str_table = (Elf64_Shdr){
      .sh_name = get_str(s, ".strtab"),
      .sh_type = SHT_STRTAB,
      .sh_offset = offset,
      .sh_size = size,
  };
}

static int write_sh(State *s, Elf64_Shdr sh) {
  strbuf_append(&s->b, (char *)&sh, sizeof(sh));
  return s->sh_count++;
}

static void write_sh_table(State *s) {
  size_t sh_table_offset = s->b.len;

  write_sh(s, (Elf64_Shdr){0}); // Null section
  write_sh(s, s->sh_text);
  write_sh(s, s->sh_ro_data);
  write_sh(s, s->sh_global_vars);
  write_sh(s, s->sh_sym_table);
  write_sh(s, s->sh_rel_table);
  int sh_str_table_index = write_sh(s, s->sh_str_table);

  s->elf_header.e_shoff = sh_table_offset;
  s->elf_header.e_shnum = s->sh_count;
  s->elf_header.e_shstrndx = sh_str_table_index;
}

StrBuf generate(const Ast *t) {
  State s = {0};
  // First symbol must be null
  strbuf_append(&s.sym_table, (char *)&(Elf64_Sym){0}, sizeof(Elf64_Sym));

  write_elf_header(&s);
  write_text_section(&s, t);
  write_ro_data_section(&s);
  write_global_vars_section(&s);
  write_sym_table(&s);
  write_rel_table(&s);
  write_str_table(&s);
  write_sh_table(&s);
  update_elf_header(&s);

  strbuf_free(s.str_table);
  strbuf_free(s.sym_table);
  strbuf_free(s.rel_table);
  return s.b;
}