Improve formatter (#1840)

- Closes #1793.

Now, if the body of a function clause does not fit in a line, the body
will start indented in the next line.

The example presented in the linked issue is now formatted thus:
```
  go n s :=
    if
      (s < n)
      (go (sub n 1) s)
      (go n (sub s n) + go (sub n 1) s);
```

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

@@ -530,7 +530,7 @@ instance (SingI s) => PrettyCode (FunctionClause s) where
       clauseOwnerFunction'
         <+?> clausePatterns'
         <+> kwAssign
-        <+> nest 2 clauseBody'
+        <+> oneLineOrNext clauseBody'
 
 instance (SingI s) => PrettyCode (AxiomDef s) where
   ppCode AxiomDef {..} = do
@@ -636,7 +636,7 @@ instance PrettyCode Application where
       let (f, args) = unfoldApplication a
       f' <- ppCode f
       args' <- mapM ppCodeAtom args
-      return $ PP.group (f' <+> nest 2 (vsep 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)
 apeHelper a alt = do

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

@@ -274,7 +274,7 @@ instance PrettyPrint (FunctionClause 'Scoped) where
     clauseFun'
       <+?> clausePatterns'
       <+> noLoc P.kwAssign
-      <+> nest clauseBody'
+      <+> oneLineOrNext clauseBody'
 
 ppPatternAtom :: forall r. (Members '[Reader Options, ExactPrint] r) => PatternArg -> Sem r ()
 ppPatternAtom pat =

src/Juvix/Data/Ape/Base.hs

@@ -5,17 +5,13 @@ import Juvix.Prelude
 class IsApe a e where
   toApe :: a -> Ape e
 
--- TODO add ApeParens
-
 -- | Abstract pretty expression
 data Ape a
   = ApeLeaf (Leaf a)
   | ApeInfix (Infix a)
   | ApeApp (App a)
   | ApePostfix (Postfix a)
 
--- TODO add CapeParens
-
 -- | Abstract pretty expressions with chains
 data Cape a
   = CapeLeaf (Leaf a)

src/Juvix/Data/Ape/Pretty.hs

@@ -38,11 +38,14 @@ ppCape = \case
   CapeAppChain c -> ppAppChain c
   CapeUChain c -> ppUChain c
 
+chain :: Doc CodeAnn -> NonEmpty (Doc CodeAnn) -> Doc CodeAnn
+chain f' args' = PP.group (nest' (vsep (f' : toList args')))
+
 ppAppChain :: forall a r. (Members '[Reader (ApeParams a)] r) => AppChain a -> Sem r (Doc CodeAnn)
 ppAppChain (AppChain f links) = do
   f' <- ppLinkExpr fx f
   args' <- mapM (ppLinkExpr fx) links
-  return $ PP.group (vsep (f' : toList args'))
+  return $ chain f' args'
   where
     fx :: Precedence
     fx = appFixity ^. fixityPrecedence
@@ -51,7 +54,7 @@ ppChain :: forall a r. (Members '[Reader (ApeParams a)] r) => Chain a -> Sem r (
 ppChain (Chain opFix f links) = do
   f' <- ppLinkExpr fx f
   args' <- mapM ppLink links
-  return $ PP.group (vsep (f' : toList args'))
+  return $ chain f' args'
   where
     fx :: Precedence
     fx = opFix ^. fixityPrecedence
@@ -72,25 +75,12 @@ ppUChain (UChain opFix f links) = do
     fx :: Precedence
     fx = opFix ^. fixityPrecedence
 
-nestIf :: Bool -> Doc a -> Doc a
-nestIf = \case
-  True -> nest 2
-  False -> id
-
 ppLinkExpr ::
   (Members '[Reader (ApeParams a)] r) => Precedence -> Cape a -> Sem r (Doc CodeAnn)
-ppLinkExpr opFix e =
-  nestIf (apeNest (atomicity e) opFix)
-    . parensCond cond
-    <$> ppCape e
+ppLinkExpr opFix e = parensCond cond <$> ppCape e
   where
     cond = apeParens (atomicity e) opFix
 
-apeNest :: Atomicity -> Precedence -> Bool
-apeNest argAtom opPrec = case argAtom of
-  Atom -> True
-  Aggregate argFix -> argFix ^. fixityPrecedence >= opPrec
-
 apeParens :: Atomicity -> Precedence -> Bool
 apeParens argAtom opPrec = case argAtom of
   Atom -> False

src/Juvix/Data/Effect/ExactPrint.hs

@@ -105,3 +105,6 @@ enclose l r p = l >> p >> r
 
 encloseSep :: (Monad m, Foldable f) => m () -> m () -> m () -> f (m ()) -> m ()
 encloseSep l r sep f = l >> sequenceWith sep f >> r
+
+oneLineOrNext :: Members '[ExactPrint] r => Sem r () -> Sem r ()
+oneLineOrNext = region P.oneLineOrNext

src/Juvix/Prelude/Pretty.hs

@@ -112,12 +112,18 @@ hsepMaybe l
   | null l = Nothing
   | otherwise = Just (hsep l)
 
+nest' :: Doc ann -> Doc ann
+nest' = nest 2
+
 indent' :: Doc ann -> Doc ann
 indent' = indent 2
 
 hang' :: Doc ann -> Doc ann
 hang' = hang 2
 
+oneLineOrNext :: Doc ann -> Doc ann
+oneLineOrNext x = PP.group (flatAlt (line <> indent' x) x)
+
 ordinal :: Int -> Doc a
 ordinal = \case
   1 -> "first"

tests/positive/Ape.juvix

@@ -1,5 +1,5 @@
 module Ape;
-  axiom String : Type;
+  builtin string axiom String : Type;
 
   infixl 7 *;
   axiom * : String → String → String;