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

-- | Simplified representation of the import list for the purposes of fixity
-- inference.
module Ormolu.Fixity.Imports
  ( FixityImport (..),
    extractFixityImports,
    applyModuleReexports,
  )
where

import Data.Bifunctor (second)
import Data.List.NonEmpty qualified as NE
import Data.Map.Strict qualified as Map
import Distribution.ModuleName (ModuleName)
import Distribution.Types.PackageName
import GHC.Data.FastString qualified as GHC
import GHC.Hs hiding (ModuleName, OpName)
import GHC.Types.Name.Occurrence
import GHC.Types.PkgQual (RawPkgQual (..))
import GHC.Types.SourceText (StringLiteral (..))
import GHC.Types.SrcLoc
import Ormolu.Fixity.Internal
import Ormolu.Utils (ghcModuleNameToCabal)

-- | Simplified info about an import.
data FixityImport = FixityImport
  { FixityImport -> Maybe PackageName
fimportPackage :: !(Maybe PackageName),
    FixityImport -> ModuleName
fimportModule :: !ModuleName,
    FixityImport -> FixityQualification
fimportQualified :: !FixityQualification,
    FixityImport -> Maybe (ImportListInterpretation, [OpName])
fimportList :: !(Maybe (ImportListInterpretation, [OpName]))
  }

-- | Extract 'FixityImport's from the AST.
extractFixityImports ::
  [LImportDecl GhcPs] ->
  [FixityImport]
extractFixityImports :: [LImportDecl GhcPs] -> [FixityImport]
extractFixityImports = (GenLocated SrcSpanAnnA (ImportDecl GhcPs) -> FixityImport)
-> [GenLocated SrcSpanAnnA (ImportDecl GhcPs)] -> [FixityImport]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ImportDecl GhcPs -> FixityImport
extractFixityImport (ImportDecl GhcPs -> FixityImport)
-> (GenLocated SrcSpanAnnA (ImportDecl GhcPs) -> ImportDecl GhcPs)
-> GenLocated SrcSpanAnnA (ImportDecl GhcPs)
-> FixityImport
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GenLocated SrcSpanAnnA (ImportDecl GhcPs) -> ImportDecl GhcPs
forall l e. GenLocated l e -> e
unLoc)

