Move comments around a bit

Also make sure that comments do not accidentally force-expand expressions

src/Nixfmt/Pretty.hs

@@ -4,7 +4,7 @@
  - SPDX-License-Identifier: MPL-2.0
  -}
 
-{-# LANGUAGE FlexibleInstances, OverloadedStrings #-}
+{-# LANGUAGE FlexibleInstances, OverloadedStrings, RankNTypes #-}
 
 module Nixfmt.Pretty where
 
@@ -16,14 +16,15 @@ import Data.Text (Text, isPrefixOf, isSuffixOf, stripPrefix)
 import qualified Data.Text as Text
   (dropEnd, empty, init, isInfixOf, last, null, strip, takeWhile)
 
+-- import Debug.Trace (traceShowId)
 import Nixfmt.Predoc
   (Doc, Pretty, base, emptyline, group, group', hardline, hardspace, hcat, line, line',
   nest, newline, pretty, sepBy, softline, softline', text, textWidth)
 import Nixfmt.Types
   (Ann(..), Binder(..), Expression(..), Item(..), Items(..), Leaf,
   ParamAttr(..), Parameter(..), Selector(..), SimpleSelector(..),
   StringPart(..), Term(..), Token(..), TrailingComment(..), Trivia, Trivium(..),
-  Whole(..), tokenText)
+  Whole(..), tokenText, mapFirstToken')
 import Nixfmt.Util (commonIndentation, isSpaces, replaceMultiple)
 
 prettyCommentLine :: Text -> Doc
@@ -260,7 +261,7 @@ isAbsorbable :: Term -> Bool
 isAbsorbable (String (Ann _ parts@(_:_:_) _))
     = not $ isSimpleString parts
 isAbsorbable (Set _ _ (Items (_:_)) _)                                   = True
-isAbsorbable (List (Ann [] _ Nothing) (Items [CommentedItem [] item]) _) = True
+isAbsorbable (List (Ann [] _ Nothing) (Items [CommentedItem [] _]) _)    = True
 isAbsorbable (Parenthesized (Ann [] _ Nothing) (Term t) _)               = isAbsorbable t
 isAbsorbable (List _ (Items (_:_:_)) _)                                  = True
 isAbsorbable _                                                           = False
@@ -299,12 +300,35 @@ instance Pretty Expression where
           <> absorbSet expr1
 
     -- Let bindings are always fully expanded (no single-line form)
-    pretty (Let let_ binders in_ expr)
+    -- We also take the comments around the `in` (trailing, leading and detached binder comments)
+    -- and move them down to the first token of the body
+    pretty (Let let_ binders (Ann leading in_ trailing) expr)
         = base $ letPart <> hardline <> inPart
         where
+          -- Convert the TrailingComment to a Trivium, if present
+          convertTrailing Nothing = []
+          convertTrailing (Just (TrailingComment t)) = [(LineComment (" " <> t))]
+
+          -- Extract detached comments at the bottom.
+          -- This uses a custom variant of span/spanJust/spanMaybe.
+          -- Note that this is a foldr which walks from the bottom, but the lists
+          -- are constructed in a way that they end up correct again.
+          (binderComments, bindersWithoutComments)
+            = foldr
+                (\item -> \(start, rest) ->
+                    case item of
+                        (DetachedComments inner) | null rest -> (inner : start, rest)
+                        _ -> (start, item : rest)
+                )
+                ([], [])
+                (unItems binders)
+
           letPart = groupWithStart let_ $ hardline <> letBody
-          inPart = groupWithStart in_ $ hardline <> pretty expr <> hardline
-          letBody = nest 2 $ prettyItems hardline binders
+          letBody = nest 2 $ prettyItems hardline (Items bindersWithoutComments)
+          inPart = groupWithStart (Ann [] in_ Nothing) $ hardline
+              -- Take our trailing and inject it between `in` and body
+              <> pretty (concat binderComments ++ leading ++ convertTrailing trailing)
+              <> pretty expr <> hardline
 
     pretty (Assert assert cond semicolon expr)
         = base (pretty assert <> hardspace
@@ -341,7 +365,7 @@ instance Pretty Expression where
     -- Secondly, the `line` between the second-to-last and last argument (marked with asterisk above) is moved into its preceding
     -- group. This allows the last argument to be multi-line without forcing the preceding arguments to be multiline.
     pretty (Application f a)
-        = let            
+        = let
             absorbApp (Application f' a') = (group $ absorbApp f') <> line <> (group a')
             absorbApp expr = pretty expr
 
@@ -350,9 +374,12 @@ instance Pretty Expression where
             absorbLast (Term (Parenthesized open expr close))
               = base $ group $ pretty open <> line' <> nest 2 (group expr) <> line' <> pretty close
             absorbLast arg = group arg
+
+            -- Extract comment before the first function and move it out, to prevent functions being force-expanded
+            (fWithoutComment, comment) = mapFirstToken' (\(Ann leading token trailing) -> (Ann [] token trailing, leading)) f
           in
-            group $
-              (group' False True $ (absorbApp f) <> line) <> (absorbLast a)
+            pretty comment <> (group $
+              (group' False True $ absorbApp fWithoutComment <> line) <> absorbLast a)
 
     -- '//' operator
     pretty (Operation a op@(Ann _ TUpdate _) b)

src/Nixfmt/Types.hs

@@ -4,7 +4,7 @@
  - SPDX-License-Identifier: MPL-2.0
  -}
 
-{-# LANGUAGE DeriveFoldable, OverloadedStrings #-}
+{-# LANGUAGE DeriveFoldable, OverloadedStrings, RankNTypes #-}
 
 module Nixfmt.Types where
 
@@ -128,6 +128,71 @@ data Whole a
 
 type File = Whole Expression
 
+-- Implemented by all AST-related types whose values are guaranteed to contain at least one (annotated) token
+class LanguageElement a where
+    -- Map the first token of that expression, no matter how deep it sits
+    -- in the AST. This is useful for modifying comments
+    mapFirstToken :: (forall b. Ann b -> Ann b) -> a -> a
+    mapFirstToken f a = fst (mapFirstToken' (\x -> (f x, ())) a)
+
+    -- Same as mapFirstToken, but the mapping function also yields a value that may be
+    -- returned. This is useful for getting/extracting values
+    mapFirstToken' :: (forall b. Ann b -> (Ann b, c)) -> a -> (a, c)
+
+instance LanguageElement Parameter where
+    mapFirstToken' f (IDParameter name)
+        = let (name', ret) = f name in (IDParameter name', ret)
+    mapFirstToken' f (SetParameter open items close)
+        = let (open', ret) = f open in (SetParameter open' items close, ret)
+    mapFirstToken' f (ContextParameter first at second)
+        = let (first', ret) = mapFirstToken' f first in ((ContextParameter first' at second), ret)
+
+instance LanguageElement Term where
+    mapFirstToken' f (Token leaf)
+        = let (leaf', ret) = (f leaf) in (Token leaf', ret)
+    mapFirstToken' f (String string)
+        = let (string', ret) = (f string) in (String string', ret)
+    mapFirstToken' f (Path path)
+        = let (path', ret) = (f path) in (Path path', ret)
+    mapFirstToken' f (List open items close)
+        = let (open', ret) = (f open) in (List open' items close, ret)
+    mapFirstToken' f (Set (Just rec) open items close)
+        = let (rec', ret) = (f rec) in (Set (Just rec') open items close, ret)
+    mapFirstToken' f (Set Nothing open items close)
+        = let (open', ret) = (f open) in (Set Nothing open' items close, ret)
+    mapFirstToken' f (Selection term selector)
+        = let (term', ret) = (mapFirstToken' f term) in (Selection term' selector, ret)
+    mapFirstToken' f (Parenthesized open expr close)
+        = let (open', ret) = (f open) in (Parenthesized open' expr close, ret)
+
+instance LanguageElement Expression where
+    mapFirstToken' f (Term term)
+        = let (term', ret) = (mapFirstToken' f term) in (Term term', ret)
+    mapFirstToken' f (With with expr0 semicolon expr1)
+        = let (with', ret) = (f with) in (With with' expr0 semicolon expr1, ret)
+    mapFirstToken' f (Let let_ items in_ body)
+        = let (let_', ret) = (f let_) in (Let let_' items in_ body, ret)
+    mapFirstToken' f (Assert assert cond semicolon body)
+        = let (assert', ret) = (f assert) in (Assert assert' cond semicolon body, ret)
+    mapFirstToken' f (If if_ expr0 then_ expr1 else_ expr2)
+        = let (if_', ret) = (f if_) in (If if_' expr0 then_ expr1 else_ expr2, ret)
+    mapFirstToken' f (Abstraction param colon body)
+        = let (param', ret) = (mapFirstToken' f param) in (Abstraction param' colon body, ret)
+    mapFirstToken' f (Application g a)
+        = let (g', ret) = (mapFirstToken' f g) in (Application g' a, ret)
+    mapFirstToken' f (Operation left op right)
+        = let (left', ret) = (mapFirstToken' f left) in (Operation left' op right, ret)
+    mapFirstToken' f (MemberCheck name dot selectors)
+        = let (name', ret) = (mapFirstToken' f name) in (MemberCheck name' dot selectors, ret)
+    mapFirstToken' f (Negation not_ expr)
+        = let (not_', ret) = (f not_) in (Negation not_' expr, ret)
+    mapFirstToken' f (Inversion tilde expr)
+        = let (tilde', ret) = (f tilde) in (Inversion tilde' expr, ret)
+
+instance LanguageElement a => LanguageElement (Whole a) where
+    mapFirstToken' f (Whole a trivia)
+        = let (a', ret) = (mapFirstToken' f a) in (Whole a' trivia, ret)
+
 data Token
     = Integer    Int
     | Float      Double

test/diff/apply/in.nix

@@ -1,4 +1,13 @@
 [
+  (
+    # Function call with comment
+    mapAttrsToStringsSep "\n" mkSection attrsOfAttrs
+  )
+  # Function call with comment
+  (mapAttrsToStringsSep "\n" mkSection attrsOfAttrs)
+  [
+    (mapAttrsToStringsSep [force long] "\n" mkSection attrsOfAttrs)
+  ]
   (a
   b)
   (

test/diff/apply/out.nix

@@ -1,4 +1,19 @@
 [
+  (
+    # Function call with comment
+    mapAttrsToStringsSep "\n" mkSection attrsOfAttrs)
+  # Function call with comment
+  (mapAttrsToStringsSep "\n" mkSection attrsOfAttrs)
+  [
+    (mapAttrsToStringsSep
+      [
+        force
+        long
+      ]
+      "\n"
+      mkSection
+      attrsOfAttrs)
+  ]
   (a b)
   ((a b) (a b)
     (a # b
@@ -91,8 +106,7 @@
         utils,
       }:
       # For each supported platform,
-      utils.lib.eachDefaultSystem
-      (system: { })
+      utils.lib.eachDefaultSystem (system: { })
       ;
   }
   {

test/diff/comment/out.nix

@@ -71,8 +71,8 @@
       #6
 
       d = 1;
-      #7
     in
+    #7
     d
   )
 

test/diff/idioms_lib_3/out.nix

@@ -173,8 +173,8 @@ rec {
         ''
         + toKeyValue { inherit mkKeyValue listsAsDuplicateKeys; } sectValues
         ;
-      # map input to ini sections
     in
+    # map input to ini sections
     mapAttrsToStringsSep "\n" mkSection attrsOfAttrs
     ;
 

test/diff/let_in/out.nix

@@ -17,7 +17,8 @@ let
   a =
     let
       c = 1;
-    in # e
+    in
+    # e
     f
     ;
   a =
@@ -30,7 +31,8 @@ let
   a =
     let
       c = 1; # d
-    in # e
+    in
+    # e
     f
     ;
   a =
@@ -42,7 +44,8 @@ let
   a =
     let # b
       c = 1;
-    in # e
+    in
+    # e
     f
     ;
   a =
@@ -54,7 +57,8 @@ let
   a =
     let # b
       c = 1; # d
-    in # e
+    in
+    # e
     f
     ;