Some more performance changes

nixfmt.cabal

@@ -1,55 +1,53 @@
-cabal-version:       2.0
-name:                nixfmt
-version:             0.6.0
-synopsis:            Official formatter for Nix code 
+cabal-version:      3.0
+name:               nixfmt
+version:            0.6.0
+synopsis:           Official formatter for Nix code 
 description:
   A formatter for Nix that ensures consistent and clear formatting by forgetting
   almost all existing formatting during parsing.
-homepage:            https://github.com/NixOS/nixfmt
-bug-reports:         https://github.com/NixOS/nixfmt/issues
-license:             MPL-2.0
-license-file:        LICENSE
-author:              Serokell <[email protected]> and nixfmt contributors
-copyright:           Serokell <[email protected]> and nixfmt contributors
-category:            Development
-build-type:          Simple
-extra-source-files:  README.md, CHANGELOG.md
+
+homepage:           https://github.com/NixOS/nixfmt
+bug-reports:        https://github.com/NixOS/nixfmt/issues
+license:            MPL-2.0
+license-file:       LICENSE
+author:             Serokell <[email protected]> and nixfmt contributors
+copyright:          Serokell <[email protected]> and nixfmt contributors
+category:           Development
+build-type:         Simple
+extra-source-files:
+  CHANGELOG.md
+  README.md
 
 source-repository head
   type:     git
   location: git://github.com/NixOS/nixfmt.git
 
 executable nixfmt
-  main-is:             Main.hs
+  main-is:          Main.hs
   other-modules:
     Paths_nixfmt
-    System.IO.Utf8
     System.IO.Atomic
-  autogen-modules:
-    Paths_nixfmt
-  other-extensions:    DeriveDataTypeable
-  hs-source-dirs:      main
+    System.IO.Utf8
+
+  autogen-modules:  Paths_nixfmt
+  hs-source-dirs:   main
   build-depends:
-      base             >= 4.12.0 && < 4.21
+    , base             >=4.12.0  && <4.21
     , bytestring
-    , cmdargs          >= 0.10.20 && < 0.11
+    , cmdargs          >=0.10.20 && <0.11
+    , directory        >=1.3.3   && <1.4
     , file-embed
+    , filepath         >=1.4.2   && <1.5
     , nixfmt
-    , unix             >= 2.7.2 && < 2.9
-    , text             >= 1.2.3 && < 2.2
+    , safe-exceptions  >=0.1.7   && <0.2
+    , text             >=1.2.3   && <2.2
+    , unix             >=2.7.2   && <2.9
 
-    -- for System.IO.Atomic
-    , directory        >= 1.3.3 && < 1.4
-    , filepath         >= 1.4.2 && < 1.5
-    , safe-exceptions  >= 0.1.7 && < 0.2
-  default-language:    Haskell2010
+  -- for System.IO.Atomic
+  default-language: Haskell2010
   ghc-options:
-    -Wall
-    -Wcompat
-    -Wincomplete-record-updates
-    -Wincomplete-uni-patterns
-    -Wredundant-constraints
-    -threaded
+    -Wall -Wcompat -Wincomplete-record-updates
+    -Wincomplete-uni-patterns -Wredundant-constraints -threaded
 
 library
   exposed-modules:
@@ -63,33 +61,42 @@ library
     Nixfmt.Util
 
   default-extensions:
-    PackageImports
-
-  other-extensions:
+    BangPatterns
+    BlockArguments
+    DataKinds
+    DeriveAnyClass
     DeriveFoldable
     DeriveFunctor
+    DeriveGeneric
+    DerivingVia
     FlexibleInstances
     LambdaCase
+    NamedFieldPuns
+    OverloadedLists
     OverloadedStrings
+    PackageImports
+    PatternSynonyms
+    RankNTypes
     ScopedTypeVariables
-    StandaloneDeriving
     TupleSections
+    TypeApplications
+    TypeOperators
 
-  hs-source-dirs:      src
+  hs-source-dirs:     src
   build-depends:
-      base             >= 4.12.0 && < 4.21
-    , megaparsec       >= 9.0.1 && < 9.6
+    , base                >=4.12.0 && <4.21
+    , containers
+    , deepseq
+    , megaparsec          >=9.0.1  && <9.6
     , mtl
-    , parser-combinators >= 1.0.3 && < 1.4
-    , scientific       >= 0.3.0 && < 0.4.0
-    , text             >= 1.2.3 && < 2.2
-    , transformers
+    , parser-combinators  >=1.0.3  && <1.4
     , pretty-simple
-  default-language:    Haskell2010
+    , scientific          >=0.3.0  && <0.4.0
+    , text                >=1.2.3  && <2.2
+    , transformers
+    , vector
+
+  default-language:   Haskell2010
   ghc-options:
-    -Wall
-    -Wcompat
-    -Wincomplete-record-updates
-    -Wincomplete-uni-patterns
-    -Wredundant-constraints
-    -Wno-orphans
+    -Wall -Wcompat -Wincomplete-record-updates
+    -Wincomplete-uni-patterns -Wredundant-constraints -Wno-orphans

