cli: Add AST pretty-printing option (--ast-dump)

Parsing can be a complex process, and it's not always easy to get a
clear picture of what's happening with the AST.  This commit adds a new
feature to the CLI that allows us to pretty-print the AST (outputs to
stdout), making it easier to visualize the tree structure and understand
how the parser is working.

The new --ast-dump option generates a human-readable representation of
the AST, including node types, values, and child relationships. This
information can be invaluable for debugging and understanding the
parser's behavior.

Signed-off-by: Johnny Richard <johnny@johnnyrichard.com>
Reviewed-by: Carlos Maniero <carlos@maniero.me>

4 files changed, 295 insertions, 3 deletions

diff --git a/src/ast.h b/src/ast.h
index bdace51..69fc918 100644
--- a/src/ast.h
+++ b/src/ast.h
@@ -54,7 +54,8 @@ typedef enum ast_binary_operation_kind_t
   AST_BINOP_ADITION,
   AST_BINOP_SUBTRACTION,
   AST_BINOP_MULTIPLICATION,
-  AST_BINOP_DIVISION
+  AST_BINOP_DIVISION,
+  AST_BINOP_N
 } ast_binary_operation_kind_t;
 
 typedef struct ast_binary_operation_t
diff --git a/src/ast_pretty_printer.c b/src/ast_pretty_printer.c
new file mode 100644
index 0000000..df17763
--- /dev/null
+++ b/src/ast_pretty_printer.c
@@ -0,0 +1,224 @@
+/*
+ * Copyright (C) 2023 Johnny Richard
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+#include "ast_pretty_printer.h"
+#include "ast.h"
+#include "string_view.h"
+
+#include <assert.h>
+#include <stdarg.h>
+#include <stdio.h>
+
+const char operation_kinds[AST_BINOP_N] = { [AST_BINOP_ADITION] = '+',
+                                            [AST_BINOP_SUBTRACTION] = '-',
+                                            [AST_BINOP_MULTIPLICATION] = '*',
+                                            [AST_BINOP_DIVISION] = '/' };
+
+void
+ast_pretty_printer_init(ast_pretty_printer_t *printer, FILE *stream);
+
+inline static void
+ast_pretty_printer_add_indentation(ast_pretty_printer_t *printer);
+
+inline static void
+ast_pretty_printer_rm_indentation(ast_pretty_printer_t *printer);
+
+static void
+ast_pretty_printer_print_indentation(ast_pretty_printer_t *printer);
+
+static int
+ast_pretty_printer_printf(ast_pretty_printer_t *printer, const char *fmt, ...);
+
+void
+ast_pretty_printer_init(ast_pretty_printer_t *printer, FILE *stream)
+{
+  assert(printer);
+  assert(stream);
+
+  printer->indentation_fmt = 0;
+  printer->indentation_lst_children = 0;
+  printer->indentation_level = 0;
+  printer->stream = stream;
+}
+
+void
+ast_pretty_printer_print_ast(ast_pretty_printer_t *printer, ast_node_t *ast)
+{
+  assert(ast);
+
+  switch (ast->kind) {
+    case AST_BINARY_OPERATION: {
+      ast_binary_operation_t binop = ast->data.binary_operation;
+      ast_pretty_printer_printf(printer, "BinaryOperation operation='%c'\n", operation_kinds[binop.kind]);
+
+      ast_pretty_printer_add_indentation(printer);
+      {
+        ast_pretty_printer_printf(printer, "left:\n");
+
+        ast_pretty_printer_add_indentation(printer);
+        {
+          printer->indentation_lst_children |= 1 << printer->indentation_level;
+          ast_pretty_printer_print_ast(printer, binop.left);
+
+          ast_pretty_printer_rm_indentation(printer);
+        }
+
+        printer->indentation_lst_children |= 1 << printer->indentation_level;
+        ast_pretty_printer_printf(printer, "right:\n");
+
+        ast_pretty_printer_add_indentation(printer);
+        {
+          printer->indentation_lst_children |= 1 << printer->indentation_level;
+          ast_pretty_printer_print_ast(printer, binop.right);
+
+          ast_pretty_printer_rm_indentation(printer);
+        }
+
+        ast_pretty_printer_rm_indentation(printer);
+      }
+      break;
+    }
+    case AST_FUNCTION_DECLARATION: {
+      ast_function_declaration_t function = ast->data.function;
+      ast_pretty_printer_printf(printer, "FunctionDecl name='" SVFMT "'\n", SVARG(&function.identifier.name));
+
+      ast_pretty_printer_add_indentation(printer);
+      for (size_t i = 0; i < function.body->size; ++i) {
+        if (i + 1 >= function.body->size) {
+          printer->indentation_lst_children |= 1 << printer->indentation_level;
+        }
+        ast_pretty_printer_print_ast(printer, vector_at(function.body, i));
+      }
+      ast_pretty_printer_rm_indentation(printer);
+      break;
+    }
+    case AST_LITERAL: {
+      ast_literal_t literal = ast->data.literal;
+
+      switch (literal.kind) {
+        case AST_LITERAL_INTEGER: {
+          printer->indentation_lst_children |= 1 << printer->indentation_level;
+          ast_pretty_printer_printf(printer, "Literal type=i32 value='%d'\n", literal.value.integer);
+          break;
+        }
+      }
+      break;
+    }
+    case AST_RETURN_STMT: {
+      ast_return_stmt_t return_stmt = ast->data.return_stmt;
+      ast_pretty_printer_printf(printer, "ReturnStmt\n");
+
+      ast_pretty_printer_add_indentation(printer);
+      {
+        printer->indentation_lst_children |= 1 << printer->indentation_level;
+        ast_pretty_printer_print_ast(printer, return_stmt.argument);
+
+        ast_pretty_printer_rm_indentation(printer);
+      }
+      break;
+    }
+    case AST_VARIABLE_DECLARATION: {
+      ast_variable_declaration_t var_decl = ast->data.variable_declaration;
+      ast_pretty_printer_printf(printer, "VariableDecl name='" SVFMT "'\n", SVARG(&var_decl.identifier.name));
+
+      ast_pretty_printer_add_indentation(printer);
+      {
+        printer->indentation_lst_children |= 1 << printer->indentation_level;
+        ast_pretty_printer_print_ast(printer, var_decl.value);
+
+        ast_pretty_printer_rm_indentation(printer);
+      }
+      break;
+    }
+    case AST_VARIABLE_ASSIGNMENT: {
+      ast_variable_assignment_t var_assign = ast->data.variable_assignment;
+      ast_pretty_printer_printf(printer, "VariableAssignment name='" SVFMT "'\n", SVARG(&var_assign.identifier->name));
+
+      ast_pretty_printer_add_indentation(printer);
+      {
+        printer->indentation_lst_children |= 1 << printer->indentation_level;
+        ast_pretty_printer_printf(printer, "expression:\n");
+
+        ast_pretty_printer_add_indentation(printer);
+        {
+          printer->indentation_lst_children |= 1 << printer->indentation_level;
+          ast_pretty_printer_print_ast(printer, var_assign.expression);
+
+          ast_pretty_printer_rm_indentation(printer);
+        }
+        ast_pretty_printer_rm_indentation(printer);
+      }
+
+      break;
+    }
+    case AST_VARIABLE: {
+      ast_variable_t var = ast->data.variable;
+      ast_pretty_printer_printf(printer, "Variable name='" SVFMT "'\n", SVARG(&var.identifier->name));
+      break;
+    }
+    case AST_UNKOWN_NODE: {
+      fprintf(printer->stream, "AST_UNKOWN_NODE\n");
+      break;
+    }
+  }
+}
+
+static int
+ast_pretty_printer_printf(ast_pretty_printer_t *printer, const char *fmt, ...)
+{
+  int ret = 0;
+  va_list args;
+
+  va_start(args, fmt);
+  ast_pretty_printer_print_indentation(printer);
+  ret = vprintf(fmt, args);
+  va_end(args);
+
+  return ret;
+}
+
+inline static void
+ast_pretty_printer_add_indentation(ast_pretty_printer_t *printer)
+{
+  printer->indentation_level++;
+  printer->indentation_fmt |= 1 << printer->indentation_level;
+}
+
+inline static void
+ast_pretty_printer_rm_indentation(ast_pretty_printer_t *printer)
+{
+  printer->indentation_level--;
+}
+
+static void
+ast_pretty_printer_print_indentation(ast_pretty_printer_t *printer)
+{
+  for (size_t i = 0; i <= printer->indentation_level; ++i) {
+    if ((printer->indentation_fmt >> i) & 1) {
+      if (i + 1 > printer->indentation_level && (printer->indentation_lst_children & (1 << i))) {
+        printer->indentation_lst_children ^= (1 << i);
+        printer->indentation_fmt ^= (1 << i);
+        fprintf(printer->stream, "└─ ");
+      } else if (i + 1 > printer->indentation_level) {
+        fprintf(printer->stream, "├─ ");
+      } else {
+        fprintf(printer->stream, "│  ");
+      }
+    } else {
+      fprintf(printer->stream, "  ");
+    }
+  }
+}
diff --git a/src/ast_pretty_printer.h b/src/ast_pretty_printer.h
new file mode 100644
index 0000000..d7e36b7
--- /dev/null
+++ b/src/ast_pretty_printer.h
@@ -0,0 +1,41 @@
+/*
+ * Copyright (C) 2023 Johnny Richard
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+#ifndef AST_PRETTY_PRINTER_H
+#define AST_PRETTY_PRINTER_H
+
+#include "ast.h"
+
+#include <stdbool.h>
+#include <stdint.h>
+#include <stdio.h>
+
+typedef struct ast_pretty_printer_t
+{
+  uint64_t indentation_fmt;
+  uint64_t indentation_lst_children;
+  uint32_t indentation_level;
+  bool last_children;
+  FILE *stream;
+} ast_pretty_printer_t;
+
+void
+ast_pretty_printer_init(ast_pretty_printer_t *printer, FILE *stream);
+
+void
+ast_pretty_printer_print_ast(ast_pretty_printer_t *printer, ast_node_t *ast);
+
+#endif /* AST_PRETTY_PRINTER_H */
diff --git a/src/main.c b/src/main.c
index e3caec1..4f68256 100644
--- a/src/main.c
+++ b/src/main.c
@@ -19,6 +19,7 @@
 #include <string.h>
 
 #include "ast.h"
