From baafa065c523e5a21953fe776c4087c71e1fc9a8 Mon Sep 17 00:00:00 2001
From: rhiannon morris <rhi@rhiannon.website>
Date: Fri, 3 Nov 2023 17:42:07 +0100
Subject: [PATCH] more tidying of outputs

---
 lib/Quox/Untyped/Erase.idr  | 28 ++++++++++++++++++++++++----
 lib/Quox/Untyped/Scheme.idr | 17 ++++++++++-------
 lib/Quox/Untyped/Syntax.idr | 34 ++++++++++++++++++++--------------
 3 files changed, 54 insertions(+), 25 deletions(-)

diff --git a/lib/Quox/Untyped/Erase.idr b/lib/Quox/Untyped/Erase.idr
index 2aa40cf..0fc24ef 100644
--- a/lib/Quox/Untyped/Erase.idr
+++ b/lib/Quox/Untyped/Erase.idr
@@ -18,6 +18,7 @@ import Language.Reflection
 %language ElabReflection
 
 %hide TT.Name
+%hide AppView.(.head)
 
 
 public export
@@ -464,10 +465,25 @@ inlineable : U.Term n -> Bool
 inlineable (F {})      = True
 inlineable (B {})      = True
 inlineable (Tag {})    = True
+inlineable (Nat {})    = True
+inlineable (Str {})    = True
 inlineable (Absurd {}) = True
 inlineable (Erased {}) = True
 inlineable _           = False
 
+export
+droppable : U.Term n -> Bool
+droppable (F {})      = True
+droppable (B {})      = True
+droppable (Fst e _)   = droppable e
+droppable (Snd e _)   = droppable e
+droppable (Tag {})    = True
+droppable (Nat {})    = True
+droppable (Str {})    = True
+droppable (Absurd {}) = True
+droppable (Erased {}) = True
+droppable _           = False
+
 export
 trimLets : U.Term n -> U.Term n
 trimLets (F x loc) = F x loc
@@ -479,8 +495,11 @@ trimLets (Fst pair loc) = Fst (trimLets pair) loc
 trimLets (Snd pair loc) = Snd (trimLets pair) loc
 trimLets (Tag tag loc) = Tag tag loc
 trimLets (CaseEnum tag cases loc) =
-  CaseEnum (trimLets tag)
-           (map (map $ \c => trimLets $ assert_smaller cases c) cases) loc
+  let tag   = trimLets tag
+      cases = map (map $ \c => trimLets $ assert_smaller cases c) cases in
+  if droppable tag && length cases == 1
+     then snd cases.head
+     else CaseEnum tag cases loc
 trimLets (Absurd loc) = Absurd loc
 trimLets (Nat n loc) = Nat n loc
 trimLets (Succ nat loc) = Succ (trimLets nat) loc
@@ -492,8 +511,9 @@ trimLets (CaseNat nat zer suc loc) =
 trimLets (Str s loc) = Str s loc
 trimLets (Let x rhs body loc) =
   let rhs'  = trimLets rhs
-      body' = trimLets body in
-  if inlineable rhs' || uses VZ body' == 1
+      body' = trimLets body
+      uses  = uses VZ body in
+  if inlineable rhs' || uses == 1 || (droppable rhs' && uses == 0)
      then sub1 rhs' body'
      else Let x rhs' body' loc
 trimLets (Erased loc) = Erased loc
diff --git a/lib/Quox/Untyped/Scheme.idr b/lib/Quox/Untyped/Scheme.idr
index a005d4c..be3ee6b 100644
--- a/lib/Quox/Untyped/Scheme.idr
+++ b/lib/Quox/Untyped/Scheme.idr
@@ -292,9 +292,11 @@ defToScheme x (SchemeDef isMain str) = do
   modifyAt AVOID $ insert x
   pure $ Just $ Define x $ Literal str
 
-orIndent : {opts : LayoutOpts} -> Doc opts -> Doc opts -> Doc opts
-orIndent a b =
-  parens $ ifMultiline (a <++> b) (a `vappend` indent 2 b)
+orIndent : {opts : LayoutOpts} -> Doc opts -> Doc opts -> Eff Pretty (Doc opts)
+orIndent a b = do
+  one <- parens $ a <++> b
+  two <- parens $ a `vappend` indent 2 b
+  pure $ ifMultiline one two
 
 export covering
 prettySexp : {opts : LayoutOpts} -> Sexp -> Eff Pretty (Doc opts)
@@ -303,7 +305,7 @@ private covering
 prettyLambda : {opts : LayoutOpts} ->
                String -> List Id -> Sexp -> Eff Pretty (Doc opts)
 prettyLambda lam xs e =
-  pure $ orIndent
+  orIndent
     (hsep [!(hl Syntax $ text lam), !(prettySexp $ L $ map V xs)])
     !(prettySexp e)
 
@@ -316,7 +318,7 @@ prettyLet : {opts : LayoutOpts} ->
             SnocList (Id, Sexp) -> Sexp -> Eff Pretty (Doc opts)
 prettyLet ps (Let x rhs body) = prettyLet (ps :< (x, rhs)) body
 prettyLet ps e =
-  pure $ orIndent
+  orIndent
     (hsep [!(hl Syntax "let*"),
            !(bracks . vsep . toList =<< traverse prettyBind ps)])
     !(prettySexp e)
