fix(grainfmt): Ensure constraints and keyword functions group properly (#2070)

Author: ospencer

compiler/src/formatting/fmt.re

@@ -157,6 +157,7 @@ let op_precedence = startsWith =>
 
 let precedence = expr => {
   switch (expr.pexp_desc) {
+  | PExpConstraint(_) => 140
   | PExpId({txt: Identifier.IdentName({txt: op})}) =>
     if (String.length(op) > 1) {
       switch (String.sub(op, 0, 2)) {
@@ -250,6 +251,8 @@ let needs_grouping = (~parent, ~side: infix_side, expr) => {
       } else {
         FormatterGrouping;
       };
+    } else if (is_keyword_function(fn1)) {
+      ParenGrouping;
     } else {
       FormatterGrouping;
     }
@@ -2483,8 +2486,17 @@ let print_expression = (fmt, ~infix_wrap=d => group(indent(d)), expr) => {
            )
            ++ hardline,
          )
-    | PExpConstraint(expr, typ) =>
-      fmt.print_expression(fmt, expr)
+    | PExpConstraint(inner_expr, typ) =>
+      (
+        switch (needs_grouping(~parent=expr, ~side=Left, inner_expr)) {
+        | ParenGrouping =>
+          parens(
+            indent(break ++ fmt.print_expression(fmt, inner_expr)) ++ break,
+          )
+        | FormatterGrouping => group(fmt.print_expression(fmt, inner_expr))
+        | None => fmt.print_expression(fmt, inner_expr)
+        }
+      )
       ++ string(":")
       ++ group(
            indent(
@@ -2493,7 +2505,7 @@ let print_expression = (fmt, ~infix_wrap=d => group(indent(d)), expr) => {
                ~none=breakable_space,
                ~lead=space,
                ~trail=breakable_space,
-               expr.pexp_loc,
+               inner_expr.pexp_loc,
                typ.ptyp_loc,
              )
              ++ fmt.print_type(fmt, typ),

compiler/test/grainfmt/binops.expected.gr

@@ -45,3 +45,5 @@ let (+++) = (+)
 let (---) = (-)
 
 let zz = 1 +++ 2 --- 3
+
+(5 + 5): Number

compiler/test/grainfmt/binops.input.gr

@@ -44,3 +44,5 @@ let (+++) = (+)
 let (---) = (-)
 
 let zz = 1 +++ 2 --- 3
+
+(5 + 5): Number

compiler/test/grainfmt/keyword_expression.expected.gr

@@ -0,0 +1,5 @@
+module KeywordExpression
+
+(fail "unimplemented") + 5
+
+(fail "unimplemented"): Number

compiler/test/grainfmt/keyword_expression.input.gr

@@ -0,0 +1,5 @@
+module KeywordExpression
+
+(fail "unimplemented") + 5
+
+(fail "unimplemented"): Number

compiler/test/suites/formatter.re

@@ -15,6 +15,7 @@ describe("formatter", ({test, testSkip}) => {
   assertFormatOutput("application2", "application2");
   assertFormatOutput("application_indenting", "application_indenting");
   assertFormatOutput("function_params", "function_params");
+  assertFormatOutput("keyword_expression", "keyword_expression");
   assertFormatOutput("variants", "variants");
   assertFormatOutput("matches", "matches");
   assertFormatOutput("includes", "includes");