Use APE mechanism to format Function expressions (#1852)

I paired with @janmasrovira on this work.

Before this change - long type signatures were formatted to contain line
breaks within applications:

```
  exampleFunction : {A : Type} -> List A -> List A -> List
        A -> List A -> List A -> List A -> List A -> Nat;
```

After this change we treat `->` and an infix operator and format like
other infix applications:

```
  exampleFunction :
    {A : Type}
      -> List A
      -> List A
      -> List A
      -> List A
      -> List A
      -> List A
      -> List A
      -> Nat;
```

* Fixes #1850
Co-authored-by: @janmasrovira

src/Juvix/Compiler/Concrete/Language.hs

@@ -993,6 +993,13 @@ instance Eq (ModuleRef'' 'S.Concrete t) where
 instance HasAtomicity (PatternAtom 'Parsed) where
   atomicity = const Atom
 
+instance SingI s => HasAtomicity (FunctionParameters s) where
+  atomicity p
+    | not (null (p ^. paramNames)) || p ^. paramImplicit == Implicit = Atom
+    | otherwise = case sing :: SStage s of
+        SParsed -> atomicity (p ^. paramType)
+        SScoped -> atomicity (p ^. paramType)
+
 deriving stock instance
   ( Show Symbol,
     Show (PatternAtom 'Parsed)
@@ -1296,7 +1303,12 @@ instance
 
 --------------------------------------------------------------------------------
 
-instance IsApe Application Expression where
+data ApeLeaf
+  = ApeLeafExpression Expression
+  | ApeLeafFunctionParams (FunctionParameters 'Scoped)
+  | ApeLeafFunctionKw KeywordRef
+
+instance IsApe Application ApeLeaf where
   toApe (Application l r) =
     ApeApp
       Ape.App
@@ -1304,7 +1316,7 @@ instance IsApe Application Expression where
           _appRight = toApe r
         }
 
-instance IsApe InfixApplication Expression where
+instance IsApe InfixApplication ApeLeaf where
   toApe i@(InfixApplication l op r) =
     ApeInfix
       Infix
@@ -1312,31 +1324,54 @@ instance IsApe InfixApplication Expression where
           _infixLeft = toApe l,
           _infixRight = toApe r,
           _infixIsComma = isDelimiterStr (prettyText (identifierName op ^. S.nameConcrete)),
-          _infixOp = ExpressionIdentifier op
+          _infixOp = ApeLeafExpression (ExpressionIdentifier op)
         }
 
-instance IsApe PostfixApplication Expression where
+instance IsApe PostfixApplication ApeLeaf where
   toApe p@(PostfixApplication l op) =
     ApePostfix
       Postfix
         { _postfixFixity = getFixity p,
           _postfixLeft = toApe l,
-          _postfixOp = ExpressionIdentifier op
+          _postfixOp = ApeLeafExpression (ExpressionIdentifier op)
         }
 
-instance IsApe Expression Expression where
+instance IsApe (Function 'Scoped) ApeLeaf where
+  toApe (Function ps kw ret) =
+    ApeInfix
+      Infix
+        { _infixFixity = funFixity,
+          _infixLeft = toApe ps,
+          _infixRight = toApe ret,
+          _infixIsComma = False,
+          _infixOp = ApeLeafFunctionKw kw
+        }
+
+instance IsApe Expression ApeLeaf where
   toApe = \case
     ExpressionApplication a -> toApe a
     ExpressionInfixApplication a -> toApe a
     ExpressionPostfixApplication a -> toApe a
+    ExpressionFunction a -> toApe a
     e ->
       ApeLeaf
         ( Leaf
             { _leafAtomicity = atomicity e,
-              _leafExpr = e
+              _leafExpr = ApeLeafExpression e
             }
         )
 
+instance IsApe (FunctionParameters 'Scoped) ApeLeaf where
+  toApe f
+    | atomicity f == Atom =
+        ApeLeaf
+          ( Leaf
+              { _leafAtomicity = Atom,
+                _leafExpr = ApeLeafFunctionParams f
+              }
+          )
+    | otherwise = toApe (f ^. paramType)
+
 instance HasAtomicity PatternArg where
   atomicity p
     | Implicit <- p ^. patternArgIsImplicit = Atom

src/Juvix/Compiler/Concrete/Pretty/Base.hs

@@ -431,20 +431,24 @@ instance (SingI s) => PrettyCode (TypeSignature s) where
     sigType' <- ppExpression _sigType
     builtin' <- traverse ppCode _sigBuiltin
     doc' <- mapM ppCode _sigDoc
-    body' :: Maybe (Doc Ann) <- fmap (kwAssign <+>) <$> mapM ppExpression _sigBody
-    return $ doc' ?<> builtin' <?+> sigTerminating' <> hang' (sigName' <+> kwColon <+> sigType' <+?> body')
+    body' :: Maybe (Doc Ann) <- fmap ((kwAssign <>) . oneLineOrNext) <$> mapM ppExpression _sigBody
+    return $ doc' ?<> builtin' <?+> sigTerminating' <> hang' (sigName' <+> kwColon <> oneLineOrNext (sigType' <+?> body'))
 
 instance (SingI s) => PrettyCode (Function s) where
-  ppCode :: forall r. (Members '[Reader Options] r) => Function s -> Sem r (Doc Ann)
-  ppCode Function {..} = do
-    funParameter' <- ppCode _funParameters
-    funReturn' <- ppRightExpression' funFixity _funReturn
-    funKw' <- ppCode _funKw
-    return $ funParameter' <+> funKw' <+> funReturn'
+  ppCode a = case sing :: SStage s of
+    SParsed -> ppCode' a
+    SScoped -> apeHelper a (ppCode' a)
     where
-      ppRightExpression' = case sing :: SStage s of
-        SParsed -> ppRightExpression
-        SScoped -> ppRightExpression
+      ppCode' :: forall r. (Members '[Reader Options] r) => Function s -> Sem r (Doc Ann)
+      ppCode' Function {..} = do
+        funParameter' <- ppCode _funParameters
+        funReturn' <- ppRightExpression' funFixity _funReturn
+        funKw' <- ppCode _funKw
+        return $ funParameter' <+> funKw' <+> funReturn'
+        where
+          ppRightExpression' = case sing :: SStage s of
+            SParsed -> ppRightExpression
+            SScoped -> ppRightExpression
 
 instance (SingI s) => PrettyCode (FunctionParameters s) where
   ppCode FunctionParameters {..} = do
@@ -628,13 +632,19 @@ instance PrettyCode Application where
       args' <- mapM ppCodeAtom args
       return $ PP.group (f' <+> nest' (vsep args'))
 
-apeHelper :: (IsApe a Expression, Members '[Reader Options] r) => a -> Sem r (Doc CodeAnn) -> Sem r (Doc CodeAnn)
+instance PrettyCode ApeLeaf where
+  ppCode = \case
+    ApeLeafExpression e -> ppCode e
+    ApeLeafFunctionParams a -> ppCode a
+    ApeLeafFunctionKw r -> return (pretty r)
+
+apeHelper :: (IsApe a ApeLeaf, Members '[Reader Options] r) => a -> Sem r (Doc CodeAnn) -> Sem r (Doc CodeAnn)
 apeHelper a alt = do
   opts <- ask @Options
   if
       | not (opts ^. optNoApe) ->
           return $
-            let params :: ApeParams Expression
+            let params :: ApeParams ApeLeaf
                 params = ApeParams (run . runReader opts . ppCode)
              in runApe params a
       | otherwise -> alt

src/Juvix/Compiler/Concrete/Print/Base.hs

@@ -189,16 +189,16 @@ instance PrettyPrint (TypeSignature 'Scoped) where
         name' = region (P.annDef _sigName) (ppCode _sigName)
         body' = case _sigBody of
           Nothing -> Nothing
-          Just body -> Just (noLoc P.kwAssign <+> ppCode body)
+          Just body -> Just (noLoc P.kwAssign <> oneLineOrNext (ppCode body))
     doc'
       ?<> builtin'
       <?+> termin'
         ?<> ( name'
                 <+> noLoc P.kwColon
-                <+> nest
-                  ( type'
-                      <+?> body'
-                  )
+                  <> oneLineOrNext
+                    ( type'
+                        <+?> body'
+                    )
             )
 
 instance PrettyPrint Pattern where

src/Juvix/Data/Fixity.hs

@@ -43,6 +43,7 @@ makeLenses ''Fixity
 data Atomicity
   = Atom
   | Aggregate Fixity
+  deriving stock (Eq)
 
 class HasAtomicity a where
   atomicity :: a -> Atomicity

tests/positive/Format.juvix

@@ -81,4 +81,41 @@ module Format;
   -- function with single wildcard parameter
   g : (_ : Type) -> Nat;
   g _ := 1;
+
+  -- grouping of type arguments
+  exampleFunction1 :
+    {A : Type}
+      -> List A
+      -> List A
+      -> List A
+      -> List A
+      -> List A
+      -> List A
+      -> List A
+      -> Nat;
+  exampleFunction1 _ _ _ _ _ _ _ := 1;
+
+  -- 200 in the body is idented with respect to the start of the chain
+  -- (at {A : Type})
+  exampleFunction2 :
+    {A : Type}
+      -> List A
+      -> List A
+      -> List A
+      -> List A
+      -> List A
+      -> List A
+      -> List A
+      -> Nat :=
+      200
+        + 100
+        + 100
+        + 100
+        + 100
+        + 100
+        + 100
+        + 100
+        + 100
+        + 100
+        + 100;
 end;