Improved formatting

5 files changed, 228 insertions, 20 deletions

diff --git a/src/ast.h b/src/ast.h
index 3fdff98c..43972371 100644
--- a/src/ast.h
+++ b/src/ast.h
@@ -65,6 +65,8 @@ typedef struct ast_list_s {
 } ast_list_t;
 
 typedef struct arg_ast_s {
+    file_t *file;
+    const char *start, *end;
     const char *name;
     type_ast_t *type;
     ast_t *value;
diff --git a/src/formatter.c b/src/formatter.c
index 83d1b7a0..2b367e3d 100644
--- a/src/formatter.c
+++ b/src/formatter.c
@@ -2,7 +2,6 @@
 
 #include <gc.h>
 #include <setjmp.h>
-#include <string.h>
 
 #include "ast.h"
 #include "formatter.h"
@@ -14,6 +13,26 @@
 #include "stdlib/tables.h"
 #include "stdlib/text.h"
 
+#define MAX_WIDTH 100
+
+#define must(expr)                                                                                                     \
+    ({                                                                                                                 \
+        OptionalText_t _expr = expr;                                                                                   \
+        if (_expr.length < 0) return NONE_TEXT;                                                                        \
+        (Text_t) _expr;                                                                                                \
+    })
+
+const Text_t single_indent = Text("    ");
+
+static void add_line(Text_t *code, Text_t line, Text_t indent) {
+    if (code->length == 0) {
+        *code = line;
+    } else {
+        if (line.length > 0) *code = Texts(*code, "\n", indent, line);
+        else *code = Texts(*code, "\n");
+    }
+}
+
 OptionalText_t next_comment(Table_t comments, const char **pos, const char *end) {
     for (const char *p = *pos; p < end; p++) {
         const char **comment_end = Table$get(comments, &p, parse_comments_info);
@@ -25,25 +44,209 @@ OptionalText_t next_comment(Table_t comments, const char **pos, const char *end)
     return NONE_TEXT;
 }
 
-Text_t format_code(ast_t *ast, Table_t comments) {
-    (void)comments;
+static bool range_has_comment(const char *start, const char *end, Table_t comments) {
+    OptionalText_t comment = next_comment(comments, &start, end);
+    return (comment.length >= 0);
+}
+
+static bool should_have_blank_line(ast_t *ast) {
     switch (ast->tag) {
+    case If:
+    case When:
+    case Repeat:
+    case While:
+    case For:
+    case Block:
+    case FunctionDef:
+    case StructDef:
+    case EnumDef:
+    case LangDef:
+    case ConvertDef:
+    case Extend: return true;
+    default: return false;
+    }
+}
+
+static OptionalText_t format_inline_type(type_ast_t *type, Table_t comments) {
+    if (range_has_comment(type->start, type->end, comments)) return NONE_TEXT;
+    switch (type->tag) {
     default: {
-        Text_t code = Text$from_strn(ast->start, (int64_t)(ast->end - ast->start));
-        return Text$replace(code, Text("\t"), Text("    "));
+        Text_t code = Text$from_strn(type->start, (int64_t)(type->end - type->start));
+        if (Text$has(code, Text("\n"))) return NONE_TEXT;
+        return Text$replace(code, Text("\t"), single_indent);
+    }
+    }
+}
+
+static Text_t format_type(type_ast_t *type, Table_t comments, Text_t indent) {
+    (void)comments, (void)indent;
+    switch (type->tag) {
+    default: {
+        OptionalText_t inline_type = format_inline_type(type, comments);
+        if (inline_type.length >= 0) return inline_type;
+        Text_t code = Text$from_strn(type->start, (int64_t)(type->end - type->start));
+        return Text$replace(code, Text("\t"), single_indent);
+    }
+    }
+}
+
+static OptionalText_t format_inline_arg(arg_ast_t *arg, Table_t comments) {
+    if (range_has_comment(arg->start, arg->end, comments)) return NONE_TEXT;
+    Text_t code = Text$from_str(arg->name);
+    if (arg->type) code = Texts(code, ":", must(format_inline_type(arg->type, comments)));
+    if (arg->value) code = Texts(code, " = ", must(format_inline_code(arg->value, comments)));
+    return code;
+}
+
+static Text_t format_arg(arg_ast_t *arg, Table_t comments, Text_t indent) {
+    (void)comments;
+    Text_t code = Text$from_str(arg->name);
+    if (arg->type) code = Texts(code, ":", format_type(arg->type, comments, indent));
+    if (arg->value) code = Texts(code, " = ", format_code(arg->value, comments, indent));
+    return code;
+}
+
+static OptionalText_t format_inline_args(arg_ast_t *args, Table_t comments) {
+    Text_t code = EMPTY_TEXT;
+    for (; args; args = args->next) {
+        if (args->next && args->type == args->next->type && args->value == args->next->value) {
+            code = Texts(code, must(Text$from_str(args->name)), ",");
+        } else {
+            code = Texts(code, must(format_inline_arg(args, comments)));
+            if (args->next) code = Texts(code, ", ");
+        }
+        if (args->next && range_has_comment(args->end, args->next->start, comments)) return NONE_TEXT;
+    }
+    return code;
+}
+
+static Text_t format_args(arg_ast_t *args, Table_t comments, Text_t indent) {
+    OptionalText_t inline_args = format_inline_args(args, comments);
+    if (inline_args.length >= 0) return inline_args;
+    Text_t code = EMPTY_TEXT;
+    for (; args; args = args->next) {
+        if (args->next && args->type == args->next->type && args->value == args->next->value) {
+            code = Texts(code, must(Text$from_str(args->name)), ",");
+        } else {
+            add_line(&code, Texts(format_arg(args, comments, indent), ","), indent);
+        }
     }
+    return code;
+}
+
+ast_t *unwrap_block(ast_t *ast) {
+    while (ast->tag == Block && Match(ast, Block)->statements && Match(ast, Block)->statements->next == NULL) {
+        ast = Match(ast, Block)->statements->ast;
     }
+    return ast;
 }
 
 OptionalText_t format_inline_code(ast_t *ast, Table_t comments) {
-    for (const char *p = ast->start; p < ast->end; p++) {
-        if (*p == '\n' || *p == '\r') return NONE_TEXT;
+    if (range_has_comment(ast->start, ast->end, comments)) return NONE_TEXT;
+    switch (ast->tag) {
+    case Block: {
+        ast_list_t *statements = Match(ast, Block)->statements;
+        if (statements == NULL) return Text("pass");
+        else if (statements->next == NULL) return format_inline_code(statements->ast, comments);
+        else return NONE_TEXT;
+    }
+    case FunctionDef: return NONE_TEXT;
+    case If: {
+        DeclareMatch(if_, ast, If);
+
+        if (if_->else_body == NULL && if_->condition->tag != Declare) {
+            ast_t *stmt = unwrap_block(if_->body);
+            switch (stmt->tag) {
+            case Return:
+            case Skip:
+            case Stop:
+                return Texts(must(format_inline_code(stmt, comments)), " if ",
+                             must(format_inline_code(if_->condition, comments)));
+            default: break;
+            }
+        }
+
+        Text_t code = Texts("if ", must(format_inline_code(if_->condition, comments)), " then ",
+                            must(format_inline_code(if_->body, comments)));
+        if (if_->else_body) code = Texts(code, " else ", must(format_inline_code(if_->else_body, comments)));
+        return code;
+    }
+    default: {
+        Text_t code = Text$from_strn(ast->start, (int64_t)(ast->end - ast->start));
+        if (Text$has(code, Text("\n"))) return NONE_TEXT;
+        return code;
     }
-    const char *pos = ast->start;
-    OptionalText_t comment = next_comment(comments, &pos, ast->end);
-    if (comment.length >= 0) return NONE_TEXT;
+    }
+}
+
+Text_t format_code(ast_t *ast, Table_t comments, Text_t indent) {
+    OptionalText_t inlined = format_inline_code(ast, comments);
+    bool inlined_fits = (inlined.length >= 0 && indent.length + inlined.length <= MAX_WIDTH);
+
     switch (ast->tag) {
-    default: return Text$from_strn(ast->start, (int64_t)(ast->end - ast->start));
+    case Block: {
+        Text_t code = EMPTY_TEXT;
+        bool gap_before_comment = false;
+        const char *comment_pos = ast->start;
+        for (ast_list_t *stmt = Match(ast, Block)->statements; stmt; stmt = stmt->next) {
+            for (OptionalText_t comment;
+                 (comment = next_comment(comments, &comment_pos, stmt->ast->start)).length > 0;) {
+                if (gap_before_comment) {
+                    add_line(&code, Text(""), indent);
+                    gap_before_comment = false;
+                }
+                add_line(&code, Texts("# ", Text$trim(comment, Text("# \t\r\n\v"), true, true)), indent);
+            }
+
+            add_line(&code, format_code(stmt->ast, comments, indent), indent);
+            comment_pos = stmt->ast->end;
+
+            if (should_have_blank_line(stmt->ast) && stmt->next) add_line(&code, Text(""), indent);
+            else gap_before_comment = true;
+        }
+
+        for (OptionalText_t comment; (comment = next_comment(comments, &comment_pos, ast->end)).length > 0;) {
+            if (gap_before_comment) {
+                add_line(&code, Text(""), indent);
+                gap_before_comment = false;
+            }
+            add_line(&code, Texts("# ", Text$trim(comment, Text("# \t\r\n\v"), true, true)), indent);
+        }
+        return code;
+    }
+    case If: {
+        DeclareMatch(if_, ast, If);
+        if (inlined_fits && if_->else_body == NULL) return inlined;
+
+        Text_t code = Texts("if ", format_code(if_->condition, comments, indent), "\n", indent, single_indent,
+                            format_code(if_->body, comments, Texts(indent, single_indent)));
+        if (if_->else_body)
+            code = Texts(code, "\n", indent, "else \n", indent, single_indent,
+                         format_code(if_->else_body, comments, Texts(indent, single_indent)));
+        return code;
+    }
+    case FunctionDef: {
+        DeclareMatch(func, ast, FunctionDef);
+        // ast_t *name;
+        // arg_ast_t *args;
+        // type_ast_t *ret_type;
+        // ast_t *body;
+        // ast_t *cache;
+        // bool is_inline;
+        Text_t code =
+            Texts("func ", format_code(func->name, comments, indent), "(", format_args(func->args, comments, indent));
+        if (func->ret_type)
+            code = Texts(code, func->args ? Text(" -> ") : Text("-> "), format_type(func->ret_type, comments, indent));
+        if (func->cache) code = Texts(code, "; cache=", format_code(func->cache, comments, indent));
+        if (func->is_inline) code = Texts(code, "; inline");
+        code = Texts(code, ")\n", single_indent, format_code(func->body, comments, Texts(indent, single_indent)));
+        return Texts(code, "\n");
+    }
+    default: {
+        if (inlined_fits) return inlined;
+        Text_t code = Text$from_strn(ast->start, (int64_t)(ast->end - ast->start));
+        return Text$replace(code, Text("\t"), single_indent);
+    }
     }
 }
 
@@ -77,11 +280,11 @@ Text_t format_file(const char *path) {
     const char *fmt_pos = file->text;
     Text_t code = EMPTY_TEXT;
     for (OptionalText_t comment; (comment = next_comment(ctx.comments, &fmt_pos, ast->start)).length > 0;) {
-        code = Texts(code, comment, "\n");
+        code = Texts(code, Texts("# ", Text$trim(comment, Text("# \t\r\n\v"), true, true)), "\n");
     }
-    code = Texts(code, format_code(ast, ctx.comments));
+    code = Texts(code, format_code(ast, ctx.comments, EMPTY_TEXT));
     for (OptionalText_t comment; (comment = next_comment(ctx.comments, &fmt_pos, ast->start)).length > 0;) {
-        code = Texts(code, comment, "\n");
+        code = Texts(code, Texts("# ", Text$trim(comment, Text("# \t\r\n\v"), true, true)), "\n");
     }
     return code;
 }
diff --git a/src/formatter.h b/src/formatter.h
index dd4c7815..bead52ec 100644
--- a/src/formatter.h
+++ b/src/formatter.h
@@ -6,6 +6,6 @@
 #include "stdlib/datatypes.h"
 
 Text_t format_file(const char *path);
-Text_t format_code(ast_t *ast, Table_t comments);
+Text_t format_code(ast_t *ast, Table_t comments, Text_t indentation);
 OptionalText_t format_inline_code(ast_t *ast, Table_t comments);
 OptionalText_t next_comment(Table_t comments, const char **pos, const char *end);
diff --git a/src/parse/functions.c b/src/parse/functions.c
index e820182b..b50519b7 100644
--- a/src/parse/functions.c
+++ b/src/parse/functions.c
@@ -26,6 +26,7 @@ arg_ast_t *parse_args(parse_ctx_t *ctx, const char **pos) {
         type_ast_t *type = NULL;
 
         typedef struct name_list_s {
+            const char *start, *end;
             const char *name;
             struct name_list_s *next;
         } name_list_t;
@@ -35,21 +36,22 @@ arg_ast_t *parse_args(parse_ctx_t *ctx, const char **pos) {
             whitespace(ctx, pos);
             const char *name = get_id(pos);
             if (!name) break;
+            const char *name_start = *pos;
             whitespace(ctx, pos);
 
             if (match(pos, ":")) {
                 type = expect(ctx, *pos - 1, pos, parse_type, "I expected a type here");
-                names = new (name_list_t, .name = name, .next = names);
                 whitespace(ctx, pos);
                 if (match(pos, "="))
                     default_val = expect(ctx, *pos - 1, pos, parse_term, "I expected a value after this '='");
+                names = new (name_list_t, .start = name_start, .end = *pos, .name = name, .next = names);
                 break;
             } else if (strncmp(*pos, "==", 2) != 0 && match(pos, "=")) {
                 default_val = expect(ctx, *pos - 1, pos, parse_term, "I expected a value after this '='");
-                names = new (name_list_t, .name = name, .next = names);
+                names = new (name_list_t, .start = name_start, .end = *pos, .name = name, .next = names);
                 break;
             } else if (name) {
-                names = new (name_list_t, .name = name, .next = names);
+                names = new (name_list_t, .start = name_start, .end = *pos, .name = name, .next = names);
                 spaces(pos);
                 if (!match(pos, ",")) break;
             } else {
@@ -64,7 +66,8 @@ arg_ast_t *parse_args(parse_ctx_t *ctx, const char **pos) {
 
         REVERSE_LIST(names);
         for (; names; names = names->next)
-            args = new (arg_ast_t, .name = names->name, .type = type, .value = default_val, .next = args);
+            args = new (arg_ast_t, .start = names->start, .end = names->end, .name = names->name, .type = type,
+                        .value = default_val, .next = args);
 
         if (!match_separator(ctx, pos)) break;
     }
diff --git a/src/parse/suffixes.c b/src/parse/suffixes.c
index 58d951c4..cb54b2f6 100644
--- a/src/parse/suffixes.c
+++ b/src/parse/suffixes.c
@@ -132,7 +132,7 @@ ast_t *parse_method_call_suffix(parse_ctx_t *ctx, ast_t *self) {
             if (name) parser_err(ctx, arg_start, pos, "I expected an argument here");
             break;
         }
-        args = new (arg_ast_t, .name = name, .value = arg, .next = args);
+        args = new (arg_ast_t, .start = arg_start, .end = arg->end, .name = name, .value = arg, .next = args);
         if (!match_separator(ctx, &pos)) break;
     }
     REVERSE_LIST(args);