src/Nixfmt.hs

@@ -1,5 +1,3 @@
-{-# LANGUAGE RankNTypes #-}
-
 module Nixfmt (
   errorBundlePretty,
   ParseErrorBundle,
@@ -28,8 +26,11 @@ type Layouter = forall a. (Pretty a, LanguageElement a) => a -> Text
 -- of @w@ columns where possible.
 format :: Layouter -> FilePath -> Text -> Either String Text
 format layout filename =
-  bimap errorBundlePretty layout
+  bimap errorBundlePretty f
     . Megaparsec.parse Parser.file filename
+  where
+    -- f !x = layout $ maybe () (error . show) (unsafeNoThunks x) `seq` x
+    f !x = layout x
 
 -- | Pretty print the internal AST for debugging
 printAst :: FilePath -> Text -> Either String a

src/Nixfmt/Lexer.hs

@@ -1,14 +1,12 @@
-{-# LANGUAGE BlockArguments #-}
-{-# LANGUAGE FlexibleContexts #-}
-{-# LANGUAGE LambdaCase #-}
-{-# LANGUAGE OverloadedStrings #-}
-
 module Nixfmt.Lexer (lexeme, pushTrivia, takeTrivia, whole) where
 
-import Control.Monad.State.Strict (MonadState, evalStateT, get, modify, put)
+import Control.DeepSeq (NFData, force)
+import Control.Monad.State.Strict (StateT, evalStateT, get, modify, put)
 import Data.Char (isSpace)
-import Data.List (dropWhileEnd)
+import Data.List (dropWhileEnd, singleton)
 import Data.Maybe (fromMaybe)
+import Data.Sequence (Seq)
+import qualified Data.Sequence as Seq
 import Data.Text as Text (
   Text,
   isPrefixOf,
@@ -26,6 +24,7 @@ import Data.Text as Text (
   unwords,
  )
 import Data.Void (Void)
+import GHC.Generics (Generic)
 import Nixfmt.Types (
   Ann (..),
   Parser,
@@ -54,12 +53,13 @@ import Text.Megaparsec (
 import Text.Megaparsec.Char (char, eol)
 
 data ParseTrivium
-  = PTNewlines Int
+  = PTNewlines {-# UNPACK #-} !Int
   | -- Track the column where the comment starts
-    PTLineComment Text Pos
+    PTLineComment {-# UNPACK #-} !Text !Pos
   | -- Track whether it is a doc comment
-    PTBlockComment Bool [Text]
-  deriving (Show)
+    PTBlockComment !Bool ![Text]
+  deriving stock (Show, Generic)
+  deriving anyclass (NFData)
 
 preLexeme :: Parser a -> Parser a
 preLexeme p = p <* manyP (\x -> isSpace x && x /= '\n' && x /= '\r')
@@ -128,8 +128,8 @@ blockComment = try $ preLexeme $ do
     commonIndentationLength = foldr (min . Text.length . Text.takeWhile (== ' '))
 
 -- This should be called with zero or one elements, as per `span isTrailing`
-convertTrailing :: [ParseTrivium] -> Maybe TrailingComment
-convertTrailing = toMaybe . join . map toText
+convertTrailing :: Seq ParseTrivium -> Maybe TrailingComment
+convertTrailing = toMaybe . join . foldMap (singleton . toText)
   where
     toText (PTLineComment c _) = strip c
     toText (PTBlockComment False [c]) = strip c
@@ -138,9 +138,9 @@ convertTrailing = toMaybe . join . map toText
     toMaybe "" = Nothing
     toMaybe c = Just $ TrailingComment c
 
-convertLeading :: [ParseTrivium] -> Trivia
+convertLeading :: Seq ParseTrivium -> Trivia
 convertLeading =
-  concatMap
+  foldMap
     ( \case
         PTNewlines 1 -> []
         PTNewlines _ -> [EmptyLine]
@@ -156,32 +156,34 @@ isTrailing (PTBlockComment False []) = True
 isTrailing (PTBlockComment False [_]) = True
 isTrailing _ = False
 
-convertTrivia :: [ParseTrivium] -> Pos -> (Maybe TrailingComment, Trivia)
+convertTrivia :: Seq ParseTrivium -> Pos -> (Maybe TrailingComment, Trivia)
 convertTrivia pts nextCol =
-  let (trailing, leading) = span isTrailing pts
+  let (trailing, leading) = Seq.spanl isTrailing pts
   in case (trailing, leading) of
       -- Special case: if the trailing comment visually forms a block with the start of the following line,
       -- then treat it like part of those comments instead of a distinct trailing comment.
       -- This happens especially often after `{` or `[` tokens, where the comment of the first item
       -- starts on the same line ase the opening token.
-      ([PTLineComment _ pos], (PTNewlines 1) : (PTLineComment _ pos') : _) | pos == pos' -> (Nothing, convertLeading pts)
+      ([PTLineComment _ pos], (PTNewlines 1) Seq.:<| (PTLineComment _ pos') Seq.:<| _) | pos == pos' -> (Nothing, convertLeading pts)
       ([PTLineComment _ pos], [PTNewlines 1]) | pos == nextCol -> (Nothing, convertLeading pts)
       _ -> (convertTrailing trailing, convertLeading leading)
 
-trivia :: Parser [ParseTrivium]
-trivia = many $ hidden $ lineComment <|> blockComment <|> newlines
+trivia :: Parser (Seq ParseTrivium)
+trivia =
+  Seq.fromList <$> do
+    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 []
+takeTrivia :: (Monad m) => StateT Trivia m Trivia
+takeTrivia = get <* put Seq.empty
 
-pushTrivia :: (MonadState Trivia m) => Trivia -> m ()
+pushTrivia :: (Monad m) => Trivia -> StateT Trivia m ()
 pushTrivia t = modify (<> t)
 {-# INLINEABLE pushTrivia #-}
 
-lexeme :: Parser a -> Parser (Ann a)
+lexeme :: (NFData a) => Parser a -> Parser (Ann a)
 lexeme p = do
   lastLeading <- takeTrivia
   SourcePos{Text.Megaparsec.sourceLine = line} <- getSourcePos
@@ -193,17 +195,18 @@ lexeme p = do
   pushTrivia nextLeading
   return $
     Ann
-      { preTrivia = lastLeading,
-        value = token,
-        Nixfmt.Types.sourceLine = line,
-        trailComment = trailing
+      { preTrivia = force lastLeading,
+        value = force token,
+        Nixfmt.Types.sourceLine = force line,
+        trailComment = force 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
 -- does not interact with the trivia state of its surroundings.
-whole :: Parser a -> Parsec Void Text (Whole a)
-whole pa = flip evalStateT [] do
+whole :: (NFData a) => Parser a -> Parsec Void Text (Whole a)
+whole pa = flip evalStateT Seq.empty do
   preLexeme $ pure ()
-  pushTrivia . convertLeading =<< trivia
-  Whole <$> pa <*> takeTrivia
+  trivia
+    >>= pushTrivia . convertLeading
+  Whole <$> (force <$> pa) <*> (force <$> takeTrivia)

src/Nixfmt/Parser.hs

@@ -1,9 +1,12 @@
 {-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE OverloadedLists #-}
 {-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE TypeApplications #-}
 
 module Nixfmt.Parser where
 
 import Control.Applicative (liftA2)
+import Control.DeepSeq (NFData, force)
 import Control.Monad (guard, liftM2)
 import Control.Monad.Combinators (sepBy)
 import qualified Control.Monad.Combinators.Expr as MPExpr (
@@ -78,7 +81,7 @@ import Prelude hiding (String)
 
 -- HELPER FUNCTIONS
 
-ann :: (a -> b) -> Parser a -> Parser (Ann b)
+ann :: (NFData b) => (a -> b) -> Parser a -> Parser (Ann b)
 ann f p = try $ lexeme $ f <$> p
 
 -- | parses a token without parsing trivia after it
@@ -370,8 +373,8 @@ term = label "term" $ do
     [] -> t
     _ -> Selection t sel def
 
-items :: Parser a -> Parser (Items a)
-items p = Items <$> many (item p) <> (toList <$> optional itemComment)
+items :: (NFData a) => Parser a -> Parser (Items a)
+items p = Items . force <$> many (item p) <> (toList <$> optional itemComment)
 
 item :: Parser a -> Parser (Item a)
 item p = itemComment <|> Item <$> p
@@ -430,23 +433,23 @@ inherit =
     <*> symbol TSemicolon
 
 assignment :: Parser Binder
-assignment =
-  Assignment
-    <$> selectorPath
-    <*> symbol TAssign
-    <*> expression
-    <*> symbol TSemicolon
+assignment = do
+  lhs <- selectorPath
+  assign <- symbol TAssign
+  expr <- expression
+  semicolon <- symbol TSemicolon
+  pure $! Assignment (force lhs) assign (force expr) semicolon
 
 binders :: Parser (Items Binder)
 binders = items (assignment <|> inherit)
 
 set :: Parser Term
-set =
-  Set
-    <$> optional (reserved KRec <|> reserved KLet)
-    <*> symbol TBraceOpen
-    <*> binders
-    <*> symbol TBraceClose
+set = do
+  mleaf <- optional (reserved KRec <|> reserved KLet)
+  lbrace <- symbol TBraceOpen
+  els <- binders
+  rbrace <- symbol TBraceClose
+  pure $! Set mleaf lbrace (force els) rbrace
 
 list :: Parser Term
 list = List <$> symbol TBrackOpen <*> items term <*> symbol TBrackClose
@@ -492,7 +495,7 @@ opCombiner (Op InfixR tok) = MPExpr.InfixR $ flip Operation <$> operator tok
 operation :: Parser Expression
 operation =
   MPExpr.makeExprParser
-    (Term <$> term <* notFollowedBy (oneOf (":@" :: [Char])))
+    (Term <$> term <* notFollowedBy (oneOf @[] ":@"))
     (map (map opCombiner) operators)
 
 -- EXPRESSIONS