{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}

module HIndent.Ast.Import.Collection
  ( ImportCollection
  , mkImportCollection
  , hasImports
  ) where

import Control.Monad.RWS
import Data.Function
import Data.List
import qualified GHC.Hs as GHC
import GHC.Stack
import qualified GHC.Types.SrcLoc as GHC
import HIndent.Ast.Import
import HIndent.Ast.NodeComments
import HIndent.Ast.WithComments
import HIndent.Config
import qualified HIndent.GhcLibParserWrapper.GHC.Hs as GHC
import HIndent.Pretty
import HIndent.Pretty.Combinators
import HIndent.Pretty.NodeComments
import HIndent.Printer

newtype ImportCollection =
  ImportCollection [[WithComments Import]]

instance CommentExtraction ImportCollection where
  nodeComments :: ImportCollection -> NodeComments
nodeComments ImportCollection {} = [LEpaComment] -> [LEpaComment] -> [LEpaComment] -> NodeComments
NodeComments [] [] []

instance Pretty ImportCollection where
  pretty' :: ImportCollection -> Printer ()
pretty' (ImportCollection [[WithComments Import]]
xs) =
    Printer [[WithComments Import]]
importDecls Printer [[WithComments Import]]
-> ([[WithComments Import]] -> Printer ()) -> Printer ()
forall a b. Printer a -> (a -> Printer b) -> Printer b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= [Printer ()] -> Printer ()
blanklined ([Printer ()] -> Printer ())
-> ([[WithComments Import]] -> [Printer ()])
-> [[WithComments Import]]
-> Printer ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([WithComments Import] -> Printer ())
-> [[WithComments Import]] -> [Printer ()]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [WithComments Import] -> Printer ()
outputImportGroup
    where
      outputImportGroup :: [WithComments Import] -> Printer ()
outputImportGroup = [Printer ()] -> Printer ()
lined ([Printer ()] -> Printer ())
-> ([WithComments Import] -> [Printer ()])
-> [WithComments Import]
-> Printer ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (WithComments Import -> Printer ())
-> [WithComments Import] -> [Printer ()]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap WithComments Import -> Printer ()
forall a. Pretty a => a -> Printer ()
pretty
      importDecls :: Printer [[WithComments Import]]
importDecls =
        (PrintState -> Bool) -> Printer Bool
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (Config -> Bool
configSortImports (Config -> Bool) -> (PrintState -> Config) -> PrintState -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrintState -> Config
psConfig) Printer Bool
-> (Bool -> Printer [[WithComments Import]])
-> Printer [[WithComments Import]]
forall a b. Printer a -> (a -> Printer b) -> Printer b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          Bool
True -> [[WithComments Import]] -> Printer [[WithComments Import]]
forall a. a -> Printer a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([[WithComments Import]] -> Printer [[WithComments Import]])
-> [[WithComments Import]] -> Printer [[WithComments Import]]
forall a b. (a -> b) -> a -> b
$ ([WithComments Import] -> [WithComments Import])
-> [[WithComments Import]] -> [[WithComments Import]]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [WithComments Import] -> [WithComments Import]
sortByName [[WithComments Import]]
xs
          Bool
False -> [[WithComments Import]] -> Printer [[WithComments Import]]
forall a. a -> Printer a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [[WithComments Import]]
xs

mkImportCollection :: GHC.HsModule' -> ImportCollection
mkImportCollection :: HsModule' -> ImportCollection
mkImportCollection GHC.HsModule {[LImportDecl GhcPs]
[LHsDecl GhcPs]
Maybe (XRec GhcPs [LIE GhcPs])
Maybe (XRec GhcPs ModuleName)
XCModule GhcPs
hsmodExt :: XCModule GhcPs
hsmodName :: Maybe (XRec GhcPs ModuleName)
hsmodExports :: Maybe (XRec GhcPs [LIE GhcPs])
hsmodImports :: [LImportDecl GhcPs]
hsmodDecls :: [LHsDecl GhcPs]
hsmodDecls :: forall p. HsModule p -> [LHsDecl p]
hsmodImports :: forall p. HsModule p -> [LImportDecl p]
hsmodExports :: forall p. HsModule p -> Maybe (XRec p [LIE p])
hsmodName :: forall p. HsModule p -> Maybe (XRec p ModuleName)
hsmodExt :: forall p. HsModule p -> XCModule p
..} =
  [[WithComments Import]] -> ImportCollection
