Better support for binops in formatter

1 files changed, 83 insertions, 7 deletions

diff --git a/src/formatter.c b/src/formatter.c
index 94c32b07..224a1489 100644
--- a/src/formatter.c
+++ b/src/formatter.c
@@ -69,6 +69,16 @@ static bool should_have_blank_line(ast_t *ast) {
     }
 }
 
+static Text_t indent_code(Text_t code) {
+    if (code.length <= 0) return code;
+    return Texts(single_indent, Text$replace(code, Text("\n"), Texts("\n", single_indent)));
+}
+
+static Text_t parenthesize(Text_t code, Text_t indent) {
+    if (Text$has(code, Text("\n"))) return Texts("(\n", indent, indent_code(code), "\n", indent, ")");
+    else return Texts("(", code, ")");
+}
+
 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) {
@@ -186,6 +196,44 @@ static Text_t format_namespace(ast_t *namespace, Table_t comments, Text_t indent
     return Texts("\n", indent, single_indent, format_code(namespace, comments, Texts(indent, single_indent)));
 }
 
+static CONSTFUNC const char *binop_tomo_operator(ast_e tag) {
+    switch (tag) {
+    case Power: return "^";
+    case PowerUpdate: return "^=";
+    case Concat: return "++";
+    case ConcatUpdate: return "++=";
+    case Multiply: return "*";
+    case MultiplyUpdate: return "*=";
+    case Divide: return "/";
+    case DivideUpdate: return "/=";
+    case Mod: return "mod";
+    case ModUpdate: return "mod=";
+    case Mod1: return "mod1";
+    case Mod1Update: return "mod1=";
+    case Plus: return "+";
+    case PlusUpdate: return "+=";
+    case Minus: return "-";
+    case MinusUpdate: return "-=";
+    case LeftShift: return "<<";
+    case LeftShiftUpdate: return "<<=";
+    case RightShift: return ">>";
+    case RightShiftUpdate: return ">>=";
+    case And: return "and";
+    case AndUpdate: return "and=";
+    case Or: return "or";
+    case OrUpdate: return "or=";
+    case Xor: return "xor";
+    case XorUpdate: return "xor=";
+    case Equals: return "==";
+    case NotEquals: return "!=";
+    case LessThan: return "<";
+    case LessThanOrEquals: return "<=";
+    case GreaterThan: return ">";
+    case GreaterThanOrEquals: return ">=";
+    default: return NULL;
+    }
+}
+
 OptionalText_t format_inline_code(ast_t *ast, Table_t comments) {
     if (range_has_comment(ast->start, ast->end, comments)) return NONE_TEXT;
     switch (ast->tag) {
@@ -327,9 +375,24 @@ OptionalText_t format_inline_code(ast_t *ast, Table_t comments) {
     }
     case BINOP_CASES: {
         binary_operands_t operands = BINARY_OPERANDS(ast);
-        const char *op = binop_operator(ast->tag);
-        return Texts("(", must(format_inline_code(operands.lhs, comments)), " ", Text$from_str(op), " ",
-                     must(format_inline_code(operands.rhs, comments)), ")");
+        const char *op = binop_tomo_operator(ast->tag);
+
+        Text_t lhs = must(format_inline_code(operands.lhs, comments));
+        Text_t rhs = must(format_inline_code(operands.rhs, comments));
+
+        if (is_update_assignment(ast)) {
+            return Texts(lhs, " ", Text$from_str(op), " ", rhs);
+        }
+
+        if (Text$has(lhs, Text("\n"))
+            || (is_binary_operation(operands.lhs) && op_tightness[operands.lhs->tag] < op_tightness[ast->tag]))
+            lhs = parenthesize(lhs, EMPTY_TEXT);
+        if (Text$has(rhs, Text("\n"))
+            || (is_binary_operation(operands.rhs) && op_tightness[operands.rhs->tag] < op_tightness[ast->tag]))
+            rhs = parenthesize(rhs, EMPTY_TEXT);
+
+        Text_t space = op_tightness[ast->tag] >= op_tightness[Multiply] ? EMPTY_TEXT : Text(" ");
+        return Texts(lhs, space, Text$from_str(binop_tomo_operator(ast->tag)), space, rhs);
     }
     default: {
         fail("Formatting not implemented for: ", ast_to_sexp(ast));
@@ -561,10 +624,23 @@ Text_t format_code(ast_t *ast, Table_t comments, Text_t indent) {
     case BINOP_CASES: {
         if (inlined_fits) return inlined;
         binary_operands_t operands = BINARY_OPERANDS(ast);
-        const char *op = binop_operator(ast->tag);
-        return Texts("(\n", indent, single_indent, format_code(operands.lhs, comments, Texts(indent, single_indent)),
-                     " ", Text$from_str(op), " ", format_code(operands.rhs, comments, Texts(indent, single_indent)),
-                     "\n", indent, ")");
+        const char *op = binop_tomo_operator(ast->tag);
+        Text_t lhs = format_code(operands.lhs, comments, Texts(indent, single_indent));
+        Text_t rhs = format_code(operands.rhs, comments, Texts(indent, single_indent));
+
+        if (is_update_assignment(ast)) {
+            return Texts(lhs, " ", Text$from_str(op), " ", rhs);
+        }
+
+        if (Text$has(lhs, Text("\n"))
+            || (is_binary_operation(operands.lhs) && op_tightness[operands.lhs->tag] < op_tightness[ast->tag]))
+            lhs = parenthesize(lhs, indent);
+        if (Text$has(rhs, Text("\n"))
+            || (is_binary_operation(operands.rhs) && op_tightness[operands.rhs->tag] < op_tightness[ast->tag]))
+            rhs = parenthesize(rhs, indent);
+
+        Text_t space = op_tightness[ast->tag] >= op_tightness[Multiply] ? EMPTY_TEXT : Text(" ");
+        return Texts(lhs, space, Text$from_str(binop_tomo_operator(ast->tag)), space, rhs);
     }
     default: {
         if (inlined_fits) return inlined;