packages feed

glean-hs-0.1.0: haskell/src/Glean/Indexer/HIE.hs

-- | HIE file indexer for glean-hs.
--
-- Reads GHC HIE (Haskell Interface Extended) files and converts
-- the semantic information into Glean facts stored via 'Storage'.
--
-- HIE files are generated by GHC when compiled with:
--   ghc -fwrite-ide-info -hiedir=.hie
-- or via cabal:
--   ghc-options: -fwrite-ide-info
--
-- We use HIE files rather than parsing Haskell source because:
--   * GHC is the authoritative Haskell parser
--   * HIE files contain fully resolved types and names
--   * Zero maintenance burden as GHC evolves
--   * More information than source parsing can produce
--
-- Reference implementation: Calligraphy library
-- https://hackage.haskell.org/package/calligraphy

module Glean.Indexer.HIE
  ( indexHieFile
  , indexHieDirectory
  , indexProject
  , IndexConfig (..)
  , defaultIndexConfig
  , IndexResult (..)
  , IndexStats (..)
  ) where

import Control.Exception (try, SomeException)
import Control.Monad (forM, forM_, when)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Builder as Builder
import Data.IORef
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Maybe (mapMaybe, fromMaybe)
import qualified Data.Set as Set
import Data.Text (Text)
import qualified Data.Text as Text
import qualified Data.Text.Encoding as Text
import Data.Word (Word64)
import System.Directory (listDirectory, doesFileExist, doesDirectoryExist)
import System.FilePath ((</>), takeExtension)

-- GHC HIE file reading (GHC 9.12 API)
import GHC.Iface.Ext.Types
  ( HieFile (..)
  , HieASTs (..)
  , HieAST (..)
  , NodeInfo (..)
  , SourcedNodeInfo (..)
  , Identifier
  , IdentifierDetails (..)
  , ContextInfo (..)

  , getAsts
  )
import GHC.Iface.Ext.Binary (readHieFile, HieFileResult(..))
import GHC.Types.Name.Cache (initNameCache)
import GHC.Types.Name (nameOccName, nameModule_maybe, occNameString)
import GHC.Types.SrcLoc
  ( RealSrcSpan
  , realSrcSpanStart, realSrcSpanEnd
  , srcLocLine, srcLocCol
  )
import GHC.Unit.Module (moduleName, moduleNameString)

import Glean.Indexer.Types
import Glean.Storage hiding (emptyStats)

-- Configuration

data IndexConfig = IndexConfig
  { cfgHieDir   :: FilePath
  , cfgVerbose  :: Bool
  , cfgMaxFiles :: Maybe Int
  } deriving (Show, Eq)

defaultIndexConfig :: IndexConfig
defaultIndexConfig = IndexConfig
  { cfgHieDir   = ".hie"
  , cfgVerbose  = False
  , cfgMaxFiles = Nothing
  }

-- Results

data IndexStats = IndexStats
  { statsFilesIndexed :: !Int
  , statsDefsFound    :: !Int
  , statsRefsFound    :: !Int
  , statsModulesFound :: !Int
  , statsImportsFound :: !Int
  , statsErrors       :: !Int
  } deriving (Show, Eq)

data IndexResult = IndexResult
  { resultStats   :: !IndexStats
  , resultModules :: ![IndexedModule]
  } deriving (Show, Eq)

emptyStats :: IndexStats
emptyStats = IndexStats 0 0 0 0 0 0

-- HIE file reading

indexHieFile :: FilePath -> IO (Maybe IndexedModule)
indexHieFile path = do
  exists <- doesFileExist path
  if not exists
    then return Nothing
    else do
      result <- try (readHieFileRaw path) :: IO (Either SomeException HieFile)
      case result of
        Left  _   -> return Nothing
        Right hie -> return $ Just (extractFacts hie)

readHieFileRaw :: FilePath -> IO HieFile
readHieFileRaw path = do
  nameCache <- initNameCache 'a' []
  result    <- readHieFile nameCache path
  return (hie_file_result result)

-- Fact extraction

extractFacts :: HieFile -> IndexedModule
extractFacts hie =
  let modName = Text.pack
              $ moduleNameString
              $ moduleName
              $ hie_module hie
      srcFile = SrcFile $ Text.pack $ hie_hs_file hie
      modFact = ModuleFact { modName = modName, modFile = srcFile }
      asts    = Map.elems $ getAsts $ hie_asts hie
      defs    = concatMap (extractDefs modName srcFile) asts
      refs    = concatMap (extractRefs modName srcFile) asts
  in IndexedModule
       { idxModule      = modFact
       , idxDefinitions = defs
       , idxReferences  = refs
       , idxImports     = []
       }

