{-# LANGUAGE CPP #-}

module StaticLS.HI (
    getDocs,
    getDocsBatch,
    renderNameDocs,
    NameDocs (..),
)
where

import qualified Data.IntMap as IntMap
import qualified Data.Map as Map
import Data.Maybe
import Data.Text as T
import qualified GHC
import qualified GHC.Plugins as GHC
import qualified GHC.Types.Unique.Map as GHC
import StaticLS.SDoc

data NameDocs = NameDocs
    { NameDocs -> Maybe [HsDoc GhcRn]
declComment :: Maybe [GHC.HsDoc GHC.GhcRn]
    , NameDocs -> IntMap (HsDoc GhcRn)
argComments :: IntMap.IntMap (GHC.HsDoc GHC.GhcRn)
    }

instance GHC.Outputable NameDocs where
    ppr :: NameDocs -> SDoc
ppr NameDocs
nameDoc = Maybe [HsDoc GhcRn] -> SDoc
forall a. Outputable a => a -> SDoc
GHC.ppr NameDocs
nameDoc.declComment

renderNameDocs :: NameDocs -> Text
renderNameDocs :: NameDocs -> Text
renderNameDocs NameDocs
nameDocs =
    Int -> Text -> Text
T.drop Int
1 (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ -- Drop the leading space from haddock comments
        Text -> ([HsDoc GhcRn] -> Text) -> Maybe [HsDoc GhcRn] -> Text
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Text
"" ([Text] -> Text
T.concat ([Text] -> Text)
-> ([HsDoc GhcRn] -> [Text]) -> [HsDoc GhcRn] -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (HsDoc GhcRn -> Text) -> [HsDoc GhcRn] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap HsDoc GhcRn -> Text
forall a. Outputable a => a -> Text
showGhc) NameDocs
nameDocs.declComment

getDocsBatch :: [GHC.Name] -> GHC.ModIface -> [NameDocs]
getDocsBatch :: [Name] -> ModIface -> [NameDocs]
getDocsBatch [Name]
names ModIface
iface =
    case ModIface -> Maybe Docs
forall (phase :: ModIfacePhase). ModIface_ phase -> Maybe Docs
GHC.mi_docs ModIface
iface of
        Maybe Docs
Nothing -> []
        Just
            GHC.Docs
                { docs_decls :: Docs -> UniqMap Name [HsDoc GhcRn]
GHC.docs_decls = UniqMap Name [HsDoc GhcRn]
decls
                , docs_args :: Docs -> UniqMap Name (IntMap (HsDoc GhcRn))
GHC.docs_args = UniqMap Name (IntMap (HsDoc GhcRn))
args
                } ->
                -- Lifted out compared to `getDocs` - probably slightly
                -- more efficient though the compiler may just optimize this
                let declMap :: Map String [HsDoc GhcRn]
declMap = UniqMap Name [HsDoc GhcRn] -> Map String [HsDoc GhcRn]
forall v. UniqMap Name v -> Map String v
uniqNameMapToMap UniqMap Name [HsDoc GhcRn]
decls
                    argsMap :: Map String (IntMap (HsDoc GhcRn))
argsMap = UniqMap Name (IntMap (HsDoc GhcRn))
-> Map String (IntMap (HsDoc GhcRn))
forall v. UniqMap Name v -> Map String v
uniqNameMapToMap UniqMap Name (IntMap (HsDoc GhcRn))
args
                 in ( \Name
name ->
                        NameDocs
                            { $sel:declComment:NameDocs :: Maybe [HsDoc GhcRn]
declComment =
                                String -> Map String [HsDoc GhcRn] -> Maybe [HsDoc GhcRn]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (Name -> String
GHC.nameStableString Name
name) Map String [HsDoc GhcRn]
declMap
                            , $sel:argComments:NameDocs :: IntMap (HsDoc GhcRn)
argComments =
                                IntMap (HsDoc GhcRn)
-> Maybe (IntMap (HsDoc GhcRn)) -> IntMap (HsDoc GhcRn)
forall a. a -> Maybe a -> a
fromMaybe IntMap (HsDoc GhcRn)
forall a. Monoid a => a
mempty (Maybe (IntMap (HsDoc GhcRn)) -> IntMap (HsDoc GhcRn))
-> Maybe (IntMap (HsDoc GhcRn)) -> IntMap (HsDoc GhcRn)
forall a b. (a -> b) -> a -> b
$
                                    String
-> Map String (IntMap (HsDoc GhcRn))
-> Maybe (IntMap (HsDoc GhcRn))
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (Name -> String
GHC.nameStableString Name
name) Map String (IntMap (HsDoc GhcRn))
argsMap
                            }
                    )
                        (Name -> NameDocs) -> [Name] -> [NameDocs]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Name]
names

getDocs :: GHC.Name -> GHC.ModIface -> Maybe NameDocs
getDocs :: Name -> ModIface -> Maybe NameDocs
getDocs Name
name ModIface
iface =
    case ModIface -> Maybe Docs
forall (phase :: ModIfacePhase). ModIface_ phase -> Maybe Docs
GHC.mi_docs ModIface
iface of
        Maybe Docs
Nothing -> Maybe NameDocs
forall a. Maybe a
Nothing
        Just
            GHC.Docs
                { docs_decls :: Docs -> UniqMap Name [HsDoc GhcRn]
GHC.docs_decls = UniqMap Name [HsDoc GhcRn]
decls
                , docs_args :: Docs -> UniqMap Name (IntMap (HsDoc GhcRn))
GHC.docs_args = UniqMap Name (IntMap (HsDoc GhcRn))
args
                } ->
                NameDocs -> Maybe NameDocs
forall a. a -> Maybe a
Just (NameDocs -> Maybe NameDocs) -> NameDocs -> Maybe NameDocs
forall a b. (a -> b) -> a -> b
$
                    NameDocs
                        { $sel:declComment:NameDocs :: Maybe [HsDoc GhcRn]
declComment = Name -> UniqMap Name [HsDoc GhcRn] -> Maybe [HsDoc GhcRn]
forall v. Name -> UniqMap Name v -> Maybe v
normalizeNameLookup Name
name UniqMap Name [HsDoc GhcRn]
decls
                        , $sel:argComments:NameDocs :: IntMap (HsDoc GhcRn)
argComments = IntMap (HsDoc GhcRn)
-> Maybe (IntMap (HsDoc GhcRn)) -> IntMap (HsDoc GhcRn)
forall a. a -> Maybe a -> a
fromMaybe IntMap (HsDoc GhcRn)
forall a. Monoid a => a
mempty (Maybe (IntMap (HsDoc GhcRn)) -> IntMap (HsDoc GhcRn))
-> Maybe (IntMap (HsDoc GhcRn)) -> IntMap (HsDoc GhcRn)
forall a b. (a -> b) -> a -> b
$ Name
-> UniqMap Name (IntMap (HsDoc GhcRn))
-> Maybe (IntMap (HsDoc GhcRn))
forall v. Name -> UniqMap Name v -> Maybe v
normalizeNameLookup Name
name UniqMap Name (IntMap (HsDoc GhcRn))
args
                        }

normalizeNameLookup :: GHC.Name -> GHC.UniqMap GHC.Name v -> Maybe v
normalizeNameLookup :: forall v. Name -> UniqMap Name v -> Maybe v
normalizeNameLookup Name
name UniqMap Name v
uMap =
    String -> Map String v -> Maybe v
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (Name -> String
GHC.nameStableString Name
name) (UniqMap Name v -> Map String v
forall v. UniqMap Name v -> Map String v
uniqNameMapToMap UniqMap Name v
uMap)

uniqNameMapToMap :: GHC.UniqMap GHC.Name v -> Map.Map String v
uniqNameMapToMap :: forall v. UniqMap Name v -> Map String v
uniqNameMapToMap =
    [(String, v)] -> Map String v
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList
        ([(String, v)] -> Map String v)
-> (UniqMap Name v -> [(String, v)])
-> UniqMap Name v
-> Map String v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Name, v) -> (String, v)) -> [(Name, v)] -> [(String, v)]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Name, v) -> (String, v)
forall {b}. (Name, b) -> (String, b)
stringifyNameKeys
        ([(Name, v)] -> [(String, v)])
