1 files changed, 410 insertions, 0 deletions
diff --git a/src/ast.h b/src/ast.h
new file mode 100644
index 0000000..6c326cc
--- /dev/null
+++ b/src/ast.h
@@ -0,0 +1,410 @@
+#ifndef INCLUDE_ast
+#define INCLUDE_ast
+
+#include "token.h"
+#include <stdlib.h>
+#include <string.h>
+#include <stdio.h>
+
+
+typedef enum { EX_NUMBER, EX_VAR, EX_BINOP, EX_CALL, EX_ADDR, EX_DEREF, EX_STRING, EX_INDEX } _EK;
+
+typedef struct _EX {
+  _EK kind;
+  union {
+    int value;  // for number
+    char *name; // for var
+    struct {
+      struct _EX *l;
+      _TK op;
+      struct _EX *r;
+    } binop;
+    struct {
+      char *func_name;
+      struct _EX **args;
+      int argc;
+    } call;
+    struct {
+      struct _EX *expr; // &expr
+    } addr;
+    struct {
+      struct _EX *expr; // *expr
+    } deref;
+    char *string; // for string literal
+    struct {
+      struct _EX *array; // array expression
+      struct _EX *index; // index expression
+    } index;
+  };
+} _EX;
+
+
+typedef enum { STK_RETURN, STK_VAR_DECL, STK_ASSIGN, STK_EXPR, STK_BLOCK, STK_IF, STK_WHILE, STK_FOR } _STK;
+
+typedef struct _STN {
+  _STK kind;
+  union {
+    struct {
+      char *name;
+      _EX *init;
+      _TY type;
+    } var_decl;
+    struct {
+      _EX *lhs; // var or *expr
+      _EX *expr;
+    } assign;
+    _EX *expr;         // STK_EXPR
+    _EX *return_expr;  // STK_RETURN
+    struct _STN *body; // STK_BLOCK
+    struct {           // STK_IF
+      _EX *cond;
+      struct _STN *thenb;
+      struct _STN *elseb; // may be NULL
+    } ifs;
+    struct {           // STK_WHILE
+      _EX *cond;
+      struct _STN *body;
+    } whl;
+    struct {           // STK_FOR
+      struct _STN *init; // may be NULL (a stmt like decl/assign)
+      _EX *cond;         // may be NULL
+      struct _STN *step; // may be NULL (an expr stmt)
+      struct _STN *body;
+    } fr;
+  };
+  struct _STN *n; // linked list
+} _STN;
+
+
+
+typedef struct _FN {
+  char *name;
+  char **params;
+  _TY *param_types;
+  int pac;
+  _STN *body;
+  struct _FN *n;
+} _FN;
+
+typedef struct NameEntry {
+    char *name;
+    struct NameEntry *next;
+} NameEntry;
+
+#define HASH_SIZE 128
+
+static unsigned long hash_str(const char *s) {
+    unsigned long h = 5381;
+    int c;
+    while ((c = *s++))
+        h = ((h << 5) + h) + (unsigned long)c;
+    return h;
+}
+
+static int name_seen(NameEntry **table, const char *name) {
+    unsigned long h = hash_str(name) % HASH_SIZE;
+    for (NameEntry *e = table[h]; e; e = e->next) {
+        if (strcmp(e->name, name) == 0)
+            return 1;
+    }
+    NameEntry *ne = malloc(sizeof(*ne));
+    if (!ne) {
+        fprintf(stderr, "Error: malloc failed\n");
+        exit(1);
+    }
+    ne->name = strdup(name);
+    ne->next = table[h];
+    table[h] = ne;
+    return 0;
+}
+
+static void free_hash_table(NameEntry **table) {
+    for (int i = 0; i < HASH_SIZE; i++) {
+        NameEntry *e = table[i];
+        while (e) {
+            NameEntry *next = e->next;
+            free(e->name);
+            free(e);
+            e = next;
+        }
+        table[i] = NULL;
+    }
+}
+
+_FN *fnlist_prepare(_FN *head) {
+    if (!head) return NULL;
+
+    NameEntry *table[HASH_SIZE] = {0};
+
+    _FN *main_node = NULL;
+    _FN *main_prev = NULL;
+    _FN *prev = NULL;
+
+    for (_FN *cur = head; cur; cur = cur->n) {
+        if (!cur->name) {
+            fprintf(stderr, "[fnlist prepare] error: function with NULL name\n");
+            free_hash_table(table);
+            exit(1);
+        }
+
+        if (name_seen(table, cur->name)) {
+            fprintf(stderr, "[fnlist prepare] error: duplicate function name '%s'\n", cur->name);
+            free_hash_table(table);
+            exit(1);
+        }
+
+        if (strcmp(cur->name, "main") == 0) {
+            if (main_node) {
+                fprintf(stderr, "[fnlist prepare] error: multiple 'main' functions found\n");
+                free_hash_table(table);
+                exit(1);
+            }
+            main_node = cur;
+            main_prev = prev;
+        }
+
+        prev = cur;
+    }
+
+    free_hash_table(table);
+
+    if (!main_node) return head;
+    if (!main_prev) return head;
+
+    main_prev->n = main_node->n;
+    main_node->n = head;
+    return main_node;
+}
+
+
+
+
+
+/* Generic alloc macros */
+#define NEW_EX(k)                                                              \
+  _EX *e = (_EX *)calloc(1, sizeof(_EX));                                      \
+  e->kind = k
+#define NEW_ST(k)                                                              \
+  _STN *s = (_STN *)calloc(1, sizeof(_STN));                                   \
+  s->kind = k
+#define NEW_FN() _FN *f = (_FN *)calloc(1, sizeof(_FN))
+
+/* Constructor declaration macros */
+#define DEFINE_EX_CONSTRUCTOR(name, kind, ...)                                 \
+  static inline _EX *ex_##name(__VA_ARGS__)
+
+#define DEFINE_ST_CONSTRUCTOR(name, kind, ...)                                 \
+  static inline _STN *st_##name(__VA_ARGS__)
+
+#define DEFINE_FN_CONSTRUCTOR(name, ...)                                       \
+  static inline _FN *fn_##name(__VA_ARGS__)
+
+
+
+DEFINE_EX_CONSTRUCTOR(number, EX_NUMBER, int value) {
+  NEW_EX(EX_NUMBER);
+  e->value = value;
+  return e;
+}
+
+DEFINE_EX_CONSTRUCTOR(var, EX_VAR, char *name) {
+  NEW_EX(EX_VAR);
+  e->name = name;
+  return e;
+}
+
+DEFINE_EX_CONSTRUCTOR(binop, EX_BINOP, _EX *l, _TK op, _EX *r) {
+  NEW_EX(EX_BINOP);
+  e->binop.l = l;
+  e->binop.r = r;
+  e->binop.op = op;
+  return e;
+}
+
+DEFINE_EX_CONSTRUCTOR(call, EX_CALL, char *func_name, _EX **args, int argc) {
+  NEW_EX(EX_CALL);
+  e->call.func_name = func_name;
+  e->call.args = args;
+  e->call.argc = argc;
+  return e;
+}
+
+DEFINE_EX_CONSTRUCTOR(string, EX_STRING, char *str) {
+  NEW_EX(EX_STRING);
+  e->string = str;
+  return e;
+}
+
+DEFINE_EX_CONSTRUCTOR(index, EX_INDEX, _EX *array, _EX *idx) {
+  NEW_EX(EX_INDEX);
+  e->index.array = array;
+  e->index.index = idx;
+  return e;
+}
+
+
+DEFINE_ST_CONSTRUCTOR(return, STK_RETURN, _EX *expr) {
+  NEW_ST(STK_RETURN);
+  s->return_expr = expr;
+  return s;
+}
+
+DEFINE_ST_CONSTRUCTOR(var_decl, STK_VAR_DECL, char *name, _EX *init) {
+  NEW_ST(STK_VAR_DECL);
+  s->var_decl.name = name;
+  s->var_decl.init = init;
+  return s;
+}
+
+DEFINE_ST_CONSTRUCTOR(assign, STK_ASSIGN, _EX *lhs, _EX *expr) {
+  NEW_ST(STK_ASSIGN);
+  s->assign.lhs = lhs;
+  s->assign.expr = expr;
+  return s;
+}
+
+DEFINE_ST_CONSTRUCTOR(expr, STK_EXPR, _EX *expr) {
+  NEW_ST(STK_EXPR);
+  s->expr = expr;
+  return s;
+}
+
+DEFINE_ST_CONSTRUCTOR(block, STK_BLOCK, _STN *body) {
+  NEW_ST(STK_BLOCK);
+  s->body = body;
+  return s;
+}
+
+DEFINE_ST_CONSTRUCTOR(if, STK_IF, _EX *cond, _STN *thenb, _STN *elseb) {
+  NEW_ST(STK_IF);
+  s->ifs.cond = cond;
+  s->ifs.thenb = thenb;
+  s->ifs.elseb = elseb;
+  return s;
+}
+
+DEFINE_ST_CONSTRUCTOR(while, STK_WHILE, _EX *cond, _STN *body) {
+  NEW_ST(STK_WHILE);
+  s->whl.cond = cond;
+  s->whl.body = body;
+  return s;
+}
+
+DEFINE_ST_CONSTRUCTOR(for, STK_FOR, _STN *init, _EX *cond, _STN *step, _STN *body) {
+  NEW_ST(STK_FOR);
+  s->fr.init = init;
+  s->fr.cond = cond;
+  s->fr.step = step;
+  s->fr.body = body;
+  return s;
+}
+
+
+
+DEFINE_FN_CONSTRUCTOR(new, char *name, char **params, _TY* params_types, int pac, _STN *body) {
+  NEW_FN();
+  f->name = name;
+  f->params = params;
+  f->param_types = params_types;
+  f->pac = pac;
+  f->body = body;
+  return f;
+}
+
+
+static inline void ex_free(_EX *e) {
+  if (!e)
+    return;
+  switch (e->kind) {
+  case EX_NUMBER:
+    break;
+  case EX_VAR:
+    free(e->name);
+    break;
+  case EX_BINOP:
+    ex_free(e->binop.l);
+    ex_free(e->binop.r);
+    break;
+  case EX_CALL:
+    for (int i = 0; i < e->call.argc; i++)
+      ex_free(e->call.args[i]);
+    free(e->call.args);
+    free(e->call.func_name);
+    break;
+  case EX_ADDR:
+    ex_free(e->addr.expr);
+    break;
+  case EX_DEREF:
+    ex_free(e->deref.expr);
+    break;
+  case EX_STRING:
+    free(e->string);
+    break;
+  case EX_INDEX:
+    ex_free(e->index.array);
+    ex_free(e->index.index);
+    break;
+  }
+  free(e);
+}
+
+static inline void st_free(_STN *s) {
+  if (!s)
+    return;
+  switch (s->kind) {
+  case STK_RETURN:
+    ex_free(s->return_expr);
+    break;
+  case STK_VAR_DECL:
+    ex_free(s->var_decl.init);
+    free(s->var_decl.name);
+    break;
+  case STK_ASSIGN:
+    ex_free(s->assign.expr);
+    break;
+  case STK_EXPR:
+    ex_free(s->expr);
+    break;
+  case STK_BLOCK:
+    st_free(s->body);
+    break;
+  case STK_IF:
+    ex_free(s->ifs.cond);
+    st_free(s->ifs.thenb);
+    st_free(s->ifs.elseb);
+    break;
+  case STK_WHILE:
+    ex_free(s->whl.cond);
+    st_free(s->whl.body);
+    break;
+  case STK_FOR:
+    st_free(s->fr.init);
+    ex_free(s->fr.cond);
+    st_free(s->fr.step);
+    st_free(s->fr.body);
+    break;
+  }
+  _STN *next = s->n;
+  free(s);
+  st_free(next);
+}
+
+static inline void fn_free(_FN *f) {
+  while (f) {
+    _FN *next = f->n;
+    if (f->name)
+      free(f->name);
+    if (f->params) {
+      for (int i = 0; i < f->pac; i++) {
+        free(f->params[i]);
+      }
+      free(f->params);
+    }
+    free(f->param_types);
+    st_free(f->body);
+    free(f);
+    f = next;
+  }
+}
+
+#endif