ImportCollection
    ([[WithComments Import]] -> ImportCollection)
-> [[WithComments Import]] -> ImportCollection
forall a b. (a -> b) -> a -> b
$ ([GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
 -> [WithComments Import])
-> [[GenLocated SrcSpanAnnA (ImportDecl GhcPs)]]
-> [[WithComments Import]]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
        ((GenLocated SrcSpanAnnA (ImportDecl GhcPs) -> WithComments Import)
-> [GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
-> [WithComments Import]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((ImportDecl GhcPs -> Import)
-> WithComments (ImportDecl GhcPs) -> WithComments Import
forall a b. (a -> b) -> WithComments a -> WithComments b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ImportDecl GhcPs -> Import
mkImport (WithComments (ImportDecl GhcPs) -> WithComments Import)
-> (GenLocated SrcSpanAnnA (ImportDecl GhcPs)
    -> WithComments (ImportDecl GhcPs))
-> GenLocated SrcSpanAnnA (ImportDecl GhcPs)
-> WithComments Import
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GenLocated SrcSpanAnnA (ImportDecl GhcPs)
-> WithComments (ImportDecl GhcPs)
forall l a. CommentExtraction l => GenLocated l a -> WithComments a
fromGenLocated))
        ([LImportDecl GhcPs] -> [[LImportDecl GhcPs]]
extractImports' [LImportDecl GhcPs]
hsmodImports)

hasImports :: ImportCollection -> Bool
hasImports :: ImportCollection -> Bool
hasImports (ImportCollection [[WithComments Import]]
xs) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [[WithComments Import]] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[WithComments Import]]
xs

-- | Extracts import declarations from the given module. Adjacent import
-- declarations are grouped as a single list.
extractImports' :: [GHC.LImportDecl GHC.GhcPs] -> [[GHC.LImportDecl GHC.GhcPs]]
extractImports' :: [LImportDecl GhcPs] -> [[LImportDecl GhcPs]]
extractImports' = [LImportDecl GhcPs] -> [[LImportDecl GhcPs]]
[GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
-> [[GenLocated SrcSpanAnnA (ImportDecl GhcPs)]]
groupImports ([GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
 -> [[GenLocated SrcSpanAnnA (ImportDecl GhcPs)]])
-> ([GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
    -> [GenLocated SrcSpanAnnA (ImportDecl GhcPs)])
-> [GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
-> [[GenLocated SrcSpanAnnA (ImportDecl GhcPs)]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [LImportDecl GhcPs] -> [LImportDecl GhcPs]
[GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
-> [GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
sortImportsByLocation

-- | Combines adjacent import declarations into a single list.
groupImports :: [GHC.LImportDecl GHC.GhcPs] -> [[GHC.LImportDecl GHC.GhcPs]]
groupImports :: [LImportDecl GhcPs] -> [[LImportDecl GhcPs]]
groupImports = [[LImportDecl GhcPs]]
-> [LImportDecl GhcPs] -> [[LImportDecl GhcPs]]
groupImports' []
  where
    groupImports' ::
         [[GHC.LImportDecl GHC.GhcPs]]
      -> [GHC.LImportDecl GHC.GhcPs]
      -> [[GHC.LImportDecl GHC.GhcPs]]
    groupImports' :: [[LImportDecl GhcPs]]
-> [LImportDecl GhcPs] -> [[LImportDecl GhcPs]]
groupImports' [[LImportDecl GhcPs]]
xs [] = [[LImportDecl GhcPs]]
xs
    groupImports' [] (LImportDecl GhcPs
x:[LImportDecl GhcPs]
xs) = [[LImportDecl GhcPs]]
-> [LImportDecl GhcPs] -> [[LImportDecl GhcPs]]
groupImports' [[LImportDecl GhcPs
x]] [LImportDecl GhcPs]
xs
    groupImports' [[]] (LImportDecl GhcPs
x:[LImportDecl GhcPs]
xs) = [[LImportDecl GhcPs]]
-> [LImportDecl GhcPs] -> [[LImportDecl GhcPs]]
groupImports' [[LImportDecl GhcPs
x]] [LImportDecl GhcPs]
xs
    groupImports' ([]:[LImportDecl GhcPs]
x:[[LImportDecl GhcPs]]
xs) (LImportDecl GhcPs
y:[LImportDecl GhcPs]
ys) = [[LImportDecl GhcPs]]
-> [LImportDecl GhcPs] -> [[LImportDecl GhcPs]]
groupImports' ([LImportDecl GhcPs
GenLocated SrcSpanAnnA (ImportDecl GhcPs)
y] [GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
-> [[GenLocated SrcSpanAnnA (ImportDecl GhcPs)]]
-> [[GenLocated SrcSpanAnnA (ImportDecl GhcPs)]]
forall a. a -> [a] -> [a]
: [LImportDecl GhcPs]
[GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
x [GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
-> [[GenLocated SrcSpanAnnA (ImportDecl GhcPs)]]
-> [[GenLocated SrcSpanAnnA (ImportDecl GhcPs)]]
forall a. a -> [a] -> [a]
: [[LImportDecl GhcPs]]
[[GenLocated SrcSpanAnnA (ImportDecl GhcPs)]]
xs) [LImportDecl GhcPs]
ys
    groupImports' ((LImportDecl GhcPs
z:[LImportDecl GhcPs]
zs):[[LImportDecl GhcPs]]
xs) (LImportDecl GhcPs
y:[LImportDecl GhcPs]
ys)
      | LImportDecl GhcPs
