1 files changed, 87 insertions, 5 deletions

diff --git a/src/ast.h b/src/ast.h
index 6c326cc..5cfe5d6 100644
--- a/src/ast.h
+++ b/src/ast.h
@@ -7,7 +7,7 @@
 #include <stdio.h>
 
 
-typedef enum { EX_NUMBER, EX_VAR, EX_BINOP, EX_CALL, EX_ADDR, EX_DEREF, EX_STRING, EX_INDEX } _EK;
+typedef enum { EX_NUMBER, EX_VAR, EX_BINOP, EX_CALL, EX_ADDR, EX_DEREF, EX_STRING, EX_INDEX, EX_TERNARY, EX_CAST } _EK;
 
 typedef struct _EX {
   _EK kind;
@@ -35,11 +35,20 @@ typedef struct _EX {
       struct _EX *array; // array expression
       struct _EX *index; // index expression
     } index;
+    struct {             // EX_TERNARY: cond ? then : else
+      struct _EX *cond;
+      struct _EX *then_expr;
+      struct _EX *else_expr;
+    } ternary;
+    struct {             // EX_CAST: (type)expr
+      _TY to;
+      struct _EX *expr;
+    } cast;
   };
 } _EX;
 
 
-typedef enum { STK_RETURN, STK_VAR_DECL, STK_ASSIGN, STK_EXPR, STK_BLOCK, STK_IF, STK_WHILE, STK_FOR } _STK;
+typedef enum { STK_RETURN, STK_VAR_DECL, STK_ASSIGN, STK_EXPR, STK_BLOCK, STK_IF, STK_WHILE, STK_FOR, STK_DOWHILE, STK_BREAK, STK_CONTINUE, STK_GLOBAL } _STK;
 
 typedef struct _STN {
   _STK kind;
@@ -71,6 +80,16 @@ typedef struct _STN {
       struct _STN *step; // may be NULL (an expr stmt)
       struct _STN *body;
     } fr;
+    struct {           // STK_DOWHILE
+      struct _STN *body;
+      _EX *cond;
+    } dowhl;
+    /* STK_BREAK and STK_CONTINUE carry no payload */
+    struct {           // STK_GLOBAL: a global variable declaration
+      char *name;
+      _EX  *init;      /* may be NULL; must be constant (EX_NUMBER/EX_STRING) */
+      _TY   type;
+    } global;
   };
   struct _STN *n; // linked list
 } _STN;
@@ -83,6 +102,7 @@ typedef struct _FN {
   _TY *param_types;
   int pac;
   _STN *body;
+  _TY ret_type;      /* return type of this function */
   struct _FN *n;
 } _FN;
 
@@ -180,7 +200,6 @@ _FN *fnlist_prepare(_FN *head) {
 
 
 
-/* Generic alloc macros */
 #define NEW_EX(k)                                                              \
   _EX *e = (_EX *)calloc(1, sizeof(_EX));                                      \
   e->kind = k
@@ -189,7 +208,6 @@ _FN *fnlist_prepare(_FN *head) {
   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__)
 
@@ -299,15 +317,60 @@ DEFINE_ST_CONSTRUCTOR(for, STK_FOR, _STN *init, _EX *cond, _STN *step, _STN *bod
   return s;
 }
 
+DEFINE_ST_CONSTRUCTOR(dowhile, STK_DOWHILE, _STN *body, _EX *cond) {
+  NEW_ST(STK_DOWHILE);
+  s->dowhl.body = body;
+  s->dowhl.cond = cond;
+  return s;
+}
+
+DEFINE_ST_CONSTRUCTOR(break, STK_BREAK) {
+  NEW_ST(STK_BREAK);
+  return s;
+}
+
+DEFINE_ST_CONSTRUCTOR(continue, STK_CONTINUE) {
+  NEW_ST(STK_CONTINUE);
+  return s;
+}
+
+DEFINE_ST_CONSTRUCTOR(global, STK_GLOBAL, char *name, _TY type, _EX *init) {
+  NEW_ST(STK_GLOBAL);
+  s->global.name = name;
+  s->global.type = type;
+  s->global.init = init;
+  return s;
+}
+
+static inline _EX *ex_ternary(_EX *cond, _EX *then_expr, _EX *else_expr) {
+  _EX *e = (_EX *)calloc(1, sizeof(_EX));
+  if (!e) { fprintf(stderr, "[AST] OOM\n"); exit(1); }
+  e->kind = EX_TERNARY;
+  e->ternary.cond      = cond;
+  e->ternary.then_expr = then_expr;
+  e->ternary.else_expr = else_expr;
+  return e;
+}
+
+static inline _EX *ex_cast(_TY to, _EX *expr) {
+  _EX *e = (_EX *)calloc(1, sizeof(_EX));
+  if (!e) { fprintf(stderr, "[AST] OOM\n"); exit(1); }
+  e->kind = EX_CAST;
+  e->cast.to   = to;
+  e->cast.expr = expr;
+  return e;
+}
 
 
-DEFINE_FN_CONSTRUCTOR(new, char *name, char **params, _TY* params_types, int pac, _STN *body) {
+
+DEFINE_FN_CONSTRUCTOR(new, char *name, char **params, _TY* params_types, int pac, _STN *body, _TY ret_type) {
   NEW_FN();
   f->name = name;
   f->params = params;
   f->param_types = params_types;
   f->pac = pac;
   f->body = body;
+  f->ret_type = ret_type;
   return f;
 }
 
@@ -344,6 +407,14 @@ static inline void ex_free(_EX *e) {
     ex_free(e->index.array);
     ex_free(e->index.index);
     break;
+  case EX_TERNARY:
+    ex_free(e->ternary.cond);
+    ex_free(e->ternary.then_expr);
+    ex_free(e->ternary.else_expr);
+    break;
+  case EX_CAST:
+    ex_free(e->cast.expr);
+    break;
   }
   free(e);
 }
@@ -383,6 +454,17 @@ static inline void st_free(_STN *s) {
     st_free(s->fr.step);
     st_free(s->fr.body);
     break;
+  case STK_DOWHILE:
+    st_free(s->dowhl.body);
+    ex_free(s->dowhl.cond);
+    break;
+  case STK_BREAK:
+  case STK_CONTINUE:
+    break;
+  case STK_GLOBAL:
+    free(s->global.name);
+    ex_free(s->global.init);
+    break;
   }
   _STN *next = s->n;
   free(s);