{-# 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)
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 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ImportDecl GhcPs -> FixityImport
extractFixityImport forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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)
ImportDeclQualifiedStyle
IsBootInterface
XRec GhcPs ModuleName
XCImportDecl GhcPs
ImportDeclPkgQual 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])
ideclImportList :: Maybe (ImportListInterpretation, XRec GhcPs [LIE GhcPs])
ideclAs :: Maybe (XRec GhcPs ModuleName)
ideclQualified :: ImportDeclQualifiedStyle
ideclSafe :: Bool
ideclSource :: IsBootInterface
ideclPkgQual :: ImportDeclPkgQual GhcPs
ideclName :: XRec GhcPs ModuleName
ideclExt :: XCImportDecl GhcPs
..} =
  FixityImport
    { fimportPackage :: Maybe PackageName
fimportPackage = case ImportDeclPkgQual GhcPs
ideclPkgQual of
        RawPkgQual
ImportDeclPkgQual GhcPs
NoRawPkgQual -> forall a. Maybe a
Nothing
        RawPkgQual StringLiteral
strLiteral ->
          forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> PackageName
mkPackageName forall b c a. (b -> c) -> (a -> b) -> a -> c
. FastString -> String
GHC.unpackFS forall b c a. (b -> c) -> (a -> b) -> a -> c
. StringLiteral -> FastString
sl_fs 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 =
        forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
          (forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap OccName -> OpName
occOpName forall b c a. (b -> c) -> (a -> b) -> a -> c
. IE GhcPs -> [OccName]
ieToOccNames forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall l e. GenLocated l e -> e
unLoc) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall l e. GenLocated l e -> e
unLoc))
          Maybe (ImportListInterpretation, XRec GhcPs [LIE GhcPs])
ideclImportList
    }
  where
    ideclName' :: ModuleName
ideclName' = ModuleName -> ModuleName
ghcModuleNameToCabal (forall l e. GenLocated l e -> e
unLoc XRec GhcPs ModuleName
ideclName)
    ideclAs' :: Maybe ModuleName
ideclAs' = ModuleName -> ModuleName
ghcModuleNameToCabal forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall l e. GenLocated l e -> e
unLoc forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (XRec GhcPs ModuleName)
ideclAs

ieToOccNames :: IE GhcPs -> [OccName]
ieToOccNames :: IE GhcPs -> [OccName]
ieToOccNames = \case
  IEVar XIEVar GhcPs
_ (L SrcSpanAnnA
_ IEWrappedName GhcPs
x) -> [forall name. HasOccName name => name -> OccName
occName IEWrappedName GhcPs
x]
  IEThingAbs XIEThingAbs GhcPs
_ (L SrcSpanAnnA
_ IEWrappedName GhcPs
x) -> [forall name. HasOccName name => name -> OccName
occName IEWrappedName GhcPs
x]
  IEThingAll XIEThingAll GhcPs
_ (L SrcSpanAnnA
_ IEWrappedName GhcPs
x) -> [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 -> forall name. HasOccName name => name -> OccName
occName IEWrappedName GhcPs
x forall a. a -> [a] -> [a]
: forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall name. HasOccName name => name -> OccName
occName forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall l e. GenLocated l e -> e
unLoc) [LIEWrappedName 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 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 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 -> 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 :: Maybe PackageName
fimportPackage = Maybe PackageName
mpackageName,
                    fimportModule :: ModuleName
fimportModule = ModuleName
mmodule
                  }
           in forall a. NonEmpty a -> [a]
NE.toList NonEmpty (Maybe PackageName, ModuleName)
exports forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= FixityImport -> [FixityImport]
expand forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Maybe PackageName, ModuleName) -> FixityImport
exportToImport