GenLocated SrcSpanAnnA (ImportDecl GhcPs)
z GenLocated SrcSpanAnnA (ImportDecl GhcPs)
-> GenLocated SrcSpanAnnA (ImportDecl GhcPs) -> Bool
forall {a} {a} {e} {e}.
(HasLoc a, HasLoc a) =>
GenLocated a e -> GenLocated a e -> Bool
`isAdjacentTo` LImportDecl GhcPs
GenLocated SrcSpanAnnA (ImportDecl GhcPs)
y = [[LImportDecl GhcPs]]
-> [LImportDecl GhcPs] -> [[LImportDecl GhcPs]]
groupImports' ((LImportDecl GhcPs
GenLocated SrcSpanAnnA (ImportDecl GhcPs)
y GenLocated SrcSpanAnnA (ImportDecl GhcPs)
-> [GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
-> [GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
forall a. a -> [a] -> [a]
: LImportDecl GhcPs
GenLocated SrcSpanAnnA (ImportDecl GhcPs)
z GenLocated SrcSpanAnnA (ImportDecl GhcPs)
-> [GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
-> [GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
forall a. a -> [a] -> [a]
: [LImportDecl GhcPs]
[GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
zs) [GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
-> [[GenLocated SrcSpanAnnA (ImportDecl GhcPs)]]
-> [[GenLocated SrcSpanAnnA (ImportDecl GhcPs)]]
forall a. a -> [a] -> [a]
: [[LImportDecl GhcPs]]
[[GenLocated SrcSpanAnnA (ImportDecl GhcPs)]]
xs) [LImportDecl GhcPs]
ys
      | Bool
otherwise = [[LImportDecl GhcPs]]
-> [LImportDecl GhcPs] -> [[LImportDecl GhcPs]]
groupImports' ([LImportDecl GhcPs
GenLocated SrcSpanAnnA (ImportDecl GhcPs)
y] [GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
-> [[GenLocated SrcSpanAnnA (ImportDecl GhcPs)]]
-> [[GenLocated SrcSpanAnnA (ImportDecl GhcPs)]]
forall a. a -> [a] -> [a]
: (LImportDecl GhcPs
GenLocated SrcSpanAnnA (ImportDecl GhcPs)
z GenLocated SrcSpanAnnA (ImportDecl GhcPs)
-> [GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
-> [GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
forall a. a -> [a] -> [a]
: [LImportDecl GhcPs]
[GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
zs) [GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
-> [[GenLocated SrcSpanAnnA (ImportDecl GhcPs)]]
-> [[GenLocated SrcSpanAnnA (ImportDecl GhcPs)]]
forall a. a -> [a] -> [a]
: [[LImportDecl GhcPs]]
[[GenLocated SrcSpanAnnA (ImportDecl GhcPs)]]
xs) [LImportDecl GhcPs]
ys
    GenLocated a e
a isAdjacentTo :: GenLocated a e -> GenLocated a e -> Bool
`isAdjacentTo` GenLocated a e
b =
      RealSrcSpan -> Int
GHC.srcSpanEndLine (GenLocated a e -> RealSrcSpan
forall {a} {e}. HasLoc a => GenLocated a e -> RealSrcSpan
sp GenLocated a e
a) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== RealSrcSpan -> Int
GHC.srcSpanStartLine (GenLocated a e -> RealSrcSpan
forall {a} {e}. HasLoc a => GenLocated a e -> RealSrcSpan
sp GenLocated a e
b)
        Bool -> Bool -> Bool
|| RealSrcSpan -> Int
GHC.srcSpanEndLine (GenLocated a e -> RealSrcSpan
forall {a} {e}. HasLoc a => GenLocated a e -> RealSrcSpan
sp GenLocated a e
b) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== RealSrcSpan -> Int
GHC.srcSpanStartLine (GenLocated a e -> RealSrcSpan
forall {a} {e}. HasLoc a => GenLocated a e -> RealSrcSpan
sp GenLocated a e
a)
    sp :: GenLocated a e -> RealSrcSpan
sp GenLocated a e
x =
      case a -> SrcSpan
forall a. HasLoc a => a -> SrcSpan
GHC.locA (a -> SrcSpan) -> a -> SrcSpan
forall a b. (a -> b) -> a -> b
$ GenLocated a e -> a
forall l e. GenLocated l e -> l
GHC.getLoc GenLocated a e
x of
        GHC.RealSrcSpan RealSrcSpan
x' Maybe BufSpan
_ -> RealSrcSpan
x'
        SrcSpan
_ -> [Char] -> RealSrcSpan
forall a. HasCallStack => [Char] -> a
error [Char]
"Src span unavailable."

-- | This function sorts imports by their start line numbers.
sortImportsByLocation ::
     [GHC.LImportDecl GHC.GhcPs] -> [GHC.LImportDecl GHC.GhcPs]
sortImportsByLocation :: [LImportDecl GhcPs] -> [LImportDecl GhcPs]
sortImportsByLocation = (GenLocated SrcSpanAnnA (ImportDecl GhcPs)
 -> GenLocated SrcSpanAnnA (ImportDecl GhcPs) -> Ordering)
-> [GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
-> [GenLocated SrcSpanAnnA (ImportDecl GhcPs)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy ((Int -> Int -> Ordering) -> Int -> Int -> Ordering
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Int -> Int -> Ordering)
-> (GenLocated SrcSpanAnnA (ImportDecl GhcPs) -> Int)
-> GenLocated SrcSpanAnnA (ImportDecl GhcPs)
-> GenLocated SrcSpanAnnA (ImportDecl GhcPs)
-> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` GenLocated SrcSpanAnnA (ImportDecl GhcPs) -> Int
forall {e}. GenLocated SrcSpanAnnA e -> Int
lineIdx)
  where
    lineIdx :: GenLocated SrcSpanAnnA e -> Int
lineIdx = HasCallStack => SrcSpan -> Int
SrcSpan -> Int
startLine (SrcSpan -> Int)
-> (GenLocated SrcSpanAnnA e -> SrcSpan)
-> GenLocated SrcSpanAnnA e
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SrcSpanAnnA -> SrcSpan
forall a. HasLoc a => a -> SrcSpan
GHC.locA (SrcSpanAnnA -> SrcSpan)
-> (GenLocated SrcSpanAnnA e -> SrcSpanAnnA)
-> GenLocated SrcSpanAnnA e
-> SrcSpan
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GenLocated SrcSpanAnnA e -> SrcSpanAnnA
forall l e. GenLocated l e -> l
GHC.getLoc

-- | This function returns the start line of the given 'SrcSpan'. If it is
-- not available, it raises an error.
startLine :: HasCallStack => GHC.SrcSpan -> Int
startLine :: HasCallStack => SrcSpan -> Int
startLine (GHC.RealSrcSpan RealSrcSpan
x Maybe BufSpan
_) = RealSrcSpan -> Int
GHC.srcSpanStartLine RealSrcSpan
x
startLine (GHC.UnhelpfulSpan UnhelpfulSpanReason
_) = [Char] -> Int
forall a. HasCallStack => [Char] -> a
error [Char]
"The src span is unavailable."