+#include "ast_pretty_printer.h"
 #include "gas_assembly_generator.h"
 #include "lexer.h"
 #include "parser.h"
@@ -34,6 +35,14 @@ generate_gas_x86_64_linux(ast_node_t *func)
 }
 
 static void
+pretty_print_ast(ast_node_t *ast)
+{
+  ast_pretty_printer_t printer;
+  ast_pretty_printer_init(&printer, stdout);
+  ast_pretty_printer_print_ast(&printer, ast);
+}
+
+static void
 print_usage(void)
 {
   fputs("pipac <filename.pipa>\n", stderr);
@@ -58,7 +67,20 @@ main(int argc, char **argv)
     return EXIT_FAILURE;
   }
 
-  char *filepath = argv[1];
+  char *filepath;
+  bool should_dump_ast = false;
+
+  // TODO: Handle command line arguments properly
+  if (argc < 3) {
+    filepath = argv[1];
+  } else {
+    if (strcmp(argv[1], "--ast-dump") != 0) {
+      print_usage();
+      return EXIT_FAILURE;
+    }
+    should_dump_ast = true;
+    filepath = argv[2];
+  }
 
   lexer_t lexer;
   lexer_init(&lexer, filepath);
@@ -74,7 +96,11 @@ main(int argc, char **argv)
     return EXIT_FAILURE;
   }
 
-  generate_gas_x86_64_linux(func);
+  if (should_dump_ast) {
+    pretty_print_ast(func);
+  } else {
+    generate_gas_x86_64_linux(func);
+  }
 
   scope_destroy(scope);
   ast_node_destroy(func);