Implement sorting of import statements

src/Language/Haskell/Brittany/Internal.hs

@@ -645,8 +645,7 @@ layoutBriDoc briDoc = do
 
   anns :: ExactPrint.Anns <- mAsk
 
-  let state = LayoutState
+  let state = LayoutState { _lstate_baseYs           = [0]
-        { _lstate_baseYs           = [0]
                           , _lstate_curYOrAddNewline = Right 0 -- important that we dont use left
                                              -- here because moveToAnn stuff
                                              -- of the first node needs to do
@@ -662,7 +661,16 @@ layoutBriDoc briDoc = do
   state' <- MultiRWSS.withMultiStateS state $ layoutBriDocM briDoc'
 
   let remainingComments =
-        extractAllComments =<< Map.elems (_lstate_comments state')
+        [ c
+        | (ExactPrint.AnnKey _ con, elemAnns) <- Map.toList
+          (_lstate_comments state')
+          -- With the new import layouter, we manually process comments
+          -- without relying on the backend to consume the comments out of
+          -- the state/map. So they will end up here, and we need to ignore
+          -- them.
+        , ExactPrint.unConName con /= "ImportDecl"
+        , c <- extractAllComments elemAnns
+        ]
   remainingComments
     `forM_` (fst .> show .> ErrorUnusedComment .> (: []) .> mTell)
 

src/Language/Haskell/Brittany/Internal/Layouters/IE.hs

@@ -2,6 +2,7 @@ module Language.Haskell.Brittany.Internal.Layouters.IE
   ( layoutIE
   , layoutLLIEs
   , layoutAnnAndSepLLIEs
+  , SortItemsFlag(..)
   )
 where
 
@@ -17,6 +18,7 @@ import           GHC                            ( unLoc
                                                 , AnnKeywordId(..)
                                                 , Located
                                                 , runGhc
+                                                , ModuleName
                                                 )
 import           HsSyn
 import           HsImpExp
@@ -126,6 +128,7 @@ layoutIE lie@(L _ ie) = docWrapNode lie $ case ie of
       | otherwise  -> name
 #endif
 
+data SortItemsFlag = ShouldSortItems | KeepItemsUnsorted
 -- Helper function to deal with Located lists of LIEs.
 -- In particular this will also associate documentation
 -- from the located list that actually belongs to the last IE.
@@ -134,8 +137,8 @@ layoutIE lie@(L _ ie) = docWrapNode lie $ case ie of
 -- handling of the resulting list. Adding parens is
 -- left to the caller since that is context sensitive
 layoutAnnAndSepLLIEs
-  :: Located [LIE GhcPs] -> ToBriDocM [ToBriDocM BriDocNumbered]
+  :: SortItemsFlag -> Located [LIE GhcPs] -> ToBriDocM [ToBriDocM BriDocNumbered]
-layoutAnnAndSepLLIEs llies@(L _ lies) = do
+layoutAnnAndSepLLIEs shouldSort llies@(L _ lies) = do
   let makeIENode ie = docSeq [docCommaSep, ie]
   let sortedLies =
         [ items
@@ -143,7 +146,9 @@ layoutAnnAndSepLLIEs llies@(L _ lies) = do
           $ List.sortOn lieToText lies
         , items <- mergeGroup group
         ]
-  let ieDocs = layoutIE <$> sortedLies
+  let ieDocs = fmap layoutIE $ case shouldSort of
+        ShouldSortItems -> sortedLies
+        KeepItemsUnsorted -> lies
   ieCommaDocs <-
     docWrapNodeRest llies $ sequence $ case splitFirstLast ieDocs of
       FirstLastEmpty        -> []
@@ -157,6 +162,7 @@ layoutAnnAndSepLLIEs llies@(L _ lies) = do
   mergeGroup items@[_] = items
   mergeGroup items     = if
     | all isProperIEThing items -> [List.foldl1' thingFolder items]
+    | all isIEVar items         -> [List.foldl1' thingFolder items]
     | otherwise                 -> items
   -- proper means that if it is a ThingWith, it does not contain a wildcard
   -- (because I don't know what a wildcard means if it is not already a
@@ -167,7 +173,12 @@ layoutAnnAndSepLLIEs llies@(L _ lies) = do
     L _ (IEThingAll _ _wn) -> True
     L _ (IEThingWith _ _wn NoIEWildcard _ _) -> True
     _ -> False
+  isIEVar :: LIE GhcPs -> Bool
+  isIEVar = \case
+    L _ IEVar{} -> True
+    _           -> False
   thingFolder :: LIE GhcPs -> LIE GhcPs -> LIE GhcPs
+  thingFolder l1@(L _ IEVar{}     ) _                     = l1
   thingFolder l1@(L _ IEThingAll{}) _                     = l1
   thingFolder _                     l2@(L _ IEThingAll{}) = l2
   thingFolder l1                    (   L _ IEThingAbs{}) = l1
@@ -198,9 +209,9 @@ layoutAnnAndSepLLIEs llies@(L _ lies) = do
 -- () -- no comments
 -- ( -- a comment
 -- )
-layoutLLIEs :: Bool -> Located [LIE GhcPs] -> ToBriDocM BriDocNumbered
+layoutLLIEs :: Bool -> SortItemsFlag -> Located [LIE GhcPs] -> ToBriDocM BriDocNumbered
-layoutLLIEs enableSingleline llies = do
+layoutLLIEs enableSingleline shouldSort llies = do
-  ieDs        <- layoutAnnAndSepLLIEs llies
+  ieDs        <- layoutAnnAndSepLLIEs shouldSort llies
   hasComments <- hasAnyCommentsBelow llies
   runFilteredAlternative $ case ieDs of
     [] -> do
@@ -240,8 +251,11 @@ lieToText = \case
   -- TODO: These _may_ appear in exports!
   -- Need to check, and either put them at the top (for module) or do some
   -- other clever thing.
-  L _ (IEModuleContents _ _  ) -> Text.pack "IEModuleContents"
+  L _ (IEModuleContents _ n) -> moduleNameToText n
-  L _ (IEGroup _ _ _         ) -> Text.pack "IEGroup"
+  L _ (IEGroup _ _ _         ) -> Text.pack "@IEGroup"
-  L _ (IEDoc      _ _        ) -> Text.pack "IEDoc"
+  L _ (IEDoc      _ _        ) -> Text.pack "@IEDoc"
-  L _ (IEDocNamed _ _        ) -> Text.pack "IEDocNamed"
+  L _ (IEDocNamed _ _        ) -> Text.pack "@IEDocNamed"
-  L _ (XIE _                 ) -> Text.pack "XIE"
+  L _ (XIE _                 ) -> Text.pack "@XIE"
+ where
+  moduleNameToText :: Located ModuleName -> Text
+  moduleNameToText (L _ name) = Text.pack ("@IEModuleContents" ++ moduleNameString name)

src/Language/Haskell/Brittany/Internal/Layouters/Import.hs

@@ -17,6 +17,7 @@ import           Name
 import           FieldLabel
 import qualified FastString
 import           BasicTypes
+import qualified Language.Haskell.GHC.ExactPrint.Types as ExactPrint.Types
 
 import           Language.Haskell.Brittany.Internal.Utils
 
@@ -41,8 +42,8 @@ prepModName :: e -> e
 prepModName = id
 #endif
 
-layoutImport :: ToBriDoc ImportDecl
+layoutImport :: ImportDecl GhcPs -> ToBriDocM BriDocNumbered
-layoutImport limportD@(L _ importD) = docWrapNode limportD $ case importD of
+layoutImport importD = case importD of
 #if MIN_VERSION_ghc(8,6,0)
   ImportDecl _ _ (L _ modName) pkg src safe q False mas mllies -> do
 #else
@@ -92,14 +93,14 @@ layoutImport limportD@(L _ importD) = docWrapNode limportD $ case importD of
           hasComments <- hasAnyCommentsBelow llies
           if compact
           then docAlt
-            [ docSeq [hidDoc, docForceSingleline $ layoutLLIEs True llies]
+            [ docSeq [hidDoc, docForceSingleline $ layoutLLIEs True ShouldSortItems llies]
             , let makeParIfHiding = if hiding
                     then docAddBaseY BrIndentRegular . docPar hidDoc
                     else id
-              in makeParIfHiding (layoutLLIEs True llies)
+              in makeParIfHiding (layoutLLIEs True ShouldSortItems llies)
             ]
           else do
-            ieDs <- layoutAnnAndSepLLIEs llies
+            ieDs <- layoutAnnAndSepLLIEs ShouldSortItems llies
             docWrapNodeRest llies
               $ docEnsureIndent (BrIndentSpecial hidDocCol)
               $ case ieDs of

src/Language/Haskell/Brittany/Internal/Layouters/Module.hs

@@ -1,3 +1,5 @@
+{-# LANGUAGE ViewPatterns #-}
+
 module Language.Haskell.Brittany.Internal.Layouters.Module (layoutModule) where
 
 #include "prelude.inc"
@@ -25,8 +27,16 @@ import           Language.Haskell.Brittany.Internal.Utils
 layoutModule :: ToBriDoc HsModule
 layoutModule lmod@(L _ mod') = case mod' of
     -- Implicit module Main
-  HsModule Nothing  _   imports _ _ _ -> docLines $ map layoutImport imports
+  HsModule Nothing  _   imports _ _ _ -> do
+    commentedImports <- transformToCommentedImport imports
+    -- groupify commentedImports `forM_` tellDebugMessShow
+    docLines (commentedImportsToDoc <$> sortCommentedImports commentedImports)
+    -- sortedImports <- sortImports imports
+    -- docLines $ [layoutImport y i | (y, i) <- sortedImports]
   HsModule (Just n) les imports _ _ _ -> do
+    commentedImports <- transformToCommentedImport imports
+    -- groupify commentedImports `forM_` tellDebugMessShow
+    -- sortedImports <- sortImports imports
     let tn = Text.pack $ moduleNameString $ unLoc n
     allowSingleLineExportList <- mAsk
       <&> _conf_layout
@@ -48,7 +58,7 @@ layoutModule lmod@(L _ mod') = case mod' of
                 , appSep $ docLit tn
                 , docWrapNode lmod $ appSep $ case les of
                   Nothing -> docEmpty
-                  Just x  -> layoutLLIEs True x
+                  Just x  -> layoutLLIEs True KeepItemsUnsorted x
                 , docSeparator
                 , docLit $ Text.pack "where"
                 ]
@@ -57,13 +67,140 @@ layoutModule lmod@(L _ mod') = case mod' of
               [ docAddBaseY BrIndentRegular $ docPar
                 (docSeq [appSep $ docLit $ Text.pack "module", docLit tn]
                 )
-                (docSeq [ docWrapNode lmod $ case les of
+                (docSeq [
+                          docWrapNode lmod $ case les of
                            Nothing -> docEmpty
-                            Just x  -> layoutLLIEs False x
+                           Just x  -> layoutLLIEs False KeepItemsUnsorted x
                         , docSeparator
                         , docLit $ Text.pack "where"
                         ]
                 )
               ]
           ]
-      : map layoutImport imports
+      : (commentedImportsToDoc <$> sortCommentedImports commentedImports) -- [layoutImport y i | (y, i) <- sortedImports]
+
+data CommentedImport
+  = EmptyLine
+  | IndependentComment (Comment, DeltaPos)
+  | ImportStatement ImportStatementRecord
+
+instance Show CommentedImport where
+  show = \case
+    EmptyLine            -> "EmptyLine"
+    IndependentComment _ -> "IndependentComment"
+    ImportStatement r ->
+      "ImportStatement " ++ show (length $ commentsBefore r) ++ " " ++ show
+        (length $ commentsAfter r)
+
+data ImportStatementRecord = ImportStatementRecord
+  { commentsBefore :: [(Comment, DeltaPos)]
+  , commentsAfter :: [(Comment, DeltaPos)]
+  , importStatement :: ImportDecl HsSyn.GhcPs
+  }
+
+instance Show ImportStatementRecord where
+  show r = "ImportStatement " ++ show (length $ commentsBefore r) ++ " " ++ show
+        (length $ commentsAfter r)
+
+transformToCommentedImport
+  :: [LImportDecl HsSyn.GhcPs] -> ToBriDocM [CommentedImport]
+transformToCommentedImport is = do
+  nodeWithAnnotations <- is `forM` \i@(L _ rawImport) -> do
+    annotionMay <- astAnn i
+    pure (annotionMay, rawImport)
+  let
+    convertComment (c, DP (y, x)) =
+      replicate (y - 1) EmptyLine ++ [IndependentComment (c, DP (1, x))]
+    accumF
+      :: [(Comment, DeltaPos)]
+      -> (Maybe Annotation, ImportDecl HsSyn.GhcPs)
+      -> ([(Comment, DeltaPos)], [CommentedImport])
+    accumF accConnectedComm (annMay, decl) = case annMay of
+      Nothing ->
+        ( []
+        , [ ImportStatement ImportStatementRecord { commentsBefore  = []
+                                                  , commentsAfter   = []
+                                                  , importStatement = decl
+                                                  }
+          ]
+        )
+      Just ann ->
+        let
+          blanksBeforeImportDecl = deltaRow (annEntryDelta ann) - 1
+          (newAccumulator, priorComments') =
+            List.span ((== 0) . deltaRow . snd) (annPriorComments ann)
+          go
+            :: [(Comment, DeltaPos)]
+            -> [(Comment, DeltaPos)]
+            -> ([CommentedImport], [(Comment, DeltaPos)], Int)
+          go acc []                       = ([], acc, 0)
+          go acc [c1@(_, DP (y, _))] = ([], c1 : acc, y - 1)
+          go acc (c1@(_, DP (1, _)) : xs) = go (c1 : acc) xs
+          go acc ((c1, DP (y, x)) : xs) =
+            ( (convertComment =<< xs) ++ replicate (y - 1) EmptyLine
+            , (c1, DP (1, x)) : acc
+            , 0
+            )
+          (convertedIndependentComments, beforeComments, initialBlanks) =
+            if blanksBeforeImportDecl /= 0
+              then (convertComment =<< priorComments', [], 0)
+              else go [] (reverse priorComments')
+        in
+          ( newAccumulator
+          , convertedIndependentComments
+          ++ replicate (blanksBeforeImportDecl + initialBlanks) EmptyLine
+          ++ [ ImportStatement ImportStatementRecord
+                 { commentsBefore  = beforeComments
+                 , commentsAfter   = accConnectedComm
+                 , importStatement = decl
+                 }
+             ]
+          )
+  let (finalAcc, finalList) = mapAccumR accumF [] nodeWithAnnotations
+  pure $ join $ (convertComment =<< finalAcc) : finalList
+
+sortCommentedImports :: [CommentedImport] -> [CommentedImport]
+sortCommentedImports =
+  unpackImports . mergeGroups . map (fmap (sortGroups)) . groupify
+ where
+  unpackImports :: [CommentedImport] -> [CommentedImport]
+  unpackImports xs = xs >>= \case
+    l@EmptyLine            -> [l]
+    l@IndependentComment{} -> [l]
+    ImportStatement r ->
+      map IndependentComment (commentsBefore r) ++ [ImportStatement r]
+  mergeGroups
+    :: [Either CommentedImport [ImportStatementRecord]] -> [CommentedImport]
+  mergeGroups xs = xs >>= \case
+    Left  x -> [x]
+    Right y -> ImportStatement <$> y
+  sortGroups :: [ImportStatementRecord] -> [ImportStatementRecord]
+  sortGroups =
+    List.sortOn (moduleNameString . unLoc . ideclName . importStatement)
+  groupify
+    :: [CommentedImport] -> [Either CommentedImport [ImportStatementRecord]]
+  groupify cs = go [] cs
+   where
+    go [] = \case
+      (l@EmptyLine            : rest) -> Left l : go [] rest
+      (l@IndependentComment{} : rest) -> Left l : go [] rest
+      (ImportStatement r      : rest) -> go [r] rest
+      []                              -> []
+    go acc = \case
+      (l@EmptyLine : rest) -> Right (reverse acc) : Left l : go [] rest
+      (l@IndependentComment{} : rest) ->
+        Left l : Right (reverse acc) : go [] rest
+      (ImportStatement r : rest) -> go (r : acc) rest
+      []                         -> [Right (reverse acc)]
+
+commentedImportsToDoc :: CommentedImport -> ToBriDocM BriDocNumbered
+commentedImportsToDoc = \case
+  EmptyLine -> docLitS ""
+  IndependentComment c -> commentToDoc c
+  ImportStatement r ->
+    docSeq
+      ( layoutImport (importStatement r)
+      : map commentToDoc (commentsAfter r)
+      )
+ where
+  commentToDoc (c, DP (_y, x)) = docLitS (replicate x ' ' ++ commentContents c)