@@ -335,8 +337,9 @@ prettySexp (L []) = hl Delim "()"
 prettySexp (L (x :: xs)) = do
   d  <- prettySexp x
   ds <- traverse prettySexp xs
-  parens $ (hsep $ d :: ds) <|> (hsep [d, vsep ds]) <|>
-           (vsep $ d :: map (indent 2) ds)
+  parens $ ifMultiline
+    (hsep $ d :: ds)
+    (hsep [d, vsep ds] <|> vsep (d :: map (indent 2) ds))
 prettySexp (Q (V x)) = hl Tag $ "'" <+> prettyId' x
 prettySexp (Q x) = pure $ hcat [!(hl Tag "'"), !(prettySexp x)]
 prettySexp (N n) = hl Tag $ pshow n
diff --git a/lib/Quox/Untyped/Syntax.idr b/lib/Quox/Untyped/Syntax.idr
index b8b6abb..1697a66 100644
--- a/lib/Quox/Untyped/Syntax.idr
+++ b/lib/Quox/Untyped/Syntax.idr
@@ -112,7 +112,7 @@ prettyArg xs arg = withPrec Arg $ prettyTerm xs arg
 export covering
 prettyAppHead : {opts : LayoutOpts} -> BContext n ->
                 Term n -> Eff Pretty (Doc opts)
-prettyAppHead xs fun = parensIfM App =<< prettyTerm xs fun
+prettyAppHead xs fun = withPrec App $ prettyTerm xs fun
 
 export
 prettyApp' : {opts : LayoutOpts} ->
@@ -120,14 +120,15 @@ prettyApp' : {opts : LayoutOpts} ->
 prettyApp' fun args = do
   d <- askAt INDENT
   let args = toList args
-  pure $ hsep (fun :: args)
-     <|> hsep [fun, vsep args]
-     <|> vsep (fun :: map (indent d) args)
+  parensIfM App $
+        hsep (fun :: args)
+    <|> hsep [fun, vsep args]
+    <|> vsep (fun :: map (indent d) args)
 
 export covering
 prettyApp : {opts : LayoutOpts} -> BContext n ->
             Doc opts -> SnocList (Term n) -> Eff Pretty (Doc opts)
-prettyApp xs fun args = prettyApp' fun =<< traverse (prettyTerm xs) args
+prettyApp xs fun args = prettyApp' fun =<< traverse (prettyArg xs) args
 
 public export
 record PrettyCaseArm a n where
@@ -159,10 +160,15 @@ sucPat : {opts : LayoutOpts} -> BindName -> Eff Pretty (Doc opts)
 sucPat x = pure $ !succD <++> !(prettyTBind x)
 
 export
-splitApp : Term b -> (Term b, SnocList (Term b))
+splitApp : Term n -> (Term n, SnocList (Term n))
 splitApp (App f x _) = mapSnd (:< x) $ splitApp f
 splitApp f           = (f, [<])
 
+export
+splitPair : Term n -> List (Term n)
+splitPair (Pair s t _) = s :: splitPair t
+splitPair t            = [t]
+
 export
 splitLam : Telescope' BindName a b -> Term b ->
            Exists $ \c => (Telescope' BindName a c, Term c)
@@ -185,9 +191,10 @@ prettyLets xs lets = sequence $ snd $ go lets where
   go [<]                = (xs, [<])
   go (lets :< (x, rhs)) =
     let (ys, docs) = go lets
-        doc = hsep <$> sequence
-              [letD, prettyTBind x, cstD, prettyTerm ys rhs, inD]
-    in
+        doc = do
+          x   <- prettyTBind x
+          rhs <- withPrec Outer $ prettyTerm ys rhs
+          hangDSingle (hsep [!letD, x, !cstD]) (hsep [rhs, !inD]) in
     (ys :< x, docs :< doc)
 
 private
@@ -210,8 +217,8 @@ prettyTerm xs (App fun arg _) =
   let (fun, args) = splitApp fun in
   prettyApp xs !(prettyAppHead xs fun) (args :< arg)
 prettyTerm xs (Pair fst snd _) =
-  parens =<< separateTight !commaD <$>
-    sequence {t = List} [prettyTerm xs fst, prettyTerm xs 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 (Tag tag _) = prettyTag tag
@@ -220,8 +227,7 @@ prettyTerm xs (CaseEnum tag cases _) =
     map (\(t, rhs) => MkPrettyCaseArm t [] rhs) $ toList cases
 prettyTerm xs (Absurd _) = hl Syntax "absurd"
 prettyTerm xs (Nat n _) = hl Tag $ pshow n
-prettyTerm xs (Succ nat _) =
-  prettyApp' !succD [< !(prettyTerm xs 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 _) =
@@ -244,7 +250,7 @@ prettyDef name ErasedDef =
 prettyDef name (KeptDef isMain rhs) = do
   name <- prettyFree name {opts}
   eq   <- cstD
-  rhs  <- prettyTerm [<] rhs
+  rhs  <- withPrec Outer $ prettyTerm [<] rhs
   let header = if isMain then text "#[main]" <++> name else name
   hangDSingle (header <++> eq) rhs
 prettyDef name (SchemeDef isMain str) = do