Implement printing of parens for rhs of await expression correctly.

Author: Maxim

src/res_parens.ml

@@ -175,7 +175,7 @@ let flattenOperandRhs parentOperator rhs =
   | _ when ParsetreeViewer.isTernaryExpr rhs -> true
   | _ -> false
 
-let lazyOrAssertExprRhs expr =
+let lazyOrAssertOrAwaitExprRhs expr =
   let optBraces, _ = ParsetreeViewer.processBracesAttr expr in
   match optBraces with
   | Some ({Location.loc = bracesLoc}, _) -> Braced bracesLoc

src/res_parens.mli

@@ -10,7 +10,7 @@ val subBinaryExprOperand : string -> string -> bool
 val rhsBinaryExprOperand : string -> Parsetree.expression -> bool
 val flattenOperandRhs : string -> Parsetree.expression -> bool
 
-val lazyOrAssertExprRhs : Parsetree.expression -> kind
+val lazyOrAssertOrAwaitExprRhs : Parsetree.expression -> kind
 
 val fieldExpr : Parsetree.expression -> kind
 

src/res_parsetree_viewer.ml

@@ -532,6 +532,8 @@ let isPrintableAttribute attr =
 
 let hasPrintableAttributes attrs = List.exists isPrintableAttribute attrs
 
+let filterPrintableAttributes attrs = List.filter isPrintableAttribute attrs
+
 let partitionPrintableAttributes attrs =
   List.partition isPrintableAttribute attrs
 

src/res_parsetree_viewer.mli

@@ -99,6 +99,7 @@ val hasOptionalAttribute : Parsetree.attributes -> bool
 val shouldIndentBinaryExpr : Parsetree.expression -> bool
 val shouldInlineRhsBinaryExpr : Parsetree.expression -> bool
 val hasPrintableAttributes : Parsetree.attributes -> bool
+val filterPrintableAttributes : Parsetree.attributes -> Parsetree.attributes
 val partitionPrintableAttributes :
   Parsetree.attributes -> Parsetree.attributes * Parsetree.attributes
 

src/res_printer.ml

@@ -3095,7 +3095,7 @@ and printExpression ~customLayout (e : Parsetree.expression) cmtTbl =
     | Pexp_assert expr ->
       let rhs =
         let doc = printExpressionWithComments ~customLayout expr cmtTbl in
-        match Parens.lazyOrAssertExprRhs expr with
+        match Parens.lazyOrAssertOrAwaitExprRhs expr with
         | Parens.Parenthesized -> addParens doc
         | Braced braces -> printBraces doc expr braces
         | Nothing -> doc
@@ -3104,7 +3104,7 @@ and printExpression ~customLayout (e : Parsetree.expression) cmtTbl =
     | Pexp_lazy expr ->
       let rhs =
         let doc = printExpressionWithComments ~customLayout expr cmtTbl in
-        match Parens.lazyOrAssertExprRhs expr with
+        match Parens.lazyOrAssertOrAwaitExprRhs expr with
         | Parens.Parenthesized -> addParens doc
         | Braced braces -> printBraces doc expr braces
         | Nothing -> doc
@@ -3282,7 +3282,24 @@ and printExpression ~customLayout (e : Parsetree.expression) cmtTbl =
   in
   let exprWithAwait =
     if ParsetreeViewer.hasAwaitAttribute e.pexp_attributes then
-      Doc.concat [Doc.text "await "; printedExpression]
+      let rhs =
+        match
+          Parens.lazyOrAssertOrAwaitExprRhs
+            {
+              e with
+              pexp_attributes =
+                List.filter
+                  (function
+                    | {Location.txt = "res.await" | "ns.braces"}, _ -> false
+                    | _ -> true)
+                  e.pexp_attributes;
+            }
+        with
+        | Parens.Parenthesized -> addParens printedExpression
+        | Braced braces -> printBraces printedExpression e braces
+        | Nothing -> printedExpression
+      in
+      Doc.concat [Doc.text "await "; rhs]
     else printedExpression
   in
   let shouldPrintItsOwnAttributes =
@@ -3680,11 +3697,7 @@ and printBinaryExpression ~customLayout (expr : Parsetree.expression) cmtTbl =
                 {
                   expr with
                   pexp_attributes =
-                    List.filter
-                      (fun attr ->
-                        match attr with
-                        | {Location.txt = "ns.braces"}, _ -> false
-                        | _ -> true)
+                    ParsetreeViewer.filterPrintableAttributes
                       expr.pexp_attributes;
                 }
             with

