export map improvements

Author: wz1000

ghcide/src/Development/IDE/Types/Exports.hs

@@ -21,7 +21,7 @@ module Development.IDE.Types.Exports
     updateExportsMapMg
     ) where
 
-import           Control.DeepSeq             (NFData (..))
+import           Control.DeepSeq             (NFData (..), force, ($!!))
 import           Control.Monad
 import           Data.Bifunctor              (Bifunctor (second))
 import           Data.Char                   (isUpper)
@@ -45,6 +45,9 @@ data ExportsMap = ExportsMap
     , getModuleExportsMap :: !(ModuleNameEnv (HashSet IdentInfo))
     }
 
+instance NFData ExportsMap where
+  rnf (ExportsMap a b) = foldOccEnv (\a b -> rnf a `seq` b) (seqEltsUFM rnf b) a
+
 instance Show ExportsMap where
   show (ExportsMap occs mods) =
     unwords [ "ExportsMap { getExportsMap ="
@@ -154,7 +157,7 @@ createExportsMap :: [ModIface] -> ExportsMap
 createExportsMap modIface = do
   let exportList = concatMap doOne modIface
   let exportsMap = mkOccEnv_C (<>) $ map (\(a,_,c) -> (a, c)) exportList
-  ExportsMap exportsMap $ buildModuleExportMap $ map (\(_,b,c) -> (b, c)) exportList
+  force $ ExportsMap exportsMap $ buildModuleExportMap $ map (\(_,b,c) -> (b, c)) exportList -- UFM is lazy, so need to seq
   where
     doOne modIFace = do
       let getModDetails = unpackAvail $ moduleName $ mi_module modIFace
@@ -164,7 +167,7 @@ createExportsMapMg :: [ModGuts] -> ExportsMap
 createExportsMapMg modGuts = do
   let exportList = concatMap doOne modGuts
   let exportsMap = mkOccEnv_C (<>) $ map (\(a,_,c) -> (a, c)) exportList
-  ExportsMap exportsMap $ buildModuleExportMap $ map (\(_,b,c) -> (b, c)) exportList
+  force $ ExportsMap exportsMap $ buildModuleExportMap $ map (\(_,b,c) -> (b, c)) exportList -- UFM is lazy, so need to seq
   where
     doOne mi = do
       let getModuleName = moduleName $ mg_module mi
@@ -188,7 +191,7 @@ createExportsMapHieDb withHieDb = do
         fmap (unwrap mn) <$> withHieDb (\hieDb -> getExportsForModule hieDb mn)
     let idents = concat idents'
     let exportsMap = mkOccEnv_C (<>) (keyWith name idents)
-    return $! ExportsMap exportsMap $ buildModuleExportMap (keyWith identModuleName idents)
+    return $!! ExportsMap exportsMap $ buildModuleExportMap (keyWith identModuleName idents) -- UFM is lazy so need to seq
   where
     unwrap m ExportRow{..} = IdentInfo exportName exportParent m
     keyWith f xs = [(f x, Set.singleton x) | x <- xs]

plugins/hls-refactor-plugin/src/Development/IDE/Plugin/CodeAction.hs

@@ -1037,19 +1037,21 @@ suggestExtendImport exportsMap (L _ HsModule {hsmodImports}) Diagnostic{_range=_
             ]
           | otherwise = []
         lookupExportMap binding mod
-          | Just match <- Map.lookup binding (getExportsMap exportsMap)
+          | let em = getExportsMap exportsMap
+                match1 = lookupOccEnv em (mkVarOrDataOcc binding)
+                match2 = lookupOccEnv em (mkTypeOcc binding)
+          , Just match <- match1 <> match2
           -- Only for the situation that data constructor name is same as type constructor name,
           -- let ident with parent be in front of the one without.
           , sortedMatch <- sortBy (\ident1 ident2 -> parent ident2 `compare` parent ident1) (Set.toList match)
           , idents <- filter (\ident -> moduleNameText ident == mod && (canUseDatacon || not (isDatacon ident))) sortedMatch
-          , (not . null) idents -- Ensure fallback while `idents` is empty
-          , ident <- head idents
+          , (ident:_) <- idents -- Ensure fallback while `idents` is empty
           = Just ident
 
             -- fallback to using GHC suggestion even though it is not always correct
           | otherwise
           = Just IdentInfo
-                { name = mkVarOccFS $ mkFastStringByteString $ T.encodeUtf8 binding
+                { name = mkVarOrDataOcc binding
                 , parent = Nothing
                 , identModuleName  = mkModuleNameFS $ mkFastStringByteString $ T.encodeUtf8 mod}
 #endif
@@ -1452,7 +1454,7 @@ suggestNewOrExtendImportForClassMethod packageExportsMap ps fileContents Diagnos
         "‘([^’]*)’ is not a \\(visible\\) method of class ‘([^’]*)’",
     idents <-
       maybe [] (Set.toList . Set.filter (\x -> fmap occNameText (parent x) == Just className)) $
-        Map.lookup methodName $ getExportsMap packageExportsMap =
+        lookupOccEnv (getExportsMap packageExportsMap) (mkVarOrDataOcc methodName) =
     mconcat $ suggest <$> idents
   | otherwise = []
   where
@@ -1507,7 +1509,7 @@ constructNewImportSuggestions
 constructNewImportSuggestions exportsMap (qual, thingMissing) notTheseModules = nubOrdOn snd
   [ suggestion
   | Just name <- [T.stripPrefix (maybe "" (<> ".") qual) $ notInScope thingMissing]
-  , identInfo <- maybe [] Set.toList $ Map.lookup name (getExportsMap exportsMap)
+  , identInfo <- maybe [] Set.toList $ (lookupOccEnv (getExportsMap exportsMap) (mkVarOrDataOcc name)) <> (lookupOccEnv (getExportsMap exportsMap) (mkTypeOcc name))
   , canUseIdent thingMissing identInfo
   , moduleNameText identInfo `notElem` fromMaybe [] notTheseModules
   , suggestion <- renderNewImport identInfo