WIP Attach leading trivia to next token

nixfmt.cabal

@@ -88,9 +88,11 @@ library
   build-depends:
       base             >= 4.12.0 && < 4.17
     , megaparsec       >= 9.0.1 && < 9.3
+    , mtl
     , parser-combinators >= 1.0.3 && < 1.4
     , scientific       >= 0.3.0 && < 0.4.0
     , text             >= 1.2.3 && < 1.3
+    , transformers
   default-language:    Haskell2010
   ghc-options:
     -Wall

src/Nixfmt/Lexer.hs

@@ -4,22 +4,25 @@
  - SPDX-License-Identifier: MPL-2.0
  -}
 
-{-# LANGUAGE LambdaCase, OverloadedStrings #-}
+{-# LANGUAGE BlockArguments, FlexibleContexts, LambdaCase, OverloadedStrings #-}
 
-module Nixfmt.Lexer (lexeme) where
+module Nixfmt.Lexer (lexeme, whole) where
 
+import Control.Monad.State (MonadState, evalStateT, get, modify, put)
 import Data.Char (isSpace)
 import Data.List (dropWhileEnd)
 import Data.Maybe (fromMaybe)
 import Data.Text as Text
-  (Text, length, lines, null, pack, replace, replicate, strip,
-  stripEnd, stripPrefix, stripStart, takeWhile, unwords)
+  (Text, length, lines, null, pack, replace, replicate, strip, stripEnd,
+  stripPrefix, stripStart, takeWhile, unwords)
+import Data.Void (Void)
 import Text.Megaparsec
-  (SourcePos(..), anySingle, chunk, getSourcePos, hidden, many, manyTill, some,
-  try, unPos, (<|>))
+  (Parsec, SourcePos(..), anySingle, chunk, getSourcePos, hidden, many,
+  manyTill, some, try, unPos, (<|>))
 import Text.Megaparsec.Char (eol)
 
-import Nixfmt.Types (Ann(..), Parser, TrailingComment(..), Trivia, Trivium(..))
+import Nixfmt.Types
+  (Ann(..), Whole(..), Parser, TrailingComment(..), Trivia, Trivium(..))
 import Nixfmt.Util (manyP)
 
 data ParseTrivium
@@ -92,8 +95,27 @@ convertTrivia pts =
 trivia :: Parser [ParseTrivium]
 trivia = many $ hidden $ lineComment <|> blockComment <|> newlines
 
+-- The following primitives to interact with the state monad that stores trivia
+-- are designed to prevent trivia from being dropped or duplicated by accident.
+
+takeTrivia :: MonadState Trivia m => m Trivia
+takeTrivia = get <* put []
+
+pushTrivia :: MonadState Trivia m => Trivia -> m ()
+pushTrivia t = modify (<>t)
+
 lexeme :: Parser a -> Parser (Ann a)
 lexeme p = do
+    lastLeading <- takeTrivia
     token <- preLexeme p
-    (trailing, leading) <- convertTrivia <$> trivia
-    return $ Ann token trailing leading
+    (trailing, nextLeading) <- convertTrivia <$> trivia
+    pushTrivia nextLeading
+    return $ Ann token trailing lastLeading
+
+-- | 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
+-- does not interact with the trivia state of its surroundings.
+whole :: Parser a -> Parsec Void Text (Whole a)
+whole pa = flip evalStateT [] do
+    pushTrivia . convertLeading =<< trivia
+    Whole <$> pa <*> takeTrivia

src/Nixfmt/Parser.hs

@@ -4,7 +4,7 @@
  - SPDX-License-Identifier: MPL-2.0
  -}
 
-{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE FlexibleInstances, OverloadedStrings #-}
 
 module Nixfmt.Parser where
 
@@ -14,23 +14,25 @@ import Control.Monad (guard, liftM2)
 import Control.Monad.Combinators (sepBy)
 import qualified Control.Monad.Combinators.Expr as MPExpr
   (Operator(..), makeExprParser)
+import Control.Monad.Trans.Class (lift)
 import Data.Char (isAlpha)
 import Data.Foldable (toList)
 import Data.Functor (($>))
 import Data.Maybe (fromMaybe, mapMaybe, maybeToList)
 import Data.Text as Text (Text, cons, empty, singleton, split, stripPrefix)
+import Data.Void (Void)
 import Text.Megaparsec
-  (anySingle, chunk, eof, label, lookAhead, many, notFollowedBy, oneOf,
+  (Parsec, anySingle, chunk, eof, label, lookAhead, many, notFollowedBy, oneOf,
   optional, satisfy, some, try, (<|>))
 import Text.Megaparsec.Char (char)
 import qualified Text.Megaparsec.Char.Lexer as L (decimal)
 
-import Nixfmt.Lexer (lexeme)
+import Nixfmt.Lexer (lexeme, whole)
+import Nixfmt.Parser.Float (floatParse)
 import Nixfmt.Types
-  (Ann, Binder(..), Expression(..), File(..), Fixity(..), Leaf, Operator(..),
+  (Ann, Binder(..), Expression(..), File, Fixity(..), Leaf, Operator(..),
   ParamAttr(..), Parameter(..), Parser, Path, Selector(..), SimpleSelector(..),
   String, StringPart(..), Term(..), Token(..), operators, tokenText)
-import Nixfmt.Parser.Float (floatParse)
 import Nixfmt.Util
   (commonIndentation, identChar, isSpaces, manyP, manyText, pathChar,
   schemeChar, someP, someText, uriChar)
@@ -79,6 +81,12 @@ identifier = ann Identifier $ do
 slash :: Parser Text
 slash = chunk "/" <* notFollowedBy (char '/')
 
+instance Semigroup a => Semigroup (Parser a) where
+  fx <> fy = do
+    x <- fx
+    y <- fy
+    pure $ x <> y
+
 envPath :: Parser (Ann Token)
 envPath = ann EnvPath $ char '<' *>
     someP pathChar <> manyText (slash <> someP pathChar)
@@ -102,7 +110,7 @@ uri = fmap (pure . pure . TextPart) $ try $
 
 interpolation :: Parser StringPart
 interpolation = Interpolation <$>
-    symbol TInterOpen <*> expression <*> rawSymbol TInterClose
+    (rawSymbol TInterOpen *> lift (whole expression) <* rawSymbol TInterClose)
 
 simpleStringPart :: Parser StringPart
 simpleStringPart = TextPart <$> someText (
@@ -342,5 +350,5 @@ expression :: Parser Expression
 expression = label "expression" $ try operation <|> abstraction <|>
     with <|> letIn <|> ifThenElse <|> assert
 
-file :: Parser File
-file = File <$> lexeme (return SOF) <*> expression <* eof
+file :: Parsec Void Text File
+file = whole (expression <* eof)

src/Nixfmt/Predoc.hs

@@ -32,7 +32,7 @@ module Nixfmt.Predoc
     ) where
 
 import Data.List (intersperse)
-import Data.Text as Text (Text, concat, length, pack, replicate)
+import Data.Text as Text (Text, concat, length, pack, replicate, strip)
 
 -- | Sequential Spacings are reduced to a single Spacing by taking the maximum.
 -- This means that e.g. a Space followed by an Emptyline results in just an
@@ -192,7 +192,7 @@ moveLinesIn (Node ann xs : ys) =
 moveLinesIn (x : xs) = x : moveLinesIn xs
 
 layout :: Pretty a => Int -> a -> Text
-layout w = layoutGreedy w . fixup . pretty
+layout w = (<>"\n") . Text.strip . layoutGreedy w . fixup . pretty
 
 -- 1. Move and merge Spacings.
 -- 2. Convert Softlines to Grouped Lines and Hardspaces to Texts.

src/Nixfmt/Pretty.hs

@@ -20,7 +20,7 @@ import Nixfmt.Predoc
   (Doc, Pretty, base, emptyline, group, hardline, hardspace, hcat, line, line',
   nest, newline, pretty, sepBy, softline, softline', text, textWidth)
 import Nixfmt.Types
-  (Ann(..), Binder(..), Expression(..), File(..), Leaf, ParamAttr(..),
+  (Ann(..), Binder(..), Expression(..), Whole(..), Leaf, ParamAttr(..),
   Parameter(..), Selector(..), SimpleSelector(..), StringPart(..), Term(..),
   Token(..), TrailingComment(..), Trivia, Trivium(..), tokenText)
 import Nixfmt.Util (commonIndentation, isSpaces, replaceMultiple)
@@ -53,7 +53,7 @@ instance Pretty [Trivium] where
 
 instance Pretty a => Pretty (Ann a) where
     pretty (Ann x trailing leading)
-        = pretty x <> pretty trailing <> pretty leading
+        = pretty leading <> pretty x <> pretty trailing
 
 instance Pretty SimpleSelector where
     pretty (IDSelector i)              = pretty i
@@ -74,7 +74,7 @@ instance Pretty Binder where
         = base $ group (pretty inherit <> softline
                  <> nest 2 (sepBy softline ids)) <> pretty semicolon
 
-    pretty (Inherit inherit source ids semicolon)
+    pretty (Inherit inherit (Just source) ids semicolon)
         = base $ group (pretty inherit <> hardspace
                  <> pretty source <> line
                  <> nest 2 (sepBy softline ids)) <> pretty semicolon
@@ -200,16 +200,11 @@ instance Pretty Expression where
           <> nest 2 (group expr0) <> pretty semicolon)
           <> absorbSet expr1
 
-    pretty (Let (Ann let_ letTrailing letLeading) binders
-                (Ann in_ inTrailing inLeading) expr)
+    pretty (Let let_ binders in_ expr)
         = base $ group letPart <> line <> group inPart
-        where letPart = pretty let_ <> pretty letTrailing <> line <> letBody
+        where letPart = pretty let_ <> line <> letBody
               inPart = pretty in_ <> hardspace <> pretty expr
-              letBody = nest 2 $
-                  pretty letLeading
-                  <> sepBy hardline binders
-                  <> pretty (toLeading inTrailing)
-                  <> pretty inLeading
+              letBody = nest 2 $ sepBy hardline binders
 
     pretty (Assert assert cond semicolon expr)
         = base (pretty assert <> hardspace
@@ -247,13 +242,9 @@ instance Pretty Expression where
     pretty (Inversion bang expr)
         = pretty bang <> pretty expr
 
-instance Pretty File where
-    pretty (File (Ann _ Nothing leading) expr)
-        = group $ hcat leading <> pretty expr <> hardline
-
-    pretty (File (Ann _ (Just (TrailingComment trailing)) leading) expr)
-        = group $ text "# " <> pretty trailing <> hardline
-                  <> hcat leading <> pretty expr <> hardline
+instance Pretty a => Pretty (Whole a) where
+    pretty (Whole x finalTrivia)
+        = group $ pretty x <> pretty finalTrivia
 
 instance Pretty Token where
     pretty = text . tokenText
@@ -331,22 +322,26 @@ isSimpleString parts
 
 instance Pretty StringPart where
     pretty (TextPart t) = text t
-    pretty (Interpolation paropen (Term t) parclose)
+    pretty (Interpolation (Whole (Term t) []))
         | isAbsorbable t
-            = group $ pretty paropen <> prettyTerm t <> pretty parclose
+            = group $ text "${" <> prettyTerm t <> text "}"
 
-    pretty (Interpolation paropen expr parclose)
+    pretty (Interpolation (Whole expr []))
         | isSimple expr
-            = pretty paropen <> pretty expr <> pretty parclose
-        | otherwise
-            = group $ pretty paropen <> line'
-                <> nest 2 (pretty expr) <> line'
-                <> pretty parclose
+            = text "${" <> pretty expr <> text "}"
+
+    pretty (Interpolation whole)
+        = group $ text "${" <> line'
+            <> nest 2 (pretty whole) <> line'
+            <> text "}"
 
 instance Pretty [StringPart] where
-    pretty [Interpolation paropen expr parclose]
-        = group $ pretty paropen <> pretty expr <> pretty parclose
+    pretty [Interpolation expr]
+        = group $ text "${" <> pretty expr <> text "}"
 
+    -- If we split a string line over multiple code lines due to large
+    -- interpolations, make sure to indent based on the indentation of the line
+    -- in the string.
     pretty (TextPart t : parts)
         = text t <> nest indentation (hcat parts)
         where indentation = textWidth $ Text.takeWhile isSpace t
@@ -370,7 +365,7 @@ prettyLine escapeText unescapeInterpol
           unescapeInterpols [] = []
           unescapeInterpols (TextPart t : TextPart u : xs)
               = unescapeInterpols (TextPart (t <> u) : xs)
-          unescapeInterpols (TextPart t : xs@(Interpolation _ _ _ : _))
+          unescapeInterpols (TextPart t : xs@(Interpolation{} : _))
               = TextPart (unescapeInterpol t) : unescapeInterpols xs
           unescapeInterpols (x : xs) = x : unescapeInterpols xs
 

src/Nixfmt/Types.hs

@@ -10,12 +10,13 @@ module Nixfmt.Types where
 
 import Prelude hiding (String)
 
+import Control.Monad.State (StateT)
 import Data.Text (Text, pack)
 import Data.Void (Void)
 import qualified Text.Megaparsec as MP (ParseErrorBundle, Parsec)
 
 -- | A @megaparsec@ @ParsecT@ specified for use with @nixfmt@.
-type Parser = MP.Parsec Void Text
+type Parser = StateT Trivia (MP.Parsec Void Text)
 
 -- | A @megaparsec@ @ParseErrorBundle@ specified for use with @nixfmt@.
 type ParseErrorBundle = MP.ParseErrorBundle Text Void
@@ -43,7 +44,7 @@ type Leaf = Ann Token
 
 data StringPart
     = TextPart Text
-    | Interpolation Leaf Expression Token
+    | Interpolation (Whole Expression)
     deriving (Eq, Show)
 
 type Path = Ann [StringPart]
@@ -101,10 +102,14 @@ data Expression
     | Inversion Leaf Expression
     deriving (Eq, Show)
 
-data File
-    = File Leaf Expression
+-- | A Whole a is an a including final trivia. It's assumed the a stores the
+-- initial trivia.
+data Whole a
+    = Whole a Trivia
     deriving (Eq, Show)
 
+type File = Whole Expression
+
 data Token
     = Integer    Int
     | Float      Double