path: root/src/HaskellCodeExplorer/ModuleInfo.hs

{-# LANGUAGE TupleSections #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE StrictData #-}

module HaskellCodeExplorer.ModuleInfo
  ( createModuleInfo
  , ModuleDependencies
  ) where

import           Control.Monad.State.Strict     ( State
                                                , evalState
                                                , execState
                                                , get
                                                , put
                                                )
import qualified Data.Aeson                    as Aeson
import           Data.Aeson.Text                ( encodeToLazyText )
import qualified Data.Generics.Uniplate.Data   as U
import qualified Data.HashMap.Strict           as HM
import qualified Data.IntMap.Strict            as IM
import qualified Data.IntervalMap.Strict       as IVM
import qualified Data.List                     as L
                                         hiding ( span )
import qualified Data.Map.Strict               as M
import           Data.Maybe                     ( fromJust
                                                , fromMaybe
                                                , mapMaybe
                                                )
import qualified Data.Set                      as S
import qualified Data.Text                     as T
import qualified Data.Text.Encoding            as TE
import           Data.Text.Lazy                 ( toStrict )
import qualified Data.Vector                   as V
import           Documentation.Haddock.Types    ( DocH )
import           GHC                            ( DynFlags
                                                , GenLocated(..)
                                                , LHsBindLR
                                                , LHsDecl
                                                , ModSummary
                                                , ModuleInfo
                                                , ModuleName
                                                , SrcSpan
                                                , TyThing(..)
                                                , Type
                                                , TypecheckedModule
                                                , getLoc
                                                , isGoodSrcSpan
                                                , modInfoExportsWithSelectors
                                                , modInfoInstances
                                                , moduleInfo
                                                , moduleNameString
                                                , ms_hspp_buf
                                                , ms_mod
                                                , renamedSource
                                                , tm_internals_
                                                , tm_typechecked_source
                                                , unLoc
                                                )
import           GHC.Core.InstEnv               ( InstEnvs(..)
                                                , is_dfun
                                                )
import           GHC.Core.TyCon                 ( isFamInstTyCon
                                                , tyConName
                                                )
import           GHC.Core.Type                  ( expandTypeSynonyms )
import           GHC.Hs.Decls                   ( ForeignDecl(..)
                                                , HsDecl(..)
                                                , HsGroup(..)
                                                , InstDecl(..)
                                                , LForeignDecl
                                                , LInstDecl
                                                , LTyClDecl
                                                , group_tyclds
                                                , hsGroupInstDecls
                                                , tcdName
                                                , tyClDeclLName
                                                )
import           GHC.Hs.Doc                     ( HsDocString )
import           GHC.Hs.Extension               ( GhcRn )
import           GHC.Hs.ImpExp                  ( IE(..)
                                                , ImportDecl(..)
                                                )
import           GHC.Hs.Utils                   ( CollectFlag(..)
                                                , collectHsBindBinders
                                                )
import           GHC.Parser.Annotation          ( getLocA
                                                , sortLocatedA
                                                )
import           GHC.Tc.Types                   ( tcVisibleOrphanMods
                                                , tcg_inst_env
                                                , tcg_rdr_env
                                                , tcg_type_env
                                                )
import           GHC.Types.Name                 ( Name
                                                , OccName
                                                , getSrcSpan
                                                , nameOccName
                                                , nameSrcSpan
                                                , nameUnique
                                                )
import           GHC.Types.Name.Reader          ( GlobalRdrEnv )
import           GHC.Types.SrcLoc               ( isOneLineSpan )
import           GHC.Types.TypeEnv              ( TypeEnv
                                                , mkTypeEnv
                                                , typeEnvElts
                                                )
import           GHC.Types.Unique               ( getKey )
import           GHC.Types.Unique.DFM           ( eltsUDFM )
import           GHC.Types.Var                  ( Id
                                                , varName
                                                , varType
                                                )
import           GHC.Types.Var.Env              ( emptyTidyEnv )
import           GHC.Unit.External              ( ExternalPackageState
                                                , eps_PTE
                                                , eps_inst_env
                                                )
import           GHC.Unit.Home.ModInfo          ( HomePackageTable
                                                , hm_details
                                                )
import           GHC.Unit.Module.ModDetails     ( md_types )
import           GHC.Unit.State                 ( UnitState )
import           GHC.Unit.Types                 ( GenModule(..) )
import           HaskellCodeExplorer.AST.RenamedSource
import           HaskellCodeExplorer.AST.TypecheckedSource
import           HaskellCodeExplorer.GhcUtils
import           HaskellCodeExplorer.Preprocessor
                                                ( createSourceCodeTransformation
                                                )
import qualified HaskellCodeExplorer.Types     as HCE
import           Prelude                 hiding ( id
                                                , span
                                                )
import qualified Text.Blaze.Html5              as H
import qualified Text.Blaze.Html5.Attributes   as A

type ModuleDependencies
  = ( HM.HashMap HCE.HaskellFilePath HCE.HaskellModulePath
    , HM.HashMap HCE.HaskellModulePath HCE.DefinitionSiteMap
    , HM.HashMap
        HCE.HaskellModuleName
        (HM.HashMap HCE.ComponentId HCE.HaskellModulePath)
    )

type ModuleGhcData
  = ( DynFlags
    , UnitState
    , TypecheckedModule
    , HomePackageTable
    , ExternalPackageState
    , ModSummary
    )

createModuleInfo
  :: ModuleDependencies -- ^ Modules that have already been indexed
  -> ModuleGhcData -- ^ Data types from GHC
  -> HCE.HaskellModulePath -- ^ Current module path
  -> HCE.PackageId -- ^ Current package id
  -> HCE.ComponentId -- ^ Current build component id
  -> (T.Text, HCE.SourceCodePreprocessing) -- ^ Source code
  -> (HCE.ModuleInfo, ModuleDependencies, [TypeError])
createModuleInfo (fileMap, defSiteMap, moduleNameMap) (flags, unitState, typecheckedModule, homePackageTable, externalPackageState, modSum) modulePath currentPackageId compId (originalSourceCode, sourceCodePreprocessing)
  = let
      globalRdrEnv = tcg_rdr_env . fst . tm_internals_ $ typecheckedModule
      modInfo = moduleInfo typecheckedModule
      (hsGroup, _, _, _) = fromJust $ renamedSource typecheckedModule
      exportedNamesSet = S.fromList $ modInfoExportsWithSelectors modInfo
      --------------------------------------------------------------------------------
      -- Preprocessed source
      --------------------------------------------------------------------------------
      (transformation, sourceCode') = prepareSourceCode
        sourceCodePreprocessing
        originalSourceCode
        modSum
        modulePath
      includedFiles         = HM.keys $ HCE.fileIndex transformation
      --------------------------------------------------------------------------------
      -- Type environment
      --------------------------------------------------------------------------------
      (tcGblEnv, _)         = tm_internals_ typecheckedModule
      currentModuleTyThings = typeEnvElts $ tcg_type_env tcGblEnv
      homePackageTyThings   = concatMap (typeEnvElts . md_types . hm_details)
        $ eltsUDFM homePackageTable
      externalPackagesTyThings = typeEnvElts $ eps_PTE externalPackageState
      typeEnv                  = mkTypeEnv
        (  currentModuleTyThings
        ++ homePackageTyThings
        ++ externalPackagesTyThings
        )
      --------------------------------------------------------------------------------
      -- Exported entities
      --------------------------------------------------------------------------------
      dataFamTyCons = mapMaybe
        (\case
          ATyCon tc | isFamInstTyCon tc -> Just $ tyConName tc
          _                             -> Nothing
        )
        currentModuleTyThings
      (defSites, allNames) = createDefinitionSiteMap flags
                                                     unitState
                                                     currentPackageId
                                                     compId
                                                     defSiteMap
                                                     fileMap
                                                     globalRdrEnv
                                                     transformation
                                                     modInfo
                                                     dataFamTyCons
                                                     hsGroup
      --------------------------------------------------------------------------------
      -- Instance environment
      --------------------------------------------------------------------------------
      homeInstEnv      = tcg_inst_env tcGblEnv
      visOrphanModules = tcVisibleOrphanMods tcGblEnv
      packageInstEnv   = eps_inst_env externalPackageState
      instEnv          = InstEnvs packageInstEnv homeInstEnv visOrphanModules
      --------------------------------------------------------------------------------
      declarations     = createDeclarations flags
                                            hsGroup
                                            typeEnv
                                            exportedNamesSet
                                            transformation
      environment = Environment { envDynFlags              = flags
                                , envUnitState             = unitState
                                , envInstEnv               = instEnv
                                , envTypeEnv               = typeEnv
                                , envTransformation        = transformation
                                , envCurrentModuleDefSites = defSites
                                , envFileMap               = fileMap
                                , envDefSiteMap            = defSiteMap
                                , envModuleNameMap         = moduleNameMap
                                , envPackageId             = currentPackageId
                                , envComponentId           = compId
                                , envExportedNames         = exportedNamesSet
                                }
      externalIds = L.foldl'
        (\acc name -> maybe
          acc
          (\id ->
            ( HCE.ExternalIdentifierInfo
              $ mkIdentifierInfo environment id (Just name)
              )
              : acc
          )
          (lookupIdInTypeEnv typeEnv name)
        )
        []
        allNames
      currentModuleName =
        (\(Module _ name) ->
            HCE.HaskellModuleName . T.pack . moduleNameString $ name
          )
          . ms_mod
          $ modSum
      SourceInfo {..} = foldAST environment typecheckedModule
    in
      ( tidyInternalIds HCE.ModuleInfo
        { id                = modulePath
        , transformation    = transformation
        , name              = currentModuleName
        , declarations      = declarations
        , exprInfoMap       = sourceInfoExprMap
        , idInfoMap         = sourceInfoIdMap
        , idOccMap          = sourceInfoIdOccMap
        , definitionSiteMap = defSites
        , source            = V.fromList . T.splitOn "\n" $ sourceCode'
        , externalIds       = externalIds
        }
      , if not $ isHsBoot modulePath
        then
          ( HM.union
            ( HM.fromList
            . (( HCE.HaskellFilePath $ HCE.getHaskellModulePath modulePath
               , modulePath
               ) :
              )
            . map (, modulePath)
            $ includedFiles
            )
            fileMap
          , HM.union (HM.singleton modulePath defSites) defSiteMap
          , HM.insertWith HM.union
                          currentModuleName
                          (HM.singleton compId modulePath)
                          moduleNameMap
          )
        else (fileMap, defSiteMap, moduleNameMap)
      , sourceInfoTypeErrors
      )

data SourceInfo = SourceInfo
  { sourceInfoExprMap    :: HCE.ExpressionInfoMap
  , sourceInfoIdMap      :: HCE.IdentifierInfoMap
  , sourceInfoIdOccMap   :: HCE.IdentifierOccurrenceMap
  , sourceInfoTypeErrors :: [TypeError]
  }
  deriving (Show, Eq)

tidyInternalIds :: HCE.ModuleInfo -> HCE.ModuleInfo
tidyInternalIds modInfo = evalState (U.transformBiM tidy modInfo) (HM.empty, 0)
 where
  tidy
    :: HCE.InternalId -> State (HM.HashMap T.Text T.Text, Int) HCE.InternalId
  tidy (HCE.InternalId text) = do
    (hmap, number) <- get
    case HM.lookup text hmap of
      Just val -> return $ HCE.InternalId val
      Nothing  -> do
        let nextInternalId = T.pack . show $ number
        put (HM.insert text nextInternalId hmap, number + 1)
        return $ HCE.InternalId nextInternalId

prepareSourceCode
  :: HCE.SourceCodePreprocessing
  -> T.Text
  -> ModSummary
  -> HCE.HaskellModulePath
  -> (HCE.SourceCodeTransformation, T.Text)
prepareSourceCode sourceCodePreprocessing originalSourceCode modSum modulePath
  = let sourceCodeAfterPreprocessing =
          case
              TE.decodeUtf8' $ maybe (error "ms_hspp_buf is Nothing")
                                     stringBufferToByteString
                                     (ms_hspp_buf modSum)
            of
              Right text -> T.replace "\t" "        " text
              Left err ->
                error
                  $  "decodeUtf8' : "
                  ++ show err
                  ++ " , file : "
                  ++ show modulePath
    in  case sourceCodePreprocessing of
          HCE.BeforePreprocessing ->
            let sourceCodeLines = T.splitOn "\n" originalSourceCode
            in  ( HCE.SourceCodeTransformation (length sourceCodeLines)
                                               modulePath
                                               S.empty
                                               HM.empty
                , originalSourceCode
                )
          HCE.AfterPreprocessing -> createSourceCodeTransformation
            modulePath
            originalSourceCode
            sourceCodeAfterPreprocessing

createDefinitionSiteMap
  :: DynFlags
  -> UnitState
  -> HCE.PackageId
  -> HCE.ComponentId
  -> HM.HashMap HCE.HaskellModulePath HCE.DefinitionSiteMap
  -> HM.HashMap HCE.HaskellFilePath HCE.HaskellModulePath
  -> GlobalRdrEnv
  -> HCE.SourceCodeTransformation
  -> ModuleInfo
  -> [Name]
  -> HsGroup GhcRn
  -> (HCE.DefinitionSiteMap, [Name])
createDefinitionSiteMap flags unitState currentPackageId compId defSiteMap fileMap globalRdrEnv transformation modInfo dataFamTyCons hsGroup
  = let
      allDecls :: [LHsDecl GhcRn]
      allDecls = sortLocatedA . ungroup $ hsGroup
      (instanceDeclsWithDocs, valueAndTypeDeclsWithDocs) =
        L.partition
            (\(L _ decl, _) -> case decl of
              InstD{} -> True
              _       -> False
            )
          $ collectDocs allDecls
      --------------------------------------------------------------------------------
      -- Instances
      --------------------------------------------------------------------------------
      -- No type instances or data instances here for now
      instanceDocMap :: M.Map SrcSpan [HsDocString]
      instanceDocMap =
        M.fromList
          . mapMaybe
              (\(L _n decl, docs) -> case decl of
                InstD _ (ClsInstD _ inst) ->
                  Just (clsInstDeclSrcSpan inst, docs)
                _ -> Nothing
              )
          $ instanceDeclsWithDocs
      nameLocation :: Maybe SrcSpan -> Name -> HCE.LocationInfo
      nameLocation = nameLocationInfo unitState
                                      currentPackageId
                                      compId
                                      transformation
                                      fileMap
                                      defSiteMap
                                      Nothing
      docHToHtml :: DocH (ModuleName, OccName) Name -> HCE.HTML
      docHToHtml = docWithNamesToHtml flags
                                      unitState
                                      currentPackageId
                                      compId
                                      transformation
                                      fileMap
                                      defSiteMap
      instancesWithDocumentation =
        HM.fromList
          . map
              (\clsInst ->
                ( instanceToText flags clsInst
                , let location =
                        nameLocation Nothing (varName . is_dfun $ clsInst)
                  in  case M.lookup (getSrcSpan clsInst) instanceDocMap of
                        Just hsDocString -> HCE.DefinitionSite
                          location
                          ( Just
                          . docHToHtml
                          . hsDocsToDocH flags globalRdrEnv
                          $ hsDocString
                          )
                        Nothing -> HCE.DefinitionSite location Nothing
                )
              )
          $ modInfoInstances modInfo -- all instances (including derived)
      --------------------------------------------------------------------------------
      -- Values and types
      --------------------------------------------------------------------------------
      mainDeclNamesWithDocumentation = concatMap
        (\(dec@(L _ decl), docs) ->
          map (, docs, getLocA dec) $ getMainDeclBinder decl
        )
        valueAndTypeDeclsWithDocs
      dataFamTyConsWithoutDocs =
        map (\name -> (name, [], nameSrcSpan name)) dataFamTyCons
      allNamesWithDocumentation =
        mainDeclNamesWithDocumentation
          ++ subordinateNamesWithDocs allDecls
          ++ dataFamTyConsWithoutDocs
      (valuesWithDocumentation, typesWithDocumentation) = L.partition
        (\(name, _doc, _srcSpan) ->
          case occNameNameSpace . nameOccName $ name of
            HCE.VarName  -> True
            HCE.DataName -> True
            _            -> False
        )
        allNamesWithDocumentation
      toHashMap
        :: [(Name, [HsDocString], SrcSpan)]
        -> HM.HashMap HCE.OccName HCE.DefinitionSite
      toHashMap =
        HM.fromListWith
            (\(HCE.DefinitionSite loc newDoc) (HCE.DefinitionSite _ oldDoc) ->
              (HCE.DefinitionSite loc $ mappend newDoc oldDoc)
            )
          . map
              (\(name, docs, srcSpan) ->
                let location = nameLocation (Just srcSpan) name
                    htmlDoc  = if not . null $ docs
                      then
                        Just
                        . docHToHtml
                        . hsDocsToDocH flags globalRdrEnv
                        $ docs
                      else Nothing
                in  ( HCE.OccName $ toText flags name
                    , HCE.DefinitionSite location htmlDoc
                    )
              )
        --------------------------------------------------------------------------------
    in
      ( HCE.DefinitionSiteMap
        { HCE.values    = toHashMap valuesWithDocumentation
        , HCE.types     = toHashMap
                          $  typesWithDocumentation
                          ++ dataFamTyConsWithoutDocs
        , HCE.instances = instancesWithDocumentation
        }
      , map (\(n, _, _) -> n) allNamesWithDocumentation
      )

occNameToHtml
  :: DynFlags
  -> HCE.PackageId
  -> HCE.ComponentId
  -> (ModuleName, OccName)
  -> H.Html
occNameToHtml flags packageId compId (modName, occName) =
  let location =
        H.textValue
          . toStrict
          . encodeToLazyText
          . Aeson.toJSON
          $ occNameLocationInfo flags packageId compId (modName, occName)
  in  (H.span H.! H.dataAttribute "location" location)
        H.! A.class_ "link"
        $   H.toHtml (toText flags occName)

nameToHtml
  :: UnitState
  -> HCE.PackageId
  -> HCE.ComponentId
  -> HCE.SourceCodeTransformation
  -> HM.HashMap HCE.HaskellFilePath HCE.HaskellModulePath
  -> HM.HashMap HCE.HaskellModulePath HCE.DefinitionSiteMap
  -> Name
  -> H.Html
nameToHtml unitState packageId compId transformation files defSiteMap name =
  let location =
        H.textValue
          . toStrict
          . encodeToLazyText
          . Aeson.toJSON
          $ nameLocationInfo unitState
                             packageId
                             compId
                             transformation
                             files
                             defSiteMap
                             Nothing
                             Nothing
                             name
  in  H.span
        H.! H.dataAttribute "location" location
        H.! A.class_ "link"
        $   H.toHtml (nameToText name)

docWithNamesToHtml
  :: DynFlags
  -> UnitState
  -> HCE.PackageId
  -> HCE.ComponentId
  -> HCE.SourceCodeTransformation
  -> HM.HashMap HCE.HaskellFilePath HCE.HaskellModulePath
  -> HM.HashMap HCE.HaskellModulePath HCE.DefinitionSiteMap
  -> DocH (ModuleName, OccName) Name
  -> HCE.HTML
docWithNamesToHtml flags unitState packageId compId transformation fileMap defSiteMap
  = HCE.docToHtml
    (occNameToHtml flags packageId compId)
    (nameToHtml unitState packageId compId transformation fileMap defSiteMap)

createDeclarations
  :: DynFlags
  -> HsGroup GhcRn
  -> TypeEnv
  -> S.Set Name
  -> HCE.SourceCodeTransformation
  -> [HCE.Declaration]
createDeclarations flags hsGroup typeEnv exportedSet transformation =
  let lineNumber :: SrcSpan -> Int
      lineNumber srcSpan =
        case srcSpanToLineAndColNumbers transformation srcSpan of
          Just (_file, (lineNum, _), (_, _)) -> lineNum
          Nothing                            -> 1
      nameType :: Name -> Maybe HCE.Type
      nameType n = case lookupIdInTypeEnv typeEnv n of
        Just i  -> Just . mkType flags . varType $ i
        Nothing -> Nothing
      -- | Top-level functions
      --------------------------------------------------------------------------------
      valToDeclarations :: LHsBindLR GhcRn GhcRn -> [HCE.Declaration]
      valToDeclarations lb@(L _ bind) =
        map
            (\name -> HCE.Declaration HCE.ValD
                                      (toText flags name)
                                      (nameType name)
                                      (S.member name exportedSet)
                                      (lineNumber (getLocA lb))
            )
          $ collectHsBindBinders CollNoDictBinders bind
      vals = concatMap valToDeclarations $ hsGroupVals hsGroup
      -- | Data, newtype, type, type family, data family or class declaration
      --------------------------------------------------------------------------------
      tyClToDeclaration :: LTyClDecl GhcRn -> HCE.Declaration
      tyClToDeclaration lt@(L _ tyClDecl) = HCE.Declaration
        HCE.TyClD
        (T.append (tyClDeclPrefix tyClDecl) (toText flags $ tcdName tyClDecl))
        (nameType $ tcdName tyClDecl)
        (S.member (unLoc $ tyClDeclLName tyClDecl) exportedSet)
        (lineNumber (getLocA lt))
      tyclds =
        map tyClToDeclaration
          . filter (isGoodSrcSpan . getLocA)
          . concatMap group_tyclds
          . hs_tyclds
          $ hsGroup
      -- | Instances
      --------------------------------------------------------------------------------
      instToDeclaration :: LInstDecl GhcRn -> HCE.Declaration
      instToDeclaration li@(L _ inst) = HCE.Declaration
        HCE.InstD
        (instanceDeclToText flags inst)
        Nothing
        True
        (lineNumber (getLocA li))
      insts =
        map instToDeclaration
          . filter (isGoodSrcSpan . getLocA)
          . hsGroupInstDecls
          $ hsGroup
      -- | Foreign functions
      --------------------------------------------------------------------------------
      foreignFunToDeclaration :: LForeignDecl GhcRn -> HCE.Declaration
      foreignFunToDeclaration lf@(L _ fd) =
        let name = unLoc $ fd_name fd
        in  HCE.Declaration HCE.ForD
                            (toText flags name)
                            (nameType name)
                            True
                            (lineNumber (getLocA lf))
      fords = map foreignFunToDeclaration $ hs_fords hsGroup
      --------------------------------------------------------------------------------
  in  L.sortOn (\HCE.Declaration { lineNumber = lineNo } -> lineNo)
        $  vals
        ++ tyclds
        ++ insts
        ++ fords

foldAST :: Environment -> TypecheckedModule -> SourceInfo
foldAST environment typecheckedModule =
  let
    renamed@(_, importDecls, mbExported, _) =
      fromJust $ renamedSource typecheckedModule
    emptyASTState =
      ASTState IVM.empty IM.empty M.empty emptyTidyEnv Nothing environment []
    ASTState {..} = execState
      (foldTypecheckedSource $ tm_typechecked_source typecheckedModule)
      emptyASTState
    -- A few things that are not in the output of the typechecker:
    --     - the export list
    --     - the imports
    --     - type signatures
    --     - type/data/newtype declarations
    --     - class declarations

    -- Both typechecked source and renamed source are used to populate
    -- 'IdentifierInfoMap' and 'IdentifierOccurrenceMap'
    (idInfoMap, idOccMap) = L.foldl'
      (addIdentifierToMaps environment astStateIdSrcSpanMap)
      (HM.empty, astStateIdOccMap)
      (namesFromRenamedSource renamed)
    unitState = envUnitState environment
    packageId = envPackageId environment
    compId    = envComponentId environment
    importedModules =
      map
          ( (\lm@(L _ modName) ->
              ( modName
              , getLocA lm
              , moduleLocationInfo unitState
                                   (envModuleNameMap environment)
                                   packageId
                                   compId
                                   modName
              )
            )
          . ideclName
          . unLoc
          )
        . filter (not . ideclImplicit . unLoc)
        $ importDecls
    exportedModules = case mbExported of
      Just lieNames -> mapMaybe
        (\(li@(L _ ie), _) -> case ie of
          IEModuleContents _ (L _ modName) -> Just
            ( modName
            , getLocA li
            , moduleLocationInfo unitState
                                 (envModuleNameMap environment)
                                 packageId
                                 compId
                                 modName
            )
          _ -> Nothing
        )
        lieNames
      Nothing -> []
    addImportedAndExportedModulesToIdOccMap
      :: HCE.IdentifierOccurrenceMap -> HCE.IdentifierOccurrenceMap
    addImportedAndExportedModulesToIdOccMap =
      IM.map (L.sortOn fst) . addModules (envTransformation environment)
                                         (importedModules ++ exportedModules)
  in
    SourceInfo
      { sourceInfoExprMap    = astStateExprInfoMap
      , sourceInfoIdMap      = idInfoMap
      , sourceInfoIdOccMap   = addImportedAndExportedModulesToIdOccMap idOccMap
      , sourceInfoTypeErrors = astStateTypeErrors
      }

-- | Updates 'IdentifierOccurrenceMap' and 'IdentifierInfoMap' using information
-- from typechecked source and renamed source
addIdentifierToMaps
  :: Environment
  -> M.Map SrcSpan (Id, Maybe (Type, [Type]))
  -> (HCE.IdentifierInfoMap, HCE.IdentifierOccurrenceMap)
  -> NameOccurrence
  -> (HCE.IdentifierInfoMap, HCE.IdentifierOccurrenceMap)
addIdentifierToMaps environment idSrcSpanMap idMaps@(idInfoMap, idOccMap) nameOcc
  | isGoodSrcSpan (getLoc $ locatedName nameOcc)
    && isOneLineSpan (getLoc $ locatedName nameOcc)
  , Just (_, (lineNumber, startCol), (_, endCol)) <-
    srcSpanToLineAndColNumbers (envTransformation environment)
    . getLoc
    . locatedName
    $ nameOcc
  = case nameOcc of
    TyLitOccurrence { kind = kind } -> addNameToMaps environment
                                                     idMaps
                                                     (Just kind)
                                                     Nothing
                                                     (description nameOcc)
                                                     lineNumber
                                                     startCol
                                                     endCol
    NameOccurrence { isBinder = isBinder } ->
      case lookupIdByNameOccurrence environment idSrcSpanMap nameOcc of
        Just (identifier, mbTypes) ->
          let name =
                fromMaybe (varName identifier) (unLoc $ locatedName nameOcc)
              identifierType         = varType identifier
              identifierTypeExpanded = expandTypeSynonyms identifierType
              tyConsAndTyVars        = map
                (, Nothing)
                (  tyConsOfType identifierType
                ++ tyVarsOfType identifierType
                ++ tyConsOfType identifierTypeExpanded
                ++ tyVarsOfType identifierTypeExpanded
                ++ maybe [] (tyConsOfType . fst) mbTypes
                ++ maybe [] (tyVarsOfType . fst) mbTypes
                )
              idInfoMap' = updateIdMap
                environment
                ((identifier, unLoc $ locatedName nameOcc) : tyConsAndTyVars)
                idInfoMap
              idOcc = mkIdentifierOccurrence environment
                                             identifier
                                             name
                                             mbTypes
                                             isBinder
                                             (description nameOcc)
              idOccMap' = IM.insertWith removeOverlappingInterval
                                        lineNumber
                                        [((startCol, endCol), idOcc)]
                                        idOccMap
          in  (idInfoMap', idOccMap')
        -- type variable or an internal identifier in a pattern synonym
        Nothing -> case unLoc $ locatedName nameOcc of
          Just name -> addNameToMaps environment
                                     idMaps
                                     Nothing
                                     (Just name)
                                     (description nameOcc)
                                     lineNumber
                                     startCol
                                     endCol
          Nothing -> idMaps
addIdentifierToMaps _ _ idMaps _ = idMaps

addNameToMaps
  :: Environment
  -> (HCE.IdentifierInfoMap, HCE.IdentifierOccurrenceMap)
  -> Maybe Type
  -> Maybe Name
  -> T.Text
  -> Int
  -> Int
  -> Int
  -> ( HCE.IdentifierInfoMap
     , HCE.IdentifierOccurrenceMap
     )
addNameToMaps environment (idInfoMap, idOccMap) mbKind mbName descr lineNumber colStart colEnd
  = let flags      = envDynFlags environment
        idInfoMap' = maybe
          idInfoMap
          (\kind -> updateIdMap
            environment
            (map (, Nothing) $ tyConsOfType kind ++ tyVarsOfType kind)
            idInfoMap
          )
          mbKind
        idOcc = HCE.IdentifierOccurrence
          { internalId                  = fmap (HCE.InternalId . nameKey) mbName
          , internalIdFromRenamedSource = HCE.InternalId
                                          .   T.pack
                                          .   show
                                          .   getKey
                                          .   nameUnique
                                          <$> mbName
          , isBinder                    = False
          , instanceResolution          = Nothing
          , idOccType                   = mkType flags <$> mbKind
          , typeArguments               = Nothing
          , description                 = descr
          , sort                        = maybe
                                            HCE.TypeId
                                            (\name -> case occNameNameSpace . nameOccName $ name of
                                              HCE.VarName  -> HCE.ValueId
                                              HCE.DataName -> HCE.ValueId
                                              _            -> HCE.TypeId
                                            )
                                            mbName
          }
        idOccMap' = IM.insertWith removeOverlappingInterval
                                  lineNumber
                                  [((colStart, colEnd), idOcc)]
                                  idOccMap
    in  (idInfoMap', idOccMap')

lookupIdByNameOccurrence
  :: Environment
  -> M.Map SrcSpan (Id, Maybe (Type, [Type]))
  -> NameOccurrence
  -> Maybe (Id, Maybe (Type, [Type]))
lookupIdByNameOccurrence environment idSrcSpanMap (NameOccurrence (L span mbName) _ _)
  = case M.lookup span idSrcSpanMap of
    Just (identifier, mbTypes) -> Just (identifier, mbTypes)
    Nothing                    -> case mbName of
      Just name -> case M.lookup (nameSrcSpan name) idSrcSpanMap of
          -- LHS of a Match
        Just (identifier, mbTypes) -> Just (identifier, mbTypes)
        Nothing ->
          -- Things that are not in the typechecked source
                   case lookupIdInTypeEnv (envTypeEnv environment) name of
          Just t  -> Just (t, Nothing)
          Nothing -> Nothing
      Nothing -> Nothing
lookupIdByNameOccurrence _ _ TyLitOccurrence{} = Nothing

updateIdMap
  :: Environment
  -> [(Id, Maybe Name)]
  -> HCE.IdentifierInfoMap
  -> HCE.IdentifierInfoMap
updateIdMap environment ids identifiersMap =
  let flags = envDynFlags environment
      update
        :: HCE.IdentifierInfoMap -> (Id, Maybe Name) -> HCE.IdentifierInfoMap
      update idMap (identifier, mbName) =
        let info = mkIdentifierInfo environment identifier mbName
        in  HM.insertWith (flip const)
                          (HCE.InternalId $ identifierKey flags identifier)
                          info
                          idMap
  in  L.foldl' update identifiersMap ids

addModules
  :: HCE.SourceCodeTransformation
  -> [(ModuleName, SrcSpan, HCE.LocationInfo)]
  -> HCE.IdentifierOccurrenceMap
  -> HCE.IdentifierOccurrenceMap
addModules transformation modules idMap =
  let update
        :: IM.IntMap [((Int, Int), HCE.IdentifierOccurrence)]
        -> (a, SrcSpan, HCE.LocationInfo)
        -> IM.IntMap [((Int, Int), HCE.IdentifierOccurrence)]
      update idOccMap (_modInfo, span, locInfo)
        | Just (_file, (lineNumber, colStart), (_, colEnd)) <-
          srcSpanToLineAndColNumbers transformation span
        = let idOcc = HCE.IdentifierOccurrence
                { internalId                  = Nothing
                , internalIdFromRenamedSource = Nothing
                , isBinder                    = False
                , instanceResolution          = Nothing
                , idOccType                   = Nothing
                , typeArguments               = Nothing
                , description                 = "ImportDecl"
                , sort                        = HCE.ModuleId locInfo
                }
          in  IM.insertWith removeOverlappingInterval
                            lineNumber
                            [((colStart, colEnd), idOcc)]
                            idOccMap
      update idOccMap _ = idOccMap
  in  L.foldl' update idMap modules