Fix add suggested import for operators (#428)

Author: pepeiborra

src/Development/IDE/GHC/Util.hs

@@ -9,6 +9,7 @@ module Development.IDE.GHC.Util(
     runGhcEnv,
     -- * GHC wrappers
     prettyPrint,
+    ParseResult(..), runParser,
     lookupPackageConfig,
     moduleImportPath,
     cgGutsToCoreModule,
@@ -47,6 +48,7 @@ import qualified Data.Text                as T
 import qualified Data.Text.Encoding       as T
 import qualified Data.Text.Encoding.Error as T
 import qualified Data.ByteString          as BS
+import Lexer
 import StringBuffer
 import System.FilePath
 
@@ -82,6 +84,13 @@ lookupPackageConfig unitId env =
 textToStringBuffer :: T.Text -> StringBuffer
 textToStringBuffer = stringToStringBuffer . T.unpack
 
+runParser :: DynFlags -> String -> P a -> ParseResult a
+runParser flags str parser = unP parser parseState
+    where
+      filename = "<interactive>"
+      location = mkRealSrcLoc (mkFastString filename) 1 1
+      buffer = stringToStringBuffer str
+      parseState = mkPState flags buffer location
 
 -- | Pretty print a GHC value using 'unsafeGlobalDynFlags '.
 prettyPrint :: Outputable a => a -> String

src/Development/IDE/Plugin/CodeAction.hs

@@ -17,6 +17,7 @@ import Development.IDE.Core.RuleTypes
 import Development.IDE.Core.Service
 import Development.IDE.Core.Shake
 import Development.IDE.GHC.Error
+import Development.IDE.GHC.Util
 import Development.IDE.LSP.Server
 import Development.IDE.Types.Location
 import Development.IDE.Types.Options
@@ -32,9 +33,13 @@ import Data.Maybe
 import Data.List.Extra
 import qualified Data.Text as T
 import Data.Tuple.Extra ((&&&))
+import HscTypes
+import OccName
+import Parser
+import RdrName
 import Text.Regex.TDFA ((=~), (=~~))
 import Text.Regex.TDFA.Text()
-import Outputable (ppr, showSDocUnsafe)
+import Outputable (showSDoc, ppr, showSDocUnsafe)
 import DynFlags (xFlags, FlagSpec(..))
 import GHC.LanguageExtensions.Type (Extension)
 
@@ -54,12 +59,15 @@ codeAction lsp state (TextDocumentIdentifier uri) _range CodeActionContext{_diag
     -- logInfo (ideLogger ide) $ T.pack $ "Code action req: " ++ show arg
     contents <- LSP.getVirtualFileFunc lsp $ toNormalizedUri uri
     let text = Rope.toText . (_text :: VirtualFile -> Rope.Rope) <$> contents
-    (ideOptions, parsedModule) <- runAction state $
-      (,) <$> getIdeOptions
-          <*> (getParsedModule . toNormalizedFilePath) `traverse` uriToFilePath uri
+        mbFile = toNormalizedFilePath <$> uriToFilePath uri
+    (ideOptions, parsedModule, env) <- runAction state $
+      (,,) <$> getIdeOptions
+           <*> getParsedModule `traverse` mbFile
+           <*> use_ GhcSession `traverse` mbFile
+    let dflags = hsc_dflags . hscEnv <$> env
     pure $ Right
         [ CACodeAction $ CodeAction title (Just CodeActionQuickFix) (Just $ List [x]) (Just edit) Nothing
-        | x <- xs, (title, tedit) <- suggestAction ideOptions ( join parsedModule ) text x
+        | x <- xs, (title, tedit) <- suggestAction dflags ideOptions ( join parsedModule ) text x
         , let edit = WorkspaceEdit (Just $ Map.singleton uri $ List tedit) Nothing
         ]
 
@@ -98,10 +106,10 @@ executeAddSignatureCommand _lsp _ideState ExecuteCommandParams{..}
     | otherwise
     = return (Null, Nothing)
 
-suggestAction  :: IdeOptions -> Maybe ParsedModule -> Maybe T.Text -> Diagnostic -> [(T.Text, [TextEdit])]
-suggestAction ideOptions parsedModule text diag = concat
+suggestAction :: Maybe DynFlags -> IdeOptions -> Maybe ParsedModule -> Maybe T.Text -> Diagnostic -> [(T.Text, [TextEdit])]
+suggestAction dflags ideOptions parsedModule text diag = concat
     [ suggestAddExtension diag
-    , suggestExtendImport text diag
+    , suggestExtendImport dflags text diag
     , suggestFillHole diag
     , suggestFillTypeWildcard diag
     , suggestFixConstructorImport text diag
@@ -268,20 +276,23 @@ suggestFillHole Diagnostic{_range=_range,..}
 
     | otherwise = []
 
-suggestExtendImport :: Maybe T.Text -> Diagnostic -> [(T.Text, [TextEdit])]
-suggestExtendImport contents Diagnostic{_range=_range,..}
+suggestExtendImport :: Maybe DynFlags -> Maybe T.Text -> Diagnostic -> [(T.Text, [TextEdit])]
+suggestExtendImport (Just dflags) contents Diagnostic{_range=_range,..}
     | Just [binding, mod, srcspan] <-
       matchRegex _message
       "Perhaps you want to add ‘([^’]*)’ to the import list in the import of ‘([^’]*)’ *\\((.*)\\).$"
     , Just c <- contents
+    , POk _ (L _ name) <- runParser dflags (T.unpack binding) parseIdentifier
     = let range = case [ x | (x,"") <- readSrcSpan (T.unpack srcspan)] of
             [s] -> let x = srcSpanToRange s
                    in x{_end = (_end x){_character = succ (_character (_end x))}}
             _ -> error "bug in srcspan parser"
           importLine = textInRange range c
+          printedName = let rn = rdrNameOcc name in showSDoc dflags $ parenSymOcc rn (ppr rn)
         in [("Add " <> binding <> " to the import list of " <> mod
-        , [TextEdit range (addBindingToImportList binding importLine)])]
+        , [TextEdit range (addBindingToImportList (T.pack printedName) importLine)])]
     | otherwise = []
+suggestExtendImport Nothing _ _ = []
 
 suggestFixConstructorImport :: Maybe T.Text -> Diagnostic -> [(T.Text, [TextEdit])]
 suggestFixConstructorImport _ Diagnostic{_range=_range,..}

test/exe/Main.hs

@@ -752,6 +752,26 @@ extendImportTests = testGroup "extend import actions"
             , "import ModuleA as A (stuffA, stuffB)"
             , "main = print (stuffA, stuffB)"
             ])
+  , testSession "extend single line import with operator" $ template
+      (T.unlines
+            [ "module ModuleA where"
+            , "(.*) :: Integer -> Integer -> Integer"
+            , "x .* y = x * y"
+            , "stuffB :: Integer"
+            , "stuffB = 123"
+            ])
+      (T.unlines
+            [ "module ModuleB where"
+            , "import ModuleA as A (stuffB)"
+            , "main = print (stuffB .* stuffB)"
+            ])
+      (Range (Position 3 17) (Position 3 18))
+      "Add .* to the import list of ModuleA"
+      (T.unlines
+            [ "module ModuleB where"
+            , "import ModuleA as A ((.*), stuffB)"
+            , "main = print (stuffB .* stuffB)"
+            ])
   , testSession "extend single line import with type" $ template
       (T.unlines
             [ "module ModuleA where"