pretty printing fixes

lib/Quox/Pretty.idr

@@ -308,9 +308,9 @@ export
 prettyApp : {opts : LayoutOpts} -> Nat -> Doc opts ->
             List (Doc opts) -> Doc opts
 prettyApp ind f args =
-      hsep (f :: args)
-  <|> hsep [f, vsep args]
-  <|> vsep (f :: map (indent ind) args)
+  ifMultiline
+    (hsep (f :: args))
+    (f <++> vsep args <|> vsep (f :: map (indent ind) args))
 
 export
 prettyAppD : {opts : LayoutOpts} -> Doc opts -> List (Doc opts) ->

lib/Quox/Syntax/Term/Pretty.idr

@@ -342,18 +342,39 @@ private covering
 prettyLets : {opts : LayoutOpts} ->
              BContext d -> BContext a -> Telescope (LetBinder d) a b ->
              Eff Pretty (SnocList (Doc opts))
-prettyLets dnames xs lets = sequence $ snd $ go lets where
-  go : forall b. Telescope (LetBinder d) a b ->
-       (BContext b, SnocList (Eff Pretty (Doc opts)))
-  go [<] = (xs, [<])
-  go (lets :< (qty, x, rhs)) =
-    let (ys, docs) = go lets
-        doc = do
+prettyLets dnames xs lets = snd <$> go lets where
+  peelAnn : forall d, n. Elim d n -> Maybe (Term d n, Term d n)
+  peelAnn (Ann tm ty _) = Just (tm, ty)
+  peelAnn e             = Nothing
+
+  letHeader : Qty -> BindName -> Eff Pretty (Doc opts)
+  letHeader qty x = do
     lett <- [|letD <+> prettyQty qty|]
     x    <- prettyTBind x
-          rhs  <- withPrec Outer $ assert_total prettyElim dnames ys rhs
-          hangDSingle (hsep [lett, x, !cstD]) (hsep [rhs, !inD]) in
-    (ys :< x, docs :< doc)
+    pure $ lett <++> x
+
+  letBody : forall n. BContext n ->
+            Doc opts -> Elim d n -> Eff Pretty (Doc opts)
+  letBody tnames hdr e = case peelAnn e of
+    Just (tm, ty) => do
+      ty <- withPrec Outer $ assert_total prettyTerm dnames tnames ty
+      tm <- withPrec Outer $ assert_total prettyTerm dnames tnames tm
+      colon <- colonD; eq <- cstD; d <- askAt INDENT
+      pure $ hangSingle d (hangSingle d hdr (colon <++> ty)) (eq <++> tm)
+    Nothing => do
+      e <- withPrec Outer $ assert_total prettyElim dnames tnames e
+      eq <- cstD; d <- askAt INDENT
+      pure $ ifMultiline
+        (hsep [hdr, eq, e])
+        (vsep [hdr, indent d $ hsep [eq, e]])
+
+  go : forall b. Telescope (LetBinder d) a b ->
+       Eff Pretty (BContext b, SnocList (Doc opts))
+  go [<] = pure (xs, [<])
+  go (lets :< (qty, x, rhs)) = do
+    (ys, docs) <- go lets
+    doc        <- letBody ys !(letHeader qty x) rhs
+    pure (ys :< x, docs :< doc)
 
 
 private
@@ -504,7 +525,10 @@ prettyTerm dnames tnames (Let qty rhs body _) = do
   let lines = toList $ heads :< body
   pure $ ifMultiline (hsep lines) (vsep lines)
 
