Add refs.is_generated to distinguis references from source/generated by ghc

Author: jhrcek

src/HieDb/Create.hs

@@ -4,6 +4,8 @@
 {-# LANGUAGE ViewPatterns #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE LambdaCase #-}
+{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}
+{-# HLINT ignore "Use camelCase" #-}
 module HieDb.Create where
 
 import Prelude hiding (mod)
@@ -37,7 +39,7 @@ import HieDb.Types
 import HieDb.Utils
 
 sCHEMA_VERSION :: Integer
-sCHEMA_VERSION = 8
+sCHEMA_VERSION = 9
 
 dB_VERSION :: Integer
 dB_VERSION = read (show sCHEMA_VERSION ++ "999" ++ show hieVersion)
@@ -117,6 +119,7 @@ initConn (getConn -> conn) = do
                 \, sc   INTEGER NOT NULL \
                 \, el   INTEGER NOT NULL \
                 \, ec   INTEGER NOT NULL \
+                \, is_generated BOOLEAN NOT NULL \
                 \, FOREIGN KEY(hieFile) REFERENCES mods(hieFile) ON UPDATE CASCADE ON DELETE CASCADE DEFERRABLE INITIALLY DEFERRED \
                 \)"
   execute_ conn "CREATE INDEX IF NOT EXISTS refs_mod ON refs(hieFile)"
@@ -331,27 +334,39 @@ addRefsFromLoaded_unsafe
       mod    = moduleName smod
       uid    = moduleUnit smod
       smod   = hie_module hf
-      refmap = generateReferencesMap $ getAsts $ hie_asts hf
+      asts   = getAsts $ hie_asts hf
+      refmapAll = generateReferencesMap asts
+      refmapSourceOnly = generateReferencesMap $ fmap (dropNodeInfos GeneratedInfo) asts
+      refmapGeneratedOnly = generateReferencesMap $ fmap (dropNodeInfos SourceInfo) asts
       (srcFile, isReal) = case sourceFile of
         RealFile f -> (Just f, True)
         FakeFile mf -> (mf, False)
       modrow = HieModuleRow path (ModuleInfo mod uid isBoot srcFile isReal hash)
 
+      -- We want to distinguish between references from source (NodeOrigin is SourceInfo)
+      -- vs. generated by compiler (NodeOrigin is GeneratedInfo).
+      -- Unfortunately generateReferencesMap throws away the info about NodeOrigin,
+      -- so we need to use this to preprocess the ASTs from which the references map is generated.
+      dropNodeInfos :: NodeOrigin -> HieAST a -> HieAST a
+      dropNodeInfos originToDrop (Node (SourcedNodeInfo sniMap) sp children) =
+        let sourceOnlyNodeInfo = SourcedNodeInfo $ M.delete originToDrop sniMap
+        in Node sourceOnlyNodeInfo sp (map (dropNodeInfos originToDrop) children)
+
   execute conn "INSERT INTO mods VALUES (?,?,?,?,?,?,?)" modrow
 
-  let AstInfo rows decls imports = genAstInfo path smod refmap
+  let AstInfo refsSrc declsSrc importsSrc = genAstInfo path smod SourceInfo refmapSourceOnly
+      AstInfo refsGen declsGen importsGen = genAstInfo path smod GeneratedInfo refmapGeneratedOnly
 
   unless (skipRefs skipOptions) $
-    executeMany conn "INSERT INTO refs  VALUES (?,?,?,?,?,?,?,?)" rows
+    executeMany conn "INSERT INTO refs VALUES (?,?,?,?,?,?,?,?,?)" (refsSrc <> refsGen)
   unless (skipDecls skipOptions) $
-    executeMany conn "INSERT INTO decls VALUES (?,?,?,?,?,?,?)" decls
+    executeMany conn "INSERT INTO decls VALUES (?,?,?,?,?,?,?)" (declsSrc <> declsGen)
   unless (skipImports skipOptions) $
-    executeMany conn "INSERT INTO imports VALUES (?,?,?,?,?,?)" imports
+    executeMany conn "INSERT INTO imports VALUES (?,?,?,?,?,?)" (importsSrc <> importsGen)
 
-  let defs = genDefRow path smod refmap
+  let defs = genDefRow path smod refmapAll
   unless (skipDefs skipOptions) $
-    forM_ defs $ \def ->
-      execute conn "INSERT INTO defs VALUES (?,?,?,?,?,?)" def
+      executeMany conn "INSERT INTO defs VALUES (?,?,?,?,?,?)" defs
 
   let exports = generateExports path $ hie_exports hf
   unless (skipExports skipOptions) $

src/HieDb/Types.hs

@@ -150,15 +150,17 @@ data RefRow
   , refSCol :: Int
   , refELine :: Int
   , refECol :: Int
+  , refIsGenerated :: Bool -- ^ True if the reference to this name is generated by GHC (NodeOrigin is GeneratedInfo)
+                           -- False if it comes from the source code (NodeOrigin is SourceInfo)
   }
 
 instance ToRow RefRow where
