From db79a178a2d9379ecbb1cf8d00c0ad74ff075ecb Mon Sep 17 00:00:00 2001
From: piegames <git@piegames.de>
Date: Tue, 18 Apr 2023 23:51:49 +0200
Subject: [PATCH] Expand let statements more

---
 src/Nixfmt/Pretty.hs | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/src/Nixfmt/Pretty.hs b/src/Nixfmt/Pretty.hs
index 16557ff8..036b8cb5 100644
--- a/src/Nixfmt/Pretty.hs
+++ b/src/Nixfmt/Pretty.hs
@@ -188,6 +188,16 @@ absorbThen :: Expression -> Doc
 absorbThen (Term t) | isAbsorbable t = hardspace <> prettyTerm t <> hardspace
 absorbThen x                         = line <> nest 2 (group x) <> line
 
+-- What is allowed to come on the same line as `in`?
+-- Absorbable terms like sets
+-- if, with, let
+absorbIn :: Expression -> Doc
+absorbIn (Term t) | isAbsorbable t = hardspace <> prettyTerm t <> hardspace
+absorbIn x@(If _ _ _ _ _ _)        = group x
+absorbIn x@(With _ _ _ _)          = group x
+absorbIn x@(Let _ _ _ _)           = group x
+absorbIn x                         = line <> nest 2 (group x) <> line
+
 absorbElse :: Expression -> Doc
 absorbElse (If if_ cond then_ expr0 else_ expr1)
     = hardspace <> pretty if_ <> hardspace <> group cond <> hardspace
@@ -212,9 +222,9 @@ instance Pretty Expression where
 
     pretty (Let (Ann let_ letTrailing letLeading) binders
                 (Ann in_ inTrailing inLeading) expr)
-        = base $ group letPart <> line <> group inPart
-        where letPart = pretty let_ <> pretty letTrailing <> line <> letBody
-              inPart = pretty in_ <> hardspace <> pretty expr
+        = base $ group letPart <> line <> inPart
+        where letPart = pretty let_ <> pretty letTrailing <> hardline <> letBody
+              inPart = pretty in_ <> hardspace <> absorbIn expr
               letBody = nest 2 $
                   pretty letLeading
                   <> sepBy hardline binders