print ast

4 files changed, 169 insertions, 6 deletions

diff --git a/README.md b/README.md
index 111d122..ca09e1b 100644
--- a/README.md
+++ b/README.md
@@ -30,9 +30,10 @@ All build artifacts are created in the build folder `./build`, including the imp
 
 ```
 Usage: imp [ARGS]
-  -i <program.imp>   interpret program and exit
   (no args)          start REPL
-  -h                 show this message
+  -i <program.imp>   interpret program and exit
+  -a <program.imp>   print ast
+  -h                 print this message
 ```
 
 ### REPL
@@ -44,6 +45,7 @@ Commands:
   %run <path/to/file.imp> interpret program
   %set <var> <val>        set variable
   %print [<var>]          print variable, or all variables
+  %proc                   print declared procedures
   %help                   show this message
 ```
 
diff --git a/include/interpreter.h b/include/interpreter.h
index b21e548..09f1e18 100644
--- a/include/interpreter.h
+++ b/include/interpreter.h
@@ -20,5 +20,7 @@ int interp_ast(context_t context, ASTNode *node);
 int interp_file (context_t context, const char *path);
 int interp_str (context_t context, const char *str);
 
+int print_ast_file (const char *path);
+
 
 #endif
\ No newline at end of file
diff --git a/src/driver.c b/src/driver.c
index c9ac098..8902a55 100644
--- a/src/driver.c
+++ b/src/driver.c
@@ -21,17 +21,22 @@ static int interpret_file(const char *path) {
 
 int main(int argc, char **argv) {
   int opt;
-  while ((opt = getopt(argc, argv, "i:h")) != -1) {
-    switch (opt){
+  while ((opt = getopt(argc, argv, "i:a:h")) != -1) {
+    switch (opt) {
     case 'i': 
       interpret_file(optarg);
       return EXIT_SUCCESS;
+    case 'a':
+      print_ast_file(optarg);
+      return EXIT_SUCCESS;
     case 'h':
     default:
       fprintf(stderr, 
-        "Usage: %s [-i program.imp]\n"
+        "Usage: %s [ARGS]\n"
+        "  (no args)          start REPL\n"
         "  -i <program.imp>   interpret program and exit\n"
-        "  (no args)          start REPL\n",
+        "  -a <program.imp>   print ast\n"
+        "  -h                 print this message\n",
         argv[0]);
       return (opt == 'h') ? EXIT_SUCCESS : EXIT_FAILURE;
     }
diff --git a/src/interpreter.c b/src/interpreter.c
index 8bcc6f4..2de5b63 100644
--- a/src/interpreter.c
+++ b/src/interpreter.c
@@ -97,6 +97,143 @@ void context_print_proc_table(Context *context) {
   hashmap_keys_iter_free(iter);
 }
 
+void ast_print(ASTNode *node, int depth) {
+  int indent = depth * 2;
+  switch (node->type) {
+    case NT_SKIP: {
+      printf("%*sSKIP\n", indent, "");
+      break;
+    }
+    case NT_ASSIGN: {
+      printf("%*sASSIGN %s=%d\n", indent, "", 
+             node->u.d_assign.var->u.d_var.name,
+             node->u.d_assign.aexp->u.d_int.val);
+      break;
+    }
+    case NT_SEQ: {
+      ast_print(node->u.d_seq.stm1, depth);
+      printf("%*sSEQ\n", indent, "");
+      ast_print(node->u.d_seq.stm2, depth);
+      break;
+    }
+    case NT_IF: {
+      printf("%*sIF (", indent, "");
+      ast_print(node->u.d_if.bexp, 0);
+      printf(")\n");
+      ast_print(node->u.d_if.stm1, depth + 1);
+      printf("%*sELSE\n", indent, "");
+      ast_print(node->u.d_if.stm2, depth + 1);
+      break;
+    }
+    case NT_WHILE: {
+      printf("%*sWHILE (", indent, "");
+      ast_print(node->u.d_while.bexp, 0);
+      printf(")\n");
+      ast_print(node->u.d_while.stm, depth + 1);
+      break;
+    }
+    case NT_INT: {
+      printf("%d", node->u.d_int.val);
+      break;
+    }
+    case NT_VAR: {
+      printf("%s", node->u.d_var.name);
+      break;
+    }
+    case NT_AOP: {
+      printf("(");
+      ast_print(node->u.d_aop.aexp1, 0);
+      switch (node->u.d_aop.aop) {
+        case AOP_ADD: printf(" + "); break;
+        case AOP_SUB: printf(" - "); break;
+        case AOP_MUL: printf(" * "); break;
+        default: assert(0);
+      }
+      ast_print(node->u.d_aop.aexp2, 0);
+      printf(")");
+      break;
+    }
+    case NT_BOP: {
+      printf("(");
+      ast_print(node->u.d_bop.bexp1, 0);
+      switch (node->u.d_bop.bop) {
+        case BOP_AND: printf(" && "); break;
+        case BOP_OR:  printf(" || "); break;
+        default: assert(0);
+      }
+      ast_print(node->u.d_bop.bexp2, 0);
+      printf(")");
+      break;
+    }
+    case NT_NOT: {
+      printf("!");
+      ast_print(node->u.d_not.bexp, 0);
+      break;
+    }
+    case NT_ROP: {
+      printf("(");
+      ast_print(node->u.d_rop.aexp1, 0);
+      switch (node->u.d_rop.rop) {
+        case ROP_EQ: printf(" == "); break;
+        case ROP_NE: printf(" != "); break;
+        case ROP_LT: printf(" < "); break;
+        case ROP_LE: printf(" <= "); break;
+        case ROP_GT: printf(" > "); break;
+        case ROP_GE: printf(" >= "); break;
+        default: assert(0);
+      }
+      ast_print(node->u.d_rop.aexp2, 0);
+      printf(")");
+      break;
+    }
+    case NT_LET: {
+      printf("%*sLET %s = ", indent, "", node->u.d_let.var->u.d_var.name);
+      ast_print(node->u.d_let.aexp, 0);
+      printf("\n");
+      ast_print(node->u.d_let.stm, depth + 1);
+      break;
+    }
+    case NT_PROCDECL: {
+      printf("%*sPROC %s(", indent, "", node->u.d_procdecl.name);
+      ASTNodeList *args = node->u.d_procdecl.args;
+      while (args) {
+        printf("%s", args->node->u.d_var.name);
+        args = args->next;
+        if (args) printf(", ");
+      }
+      printf("; ");
+      ASTNodeList *vargs = node->u.d_procdecl.vargs;
+      while (vargs) {
+        printf("%s", vargs->node->u.d_var.name);
+        vargs = vargs->next;
+        if (vargs) printf(", ");
+      }
+      printf(")\n");
+      ast_print(node->u.d_procdecl.stm, depth + 1);
+      break;
+    }
+    case NT_PROCCALL: {
+      printf("%*sCALL %s(", indent, "", node->u.d_proccall.name);
+      ASTNodeList *args = node->u.d_proccall.args;
+      while (args) {
+        printf("%s", args->node->u.d_var.name);
+        args = args->next;
+        if (args) printf(", ");
+      }
+      printf("; ");
+      ASTNodeList *vargs = node->u.d_proccall.vargs;
+      while (vargs) {
+        printf("%s", vargs->node->u.d_var.name);
+        vargs = vargs->next;
+        if (vargs) printf(", ");
+      }
+      printf(")\n");
+      break;
+    }
+    default: assert(0);
+  }
+}
+
 static int eval_aexpr(Context *context, ASTNode *node) {
   switch (node->type) {
     case NT_INT: return node->u.d_int.val;
@@ -278,4 +415,21 @@ int interp_str (Context *context, const char *str) {
   }
   yy_delete_buffer(buf);
   return 0;
+}
+
+int print_ast_file (const char *path) {
+  yyin = fopen(path, "r");
+  if (!yyin) {
+    return -1;
+  }
+  yyrestart(yyin);
+  if (yyparse()) {
+    ast_free(ast_root);
+    fclose(yyin);
+    return -1;
+  }
+  ast_print(ast_root, 0);
+  ast_free(ast_root);
+  fclose(yyin);
+  return 0;
 }
\ No newline at end of file