WIP: Maintain expanded attrs

src/Nixfmt/Lexer.hs

@@ -183,13 +183,14 @@ pushTrivia t = modify (<> t)
 lexeme :: Parser a -> Parser (Ann a)
 lexeme p = do
   lastLeading <- takeTrivia
+  SourcePos{sourceLine = line} <- getSourcePos
   token <- preLexeme p
   parsedTrivia <- trivia
   -- This is the position of the next lexeme after the currently parsed one
   SourcePos{sourceColumn = col} <- getSourcePos
   let (trailing, nextLeading) = convertTrivia parsedTrivia col
   pushTrivia nextLeading
-  return $ Ann lastLeading token trailing
+  return $ Ann lastLeading line token trailing
 
 -- | Tokens normally have only leading trivia and one trailing comment on the same
 -- line. A whole x also parses and stores final trivia after the x. A whole also

src/Nixfmt/Pretty.hs

@@ -60,14 +60,15 @@ import Nixfmt.Types (
   mapLastToken',
   tokenText,
  )
+import Text.Megaparsec (pos1)
 import Prelude hiding (String)
 
 toLineComment :: TrailingComment -> Trivium
 toLineComment (TrailingComment c) = LineComment $ " " <> c
 
 -- If the token has some trailing comment after it, move that in front of the token
 moveTrailingCommentUp :: Ann a -> Ann a
-moveTrailingCommentUp (Ann pre a (Just post)) = Ann (pre ++ [toLineComment post]) a Nothing
+moveTrailingCommentUp (Ann pre pos a (Just post)) = Ann (pre ++ [toLineComment post]) pos a Nothing
 moveTrailingCommentUp a = a
 
 instance Pretty TrailingComment where
@@ -103,13 +104,13 @@ instance Pretty [Trivium] where
   pretty trivia = hardline <> hcat trivia
 
 instance (Pretty a) => Pretty (Ann a) where