-- | Get merged identifiers from a SourcedNodeInfo (GHC 9.12).
-- SourcedNodeInfo is a Map NodeOrigin (NodeInfo a) in GHC 9.12.
getNodeIdentifiers :: SourcedNodeInfo a -> Map Identifier (IdentifierDetails a)
getNodeIdentifiers (SourcedNodeInfo nodeMap) =
  Map.foldl' (\acc ni -> Map.unionWith mergeDetails acc (nodeIdentifiers ni))
             Map.empty
             nodeMap

mergeDetails :: IdentifierDetails a -> IdentifierDetails a -> IdentifierDetails a
mergeDetails d1 d2 = d1
  { identInfo = Set.union (identInfo d1) (identInfo d2) }

extractDefs :: Text -> SrcFile -> HieAST a -> [DefinitionFact]
extractDefs modName srcFile node =
  let glSpan   = convertSpan srcFile (nodeSpan node)
      idents   = getNodeIdentifiers (sourcedNodeInfo node)
      nodeDefs = mapMaybe (extractDef modName glSpan) (Map.toList idents)
      children = concatMap (extractDefs modName srcFile) (nodeChildren node)
  in nodeDefs ++ children

extractDef :: Text -> SrcSpan -> (Identifier, IdentifierDetails a)
           -> Maybe DefinitionFact
extractDef modName span (ident, details) =
  case ident of
    Left  _ -> Nothing
    Right name ->
      if any isDefinition (Set.toList (identInfo details))
        then Just DefinitionFact
               { defName   = Text.pack $ occNameString $ nameOccName name
               , defModule = modName
               , defSpan   = span
               , defType   = Nothing
               }
        else Nothing

isDefinition :: ContextInfo -> Bool
isDefinition (ValBind _ _ _)     = True
isDefinition (PatternBind _ _ _) = True
isDefinition (Decl _ _)          = True
isDefinition TyDecl               = True
isDefinition (ClassTyDecl _)      = True
isDefinition _                    = False

extractRefs :: Text -> SrcFile -> HieAST a -> [ReferenceFact]
extractRefs modName srcFile node =
  let glSpan   = convertSpan srcFile (nodeSpan node)
      idents   = getNodeIdentifiers (sourcedNodeInfo node)
      nodeRefs = mapMaybe (extractRef modName glSpan) (Map.toList idents)
      children = concatMap (extractRefs modName srcFile) (nodeChildren node)
  in nodeRefs ++ children

extractRef :: Text -> SrcSpan -> (Identifier, IdentifierDetails a)
           -> Maybe ReferenceFact
extractRef modName span (ident, details) =
  case ident of
    Left  _ -> Nothing
    Right name ->
      if any isReference (Set.toList (identInfo details))
        then Just ReferenceFact
               { refName   = Text.pack $ occNameString $ nameOccName name
               , refModule = modName
               , refSpan   = span
               , refTarget = fmap (Text.pack . moduleNameString . moduleName)
                           $ nameModule_maybe name
               }
        else Nothing

isReference :: ContextInfo -> Bool
isReference Use          = True
isReference (IEThing _) = True
isReference _            = False

convertSpan :: SrcFile -> RealSrcSpan -> SrcSpan
convertSpan file rss = SrcSpan
  { spanFile  = file
  , spanStart = SrcPos
      { posLine = srcLocLine $ realSrcSpanStart rss
      , posCol  = srcLocCol  $ realSrcSpanStart rss
      }
  , spanEnd   = SrcPos
      { posLine = srcLocLine $ realSrcSpanEnd rss
      , posCol  = srcLocCol  $ realSrcSpanEnd rss
      }
  }

-- Directory indexing