tests/printer/expr/asyncAwait.res

@@ -34,4 +34,50 @@ user.data = await fetch()
 let inBinaryExpression = await x->Js.Promise.resolve + 1
 let inBinaryExpression = await x->Js.Promise.resolve + await y->Js.Promise.resolve
 
+let withCallback = async (. ()) => {
+    async (. x) => await (x->Js.promise.resolve) + 1
+}
+
 let () = await (await fetch(url))->(await resolve)
+
+let _ = await (1 + 1)
+let _ = await 1 + 1
+let _ = await (-1)
+let _ = await - 1
+let _ = await !ref 
+let _ = await f
+let _ = await %extension 
+let _ = await "test"
+let _ = await ((a, b) => a + b)
+let _ = await (async (a, b) => a + b)
+let _ = await (switch x { | A => () | B => ()})
+let _ = await [1, 2, 3]
+let _ = await (1, 2, 3)
+let _ = await {name: "Steve", age: 32}
+let _ = await (user.name = "Steve")
+let _ = await (if 20 { true } else {false})
+let _ = await (condition() ? true : false)
+let _ = await f(x)
+let _ = await f(. x)
+let _ = await (f(x) : Js.Promise.t<unit>)
+let _ = await (while true { infiniteLoop() })
+let _ = await (try ok() catch { | _  =>  logError() })
+let _ = await (for i in 0 to 10 { sideEffect()})
+let _ = await (lazy x)
+let _ = await (assert x)
+let _ = await promises[0]
+let _ = await promises["resolved"]
+let _ = await (promises["resolved"] = sideEffect())
+let _ = await (@attr expr)
+let _ = await module(Foo)
+let _ = await module(Foo: Bar)
+let _ = await Promise
+let _ = await Promise(status)
+
+// braces are being dropped, because the ast only has space to record braces surrounding the await
+let _ = await {x}
+// here braces stay, because of precedence
+let _ = await {
+    let x = 1
+    Js.Promise.resolve(x)
+}

tests/printer/expr/expected/asyncAwait.res.txt

@@ -33,4 +33,73 @@ user.data = await fetch()
 let inBinaryExpression = (await x)->Js.Promise.resolve + 1
 let inBinaryExpression = (await x)->Js.Promise.resolve + (await y)->Js.Promise.resolve
 
+let withCallback = async (. ()) => {
+  async (. x) => await (x->Js.promise.resolve) + 1
+}
+
 let () = (await fetch(url))->(await resolve)
+
+let _ = await (1 + 1)
+let _ = (await 1) + 1
+let _ = await -1
+let _ = await -1
+let _ = await !ref
+let _ = await f
+let _ = await %extension
+let _ = await "test"
+let _ = await ((a, b) => a + b)
+let _ = await (async (a, b) => a + b)
+let _ = await (
+  switch x {
+  | A => ()
+  | B => ()
+  }
+)
+let _ = await [1, 2, 3]
+let _ = await (1, 2, 3)
+let _ = await {name: "Steve", age: 32}
+let _ = await (user.name = "Steve")
+let _ = await (
+  if 20 {
+    true
+  } else {
+    false
+  }
+)
+let _ = await (condition() ? true : false)
+let _ = await f(x)
+let _ = await f(. x)
+let _ = (await (f(x): Js.Promise.t<unit>))
+let _ = await (
+  while true {
+    infiniteLoop()
+  }
+)
+let _ = await (
+  try ok() catch {
+  | _ => logError()
+  }
+)
+let _ = await (
+  for i in 0 to 10 {
+    sideEffect()
+  }
+)
+let _ = await (lazy x)
+let _ = await (assert x)
+let _ = await promises[0]
+let _ = await promises["resolved"]
+let _ = await promises["resolved"] = sideEffect()
+let _ = @attr await (expr)
+let _ = await module(Foo)
+let _ = await module(Foo: Bar)
+let _ = await Promise
+let _ = await Promise(status)
+
+// braces are being dropped, because the ast only has space to record braces surrounding the await
+let _ = await x
+// here braces stay, because of precedence
+let _ = await {
+  let x = 1
+  Js.Promise.resolve(x)
+}