From d3e862fc5179b1f8aac4b408346e65c0522fea70 Mon Sep 17 00:00:00 2001
From: Jan Mas Rovira <janmasrovira@gmail.com>
Date: Thu, 7 Dec 2023 11:26:48 +0100
Subject: [PATCH] Improve formatting of function definition arguments (#2551)

- Closes #2534
---
 examples/milestone/Bank/Bank.juvix            | 12 ++++++++---
 src/Juvix/Compiler/Concrete/Print/Base.hs     | 18 +++++++++--------
 .../Internal/DefaultArgCycleArity.juvix       |  8 +++++---
 tests/positive/ConstructorWildcard.juvix      | 13 ++++++------
 tests/positive/DefaultValues.juvix            | 14 +++++++++++--
 tests/positive/Format.juvix                   | 17 ++++++++++++++++
 tests/positive/Syntax.juvix                   |  5 +++--
 tests/positive/Traits2.juvix                  | 20 +++++++++++--------
 8 files changed, 75 insertions(+), 32 deletions(-)

diff --git a/examples/milestone/Bank/Bank.juvix b/examples/milestone/Bank/Bank.juvix
index 6c7cb2bf68..a6af6a0c7d 100644
--- a/examples/milestone/Bank/Bank.juvix
+++ b/examples/milestone/Bank/Bank.juvix
@@ -100,8 +100,8 @@ issue : Address -> Address -> Nat -> Token
     assert (caller == bankAddress) (mkToken owner 0 amount);
 
 --- Deposits some amount of money into the bank.
-deposit (bal : Balances) (token : Token) (amount : Nat)
-  : Token :=
+deposit
+  (bal : Balances) (token : Token) (amount : Nat) : Token :=
   let
     difference : Nat := sub (getAmount token) amount;
     remaining : Token :=
@@ -111,7 +111,13 @@ deposit (bal : Balances) (token : Token) (amount : Nat)
   in runOnChain (commitBalances bal') remaining;
 
 --- Returns a new ;Token; containing the amount of money withdrawn.
-withdraw (bal : Balances) (caller : Address) (recipient : Address) (amount : Nat) (rate : Nat) (periods : Nat)
+withdraw
+  (bal : Balances)
+  (caller : Address)
+  (recipient : Address)
+  (amount : Nat)
+  (rate : Nat)
+  (periods : Nat)
   : Token :=
   assert
     (caller == bankAddress)
diff --git a/src/Juvix/Compiler/Concrete/Print/Base.hs b/src/Juvix/Compiler/Concrete/Print/Base.hs
index 1769abc658..d77571ac68 100644
--- a/src/Juvix/Compiler/Concrete/Print/Base.hs
+++ b/src/Juvix/Compiler/Concrete/Print/Base.hs
@@ -946,20 +946,22 @@ ppFunctionSignature FunctionDef {..} = do
   let termin' = (<> line) . ppCode <$> _signTerminating
       coercion' = (<> if isJust instance' then space else line) . ppCode <$> _signCoercion
       instance' = (<> line) . ppCode <$> _signInstance
-      args' = hsep . fmap ppCode <$> nonEmpty _signArgs
       builtin' = (<> line) . ppCode <$> _signBuiltin
-      type' = case _signColonKw ^. unIrrelevant of
-        Just col -> oneLineOrNext (ppCode col <+> ppExpressionType (fromJust _signRetType))
-        Nothing -> mempty
+      margs' = fmap ppCode <$> nonEmpty _signArgs
+      mtype' = case _signColonKw ^. unIrrelevant of
+        Just col -> Just (ppCode col <+> ppExpressionType (fromJust _signRetType))
+        Nothing -> Nothing
+      argsAndType' = case mtype' of
+        Nothing -> margs'
+        Just ty' -> case margs' of
+          Nothing -> Just (pure ty')
+          Just args' -> Just (args' <> pure ty')
       name' = annDef _signName (ppSymbolType _signName)
    in builtin'
         ?<> termin'
         ?<> coercion'
         ?<> instance'
-        ?<> ( name'
-                <+?> args'
-                  <> type'
-            )
+        ?<> (name' <>? (oneLineOrNext . sep <$> argsAndType'))
 
 instance (SingI s) => PrettyPrint (FunctionDef s) where
   ppCode :: forall r. (Members '[ExactPrint, Reader Options] r) => FunctionDef s -> Sem r ()
diff --git a/tests/negative/Internal/DefaultArgCycleArity.juvix b/tests/negative/Internal/DefaultArgCycleArity.juvix
index 08d56b505b..f2bcdfe8c1 100644
--- a/tests/negative/Internal/DefaultArgCycleArity.juvix
+++ b/tests/negative/Internal/DefaultArgCycleArity.juvix
@@ -2,6 +2,8 @@ module DefaultArgCycleArity;
 
 import Stdlib.Data.Nat open;
 
-fun {a : Nat := 1} {b : Nat := fun + a + 1} {c : Nat := b
-  + a
-  + 1} : Nat := a * b + c;
+fun
+  {a : Nat := 1}
+  {b : Nat := fun + a + 1}
+  {c : Nat := b + a + 1}
+  : Nat := a * b + c;
diff --git a/tests/positive/ConstructorWildcard.juvix b/tests/positive/ConstructorWildcard.juvix
index 1698e294ea..3480959189 100644
--- a/tests/positive/ConstructorWildcard.juvix
+++ b/tests/positive/ConstructorWildcard.juvix
@@ -1,16 +1,17 @@
 module ConstructorWildcard;
 
 type Bool :=
- false | true;
+  | false
+  | true;
 
 type either (A B : Type) :=
-  left A
+  | left A
   | right B;
 
 isLeft {A B} : either A B → Bool
- | left@{} := true
- | right@{} := false;
+  | left@{} := true
+  | right@{} := false;
 
 not : Bool → Bool
- | false@{} := true
- | true := false;
+  | false@{} := true
+  | true := false;
diff --git a/tests/positive/DefaultValues.juvix b/tests/positive/DefaultValues.juvix
index b2b1e90c88..bb3bacfae9 100644
--- a/tests/positive/DefaultValues.juvix
+++ b/tests/positive/DefaultValues.juvix
@@ -18,12 +18,22 @@ axiom c : C;
 
 axiom d : D;
 
-mk {f1 : A := a} {f2 : Type} {f3 : C := c} (x : f2)
+mk
+  {f1 : A := a}
+  {f2 : Type}
+  {f3 : C := c}
+  (x : f2)
   : A × f2 × C := f1, x, f3;
 
 x1 : A × B × C := mk (x := b);
 
-mk2 {f1 : A := a} {f2 : Type} {f3 : C := c} (x : f2) {f4 : D := d} (y : f2)
+mk2
+  {f1 : A := a}
+  {f2 : Type}
+  {f3 : C := c}
+  (x : f2)
+  {f4 : D := d}
+  (y : f2)
   : A × f2 × C := f1, x, f3;
 
 x2 : A × B × C := mk2 (x := b) (y := b);
diff --git a/tests/positive/Format.juvix b/tests/positive/Format.juvix
index 896e6e9541..ae634e38a1 100644
--- a/tests/positive/Format.juvix
+++ b/tests/positive/Format.juvix
@@ -389,4 +389,21 @@ i2479' : {_ : Nat} -> Nat
   | {suc x} := x
   | {_} := zero;
 
+-- formatting arguments that do not fit in a line
+fffffffffffffffffffffffffffffffffffffffffffffffff
+  {f : List Nat} : Nat := zero;
+
+-- formatting arguments that do not fit in a line
+fff
+  {f : List Nat}
+  {f0 : List Nat}
+  {f1 : List Nat}
+  {f2 : List Nat}
+  {f3 : List Nat}
+  {f4 : List Nat}
+  {f5 : List Nat}
+  {f6 : List Nat}
+  {f7 : List Nat}
+  : Nat := zero;
+
 -- Comment at the end of a module
diff --git a/tests/positive/Syntax.juvix b/tests/positive/Syntax.juvix
index 83b5719238..5c6a868d93 100644
--- a/tests/positive/Syntax.juvix
+++ b/tests/positive/Syntax.juvix
@@ -1,7 +1,8 @@
 module Syntax;
 
-compose {A B C : Type} (f : B -> C) (g : A -> B) (x : A)
-  : C := f (g x);
+compose
+  {A B C : Type} (f : B -> C) (g : A -> B) (x : A) : C :=
+  f (g x);
 
 compose' {A B C : Type} (f : B -> C) (g : A -> B) : A -> C
   | x := f (g x);
diff --git a/tests/positive/Traits2.juvix b/tests/positive/Traits2.juvix
index 34ca5b955e..6374b7d918 100644
--- a/tests/positive/Traits2.juvix
+++ b/tests/positive/Traits2.juvix
@@ -51,8 +51,8 @@ syntax operator >>= bind;
     -> (A -> M B)
     -> M B := Monad.bind;
 
-Maybe-bind {A B : Type}
-  : Maybe A -> (A -> Maybe B) -> Maybe B
+Maybe-bind
+  {A B : Type} : Maybe A -> (A -> Maybe B) -> Maybe B
   | nothing _ := nothing
   | (just a) f := f a;
 
@@ -74,13 +74,16 @@ import Stdlib.Data.Product open;
 ×-Functor : {A : Type} -> Functor ((×) A) :=
   mkFunctor ×-fmap;
 
-×-bind {A : Type} {{Semigroup A}} {B C : Type}
+×-bind
+  {A : Type}
+  {{Semigroup A}}
+  {B C : Type}
   : A × B -> (B -> A × C) -> A × C
   | (a, b) f := case f b of {a', b' := a <> a', b'};
 
 instance
-×-Monad {A : Type} {{Semigroup A}} {{Monoid A}}
-  : Monad ((×) A) :=
+×-Monad
+  {A : Type} {{Semigroup A}} {{Monoid A}} : Monad ((×) A) :=
   mkMonad@{
     bind := ×-bind;
     return := λ {b := Monoid.mempty, b};
@@ -89,8 +92,8 @@ instance
 
 type Reader (r a : Type) := mkReader {runReader : r -> a};
 
-Reader-fmap {R A B : Type} (f : A -> B)
-  : Reader R A -> Reader R B
+Reader-fmap
+  {R A B : Type} (f : A -> B) : Reader R A -> Reader R B
   | (mkReader ra) := mkReader (f ∘ ra);
 
 Reader-Functor-NoNamed {R : Type} : Functor (Reader R) :=
@@ -109,7 +112,8 @@ Reader-Monad {R : Type} : Monad (Reader R) :=
     functor := Reader-Functor;
     return {A : Type} (a : A) : Reader R A :=
       mkReader (const a);
-    bind {A B : Type}
+    bind
+      {A B : Type}
       : Reader R A -> (A -> Reader R B) -> Reader R B
       | (mkReader ra) arb :=
         let