3 files changed, 31 insertions, 13 deletions

diff --git a/src/formatter/formatter.c b/src/formatter/formatter.c
index 8b8aeb1b..98f0f59b 100644
--- a/src/formatter/formatter.c
+++ b/src/formatter/formatter.c
@@ -414,9 +414,6 @@ Text_t format_code(ast_t *ast, Table_t comments, Text_t indent) {
         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 (int blanks = suggested_blank_lines(stmt->ast); blanks > 0; blanks--)
-                add_line(&code, Text(""), indent);
-
             for (OptionalText_t comment;
                  (comment = next_comment(comments, &comment_pos, stmt->ast->start)).length > 0;) {
                 if (gap_before_comment) {
@@ -430,8 +427,7 @@ Text_t format_code(ast_t *ast, Table_t comments, Text_t indent) {
             comment_pos = stmt->ast->end;
 
             if (stmt->next) {
-                for (int blanks = suggested_blank_lines(stmt->ast) - suggested_blank_lines(stmt->next->ast); blanks > 0;
-                     blanks--)
+                for (int blanks = suggested_blank_lines(stmt->ast, stmt->next->ast); blanks > 0; blanks--)
                     add_line(&code, Text(""), indent);
             } else gap_before_comment = true;
         }
diff --git a/src/formatter/utils.c b/src/formatter/utils.c
index 1f4e64d8..5619e6ce 100644
--- a/src/formatter/utils.c
+++ b/src/formatter/utils.c
@@ -38,24 +38,45 @@ bool range_has_comment(const char *start, const char *end, Table_t comments) {
     return (comment.length >= 0);
 }
 
-CONSTFUNC int suggested_blank_lines(ast_t *ast) {
-    switch (ast->tag) {
+CONSTFUNC int suggested_blank_lines(ast_t *first, ast_t *second) {
+    if (second == NULL) return 0;
+
+    if (first->tag == Use && second->tag != Use) return 1;
+
+    switch (first->tag) {
+    case If:
+    case When:
+    case Repeat:
+    case While:
+    case For:
+    case Block:
+    case Defer:
+    case ConvertDef:
+    case FunctionDef:
+    case StructDef:
+    case EnumDef:
+    case LangDef:
+    case Extend: return 1;
+    default: break;
+    }
+
+    switch (second->tag) {
     case If:
     case When:
     case Repeat:
     case While:
     case For:
     case Block:
-    case Defer: return 1;
-    case Use: return 1;
+    case Defer:
     case ConvertDef:
-    case FunctionDef: return 1;
+    case FunctionDef:
     case StructDef:
     case EnumDef:
     case LangDef:
-    case Extend: return 2;
-    default: return 0;
+    case Extend: return 1;
+    default: break;
     }
+    return 0;
 }
 
 Text_t indent_code(Text_t code) {
@@ -73,6 +94,7 @@ CONSTFUNC 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;
     }
+    if (ast->tag == Block && Match(ast, Block)->statements == NULL) return NULL;
     return ast;
 }
 
diff --git a/src/formatter/utils.h b/src/formatter/utils.h
index 0a00be6c..880da0a9 100644
--- a/src/formatter/utils.h
+++ b/src/formatter/utils.h
@@ -22,7 +22,7 @@ extern const Text_t single_indent;
 void add_line(Text_t *code, Text_t line, Text_t indent);
 OptionalText_t next_comment(Table_t comments, const char **pos, const char *end);
 bool range_has_comment(const char *start, const char *end, Table_t comments);
-CONSTFUNC int suggested_blank_lines(ast_t *ast);
+CONSTFUNC int suggested_blank_lines(ast_t *first, ast_t *second);
 Text_t indent_code(Text_t code);
 Text_t parenthesize(Text_t code, Text_t indent);
 CONSTFUNC ast_t *unwrap_block(ast_t *ast);