Introduce support for optional trailing commas

We want trailing commas, but only in expanded argument lists.
This change allows us to do that.

src/Nixfmt/Predoc.hs

@@ -11,6 +11,7 @@
 module Nixfmt.Predoc
     ( text
     , comment
+    , trailing
     , sepBy
     , surroundWith
     , hcat
@@ -91,11 +92,15 @@ data DocAnn
     | Base
     deriving (Show, Eq)
 
+-- Comments do not count towards some line length limits
+-- Trailing tokens have the property that they will only exist in expanded groups, and "swallowed" in compact groups
+data TextAnn = Regular | Comment | Trailing
+    deriving (Show, Eq)
+
 -- | Single document element. Documents are modeled as lists of these elements
 -- in order to make concatenation simple.
 data DocE
-    -- Mark comments with a flag, to not count them to line length limits
-    = Text Bool Text
+    = Text TextAnn Text
     | Spacing Spacing
     | Node DocAnn Doc
     deriving (Show, Eq)
@@ -126,11 +131,16 @@ instance (Pretty a, Pretty b, Pretty c) => Pretty (a, b, c) where
 
 text :: Text -> Doc
 text "" = []
-text t  = [Text False t]
+text t  = [Text Regular t]
 
 comment :: Text -> Doc
 comment "" = []
-comment t  = [Text True t]
+comment t  = [Text Comment t]
+
+-- Text tokens that are only needed in expanded groups
+trailing :: Text -> Doc
+trailing "" = []
+trailing t = [Text Trailing t]
 
 -- | Group document elements together (see Node Group documentation)
 -- Must not contain non-hard whitespace (e.g. line, softline' etc.) at the start of the end.
@@ -276,8 +286,8 @@ mergeSpacings _            y            = y
 mergeLines :: Doc -> Doc
 mergeLines []                           = []
 mergeLines (Spacing a : Spacing b : xs) = mergeLines $ Spacing (mergeSpacings a b) : xs
-mergeLines (Text isComment a : Text isComment' b : xs) | isComment == isComment'
-    = mergeLines $ Text isComment (a <> b) : xs
+mergeLines (Text ann a : Text ann' b : xs) | ann == ann'
+    = mergeLines $ Text ann (a <> b) : xs
 mergeLines (Node ann xs : ys)           = Node ann (mergeLines xs) : mergeLines ys
 mergeLines (x : xs)                     = x : mergeLines xs
 
@@ -329,8 +339,9 @@ fits :: Int -> Doc -> Maybe Text
 fits c _ | c < 0 = Nothing
 fits _ [] = Just ""
 fits c (x:xs) = case x of
-    Text False t         -> (t<>) <$> fits (c - textWidth t) xs
-    Text True t          -> (t<>) <$> fits c xs
+    Text Regular t       -> (t<>) <$> fits (c - textWidth t) xs
+    Text Comment t       -> (t<>) <$> fits c xs
+    Text Trailing _      -> fits c xs
     Spacing Softbreak    -> fits c xs
     Spacing Break        -> fits c xs
     Spacing Softspace    -> (" "<>) <$> fits (c - 1) xs
@@ -345,8 +356,8 @@ fits c (x:xs) = case x of
 -- width 0, which always forces line breaks when possible.
 firstLineWidth :: Doc -> Int
 firstLineWidth []                       = 0
-firstLineWidth (Text False t : xs)      = textWidth t + firstLineWidth xs
-firstLineWidth (Text True _ : xs)       = firstLineWidth xs
+firstLineWidth (Text Comment _ : xs)    = firstLineWidth xs
+firstLineWidth (Text _ t : xs)          = textWidth t + firstLineWidth xs
 firstLineWidth (Spacing Hardspace : xs) = 1 + firstLineWidth xs
 firstLineWidth (Spacing _ : _)          = 0
 firstLineWidth (Node _ xs : ys)         = firstLineWidth (xs ++ ys)
@@ -357,8 +368,8 @@ firstLineFits :: Int -> Int -> Doc -> Bool
 firstLineFits targetWidth maxWidth docs = go maxWidth docs
     where go c _ | c < 0                = False
           go c []                       = maxWidth - c <= targetWidth
-          go c (Text False t : xs)      = go (c - textWidth t) xs
-          go c (Text True _ : xs)       = go c xs
+          go c (Text Regular t : xs)    = go (c - textWidth t) xs
+          go c (Text _ _ : xs)          = go c xs
           go c (Spacing Hardspace : xs) = go (c - 1) xs
           go c (Spacing _ : _)          = maxWidth - c <= targetWidth
           go c (Node (Group _) ys : xs)     =
@@ -422,8 +433,7 @@ layoutGreedy tw doc = Text.concat $ evalState (go [Chunk 0 $ Node (Group False)
             putNL = put (0, ti)
         in
         case x of
-        Text False t         -> putCC (nc + textWidth t) $> [lineStart, t]
-        Text True t          -> putCC (nc + textWidth t) $> [lineStart, t]
+        Text _ t             -> putCC (nc + textWidth t) $> [lineStart, t]
 
         -- This code treats whitespace as "expanded"
         -- A new line resets the column counter and sets the target indentation as current indentation

src/Nixfmt/Pretty.hs

@@ -19,7 +19,7 @@ import qualified Data.Text as Text
 -- import Debug.Trace (traceShowId)
 import Nixfmt.Predoc
   (Doc, Pretty, base, emptyline, group, group', hardline, hardspace, hcat, line, line',
-  nest, newline, pretty, sepBy, surroundWith, softline, softline', text, comment, textWidth)
+  nest, newline, pretty, sepBy, surroundWith, softline, softline', text, comment, trailing, textWidth)
 import Nixfmt.Types
   (Ann(..), Binder(..), Expression(..), Item(..), Items(..), Leaf,
   ParamAttr(..), Parameter(..), Selector(..), SimpleSelector(..),
@@ -50,8 +50,8 @@ moveTrailingCommentUp a = a
 -- if the first argument has some leading comments they will be put before
 -- the group
 groupWithStart :: HasCallStack => Pretty a => Ann a -> Doc -> Doc
-groupWithStart (Ann leading a trailing) b
-    = pretty leading <> group (pretty a <> pretty trailing <> b)
+groupWithStart (Ann leading a trailing') b
+    = pretty leading <> group (pretty a <> pretty trailing' <> b)
 
 instance Pretty TrailingComment where
     pretty (TrailingComment c)
@@ -88,14 +88,14 @@ instance Pretty [Trivium] where
     pretty trivia = hardline <> hcat trivia
 
 instance Pretty a => Pretty (Ann a) where
-    pretty (Ann leading x trailing)
-        = pretty leading <> pretty x <> pretty 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 leading <> prettySimpleString s <> pretty trailing
+    pretty (StringSelector (Ann leading s trailing'))
+        = pretty leading <> prettySimpleString s <> pretty trailing'
 
 instance Pretty Selector where
     pretty (Selector dot sel Nothing)
@@ -207,8 +207,8 @@ prettyTerm (Selection term@(Parenthesized _ _ _) selectors) = pretty term <> sof
 prettyTerm (Selection term selectors) = pretty term <> line' <> hcat selectors
 
 -- Empty list
-prettyTerm (List (Ann leading paropen Nothing) (Items []) (Ann [] parclose trailing))
-    = pretty leading <> pretty paropen <> hardspace <> pretty parclose <> pretty trailing
+prettyTerm (List (Ann leading paropen Nothing) (Items []) (Ann [] parclose trailing'))
+    = pretty leading <> pretty paropen <> hardspace <> pretty parclose <> pretty trailing'
 
 -- Singleton list
 -- Expand unless absorbable term or single line
@@ -350,7 +350,7 @@ prettyApp commentPre pre post commentPost f a
         absorbLast arg = group' False $ nest 2 $ pretty 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
+        (fWithoutComment, comment') = mapFirstToken' (\(Ann leading token trailing') -> (Ann [] token trailing', leading)) f
         in
         (if null comment' then mempty else commentPre)
         <> pretty comment' <> (group' False $
@@ -408,7 +408,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 in_ trailing') expr)
         = base $ letPart <> hardline <> inPart
         where
           -- Convert the TrailingComment to a Trivium, if present
@@ -433,7 +433,7 @@ instance Pretty Expression where
           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 (concat binderComments ++ leading ++ convertTrailing trailing')
               <> pretty expr <> hardline
 
     pretty (Assert assert cond semicolon expr)
@@ -494,7 +494,7 @@ instance Pretty Expression where
                 line <> pretty (moveTrailingCommentUp op') <> nest 2 (absorbOperation expr)
 
             -- Extract comment before the first operand and move it out, to prevent force-expanding the expression
-            (operationWithoutComment, comment') = mapFirstToken' (\(Ann leading token trailing) -> (Ann [] token trailing, leading)) operation
+            (operationWithoutComment, comment') = mapFirstToken' (\(Ann leading token trailing') -> (Ann [] token trailing', leading)) operation
           in
             pretty comment' <> (group $
                 (concat . map prettyOperation . (flatten Nothing)) operationWithoutComment)