From be68e2b486eab28ebc012c2c7a05732ad271a993 Mon Sep 17 00:00:00 2001
From: Ray Shih <rayshih@fb.com>
Date: Wed, 5 May 2021 12:31:00 +0100
Subject: [PATCH] fix unwanted import refinement

---
 .../src/Ide/Plugin/RefineImports.hs           | 25 +++++++++++++------
 .../test/testdata/F.hs                        |  7 ++++++
 .../test/testdata/G.hs                        |  4 +++
 .../test/testdata/WithOverride.expected.hs    |  3 ++-
 .../test/testdata/WithOverride.hs             |  3 ++-
 .../test/testdata/hie.yaml                    |  4 ++-
 6 files changed, 36 insertions(+), 10 deletions(-)
 create mode 100644 plugins/hls-refine-imports-plugin/test/testdata/F.hs
 create mode 100644 plugins/hls-refine-imports-plugin/test/testdata/G.hs

diff --git a/plugins/hls-refine-imports-plugin/src/Ide/Plugin/RefineImports.hs b/plugins/hls-refine-imports-plugin/src/Ide/Plugin/RefineImports.hs
index d5528e7e84..283f569fdb 100644
--- a/plugins/hls-refine-imports-plugin/src/Ide/Plugin/RefineImports.hs
+++ b/plugins/hls-refine-imports-plugin/src/Ide/Plugin/RefineImports.hs
@@ -22,6 +22,7 @@ import           Data.IORef                           (readIORef)
 import           Data.List                            (intercalate)
 import qualified Data.Map.Strict                      as Map
 import           Data.Maybe                           (catMaybes, fromMaybe)
+import qualified Data.Set                             as S
 import qualified Data.Text                            as T
 import           Data.Traversable                     (forM)
 import           Development.IDE
@@ -197,13 +198,23 @@ refineImportsRule = define $ \RefineImports nfp -> do
   let filterByImport
         :: LImportDecl GhcRn
         -> Map.Map ModuleName [AvailInfo]
-        -> Map.Map ModuleName [AvailInfo]
+        -> Maybe (Map.Map ModuleName [AvailInfo])
       filterByImport (L _ ImportDecl{ideclHiding = Just (_, L _ names)}) avails =
-        let importedNames = map (ieName . unLoc) names
-        in flip Map.filter avails $ \a ->
-              any (`elem` importedNames)
-                $ concatMap availNamesWithSelectors a
-      filterByImport _ _ = mempty
+        let importedNames = S.fromList $ map (ieName . unLoc) names
+            res = flip Map.filter avails $ \a ->
+                    any (`S.member` importedNames)
+                      $ concatMap availNamesWithSelectors a
+            allFilteredAvailsNames = S.fromList
+              $ concatMap availNamesWithSelectors
+              $ mconcat
+              $ Map.elems res
+            -- if there is a function defined in the current module and is used
+            -- i.e. if a function is not reexported but defined in current
+            -- module then this import cannot be refined
+        in if importedNames `S.isSubsetOf` allFilteredAvailsNames
+              then Just res
+              else Nothing
+      filterByImport _ _ = Nothing
   let constructImport
         :: LImportDecl GhcRn
         -> (ModuleName, [AvailInfo])
@@ -229,7 +240,7 @@ refineImportsRule = define $ \RefineImports nfp -> do
         -- we check for the inner imports
         , Just innerImports <- [Map.lookup mn import2Map]
         -- and only get those symbols used
-        , filteredInnerImports <- [filterByImport i innerImports]
+        , Just filteredInnerImports <- [filterByImport i innerImports]
         -- if no symbols from this modules then don't need to generate new import
         , not $ null filteredInnerImports
         ]
diff --git a/plugins/hls-refine-imports-plugin/test/testdata/F.hs b/plugins/hls-refine-imports-plugin/test/testdata/F.hs
new file mode 100644
index 0000000000..9fab4ee9b8
--- /dev/null
+++ b/plugins/hls-refine-imports-plugin/test/testdata/F.hs
@@ -0,0 +1,7 @@
+module F (module F, module G) where
+
+import G
+
+f1 :: String 
+f1 = "f1"
+
diff --git a/plugins/hls-refine-imports-plugin/test/testdata/G.hs b/plugins/hls-refine-imports-plugin/test/testdata/G.hs
new file mode 100644
index 0000000000..634a7f3ed8
--- /dev/null
+++ b/plugins/hls-refine-imports-plugin/test/testdata/G.hs
@@ -0,0 +1,4 @@
+module G where
+
+g1 :: String 
+g1 = "g1"
\ No newline at end of file
diff --git a/plugins/hls-refine-imports-plugin/test/testdata/WithOverride.expected.hs b/plugins/hls-refine-imports-plugin/test/testdata/WithOverride.expected.hs
index c743d4d110..125b6b123d 100644
--- a/plugins/hls-refine-imports-plugin/test/testdata/WithOverride.expected.hs
+++ b/plugins/hls-refine-imports-plugin/test/testdata/WithOverride.expected.hs
@@ -3,9 +3,10 @@ module Main where
 import B ( b1 )
 import C ( c1 )
 import D
+import F
 import Data.List (intercalate)
 
 main :: IO ()
 main = putStrLn 
      $ "hello " 
-    <> intercalate ", " [b1, c1, e1]
+    <> intercalate ", " [b1, c1, e1, f1, g1]
diff --git a/plugins/hls-refine-imports-plugin/test/testdata/WithOverride.hs b/plugins/hls-refine-imports-plugin/test/testdata/WithOverride.hs
index e25fa41bea..66d62989cd 100644
--- a/plugins/hls-refine-imports-plugin/test/testdata/WithOverride.hs
+++ b/plugins/hls-refine-imports-plugin/test/testdata/WithOverride.hs
@@ -2,9 +2,10 @@ module Main where
 
 import A
 import D
+import F
 import Data.List (intercalate)
 
 main :: IO ()
 main = putStrLn 
      $ "hello " 
-    <> intercalate ", " [b1, c1, e1]
+    <> intercalate ", " [b1, c1, e1, f1, g1]
diff --git a/plugins/hls-refine-imports-plugin/test/testdata/hie.yaml b/plugins/hls-refine-imports-plugin/test/testdata/hie.yaml
index 0d1383c68f..4770978cb2 100644
--- a/plugins/hls-refine-imports-plugin/test/testdata/hie.yaml
+++ b/plugins/hls-refine-imports-plugin/test/testdata/hie.yaml
@@ -7,4 +7,6 @@ cradle:
     - B.hs
     - C.hs
     - D.hs
-    - E.hs
\ No newline at end of file
+    - E.hs
+    - F.hs
+    - G.hs
\ No newline at end of file