indexHieDirectory :: IndexConfig -> IO IndexResult
indexHieDirectory config = do
  files <- findHieFiles (cfgHieDir config)
  let files' = maybe files (`take` files) (cfgMaxFiles config)
  when (cfgVerbose config) $
    putStrLn $ "Found " ++ show (length files') ++ " HIE files"
  statsRef   <- newIORef emptyStats
  modulesRef <- newIORef []
  forM_ files' $ \f -> do
    when (cfgVerbose config) $ putStrLn $ "Indexing: " ++ f
    result <- indexHieFile f
    case result of
      Nothing -> modifyIORef' statsRef $ \s ->
                   s { statsErrors = statsErrors s + 1 }
      Just m  -> do
        modifyIORef' statsRef $ \s -> s
          { statsFilesIndexed = statsFilesIndexed s + 1
          , statsDefsFound    = statsDefsFound    s + length (idxDefinitions m)
          , statsRefsFound    = statsRefsFound    s + length (idxReferences  m)
          , statsModulesFound = statsModulesFound s + 1
          , statsImportsFound = statsImportsFound s + length (idxImports     m)
          }
        modifyIORef' modulesRef (m :)
  stats   <- readIORef statsRef
  modules <- readIORef modulesRef
  return IndexResult
    { resultStats   = stats
    , resultModules = reverse modules
    }

findHieFiles :: FilePath -> IO [FilePath]
findHieFiles dir = do
  exists <- doesDirectoryExist dir
  if not exists
    then return []
    else do
      entries <- listDirectory dir
      results <- forM entries $ \entry -> do
        let path = dir </> entry
        isDir  <- doesDirectoryExist path
        isFile <- doesFileExist path
        if isDir
          then findHieFiles path
          else if isFile && takeExtension path == ".hie"
               then return [path]
               else return []
      return $ concat results

-- Project indexing

indexProject :: Storage s => s -> IndexConfig -> IO IndexStats
indexProject db config = do
  result <- indexHieDirectory config
  when (cfgVerbose config) $ do
    let s = resultStats result
    putStrLn $ "Indexing complete:"
    putStrLn $ "  Files:       " ++ show (statsFilesIndexed s)
    putStrLn $ "  Definitions: " ++ show (statsDefsFound    s)
    putStrLn $ "  References:  " ++ show (statsRefsFound    s)
    putStrLn $ "  Modules:     " ++ show (statsModulesFound s)
    putStrLn $ "  Errors:      " ++ show (statsErrors       s)
  forM_ (resultModules result) $ \m ->
    Glean.Storage.store db (serializeModule m)
  return (resultStats result)

-- Serialization

serializeModule :: IndexedModule -> FactBatch
serializeModule m =
  let defBytes = foldMap serializeDef (idxDefinitions m)
      refBytes = foldMap serializeRef (idxReferences  m)
      modBytes = serializeMod         (idxModule       m)
      impBytes = foldMap serializeImp (idxImports      m)
      allBytes = Builder.toLazyByteString
               $ defBytes <> refBytes <> modBytes <> impBytes
      count    = length (idxDefinitions m)
               + length (idxReferences  m)
               + 1
               + length (idxImports     m)
  in FactBatch
       { batchData       = BS.toStrict allBytes
       , batchFirstId    = 1024
       , batchCount      = count
       , batchPredicates = Map.fromList
           [ (pidDefinition, length (idxDefinitions m))
           , (pidReference,  length (idxReferences  m))
           , (pidModule,     1)
           , (pidImport,     length (idxImports      m))
           ]
       }

-- | Wrap a fact body with pid + length header.
-- Format: word64LE(pid) + word32LE(body_len) + body_bytes
-- This allows the deserializer to find fact boundaries.
wrapFact :: Word64 -> Builder.Builder -> Builder.Builder
wrapFact pid body =
  let bodyBytes = BS.toStrict $ Builder.toLazyByteString body
  in Builder.word64LE pid <>
     Builder.word32LE (fromIntegral (BS.length bodyBytes)) <>
     Builder.byteString bodyBytes

serializeDef :: DefinitionFact -> Builder.Builder
serializeDef def = wrapFact pidDefinition $
  encodeText (defName   def) <>
  encodeText (defModule def) <>
  encodeSpan (defSpan   def)

serializeRef :: ReferenceFact -> Builder.Builder
serializeRef ref = wrapFact pidReference $
  encodeText (refName   ref) <>
  encodeText (refModule ref) <>
  encodeSpan (refSpan   ref) <>
  encodeText (fromMaybe Text.empty (refTarget ref))

serializeMod :: ModuleFact -> Builder.Builder
serializeMod m = wrapFact pidModule $
  encodeText (modName m) <>
  encodeText (srcFilePath (modFile m))

serializeImp :: ImportFact -> Builder.Builder
serializeImp imp = wrapFact pidImport $
  encodeText (impFrom   imp) <>
  encodeText (impTarget imp) <>
  Builder.word8 (if impQualified imp then 1 else 0) <>
  encodeText (fromMaybe Text.empty (impAlias imp))

encodeText :: Text -> Builder.Builder
encodeText t =
  let bs = Text.encodeUtf8 t
  in Builder.word32LE (fromIntegral (BS.length bs)) <>
     Builder.byteString bs

encodeSpan :: SrcSpan -> Builder.Builder
encodeSpan span =
  encodeText (srcFilePath (spanFile span)) <>
  Builder.word32LE (fromIntegral (posLine (spanStart span))) <>
  Builder.word32LE (fromIntegral (posCol  (spanStart span))) <>
  Builder.word32LE (fromIntegral (posLine (spanEnd   span))) <>
  Builder.word32LE (fromIntegral (posCol  (spanEnd   span)))