-- | Extract an individual 'FixityImport'.
extractFixityImport :: ImportDecl GhcPs -> FixityImport
extractFixityImport :: ImportDecl GhcPs -> FixityImport
extractFixityImport ImportDecl {Bool
Maybe (ImportListInterpretation, XRec GhcPs [LIE GhcPs])
Maybe (XRec GhcPs ModuleName)
ImportDeclPkgQual GhcPs
XCImportDecl GhcPs
XRec GhcPs ModuleName
IsBootInterface
ImportDeclQualifiedStyle
ideclExt :: XCImportDecl GhcPs
ideclName :: XRec GhcPs ModuleName
ideclPkgQual :: ImportDeclPkgQual GhcPs
ideclSource :: IsBootInterface
ideclSafe :: Bool
ideclQualified :: ImportDeclQualifiedStyle
ideclAs :: Maybe (XRec GhcPs ModuleName)
ideclImportList :: Maybe (ImportListInterpretation, XRec GhcPs [LIE GhcPs])
ideclExt :: forall pass. ImportDecl pass -> XCImportDecl pass
ideclName :: forall pass. ImportDecl pass -> XRec pass ModuleName
ideclPkgQual :: forall pass. ImportDecl pass -> ImportDeclPkgQual pass
ideclSource :: forall pass. ImportDecl pass -> IsBootInterface
ideclSafe :: forall pass. ImportDecl pass -> Bool
ideclQualified :: forall pass. ImportDecl pass -> ImportDeclQualifiedStyle
ideclAs :: forall pass. ImportDecl pass -> Maybe (XRec pass ModuleName)
ideclImportList :: forall pass.
ImportDecl pass
-> Maybe (ImportListInterpretation, XRec pass [LIE pass])
..} =
  FixityImport
    { fimportPackage :: Maybe PackageName
fimportPackage = case ImportDeclPkgQual GhcPs
ideclPkgQual of
        ImportDeclPkgQual GhcPs
RawPkgQual
NoRawPkgQual -> Maybe PackageName
forall a. Maybe a
Nothing
        RawPkgQual StringLiteral
strLiteral ->
          PackageName -> Maybe PackageName
forall a. a -> Maybe a
Just (PackageName -> Maybe PackageName)
-> (StringLiteral -> PackageName)
-> StringLiteral
-> Maybe PackageName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> PackageName
mkPackageName (String -> PackageName)
-> (StringLiteral -> String) -> StringLiteral -> PackageName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FastString -> String
GHC.unpackFS (FastString -> String)
-> (StringLiteral -> FastString) -> StringLiteral -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StringLiteral -> FastString
sl_fs (StringLiteral -> Maybe PackageName)
-> StringLiteral -> Maybe PackageName
forall a b. (a -> b) -> a -> b
$ StringLiteral
strLiteral,
      fimportModule :: ModuleName
fimportModule = ModuleName
ideclName',
      fimportQualified :: FixityQualification
fimportQualified = case (ImportDeclQualifiedStyle
ideclQualified, Maybe ModuleName
ideclAs') of
        (ImportDeclQualifiedStyle
QualifiedPre, Maybe ModuleName
Nothing) ->
          ModuleName -> FixityQualification
OnlyQualified ModuleName
ideclName'
        (ImportDeclQualifiedStyle
QualifiedPost, Maybe ModuleName
Nothing) ->
          ModuleName -> FixityQualification
OnlyQualified ModuleName
ideclName'
        (ImportDeclQualifiedStyle
QualifiedPre, Just ModuleName
m) -> ModuleName -> FixityQualification
OnlyQualified ModuleName
m
        (ImportDeclQualifiedStyle
QualifiedPost, Just ModuleName
m) -> ModuleName -> FixityQualification
OnlyQualified ModuleName
m
        (ImportDeclQualifiedStyle
NotQualified, Maybe ModuleName
Nothing) ->
          ModuleName -> FixityQualification
UnqualifiedAndQualified ModuleName
ideclName'
        (ImportDeclQualifiedStyle
NotQualified, Just ModuleName
m) ->
          ModuleName -> FixityQualification
UnqualifiedAndQualified ModuleName
m,
      fimportList :: Maybe (ImportListInterpretation, [OpName])
fimportList =
        ((ImportListInterpretation,
  GenLocated SrcSpanAnnL [GenLocated SrcSpanAnnA (IE GhcPs)])
 -> (ImportListInterpretation, [OpName]))
-> Maybe
     (ImportListInterpretation,
      GenLocated SrcSpanAnnL [GenLocated SrcSpanAnnA (IE GhcPs)])
-> Maybe (ImportListInterpretation, [OpName])
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
          ((GenLocated SrcSpanAnnL [GenLocated SrcSpanAnnA (IE GhcPs)]
 -> [OpName])
-> (ImportListInterpretation,
    GenLocated SrcSpanAnnL [GenLocated SrcSpanAnnA (IE GhcPs)])
-> (ImportListInterpretation, [OpName])
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ((GenLocated SrcSpanAnnA (IE GhcPs) -> [OpName])
-> [GenLocated SrcSpanAnnA (IE GhcPs)] -> [OpName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((OccName -> OpName) -> [OccName] -> [OpName]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap OccName -> OpName
occOpName ([OccName] -> [OpName])
-> (GenLocated SrcSpanAnnA (IE GhcPs) -> [OccName])
-> GenLocated SrcSpanAnnA (IE GhcPs)
-> [OpName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IE GhcPs -> [OccName]
ieToOccNames (IE GhcPs -> [OccName])
-> (GenLocated SrcSpanAnnA (IE GhcPs) -> IE GhcPs)
-> GenLocated SrcSpanAnnA (IE GhcPs)
-> [OccName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GenLocated SrcSpanAnnA (IE GhcPs) -> IE GhcPs
forall l e. GenLocated l e -> e
unLoc) ([GenLocated SrcSpanAnnA (IE GhcPs)] -> [OpName])
-> (GenLocated SrcSpanAnnL [GenLocated SrcSpanAnnA (IE GhcPs)]
    -> [GenLocated SrcSpanAnnA (IE GhcPs)])
-> GenLocated SrcSpanAnnL [GenLocated SrcSpanAnnA (IE GhcPs)]
-> [OpName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GenLocated SrcSpanAnnL [GenLocated SrcSpanAnnA (IE GhcPs)]
-> [GenLocated SrcSpanAnnA (IE GhcPs)]
forall l e. GenLocated l e -> e
unLoc))
          Maybe (ImportListInterpretation, XRec GhcPs [LIE GhcPs])
Maybe
  (ImportListInterpretation,
   GenLocated SrcSpanAnnL [GenLocated SrcSpanAnnA (IE GhcPs)])
ideclImportList
    }
  where
    ideclName' :: ModuleName
ideclName' = ModuleName -> ModuleName
ghcModuleNameToCabal (GenLocated SrcSpanAnnA ModuleName -> ModuleName
forall l e. GenLocated l e -> e
unLoc XRec GhcPs ModuleName
GenLocated SrcSpanAnnA ModuleName
ideclName)
    ideclAs' :: Maybe ModuleName
ideclAs' = ModuleName -> ModuleName
ghcModuleNameToCabal (ModuleName -> ModuleName)
-> (GenLocated SrcSpanAnnA ModuleName -> ModuleName)
-> GenLocated SrcSpanAnnA ModuleName
-> ModuleName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GenLocated SrcSpanAnnA ModuleName -> ModuleName
forall l e. GenLocated l e -> e
unLoc (GenLocated SrcSpanAnnA ModuleName -> ModuleName)
-> Maybe (GenLocated SrcSpanAnnA ModuleName) -> Maybe ModuleName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (XRec GhcPs ModuleName)
Maybe (GenLocated SrcSpanAnnA ModuleName)
ideclAs

ieToOccNames :: IE GhcPs -> [OccName]
ieToOccNames :: IE GhcPs -> [OccName]
ieToOccNames = \case
  IEVar XIEVar GhcPs
_ (L SrcSpanAnnA
_ IEWrappedName GhcPs
x) -> [IEWrappedName GhcPs -> OccName
forall name. HasOccName name => name -> OccName
occName IEWrappedName GhcPs
x]
  IEThingAbs XIEThingAbs GhcPs
_ (L SrcSpanAnnA
_ IEWrappedName GhcPs
x) -> [IEWrappedName GhcPs -> OccName
forall name. HasOccName name => name -> OccName
occName IEWrappedName GhcPs
x]
  IEThingAll XIEThingAll GhcPs
_ (L SrcSpanAnnA
_ IEWrappedName GhcPs
x) -> [IEWrappedName GhcPs -> OccName
forall name. HasOccName name => name -> OccName
occName IEWrappedName GhcPs
x] -- TODO not quite correct, but how to do better?
  IEThingWith XIEThingWith GhcPs
_ (L SrcSpanAnnA
_ IEWrappedName GhcPs
x) IEWildcard
_ [LIEWrappedName GhcPs]
xs -> IEWrappedName GhcPs -> OccName
forall name. HasOccName name => name -> OccName
occName IEWrappedName GhcPs
x OccName -> [OccName] -> [OccName]
forall a. a -> [a] -> [a]
: (GenLocated SrcSpanAnnA (IEWrappedName GhcPs) -> OccName)
-> [GenLocated SrcSpanAnnA (IEWrappedName GhcPs)] -> [OccName]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (IEWrappedName GhcPs -> OccName
forall name. HasOccName name => name -> OccName
occName (IEWrappedName GhcPs -> OccName)
-> (GenLocated SrcSpanAnnA (IEWrappedName GhcPs)
    -> IEWrappedName GhcPs)
-> GenLocated SrcSpanAnnA (IEWrappedName GhcPs)
-> OccName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GenLocated SrcSpanAnnA (IEWrappedName GhcPs) -> IEWrappedName GhcPs
forall l e. GenLocated l e -> e
unLoc) [LIEWrappedName GhcPs]
[GenLocated SrcSpanAnnA (IEWrappedName GhcPs)]
xs
  IE GhcPs
_ -> []

-- | Apply given module re-exports.
applyModuleReexports :: ModuleReexports -> [FixityImport] -> [FixityImport]
applyModuleReexports :: ModuleReexports -> [FixityImport] -> [FixityImport]
applyModuleReexports (ModuleReexports Map ModuleName (NonEmpty (Maybe PackageName, ModuleName))
reexports) [FixityImport]
imports = [FixityImport]
imports [FixityImport]
-> (FixityImport -> [FixityImport]) -> [FixityImport]
forall a b. [a] -> (a -> [b]) -> [b]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= FixityImport -> [FixityImport]
expand
  where
    expand :: FixityImport -> [FixityImport]
expand FixityImport
i = do
      case ModuleName
-> Map ModuleName (NonEmpty (Maybe PackageName, ModuleName))
-> Maybe (NonEmpty (Maybe PackageName, ModuleName))
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (FixityImport -> ModuleName
fimportModule FixityImport
i) Map ModuleName (NonEmpty (Maybe PackageName, ModuleName))
reexports of
        Maybe (NonEmpty (Maybe PackageName, ModuleName))
Nothing -> FixityImport -> [FixityImport]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure FixityImport
i
        Just NonEmpty (Maybe PackageName, ModuleName)
exports ->
          let exportToImport :: (Maybe PackageName, ModuleName) -> FixityImport
exportToImport (Maybe PackageName
mpackageName, ModuleName
mmodule) =
                FixityImport
i
                  { fimportPackage = mpackageName,
                    fimportModule = mmodule
                  }
           in NonEmpty (Maybe PackageName, ModuleName)
-> [(Maybe PackageName, ModuleName)]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty (Maybe PackageName, ModuleName)
exports [(Maybe PackageName, ModuleName)]
-> ((Maybe PackageName, ModuleName) -> [FixityImport])
-> [FixityImport]
forall a b. [a] -> (a -> [b]) -> [b]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= FixityImport -> [FixityImport]
expand (FixityImport -> [FixityImport])
-> ((Maybe PackageName, ModuleName) -> FixityImport)
-> (Maybe PackageName, ModuleName)
-> [FixityImport]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Maybe PackageName, ModuleName) -> FixityImport
exportToImport