-  pretty (Ann leading x trailing') =
+  pretty (Ann leading _ x trailing') =
     pretty leading <> pretty x <> pretty trailing'
 
 instance Pretty SimpleSelector where
   pretty (IDSelector i) = pretty i
   pretty (InterpolSelector interpol) = pretty interpol
-  pretty (StringSelector (Ann leading s trailing')) =
+  pretty (StringSelector (Ann leading _ s trailing')) =
     pretty leading <> prettySimpleString s <> pretty trailing'
 
 instance Pretty Selector where
@@ -152,17 +153,19 @@ instance Pretty Binder where
 -- be even more eager at expanding, except for empty sets and inherit statements.
 prettySet :: Bool -> (Maybe Leaf, Leaf, Items Binder, Leaf) -> Doc
 -- Empty attribute set
-prettySet _ (krec, Ann [] paropen Nothing, Items [], parclose@(Ann [] _ _)) =
-  pretty (fmap (,hardspace) krec) <> pretty paropen <> hardspace <> pretty parclose
+prettySet _ (krec, Ann [] pos paropen Nothing, Items [], parclose@(Ann [] pos' _ _)) =
+  pretty (fmap (,hardspace) krec) <> pretty paropen <> sep <> pretty parclose
+  where
+    sep = if pos /= pos' then hardline else hardspace
 -- Singleton sets are allowed to fit onto one line,
 -- but apart from that always expand.
-prettySet wide (krec, Ann pre paropen post, binders, parclose) =
+prettySet wide (krec, Ann pre pos paropen post, binders, parclose@(Ann _ pos' _ _)) =
   pretty (fmap (,hardspace) krec)
-    <> pretty (Ann pre paropen Nothing)
+    <> pretty (Ann pre pos paropen Nothing)
     <> surroundWith sep (nest $ pretty post <> prettyItems binders)
     <> pretty parclose
   where
-    sep = if wide && not (null (unItems binders)) then hardline else line
+    sep = if (wide && not (null (unItems binders))) || pos /= pos' then hardline else line
 
 prettyTermWide :: Term -> Doc
 prettyTermWide (Set krec paropen items parclose) = prettySet True (krec, paropen, items, parclose)
@@ -171,8 +174,8 @@ prettyTermWide t = prettyTerm t
 -- | Pretty print a term without wrapping it in a group.
 prettyTerm :: Term -> Doc
 prettyTerm (Token t) = pretty t
-prettyTerm (SimpleString (Ann leading s trailing')) = pretty leading <> prettySimpleString s <> pretty trailing'
-prettyTerm (IndentedString (Ann leading s trailing')) = pretty leading <> prettyIndentedString s <> pretty trailing'
+prettyTerm (SimpleString (Ann leading _ s trailing')) = pretty leading <> prettySimpleString s <> pretty trailing'
+prettyTerm (IndentedString (Ann leading _ s trailing')) = pretty leading <> prettyIndentedString s <> pretty trailing'
 prettyTerm (Path p) = pretty p
 prettyTerm (Selection term selectors rest) =
   pretty term
@@ -190,21 +193,21 @@ prettyTerm (Selection term selectors rest) =
       _ -> line'
 
 -- Empty list
-prettyTerm (List (Ann leading paropen Nothing) (Items []) (Ann [] parclose trailing')) =
+prettyTerm (List (Ann leading _ paropen Nothing) (Items []) (Ann [] _ parclose trailing')) =
   pretty leading <> pretty paropen <> hardspace <> pretty parclose <> pretty trailing'
 -- General list
 -- Always expand if len > 1
-prettyTerm (List (Ann pre paropen post) items parclose) =
-  pretty (Ann pre paropen Nothing)
+prettyTerm (List (Ann pre pos paropen post) items parclose) =
+  pretty (Ann pre pos paropen Nothing)
     <> surroundWith line (nest $ pretty post <> prettyItems items)
     <> pretty parclose
 prettyTerm (Set krec paropen items parclose) = prettySet False (krec, paropen, items, parclose)
 -- Parentheses
-prettyTerm (Parenthesized paropen expr (Ann closePre parclose closePost)) =
+prettyTerm (Parenthesized paropen expr (Ann closePre pos parclose closePost)) =
   group $
     pretty (moveTrailingCommentUp paropen)
       <> nest (inner <> pretty closePre)
-      <> pretty (Ann [] parclose closePost)
+      <> pretty (Ann [] pos parclose closePost)
   where
     inner =
       case expr of
@@ -244,17 +247,17 @@ instance Pretty ParamAttr where
 -- This assumes that all items already have a trailing comma from earlier pre-processing
 moveParamAttrComment :: ParamAttr -> ParamAttr
 -- Simple parameter
-moveParamAttrComment (ParamAttr (Ann trivia name (Just comment')) Nothing (Just (Ann [] comma Nothing))) =
-  ParamAttr (Ann trivia name Nothing) Nothing (Just (Ann [] comma (Just comment')))
+moveParamAttrComment (ParamAttr (Ann trivia pos name (Just comment')) Nothing (Just (Ann [] pos' comma Nothing))) =
+  ParamAttr (Ann trivia pos name Nothing) Nothing (Just (Ann [] pos' comma (Just comment')))
 -- Parameter with default value
-moveParamAttrComment (ParamAttr name (Just (qmark, def)) (Just (Ann [] comma Nothing))) =
-  ParamAttr name (Just (qmark, def')) (Just (Ann [] comma comment'))
+moveParamAttrComment (ParamAttr name (Just (qmark, def)) (Just (Ann [] pos comma Nothing))) =
+  ParamAttr name (Just (qmark, def')) (Just (Ann [] pos comma comment'))
   where
     -- Extract comment at the end of the line
     (def', comment') =
       mapLastToken'
         ( \case
-            (Ann trivia t (Just comment'')) -> (Ann trivia t Nothing, Just comment'')
+            (Ann trivia pos' t (Just comment'')) -> (Ann trivia pos' t Nothing, Just comment'')
             ann -> (ann, Nothing)
         )
         def
@@ -268,25 +271,25 @@ moveParamsComments
   -- , name1
   -- # comment
   -- , name2
-  ( (ParamAttr name maybeDefault (Just (Ann trivia comma Nothing)))
-      : (ParamAttr (Ann trivia' name' trailing') maybeDefault' maybeComma')
+  ( (ParamAttr name maybeDefault (Just (Ann trivia pos comma Nothing)))
+      : (ParamAttr (Ann trivia' pos' name' trailing') maybeDefault' maybeComma')
       : xs
     ) =
-    ParamAttr name maybeDefault (Just (Ann [] comma Nothing))
-      : moveParamsComments (ParamAttr (Ann (trivia ++ trivia') name' trailing') maybeDefault' maybeComma' : xs)
+    ParamAttr name maybeDefault (Just (Ann [] pos comma Nothing))
+      : moveParamsComments (ParamAttr (Ann (trivia ++ trivia') pos' name' trailing') maybeDefault' maybeComma' : xs)
 -- This may seem like a nonsensical case, but keep in mind that blank lines also count as comments (trivia)
 moveParamsComments
   -- , name
   -- # comment
   -- ellipsis
-  [ ParamAttr name maybeDefault (Just (Ann trivia comma Nothing)),
-    ParamEllipsis (Ann trivia' name' trailing')
+  [ ParamAttr name maybeDefault (Just (Ann trivia pos comma Nothing)),
+    ParamEllipsis (Ann trivia' pos' name' trailing')
     ] =
-    [ ParamAttr name maybeDefault (Just (Ann [] comma Nothing)),
-      ParamEllipsis (Ann (trivia ++ trivia') name' trailing')
+    [ ParamAttr name maybeDefault (Just (Ann [] pos comma Nothing)),
+      ParamEllipsis (Ann (trivia ++ trivia') pos' name' trailing')
     ]
 -- Inject a trailing comma on the last element if nessecary
-moveParamsComments [ParamAttr name def Nothing] = [ParamAttr name def (Just (Ann [] TComma Nothing))]
+moveParamsComments [ParamAttr name def Nothing] = [ParamAttr name def (Just (Ann [] pos1 TComma Nothing))]
 moveParamsComments (x : xs) = x : moveParamsComments xs
 moveParamsComments [] = []
 
@@ -308,7 +311,7 @@ instance Pretty Parameter where
       handleTrailingComma :: [ParamAttr] -> [Doc]
       handleTrailingComma [] = []
       -- That's the case we're interested in
-      handleTrailingComma [ParamAttr name maybeDefault (Just (Ann [] TComma Nothing))] =
+      handleTrailingComma [ParamAttr name maybeDefault (Just (Ann [] _ TComma Nothing))] =
         [pretty (ParamAttr name maybeDefault Nothing) <> trailing ","]
       handleTrailingComma (x : xs) = pretty x : handleTrailingComma xs
 
@@ -378,7 +381,7 @@ prettyApp indentFunction pre hasPost f a =
         ( Term
             ( Parenthesized
                 open
-                (Application (Term (Token ident@(Ann _ fn@(Identifier _) _))) (Term body))
+                (Application (Term (Token ident@(Ann _ _ fn@(Identifier _) _))) (Term body))
                 close
               )
           )
@@ -397,7 +400,7 @@ prettyApp indentFunction pre hasPost f a =
       -- 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)) . moveTrailingCommentUp)
+          ((\(Ann leading pos token trailing') -> (Ann [] pos token trailing', leading)) . moveTrailingCommentUp)
           f
 
       renderedF = pre <> group' Transparent (absorbApp fWithoutComment)
@@ -446,36 +449,36 @@ isAbsorbableExpr expr = case expr of
 
 isAbsorbable :: Term -> Bool
 -- Multi-line indented string
-isAbsorbable (IndentedString (Ann _ (_ : _ : _) _)) = True
+isAbsorbable (IndentedString (Ann _ _ (_ : _ : _) _)) = True
 isAbsorbable (Path _) = True
 -- Non-empty sets and lists
 isAbsorbable (Set _ _ (Items (_ : _)) _) = True
 isAbsorbable (List _ (Items (_ : _)) _) = True
-isAbsorbable (Parenthesized (Ann [] _ Nothing) (Term t) _) = isAbsorbable t
+isAbsorbable (Parenthesized (Ann [] _ _ Nothing) (Term t) _) = isAbsorbable t
 isAbsorbable _ = False
 
 isAbsorbableTerm :: Term -> Bool
 isAbsorbableTerm = isAbsorbable
 
 absorbParen :: Ann Token -> Expression -> Ann Token -> Doc
-absorbParen (Ann pre' open post') expr (Ann pre'' close post'') =
+absorbParen (Ann pre' pos open post') expr (Ann pre'' pos' close post'') =
   group' Priority $
     nest $
-      pretty (Ann pre' open Nothing)
+      pretty (Ann pre' pos open Nothing)
         -- Move any trailing comments on the opening parenthesis down into the body
         <> surroundWith
           line'
           ( group' RegularG $
               nest $
                 pretty
                   ( mapFirstToken
-                      (\(Ann leading token trailing') -> Ann (maybeToList (toLineComment <$> post') ++ leading) token trailing')
+                      (\(Ann leading pos'' token trailing') -> Ann (maybeToList (toLineComment <$> post') ++ leading) pos'' token trailing')
                       expr
                   )
                   -- Move any leading comments on the closing parenthesis up into the nest
                   <> pretty pre''
           )
-        <> pretty (Ann [] close post'')
+        <> pretty (Ann [] pos' close post'')
 
 -- Note that unlike for absorbable terms which can be force-absorbed, some expressions
 -- may turn out to not be absorbable. In that case, they should start with a line' so that
@@ -509,15 +512,15 @@ absorbRHS expr = case expr of
   (With{}) -> group' RegularG $ line <> pretty expr
   -- Special case `//` and `++` operations to be more compact in some cases
   -- Case 1: two arguments, LHS is absorbable term, RHS fits onto the last line
-  (Operation (Term t) (Ann [] op Nothing) b)
+  (Operation (Term t) (Ann [] _ op Nothing) b)
     | isAbsorbable t && isUpdateOrConcat op ->
         group' RegularG $ line <> group' Priority (prettyTermWide t) <> line <> pretty op <> hardspace <> pretty b
   -- Case 2a: LHS fits onto first line, RHS is an absorbable term
-  (Operation l (Ann [] op Nothing) (Term t))
+  (Operation l (Ann [] _ op Nothing) (Term t))
     | isAbsorbable t && isUpdateOrConcat op ->
         group' RegularG $ line <> pretty l <> line <> group' Transparent (pretty op <> hardspace <> group' Priority (prettyTermWide t))
   -- Case 2b: LHS fits onto first line, RHS is a function application
-  (Operation l (Ann [] op Nothing) (Application f a))
+  (Operation l (Ann [] _ op Nothing) (Application f a))
     | isUpdateOrConcat op ->
         line <> group l <> line <> prettyApp False (pretty op <> hardspace) False f a
   -- Everything else:
@@ -536,7 +539,7 @@ instance Pretty Expression where
   -- Let bindings are always fully expanded (no single-line form)
   -- 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) =
+  pretty (Let let_ binders (Ann leading pos in_ trailing') expr) =
     letPart <> hardline <> inPart
     where
       -- Convert the TrailingComment to a Trivium, if present
@@ -564,7 +567,7 @@ instance Pretty Expression where
       letBody = nest $ prettyItems (Items bindersWithoutComments)
       inPart =
         group $
-          pretty (Ann [] in_ Nothing)
+          pretty (Ann [] pos in_ Nothing)
             <> hardline
             -- Take our trailing and inject it between `in` and body
             <> pretty (concat binderComments ++ leading ++ convertTrailing trailing')
@@ -623,7 +626,7 @@ instance Pretty Expression where
   pretty (Application f a) =
     prettyApp False mempty False f a
   -- not chainable binary operators: <, >, <=, >=, ==, !=
-  pretty (Operation a op@(Ann _ op' _) b)
+  pretty (Operation a op@(Ann _ _ op' _) b)
     | op' == TLess || op' == TGreater || op' == TLessEqual || op' == TGreaterEqual || op' == TEqual || op' == TUnequal =
         pretty a <> softline <> pretty op <> hardspace <> pretty b
   -- all other operators
@@ -675,13 +678,13 @@ isSimpleSelector (Selector _ (IDSelector _)) = 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 (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 def)) =
   isSimple (Term t) && all isSimpleSelector selectors && isNothing def
-isSimple (Term (Parenthesized (Ann [] _ Nothing) e (Ann [] _ Nothing))) = isSimple e
+isSimple (Term (Parenthesized (Ann [] _ _ Nothing) e (Ann [] _ _ Nothing))) = isSimple e
 -- Function applications of simple terms are simple up to two arguments
 isSimple (Application (Application (Application _ _) _) _) = False
 isSimple (Application f a) = isSimple f && isSimple a

src/Nixfmt/Types.hs

@@ -14,6 +14,7 @@ import Data.List.NonEmpty as NonEmpty
 import Data.Maybe (maybeToList)
 import Data.Text (Text, pack)
 import Data.Void (Void)
+import Text.Megaparsec (Pos, pos1)
 import qualified Text.Megaparsec as MP (ParseErrorBundle, Parsec)
 import Prelude hiding (String)
 
@@ -38,20 +39,20 @@ type Trivia = [Trivium]
 newtype TrailingComment = TrailingComment Text deriving (Eq, Show)
 
 data Ann a
-  = Ann Trivia a (Maybe TrailingComment)
+  = Ann Trivia Pos a (Maybe TrailingComment)
   deriving (Show)
 
 hasTrivia :: Ann a -> Bool
-hasTrivia (Ann [] _ Nothing) = False
+hasTrivia (Ann [] _ _ Nothing) = False
 hasTrivia _ = True
 
 ann :: a -> Ann a
-ann a = Ann [] a Nothing
+ann a = Ann [] pos1 a Nothing
 
 -- | Equality of annotated syntax is defined as equality of their corresponding
 -- semantics, thus ignoring the annotations.
 instance (Eq a) => Eq (Ann a) where
-  Ann _ x _ == Ann _ y _ = x == y
+  Ann _ _ x _ == Ann _ _ y _ = x == y
 
 -- Trivia is ignored for Eq, so also don't show
 -- instance Show a => Show (Ann a) where
@@ -212,7 +213,7 @@ instance LanguageElement SimpleSelector where
 
   walkSubprograms = \case
     (IDSelector name) -> [Term (Token name)]
-    (InterpolSelector (Ann _ str _)) -> pure $ Term $ SimpleString $ Ann [] [[str]] Nothing
+    (InterpolSelector (Ann _ pos str _)) -> pure $ Term $ SimpleString $ Ann [] pos [[str]] Nothing
     (StringSelector str) -> [Term (SimpleString str)]
 
 instance LanguageElement Selector where

test/diff/apply/out.nix

@@ -113,17 +113,25 @@
     name3 = function arg { asdf = 1; } { qwer = 12345; } argument;
   }
   {
-    name1 = function arg { asdf = 1; };
+    name1 = function arg {
+      asdf = 1;
+    };
 
     name2 = function arg {
       asdf = 1;
       # multiline
     } argument;
 
-    name3 = function arg {
-      asdf = 1;
-      # multiline
-    } { qwer = 12345; } argument;
+    name3 =
+      function arg
+        {
+          asdf = 1;
+          # multiline
+        }
+        {
+          qwer = 12345;
+        }
+        argument;
   }
   {
     name4 = function arg { asdf = 1; } {

test/diff/attr_set/in.nix

@@ -7,6 +7,8 @@
 
   {
 
+  }
+  {
   }
 
   { a = {