From 3cf79faafb5fd651e6bfad2b9e0ba3c843edefae Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=81ukasz=20Czajka?=
 <62751+lukaszcz@users.noreply.github.com>
Date: Fri, 25 Oct 2024 11:42:01 +0200
Subject: [PATCH] Formatter: add braces when the iterator body is not enclosed
 (#3122)

* Closes #3091
* Formatter adds braces when the body is not enclosed in braces or
parentheses. Braces-enclosed body is always printed as a block on a new
line:
```
for (acc := 0) (x in lst) {
  x + acc
}
```
* If the body is enclosed in ordinary parentheses, then they are
preserved and the iterator is printed on a single line, if possible:
```
for (acc := 0) (x in lst) (x + acc)
```
This is sometimes useful when you want iterator application as an
argument to something.
---
 src/Juvix/Compiler/Concrete/Print/Base.hs     |  8 +--
 .../FromParsed/Analysis/Scoping.hs            |  7 +--
 .../Concrete/Translation/FromSource.hs        |  3 +-
 .../Anoma/Compilation/positive/test012.juvix  |  9 ++--
 .../Anoma/Compilation/positive/test041.juvix  |  2 +-
 .../Anoma/Compilation/positive/test054.juvix  | 53 ++++++++-----------
 .../Anoma/Compilation/positive/test056.juvix  |  2 +-
 .../Anoma/Compilation/positive/test058.juvix  |  7 +--
 .../Anoma/Compilation/positive/test061.juvix  |  8 +--
 tests/positive/Iterators.juvix                | 12 +++--
 tests/positive/RecordIterator.juvix           |  5 +-
 11 files changed, 60 insertions(+), 56 deletions(-)

diff --git a/src/Juvix/Compiler/Concrete/Print/Base.hs b/src/Juvix/Compiler/Concrete/Print/Base.hs
index 177ebd2d9..8143c6974 100644
--- a/src/Juvix/Compiler/Concrete/Print/Base.hs
+++ b/src/Juvix/Compiler/Concrete/Print/Base.hs
@@ -304,17 +304,17 @@ instance (SingI s) => PrettyPrint (Range s) where
     n <+> ppCode _rangeInKw <+> e
 
 ppIterator :: forall r s. (Members '[ExactPrint, Reader Options] r, SingI s) => IsTop -> Iterator s -> Sem r ()
-ppIterator isTop Iterator {..} = do
+ppIterator _isTop Iterator {..} = do
   let n = ppIdentifierType _iteratorName
       is = ppCode <$> _iteratorInitializers
       rngs = ppCode <$> _iteratorRanges
       is' = parens . hsepSemicolon <$> nonEmpty is
       rngs' = parens . hsepSemicolon <$> nonEmpty rngs
       b
-        | _iteratorBodyBraces = braces (oneLineOrNextNoIndent (ppTopExpressionType _iteratorBody))
-        | otherwise = line <> ppMaybeTopExpression isTop _iteratorBody
+        | _iteratorBodyBraces = space <> braces (blockIndent (ppTopExpressionType _iteratorBody))
+        | otherwise = parens (oneLineOrNextNoIndent (ppTopExpressionType _iteratorBody))
   parensIf _iteratorParens $
-    hang (n <+?> is' <+?> rngs' <> b)
+    n <+?> is' <+?> rngs' <> b
 
 instance PrettyPrint S.AName where
   ppCode n = annotated (AnnKind (S.getNameKind n)) (noLoc (pretty (n ^. S.anameVerbatim)))
diff --git a/src/Juvix/Compiler/Concrete/Translation/FromParsed/Analysis/Scoping.hs b/src/Juvix/Compiler/Concrete/Translation/FromParsed/Analysis/Scoping.hs
index 00cb69e4c..d9dd7b276 100644
--- a/src/Juvix/Compiler/Concrete/Translation/FromParsed/Analysis/Scoping.hs
+++ b/src/Juvix/Compiler/Concrete/Translation/FromParsed/Analysis/Scoping.hs
@@ -1615,9 +1615,10 @@ checkSections sec = topBindings helper
                               fs <-
                                 failMaybe $
                                   mkRec
-                                    ^? constructorRhs
-                                      . _ConstructorRhsRecord
-                                      . to mkRecordNameSignature
+                                    ^? ( constructorRhs
+                                           . _ConstructorRhsRecord
+                                           . to mkRecordNameSignature
+                                       )
                               let info =
                                     RecordInfo
                                       { _recordInfoSignature = fs,
diff --git a/src/Juvix/Compiler/Concrete/Translation/FromSource.hs b/src/Juvix/Compiler/Concrete/Translation/FromSource.hs
index f250b49e0..b64395c69 100644
--- a/src/Juvix/Compiler/Concrete/Translation/FromSource.hs
+++ b/src/Juvix/Compiler/Concrete/Translation/FromSource.hs
@@ -931,7 +931,8 @@ iterator = do
   do
     (_iteratorBody, _iteratorBodyBraces) <-
       (,True) <$> braces parseExpressionAtoms
-        <|> (,False) <$> parseExpressionAtoms
+        <|> (,False) <$> parens parseExpressionAtoms
+        <|> (,True) <$> parseExpressionAtoms
     let _iteratorParens = False
     return $ Iterator {..}
   where
diff --git a/tests/Anoma/Compilation/positive/test012.juvix b/tests/Anoma/Compilation/positive/test012.juvix
index a8e55613e..75b2c4284 100644
--- a/tests/Anoma/Compilation/positive/test012.juvix
+++ b/tests/Anoma/Compilation/positive/test012.juvix
@@ -24,11 +24,14 @@ gen : Nat → Tree
 
 preorder : Tree → List Nat
   | leaf := 0 :: nil
-  | (node1 c) :=  1 :: preorder c
-  | (node2 l r) :=  2 :: preorder l ++ preorder r
+  | (node1 c) := 1 :: preorder c
+  | (node2 l r) := 2 :: preorder l ++ preorder r
   | (node3 l m r) := 3 :: preorder l ++ preorder m ++ preorder r;
 
-combineDigits (xs : List Nat) : Nat := for (acc := 0) (x in xs) acc * 10 + x;
+combineDigits (xs : List Nat) : Nat :=
+  for (acc := 0) (x in xs) {
+    acc * 10 + x
+  };
 
 main : Nat :=
   trace (combineDigits (preorder (gen 3)))
diff --git a/tests/Anoma/Compilation/positive/test041.juvix b/tests/Anoma/Compilation/positive/test041.juvix
index c0d632755..3b140a96c 100644
--- a/tests/Anoma/Compilation/positive/test041.juvix
+++ b/tests/Anoma/Compilation/positive/test041.juvix
@@ -7,6 +7,6 @@ type BoxedList :=
   | boxed (List Nat);
 
 sumBoxedList : BoxedList -> Nat
-  | (boxed l) := for (acc := 0) (x in l) acc + x;
+  | (boxed l) := for (acc := 0) (x in l) (acc + x);
 
 main : Nat := sumBoxedList (boxed [1;2;3]);
diff --git a/tests/Anoma/Compilation/positive/test054.juvix b/tests/Anoma/Compilation/positive/test054.juvix
index c764cd513..bad06886f 100644
--- a/tests/Anoma/Compilation/positive/test054.juvix
+++ b/tests/Anoma/Compilation/positive/test054.juvix
@@ -4,8 +4,7 @@ module test054;
 import Stdlib.Prelude open;
 
 syntax iterator myfor;
-myfor : {A B : Type} → (A → B → A) → A → List B → A :=
-  foldl {_} {_} {{_}} {_};
+myfor : {A B : Type} → (A → B → A) → A → List B → A := foldl {_} {_} {{_}} {_};
 
 syntax iterator mymap {init := 0};
 mymap : {A B : Type} → (A → B) → List A → List B
@@ -13,7 +12,10 @@ mymap : {A B : Type} → (A → B) → List A → List B
   | f (x :: xs) := f x :: mymap f xs;
 
 sum : List Nat → Nat
-  | xs := myfor (acc := 0) (x in xs) {acc + x};
+  | xs :=
+    myfor (acc := 0) (x in xs) {
+      acc + x
+    };
 
 sum' : List Nat → Nat
   | xs := myfor λ {acc x := acc + x} 0 xs;
@@ -21,39 +23,30 @@ sum' : List Nat → Nat
 lst : List Nat := 1 :: 2 :: 3 :: 4 :: 5 :: nil;
 
 syntax iterator myfor2 {init := 1; range := 2};
-myfor2
-  : {A B C : Type}
-    → (A → B → C → A)
-    → A
-    → List B
-    → List C
-    → A
+myfor2 : {A B C : Type} → (A → B → C → A) → A → List B → List C → A
   | f acc xs ys :=
-    myfor (acc' := acc) (x in xs)
-      myfor (acc'' := acc') (y in ys)
-        f acc'' x y;
+    myfor (acc' := acc) (x in xs) {
+      myfor (acc'' := acc') (y in ys) {
+        f acc'' x y
+      }
+    };
 
 syntax iterator myzip2 {init := 2; range := 2};
-myzip2
-  : {A A' B C : Type}
-    → (A → A' → B → C → Pair A A')
-    → A
-    → A'
-    → List B
-    → List C
-    → Pair A A'
+myzip2 : {A A' B C : Type} → (A → A' → B → C → Pair A A') → A → A' → List B → List C → Pair A A'
   | f a b xs ys :=
-    myfor (acc, acc' := a, b) (x, y in zip xs ys)
-      f acc acc' x y;
+    myfor (acc, acc' := a, b) (x, y in zip xs ys) {
+      f acc acc' x y
+    };
 
 main : Nat :=
   sum lst
     + sum' lst
-    + fst (myfor (a, b := 0, 0) (x in lst) b + x, a)
-    + (myfor2 (acc := 0) (x in lst; y in 1 :: 2 :: nil)
-        acc + x + y)
+    + fst (myfor (a, b := 0, 0) (x in lst) (b + x, a))
+    + myfor2 (acc := 0) (x in lst; y in 1 :: 2 :: nil) (acc + x + y)
     + fst
-      (myzip2 (a := 0; b := 0) (x in lst; y in reverse lst)
-        a + x * y, b + y)
-    + myfor (a := 0) (x, y in mymap (x in lst) {x, x + 1})
-        a + x * y;
+      (myzip2 (a := 0; b := 0) (x in lst; y in reverse lst) {
+        a + x * y, b + y
+      })
+    + myfor (a := 0) (x, y in mymap (x in lst) (x, x + 1)) {
+      a + x * y
+    };
diff --git a/tests/Anoma/Compilation/positive/test056.juvix b/tests/Anoma/Compilation/positive/test056.juvix
index 2981ae3b5..012019210 100644
--- a/tests/Anoma/Compilation/positive/test056.juvix
+++ b/tests/Anoma/Compilation/positive/test056.juvix
@@ -26,7 +26,7 @@ myf'
   | a0 f a b := myf a0 (f a0) a b true;
 
 sum : List Nat -> Nat
-  | xs := for (acc := 0) (x in xs) {x + acc};
+  | xs := for (acc := 0) (x in xs) (x + acc);
 
 funa : {A : Type} -> (A -> A) -> A -> A
   | {A} f a :=
diff --git a/tests/Anoma/Compilation/positive/test058.juvix b/tests/Anoma/Compilation/positive/test058.juvix
index 25aca0ebb..d1cec5a99 100644
--- a/tests/Anoma/Compilation/positive/test058.juvix
+++ b/tests/Anoma/Compilation/positive/test058.juvix
@@ -3,10 +3,11 @@ module test058;
 
 import Stdlib.Prelude open;
 
-sum (x : Nat) : Nat :=
-  for (acc := 0) (n in 1 to x) {acc + n};
+sum (x : Nat) : Nat := for (acc := 0) (n in 1 to x) (acc + n);
 
 sum' (x : Nat) : Nat :=
-  for (acc := 0) (n in 1 to x step 2) {acc + n};
+  for (acc := 0) (n in 1 to x step 2) {
+    acc + n
+  };
 
 main : Nat := sum 100 + sum' 100;
diff --git a/tests/Anoma/Compilation/positive/test061.juvix b/tests/Anoma/Compilation/positive/test061.juvix
index f373e2242..a0ffec990 100644
--- a/tests/Anoma/Compilation/positive/test061.juvix
+++ b/tests/Anoma/Compilation/positive/test061.juvix
@@ -5,7 +5,7 @@ import Stdlib.Prelude open;
 import Stdlib.Debug.Trace open;
 
 trait
-type T A := mkT {toNat : A -> Nat};
+type T A := mkT@{toNat : A -> Nat};
 
 instance
 tNatI : T Nat := mkT id;
@@ -29,16 +29,16 @@ tMaybeI {A} {{T A}} : T (Maybe A) :=
 instance
 tUnitFunI {A} {{T A}} : T (Unit -> A) :=
   mkT@{
-    toNat {A} {{T A}} (f : Unit -> A) : Nat :=
-      T.toNat (f unit)
+    toNat {A} {{T A}} (f : Unit -> A) : Nat := T.toNat (f unit)
   };
 
 instance
 tListI {A} {{T A}} : T (List A) :=
   mkT@{
     toNat {A} {{T A}} (xs : List A) : Nat :=
-      for (acc := 0) (x in xs)
+      for (acc := 0) (x in xs) {
         acc + T.toNat {A} x
+      }
   };
 
 g : {A : Type} → {{T A}} → Nat := 5;
diff --git a/tests/positive/Iterators.juvix b/tests/positive/Iterators.juvix
index 49b83808d..4336fcf17 100644
--- a/tests/positive/Iterators.juvix
+++ b/tests/positive/Iterators.juvix
@@ -15,8 +15,10 @@ type Bool :=
 main : Bool :=
   let
     z : Bool := false;
-  in itconst (a := true; b := false) (c in false; d in false)
-       for (x := true) (y in false)
-         case x of
-           | true := y
-           | false := z;
+  in itconst (a := true; b := false) (c in false; d in false) {
+    for (x := true) (y in false) {
+      case x of
+        | true := y
+        | false := z
+    }
+  };
diff --git a/tests/positive/RecordIterator.juvix b/tests/positive/RecordIterator.juvix
index c504a6069..f8ae6d821 100644
--- a/tests/positive/RecordIterator.juvix
+++ b/tests/positive/RecordIterator.juvix
@@ -19,4 +19,7 @@ foldl
   (g : B -> elem -> B)
   (ini : B)
   (ls : container)
-  : B := for (acc := ini) (x in ls) {g acc x};
+  : B :=
+  for (acc := ini) (x in ls) {
+    g acc x
+  };