-> (UniqMap Name v -> [(Name, v)])
-> UniqMap Name v
-> [(String, v)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap (Name, v) -> [(Name, v)]
forall a. IntMap a -> [a]
IntMap.elems
        (IntMap (Name, v) -> [(Name, v)])
-> (UniqMap Name v -> IntMap (Name, v))
-> UniqMap Name v
-> [(Name, v)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UniqFM Name (Name, v) -> IntMap (Name, v)
forall key elt. UniqFM key elt -> IntMap elt
GHC.ufmToIntMap
        (UniqFM Name (Name, v) -> IntMap (Name, v))
-> (UniqMap Name v -> UniqFM Name (Name, v))
-> UniqMap Name v
-> IntMap (Name, v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UniqMap Name v -> UniqFM Name (Name, v)
forall k a. UniqMap k a -> UniqFM k (k, a)
getUniqMap
  where
    stringifyNameKeys :: (Name, b) -> (String, b)
stringifyNameKeys (Name
nameKey, b
v) = (Name -> String
GHC.nameStableString Name
nameKey, b
v)

getUniqMap :: GHC.UniqMap k a -> GHC.UniqFM k (k, a)
#if MIN_VERSION_base(4,18,0)
getUniqMap :: forall k a. UniqMap k a -> UniqFM k (k, a)
getUniqMap = UniqMap k a -> UniqFM k (k, a)
forall k a. UniqMap k a -> UniqFM k (k, a)
GHC.getUniqMap
#else
getUniqMap (GHC.UniqMap m) = m
#endif