From f5bd2d443e5c21a2e54d7f1cdb821fabbfa6bf75 Mon Sep 17 00:00:00 2001
From: piegames <git@piegames.de>
Date: Tue, 6 Feb 2024 00:29:28 +0100
Subject: [PATCH] foo

---
 src/Nixfmt/Pretty.hs           | 46 +++++++++++++++++++++++++++-------
 test/diff/idioms_lib_5/out.nix |  4 +--
 2 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/src/Nixfmt/Pretty.hs b/src/Nixfmt/Pretty.hs
index 3d12666f..1f20a411 100644
--- a/src/Nixfmt/Pretty.hs
+++ b/src/Nixfmt/Pretty.hs
@@ -295,6 +295,7 @@ instance Pretty Parameter where
 -- `f g a` -> pre [f line g] line a post
 -- `f g h a` -> pre [[f line g] line h] line a post
 -- `f g h i a` -> pre [[[f line g] line h] line i] line a post
+--
 -- As you can see, it separates the elements by `line` whitespace. However, there are three tricks to make it look good:
 -- First, for each function call (imagine the fully parenthesised Nix code), we group it. Due to the greedy expansion
 -- of groups this means that it will place as many function calls on the first line as possible, but then all the remaining
@@ -305,8 +306,14 @@ instance Pretty Parameter where
 -- Third, callers may inject `pre` and `post` tokens (mostly newlines) into the inside of the group.
 -- This means that callers can say "try to be compact first, but if more than the last argument does not fit onto the line,
 -- then start on a new line instead".
--- Out of necessity, callers may also inject `commentPre` and `commentPost`, which will be added before/after the entire
--- thing if the function has a comment associated with its first token
+--
+-- When the function is "simple", we force it onto a single line
+--
+-- When a non-last argument is multiline, then this will usually result in it not being absorbed and in starting
+-- on a new line with indentation, same for all following arguments.
+-- When what follows the function is "simple" though and fits onto a single line, we can instead absorb that multiline
+-- argument and put the rest onto its closing line.
+-- While the code only checks this for the last argument, due to recursion it applies to any number of single-line simple trailing arguments.
 prettyApp :: Bool -> Doc -> Bool -> Expression -> Expression -> Doc
 prettyApp indentFunction pre hasPost f a
     = let
@@ -578,19 +585,40 @@ instance Pretty a => Pretty (Whole a) where
 instance Pretty Token where
     pretty = text . tokenText
 
+isSimpleLeaf :: Leaf -> Bool
+isSimpleLeaf (Ann [] (Identifier _) Nothing) = True
+isSimpleLeaf _ = False
+
+isSimpleItem :: (a -> Bool) -> Item a -> Bool
+isSimpleItem p (CommentedItem [] a) = p a
+isSimpleItem _ _ = False
 
 isSimpleSelector :: Selector -> Bool
 isSimpleSelector (Selector _ (IDSelector _) Nothing) = True
 isSimpleSelector _                                   = False
 
-isSimple :: Expression -> Bool
-isSimple (Term (SimpleString (Ann [] _ Nothing))) = True
-isSimple (Term (IndentedString (Ann [] _ Nothing))) = True
-isSimple (Term (Path (Ann [] _ Nothing))) = True
-isSimple (Term (Token (Ann [] (Identifier _) Nothing))) = True
-isSimple (Term (Selection t selectors))
+isSimpleBinder :: Binder -> Bool
+isSimpleBinder (Inherit (Ann [] _ Nothing) Nothing leaves (Ann [] _ Nothing)) = all isSimpleLeaf leaves
+isSimpleBinder (Inherit (Ann [] _ Nothing) (Just from) leaves (Ann [] _ Nothing)) = isSimple (Term from) && all isSimpleLeaf leaves
+isSimpleBinder (Assignment sels (Ann [] _ Nothing) rhs (Ann [] _ Nothing)) =
+    all isSimpleSelector sels && isSimple rhs
+isSimpleBinder _ = False
+
+isSimpleTerm :: Term -> Bool
+isSimpleTerm (SimpleString (Ann [] _ Nothing)) = True
+isSimpleTerm (IndentedString (Ann [] _ Nothing)) = True
+isSimpleTerm (Path (Ann [] _ Nothing)) = True
+isSimpleTerm (Token leaf) = isSimpleLeaf leaf
+isSimpleTerm (Set Nothing (Ann [] _ Nothing) items (Ann [] _ Nothing)) = all (isSimpleItem isSimpleBinder) (unItems items)
+isSimpleTerm (Set (Just (Ann [] _ Nothing)) (Ann [] _ Nothing) items (Ann [] _ Nothing)) = all (isSimpleItem isSimpleBinder) (unItems items)
+isSimpleTerm (List (Ann [] _ Nothing) items (Ann [] _ Nothing)) = all (isSimpleItem isSimpleTerm) (unItems items)
+isSimpleTerm (Selection t selectors)
     = isSimple (Term t) && all isSimpleSelector selectors
-isSimple (Term (Parenthesized (Ann [] _ Nothing) e (Ann [] _ Nothing))) = isSimple e
+isSimpleTerm (Parenthesized (Ann [] _ Nothing) e (Ann [] _ Nothing)) = isSimple e
+isSimpleTerm _ = False
+
+isSimple :: Expression -> Bool
+isSimple (Term t) = isSimpleTerm t
 -- Function applications of simple terms are simple up to two arguments
 isSimple (Application (Application (Application _ _) _) _) = False
 isSimple (Application f a) = isSimple f && isSimple a
diff --git a/test/diff/idioms_lib_5/out.nix b/test/diff/idioms_lib_5/out.nix
index c98343fa..216b237c 100644
--- a/test/diff/idioms_lib_5/out.nix
+++ b/test/diff/idioms_lib_5/out.nix
@@ -383,9 +383,7 @@ let
       if typeCheck metaTypes.${k} v then
         null
       else
-        "key 'meta.${k}' has invalid value; expected ${metaTypes.${k}.description}, got\n    ${
-          lib.generators.toPretty { indent = "    "; } v
-        }"
+        "key 'meta.${k}' has invalid value; expected ${metaTypes.${k}.description}, got\n    ${lib.generators.toPretty { indent = "    "; } v}"
     else
       "key 'meta.${k}' is unrecognized; expected one of: \n  [${
         lib.concatMapStringsSep ", " (x: "'${x}'") (lib.attrNames metaTypes)