Add --prefix and --unit options (facebookincubator#626)

Summary:
Most indexers have a `--prefix` for getting the file paths right. `--unit` is useful for stacking DBs when the unit names don't match exactly... I'll leave this as undocumented for now until the workflow is more polished.

Pull Request resolved: facebookincubator#626

Reviewed By: jjuliamolin

Differential Revision: D88849435

Pulled By: phlalx

fbshipit-source-id: 28a0d38803fc56b7c35960de5341557024a7f102

Author: simonmar

glean/lang/haskell/HieIndexer/Index.hs

@@ -58,7 +58,8 @@ import qualified GHC.Types.Var as GHC (ForAllTyFlag(..), Specificity(..))
 #else
 import qualified GHC.Types.Var as GHC (ArgFlag(..), Specificity(..))
 #endif
-import GHC.Unit.Types (unitFS)
+import qualified GHC.Unit.Types as GHC
+import qualified GHC.Types.Unique.Set as GHC (mkUniqSet, elementOfUniqSet)
 #if !MIN_VERSION_ghc(9,6,0)
 import qualified GHC.Unit.Module.Name as GHC (moduleNameFS)
 #endif
@@ -71,6 +72,7 @@ import Glean.Impl.ConfigProvider ()
 import qualified Glean.Schema.Hs.Types as Hs
 import qualified Glean.Schema.Src.Types as Src
 import Glean.Util.Range
+import HieIndexer.Options
 
 {- TODO
 
@@ -96,6 +98,12 @@ import Glean.Util.Range
     Haddock's Interface type has all the ASTs for the declarations
     in addition to the Hie.
 
+- imports
+  - import refs contain only ModuleName, not Module. This means we can't
+    resolve imports to the correct file for module names that occur in
+    multiple packages.
+  - modules in the export list should be refs
+
 - map Name to exportedness?
 
 - exclude generated names in a cleaner way
@@ -108,15 +116,31 @@ import Glean.Util.Range
   - Haddock docs for symbol
 -}
 
-mkModule :: Glean.NewFact m => GHC.Module -> m Hs.Module
-mkModule mod = do
+mkModule :: Glean.NewFact m => GHC.Module -> UnitName -> m Hs.Module
+mkModule mod unit = do
   modname <- Glean.makeFact @Hs.ModuleName $
     fsToText (GHC.moduleNameFS (GHC.moduleName mod))
   unitname <- Glean.makeFact @Hs.UnitName $
-    fsToText (unitFS (GHC.moduleUnit mod))
+    toUnitName unit $ GHC.moduleUnit mod
   Glean.makeFact @Hs.Module $
     Hs.Module_key modname unitname
 
+toUnitName :: UnitName -> GHC.Unit -> Text
+toUnitName unitName u = case unitName of
+  UnitKey -> t
+  UnitId
+    | isWiredInUnit u -> t
+    | (id, _) <- Text.breakOnEnd "-" t, not (Text.null id) -> Text.init id
+    | otherwise -> t
+    -- TODO: this is not right for local libraries, which have names like
+    -- aeson-pretty-0.8.9-inplace-aeson
+  where t = fsToText $ GHC.unitFS u
+
+isWiredInUnit :: GHC.Unit -> Bool
+isWiredInUnit = \u -> GHC.toUnitId u `GHC.elementOfUniqSet` wiredInUnitIds
+  where
+  wiredInUnitIds = GHC.mkUniqSet GHC.wiredInUnitIds
+
 mkName :: Glean.NewFact m => GHC.Name -> Hs.Module -> Hs.NameSort -> m Hs.Name
 mkName name mod sort = do
   let occ = nameOccName name
@@ -254,9 +278,10 @@ nat = Glean.toNat . fromIntegral
 
 indexTypes
   :: forall m . (MonadFail m, Monad m, Glean.NewFact m)
-  => A.Array TypeIndex HieTypeFlat
+  => UnitName
+  -> A.Array TypeIndex HieTypeFlat
   -> m (IntMap Hs.Type)
-indexTypes typeArr = foldM go IntMap.empty (A.assocs typeArr)
+indexTypes unit typeArr = foldM go IntMap.empty (A.assocs typeArr)
   where
   go tymap (n,ty) = do
     fact <- mkTy ty
@@ -314,7 +339,7 @@ indexTypes typeArr = foldM go IntMap.empty (A.assocs typeArr)
         tcname <- case nameModule_maybe name of
           Nothing -> fail "HTyConApp: internal name"
           Just mod -> do
-            namemod <- mkModule mod
+            namemod <- mkModule mod unit
             mkName name namemod (Hs.NameSort_external def)
         let sort = case GHC.ifaceTyConSort info of
               GHC.IfaceNormalTyCon -> Hs.TyConSort_normal def
@@ -347,25 +372,28 @@ indexTypes typeArr = foldM go IntMap.empty (A.assocs typeArr)
 indexHieFile
   :: Glean.Writer
   -> NonEmpty Text
+  -> Maybe Text
+  -> UnitName
   -> FilePath
   -> HieFile
   -> IO ()
-indexHieFile writer srcPaths path hie = do
+indexHieFile writer srcPaths srcPrefix unit path hie = do
   srcFile <- findSourceFile srcPaths (hie_module hie) (hie_hs_file hie)
   logInfo $ "Indexing: " <> path <> " (" <> srcFile <> ")"
   Glean.writeFacts writer $ do
-    modfact <- mkModule smod
+    modfact <- mkModule smod unit
 
     let offs = getLineOffsets (hie_hs_src hie)
     let hsFileFS = GHC.mkFastString $ hie_hs_file hie
-    filefact <- Glean.makeFact @Src.File (Text.pack srcFile)
+    let file = maybe id (\p f -> (p <> "/" <> f)) srcPrefix (Text.pack srcFile)
+    filefact <- Glean.makeFact @Src.File file
     let fileLines = mkFileLines filefact offs
     Glean.makeFact_ @Src.FileLines fileLines
 
     Glean.makeFact_ @Hs.ModuleSource $
       Hs.ModuleSource_key modfact filefact
 
-    typeMap <- indexTypes (hie_types hie)
+    typeMap <- indexTypes unit (hie_types hie)
 
     let toByteRange = srcRangeToByteRange fileLines (hie_hs_src hie)
         toByteSpan sp
@@ -413,7 +441,7 @@ indexHieFile writer srcPaths path hie = do
               -- But it does happen, so let's not crash.
               return Nothing
             Just m -> do
-              mod <- if m == smod then return modfact else mkModule m
+              mod <- if m == smod then return modfact else mkModule m unit
               Just <$> mkName name mod (Hs.NameSort_external def)
 
     sigMap <- produceDeclInfo filefact toByteSpan getName declInfo
@@ -512,7 +540,7 @@ findSourceFile srcPaths mod src = do
       return src
     Just f -> return f
   where
-  pkg = fsToText (unitFS (GHC.moduleUnit mod))
+  pkg = fsToText (GHC.unitFS (GHC.moduleUnit mod))
   isVer = Text.all (\c -> isDigit c || c == '.')
   pkgNameAndVersion = case break isVer (Text.splitOn "-" pkg) of
     (before, after) -> Text.intercalate "-" (before <> take 1 after)

glean/lang/haskell/HieIndexer/Main.hs

@@ -87,7 +87,7 @@ defaultMain cfg repo backend = do
 
   Glean.withSender backend repo allPredicates def $ \sender -> do
     Glean.withWriter sender def $
-      indexHieFiles paths (srcPaths cfg)
+      indexHieFiles paths (srcPaths cfg) (srcPrefix cfg) (unitName cfg)
 
   predicates <-
     Glean.schemaInfo_predicateIds
@@ -119,20 +119,22 @@ outputMain cfg out schema_id backend = do
   paths <- getHieFilesIn (NonEmpty.toList (hiePaths cfg))
   ((), batch) <-
     Glean.withBatchWriter backend schema_id Nothing def $
-      indexHieFiles paths (srcPaths cfg)
+      indexHieFiles paths (srcPaths cfg) (srcPrefix cfg) (unitName cfg)
   BS.writeFile out (Thrift.Protocol.Compact.serializeCompact batch)
 
 indexHieFiles
   :: HashSet.HashSet FilePath
   -> NonEmpty Text
+  -> Maybe Text
+  -> UnitName
   -> Glean.Writer
   -> IO ()
-indexHieFiles paths srcs writer =
+indexHieFiles paths srcs srcPrefix unitName writer =
   forM_ (HashSet.toList paths) $ \f -> do
     nc <- newIORef =<< makeNc
     runDbM nc $ do
       withHieFile f $ \h ->
-        liftIO $ indexHieFile writer srcs f h
+        liftIO $ indexHieFile writer srcs srcPrefix unitName f h
 
 {- | Recursively search for @.hie@ and @.hie-boot@  files in given directory
    avoiding loops due to symlinks

glean/lang/haskell/HieIndexer/Options.hs

@@ -11,6 +11,7 @@
 module HieIndexer.Options (
     HieIndexerOptions(..),
     Mode(..),
+    UnitName(..),
     options,
   ) where
 
@@ -31,9 +32,24 @@ data HieIndexerOptions = HieIndexerOptions
       -- ^ Paths to look for source files. May include the string
       -- @$PACKAGE@ which will be replaced by the package name
       -- (@<pkg>-<version>@).
+  , srcPrefix :: Maybe Text
+      -- ^ Prefix to add to source paths
+  , unitName :: UnitName
+      -- ^ How to handle unit names
   , verbosity :: Int
   }
 
+data UnitName = UnitKey | UnitId
+
+parseUnit :: String -> Maybe UnitName
+parseUnit "key" = Just UnitKey
+parseUnit "id" = Just UnitId
+parseUnit _ = Nothing
+
+showUnit :: UnitName -> String
+showUnit UnitKey = "key"
+showUnit UnitId = "id"
+
 data Mode
   = WriteMode {
       repo :: Repo
@@ -61,6 +77,18 @@ options = info (helper <*> parser) fullDesc
             "Path to search for source files. The string \"$PACKAGE\" is " <>
             "replaced by the package name, e.g. text-2.0.2")))
 
+      srcPrefix <- optional (textOption (
+        long "prefix" <>
+        metavar "PATH" <>
+        help "Prefix to add to source paths in the DB"))
+
+      unitName <- option (maybeReader parseUnit) (
+        long "unit" <>
+        metavar "key|id" <>
+        value UnitKey <>
+        showDefaultWith showUnit <>
+        help "Whether to store the full unit key (key) or just name-version (id)")
+
       verbosity <-
         option
           auto