-  toRow (RefRow a b c d e f g h) = toRow ((a,b,c):.(d,e,f):.(g,h))
+  toRow (RefRow a b c d e f g h i) = toRow ((a,b,c):.(d,e,f):.(g,h,i))
 
 instance FromRow RefRow where
   fromRow = RefRow <$> field <*> field <*> field
                    <*> field <*> field <*> field
-                   <*> field <*> field
+                   <*> field <*> field <*> field
 
 data DeclRow
   = DeclRow
@@ -178,23 +180,23 @@ instance FromRow DeclRow where
   fromRow = DeclRow <$> field <*> field <*> field <*> field
                     <*> field <*> field <*> field
 
-data ImportRow 
-  = ImportRow 
+data ImportRow
+  = ImportRow
     { importSrc :: FilePath
     , importModuleName :: ModuleName
-    , importSLine :: Int 
-    , importSCol :: Int 
-    , importELine :: Int 
-    , importECol :: Int 
+    , importSLine :: Int
+    , importSCol :: Int
+    , importELine :: Int
+    , importECol :: Int
     }
 
-instance FromRow ImportRow where 
-  fromRow = 
-    ImportRow 
-      <$> field <*> field <*> field <*> field 
+instance FromRow ImportRow where
+  fromRow =
+    ImportRow
+      <$> field <*> field <*> field <*> field
       <*> field <*> field
 
-instance ToRow ImportRow where 
+instance ToRow ImportRow where
   toRow (ImportRow a b c d e f) = toRow ((a,b,c,d):.(e,f))
 
 data TypeName = TypeName

src/HieDb/Utils.hs

@@ -179,18 +179,18 @@ instance Semigroup AstInfo where
 instance Monoid AstInfo where
   mempty = AstInfo [] [] []
 
-genAstInfo :: FilePath -> Module -> M.Map Identifier [(Span, IdentifierDetails a)] -> AstInfo
-genAstInfo path smdl refmap = genRows $ flat $ M.toList refmap
+genAstInfo :: FilePath -> Module -> NodeOrigin -> M.Map Identifier [(Span, IdentifierDetails a)] -> AstInfo
+genAstInfo path smdl nodeOrigin refmap = genRows $ flat $ M.toList refmap
   where
+    isGenerated = nodeOrigin == GeneratedInfo
     flat = concatMap (\(a,xs) -> map (a,) xs)
-    genRows = foldMap go
-    go = mkAstInfo
+    genRows = foldMap mkAstInfo
 
     mkAstInfo x = AstInfo (maybeToList $ goRef x) (maybeToList $ goDec x) (maybeToList $ goImport x)
 
     goRef (Right name, (sp,_))
       | Just mod <- nameModule_maybe name = Just $
-          RefRow path occ (moduleName mod) (moduleUnit mod) sl sc el ec
+          RefRow path occ (moduleName mod) (moduleUnit mod) sl sc el ec isGenerated
           where
             occ = nameOccName name
             sl = srcSpanStartLine sp