-prettyTerm dnames tnames (E e) = prettyElim dnames tnames e
+prettyTerm dnames tnames (E e) =
+  case the (Elim d n) (pushSubsts' e) of
+    Ann tm _ _ => assert_total prettyTerm dnames tnames tm
+    _          => assert_total prettyElim dnames tnames e
 
 prettyTerm dnames tnames t0@(CloT (Sub t ph)) =
   prettyTerm dnames tnames $ assert_smaller t0 $ pushSubstsWith' id ph t
@@ -567,9 +591,12 @@ prettyElim dnames tnames e@(DApp {}) =
   prettyDTApps dnames tnames f xs
 
 prettyElim dnames tnames (Ann tm ty _) =
-  parensIfM Outer =<<
-    hangDSingle !(withPrec AnnL  [|prettyTerm dnames tnames tm <++> annD|])
-                !(withPrec Outer (prettyTerm dnames tnames ty))
+  case the (Term d n) (pushSubsts' tm) of
+    E e => assert_total prettyElim dnames tnames e
+    _   => do
+      tm <- withPrec AnnL  $ assert_total prettyTerm dnames tnames tm
+      ty <- withPrec Outer $ assert_total prettyTerm dnames tnames ty
+      parensIfM Outer =<< hangDSingle (tm <++> !annD) ty
 
 prettyElim dnames tnames (Coe ty p q val _) =
   parensIfM App =<< do

lib/Quox/Typing/Context.idr

@@ -340,8 +340,10 @@ export
 prettyTContext : {opts : _} ->
                  BContext d -> QContext n -> BContext n ->
                  TContext d n -> Eff Pretty (Doc opts)
-prettyTContext dnames qtys tnames tys =
-  separateTight !commaD <$> prettyTContext' dnames qtys tnames tys
+prettyTContext dnames qtys tnames tys = do
+  comma <- commaD
+  sepSingle . exceptLast (<+> comma) . toList <$>
+    prettyTContext' dnames qtys tnames tys
 
 export
 prettyTyContext : {opts : _} -> TyContext d n -> Eff Pretty (Doc opts)
@@ -349,7 +351,7 @@ prettyTyContext (MkTyContext dctx dnames tctx tnames qtys) =
   case dctx of
     C [<] => prettyTContext dnames qtys tnames tctx
     _     => pure $
-      sep [!(prettyDimEq dnames dctx) <++> !pipeD,
+      sepSingle [!(prettyDimEq dnames dctx) <++> !pipeD,
                  !(prettyTContext dnames qtys tnames tctx)]
 
 export

lib/Quox/Typing/Error.idr

@@ -256,7 +256,7 @@ parameters {opts : LayoutOpts} (showContext : Bool)
                Doc opts -> Eff Pretty (Doc opts)
   inContext' null ctx f doc =
     if showContext && not null then
-      pure $ vappend doc (sep ["in context", !(f ctx)])
+      vappend doc <$> hangDSingle "in context" !(f ctx)
     else pure doc
 
   export %inline
@@ -416,5 +416,6 @@ parameters {opts : LayoutOpts} (showContext : Bool)
 
   export
   prettyError : Error -> Eff Pretty (Doc opts)
-  prettyError err = sep <$> sequence
-    [prettyLoc err.loc, indentD =<< prettyErrorNoLoc err]
+  prettyError err = hangDSingle
+    !(prettyLoc err.loc)
+    !(indentD =<< prettyErrorNoLoc err)

lib/Quox/Untyped/Syntax.idr

@@ -105,25 +105,17 @@ prettyArg : {opts : LayoutOpts} -> BContext n -> Term n -> Eff Pretty (Doc opts)
 prettyArg xs arg = withPrec Arg $ prettyTerm xs arg
 
 export covering
-prettyAppHead : {opts : LayoutOpts} -> BContext n ->
-                Term n -> Eff Pretty (Doc opts)
-prettyAppHead xs fun = withPrec App $ prettyTerm xs fun
-
-export
-prettyApp' : {opts : LayoutOpts} ->
-             Doc opts -> SnocList (Doc opts) -> Eff Pretty (Doc opts)
-prettyApp' fun args = do
-  d <- askAt INDENT
-  let args = toList args
-  parensIfM App $
-        hsep (fun :: args)
-    <|> hsep [fun, vsep args]
-    <|> vsep (fun :: map (indent d) args)
-
-export covering
-prettyApp : {opts : LayoutOpts} -> BContext n ->
+prettyApp_ : {opts : LayoutOpts} -> BContext n ->
              Doc opts -> SnocList (Term n) -> Eff Pretty (Doc opts)
-prettyApp xs fun args = prettyApp' fun =<< traverse (prettyArg xs) args
+prettyApp_ xs fun args =
+  parensIfM App =<<
+    prettyAppD fun (toList !(traverse (prettyArg xs) args))
+
+export covering %inline
+prettyApp : {opts : LayoutOpts} -> BContext n ->
+            Term n -> SnocList (Term n) -> Eff Pretty (Doc opts)
+prettyApp xs fun args =
+  prettyApp_ xs !(prettyArg xs fun) args
 
 public export
 record PrettyCaseArm a n where
@@ -208,21 +200,21 @@ prettyTerm xs (Lam x body _) =
     vars <- hsep . toList' <$> traverse prettyTBind ys
     body <- withPrec Outer $ prettyTerm (xs . ys) body
     hangDSingle (hsep [!lamD, vars, !darrowD]) body
-prettyTerm xs (App fun arg _) =
-  let (fun, args) = splitApp fun in
-  prettyApp xs !(prettyAppHead xs fun) (args :< arg)
+prettyTerm xs (App fun arg _) = do
+  let (fun, args) = splitApp fun
+  prettyApp xs fun (args :< arg)
 prettyTerm xs (Pair fst snd _) =
   parens . separateTight !commaD =<<
     traverse (withPrec Outer . prettyTerm xs) (fst :: splitPair snd)
-prettyTerm xs (Fst pair _) = prettyApp xs !fstD [< pair]
-prettyTerm xs (Snd pair _) = prettyApp xs !sndD [< pair]
+prettyTerm xs (Fst pair _) = prettyApp_ xs !fstD [< pair]
+prettyTerm xs (Snd pair _) = prettyApp_ xs !sndD [< pair]
 prettyTerm xs (Tag tag _) = prettyTag tag
 prettyTerm xs (CaseEnum tag cases _) =
   prettyCase xs prettyTag tag $
     map (\(t, rhs) => MkPrettyCaseArm t [] rhs) $ toList cases
 prettyTerm xs (Absurd _) = hl Syntax "absurd"
 prettyTerm xs (Nat n _) = hl Constant $ pshow n
-prettyTerm xs (Succ nat _) = prettyApp xs !succD [< nat]
+prettyTerm xs (Succ nat _) = prettyApp_ xs !succD [< nat]
 prettyTerm xs (CaseNat nat zer suc _) =
   prettyCase xs pure nat [MkPrettyCaseArm !zeroD [] zer, !(sucCaseArm suc)]
 prettyTerm xs (Str s _) =
@@ -235,7 +227,7 @@ prettyTerm xs (Let x rhs body _) =
     let lines = toList $ heads :< body
     pure $ ifMultiline (hsep lines) (vsep lines)
 prettyTerm _ (Erased _) =
-  hl Syntax =<< ifUnicode "⌷" "[]"
+  hl Syntax =<< ifUnicode "□" "[]"
 
 export covering
 prettyDef : {opts : LayoutOpts} -> Name ->