Use exact print to extend import lists

ghcide/src/Development/IDE/GHC/ExactPrint.hs

@@ -74,10 +74,10 @@ instance NFData GetAnnotatedParsedSource
 instance Binary GetAnnotatedParsedSource
 type instance RuleResult GetAnnotatedParsedSource = Annotated ParsedSource
 
--- | Get the latest version of the annotated parse source.
+-- | Get the latest version of the annotated parse source with comments.
 getAnnotatedParsedSourceRule :: Rules ()
 getAnnotatedParsedSourceRule = define $ \GetAnnotatedParsedSource nfp -> do
-  pm <- use GetParsedModule nfp
+  pm <- use GetParsedModuleWithComments nfp
   return ([], fmap annotateParsedSource pm)
 
 annotateParsedSource :: ParsedModule -> Annotated ParsedSource
@@ -314,6 +314,10 @@ instance p ~ GhcPs => ASTElement (HsDecl p) where
     parseAST = parseDecl
     maybeParensAST = id
 
+instance p ~ GhcPs => ASTElement (ImportDecl p) where
+    parseAST = parseImport
+    maybeParensAST = id
+
 instance ASTElement RdrName where
     parseAST df fp = parseWith df fp parseIdentifier
     maybeParensAST = id

ghcide/src/Development/IDE/Plugin/CodeAction.hs

@@ -116,7 +116,7 @@ codeAction lsp state (TextDocumentIdentifier uri) _range CodeActionContext{_diag
           | x <- xs
           , Just ps <- [annotatedPS]
           , Just dynflags <- [df]
-          , (title, graft) <- suggestExactAction dynflags ps x
+          , (title, graft) <- suggestExactAction exportsMap dynflags ps x
           , let edit = either error id $
                         rewriteToEdit dynflags uri (annsA ps) graft
           ]
@@ -173,14 +173,16 @@ commandHandler lsp _ideState ExecuteCommandParams{..}
     = return (Right Null, Nothing)
 
 suggestExactAction ::
+  ExportsMap ->
   DynFlags ->
   Annotated ParsedSource ->
   Diagnostic ->
   [(T.Text, Rewrite)]
-suggestExactAction df ps x =
+suggestExactAction exportsMap df ps x =
   concat
     [ suggestConstraint df (astA ps) x
     , suggestImplicitParameter (astA ps) x
+    , suggestExtendImport exportsMap (astA ps) x
     ]
 
 suggestAction
@@ -193,7 +195,6 @@ suggestAction
 suggestAction packageExports ideOptions parsedModule text diag = concat
    -- Order these suggestions by priority
     [ suggestSignature True diag
-    , suggestExtendImport packageExports text diag
     , suggestFillTypeWildcard diag
     , suggestFixConstructorImport text diag
     , suggestModuleTypo diag
@@ -725,32 +726,31 @@ getIndentedGroupsBy pred inp = case dropWhile (not.pred) inp of
 indentation :: T.Text -> Int
 indentation = T.length . T.takeWhile isSpace
 
-suggestExtendImport :: ExportsMap -> Maybe T.Text -> Diagnostic -> [(T.Text, [TextEdit])]
-suggestExtendImport exportsMap contents Diagnostic{_range=_range,..}
+suggestExtendImport :: ExportsMap -> ParsedSource -> Diagnostic -> [(T.Text, Rewrite)]
+suggestExtendImport exportsMap (L _ HsModule {hsmodImports}) Diagnostic{_range=_range,..}
     | Just [binding, mod, srcspan] <-
       matchRegexUnifySpaces _message
       "Perhaps you want to add ‘([^’]*)’ to the import list in the import of ‘([^’]*)’ *\\((.*)\\).$"
-    , Just c <- contents
-    = suggestions c binding mod srcspan
+    = suggestions hsmodImports binding mod srcspan
     | Just (binding, mod_srcspan) <-
       matchRegExMultipleImports _message
-    , Just c <- contents
-    = mod_srcspan >>= (\(x, y) -> suggestions c binding x y)
+    = mod_srcspan >>= uncurry (suggestions hsmodImports binding)
     | otherwise = []
     where
-        suggestions c binding mod srcspan
+        unImportStyle (ImportTopLevel x) = (Nothing, T.unpack x)
+        unImportStyle (ImportViaParent x y) = (Just $ T.unpack y, T.unpack x)
+        suggestions decls binding mod srcspan
           |  range <- case [ x | (x,"") <- readSrcSpan (T.unpack srcspan)] of
                 [s] -> let x = realSrcSpanToRange s
                    in x{_end = (_end x){_character = succ (_character (_end x))}}
                 _ -> error "bug in srcspan parser",
-            importLine <- textInRange range c,
+            Just decl <- findImportDeclByRange decls range,
             Just ident <- lookupExportMap binding mod
-          = [ ( "Add " <> rendered <> " to the import list of " <> mod
-              , [TextEdit range result]
+          = [ ( "Add " <> renderImportStyle importStyle <> " to the import list of " <> mod
+              , uncurry extendImport (unImportStyle importStyle) decl
               )
             | importStyle <- NE.toList $ importStyles ident
-            , let rendered = renderImportStyle importStyle
-            , result <- maybeToList $ addBindingToImportList importStyle importLine]
+            ]
           | otherwise = []
         lookupExportMap binding mod
           | Just match <- Map.lookup binding (getExportsMap exportsMap)
@@ -765,6 +765,9 @@ suggestExtendImport exportsMap contents Diagnostic{_range=_range,..}
                 , parent = Nothing
                 , isDatacon = False}
 
+findImportDeclByRange :: [LImportDecl GhcPs] -> Range -> Maybe (LImportDecl GhcPs)
+findImportDeclByRange xs range = find (\(L l _)-> srcSpanToRange l == Just range) xs
+
 suggestFixConstructorImport :: Maybe T.Text -> Diagnostic -> [(T.Text, [TextEdit])]
 suggestFixConstructorImport _ Diagnostic{_range=_range,..}
     -- ‘Success’ is a data constructor of ‘Result’
@@ -1187,49 +1190,6 @@ rangesForBinding' b (L l (IEThingWith _ thing _  inners labels))
         [ l' | L l' x <- labels, showSDocUnsafe (ppr x) == b]
 rangesForBinding' _ _ = []
 
--- | Extends an import list with a new binding.
---   Assumes an import statement of the form:
---       import (qualified) A (..) ..
---   Places the new binding first, preserving whitespace.
---   Copes with multi-line import lists
-addBindingToImportList :: ImportStyle -> T.Text -> Maybe T.Text
-addBindingToImportList importStyle importLine =
-  case T.breakOn "(" importLine of
-    (pre, T.uncons -> Just (_, rest)) ->
-      case importStyle of
-        ImportTopLevel rendered ->
-          -- the binding has no parent, add it to the head of import list
-          Just $ T.concat [pre, "(", rendered, addCommaIfNeeds rest]
-        ImportViaParent rendered parent -> case T.breakOn parent rest of
-          -- the binding has a parent, and the current import list contains the
-          -- parent
-          --
-          -- `rest'` could be 1. `,...)`
-          --               or 2. `(),...)`
-          --               or 3. `(ConsA),...)`
-          --               or 4. `)`
-          (leading, T.stripPrefix parent -> Just rest') -> case T.uncons (T.stripStart rest') of
-            -- case 1: no children and parentheses, e.g. `import A(Foo,...)` --> `import A(Foo(Cons), ...)`
-            Just (',', rest'') -> Just $ T.concat [pre, "(", leading, parent, "(", rendered, ")", addCommaIfNeeds rest'']
-            -- case 2: no children but parentheses, e.g. `import A(Foo(),...)` --> `import A(Foo(Cons), ...)`
-            Just ('(', T.uncons -> Just (')', rest'')) -> Just $ T.concat [pre, "(", leading, parent, "(", rendered, ")", rest'']
-            -- case 3: children with parentheses, e.g. `import A(Foo(ConsA),...)` --> `import A(Foo(Cons, ConsA), ...)`
-            Just ('(', T.breakOn ")" -> (children, rest''))
-              | not (T.null children),
-                -- ignore A(Foo({-...-}), ...)
-                not $ "{-" `T.isPrefixOf` T.stripStart children
-              -> Just $ T.concat [pre, "(", leading, parent, "(", rendered, ", ", children, rest'']
-            -- case 4: no trailing, e.g. `import A(..., Foo)` --> `import A(..., Foo(Cons))`
-            Just (')', _) -> Just $ T.concat [pre, "(", leading, parent, "(", rendered, ")", rest']
-            _ -> Nothing
-          -- current import list does not contain the parent, e.g. `import A(...)` --> `import A(Foo(Cons), ...)`
-          _ -> Just $ T.concat [pre, "(", parent, "(", rendered, ")", addCommaIfNeeds rest]
-    _ -> Nothing
-  where
-    addCommaIfNeeds r = case T.uncons (T.stripStart r) of
-      Just (')', _) -> r
-      _ -> ", " <> r
-
 -- | 'matchRegex' combined with 'unifySpaces'
 matchRegexUnifySpaces :: T.Text -> T.Text -> Maybe [T.Text]
 matchRegexUnifySpaces message = matchRegex (unifySpaces message)
@@ -1321,6 +1281,7 @@ data ImportStyle
       --
       -- @P@ and @?@ can be a data type and a constructor, a class and a method,
       -- a class and an associated type/data family, etc.
+  deriving Show
 
 importStyles :: IdentInfo -> NonEmpty ImportStyle
 importStyles IdentInfo {parent, rendered, isDatacon}

ghcide/src/Development/IDE/Plugin/CodeAction/ExactPrint.hs

@@ -9,6 +9,7 @@ module Development.IDE.Plugin.CodeAction.ExactPrint
 
     -- * Utilities
     appendConstraint,
+    extendImport,
   )
 where
 
@@ -28,6 +29,8 @@ import GhcPlugins (realSrcSpanEnd, realSrcSpanStart, sigPrec)
 import Language.Haskell.GHC.ExactPrint
 import Language.Haskell.GHC.ExactPrint.Types (DeltaPos (DP), KeywordId (G), mkAnnKey)
 import Language.Haskell.LSP.Types
+import OccName
+import Outputable (ppr, showSDocUnsafe)
 
 ------------------------------------------------------------------------------
 
@@ -58,7 +61,7 @@ rewriteToEdit dflags uri anns (Rewrite dst f) = do
           [ ( uri,
               List
                 [ TextEdit (fromJust $ srcSpanToRange dst) $
-                      T.pack $ tail $ exactPrint ast anns
+                    T.pack $ tail $ exactPrint ast anns
                 ]
             )
           ]
@@ -173,3 +176,125 @@ headMaybe (a : _) = Just a
 lastMaybe :: [a] -> Maybe a
 lastMaybe [] = Nothing
 lastMaybe other = Just $ last other
+
+liftMaybe :: String -> Maybe a -> TransformT (Either String) a
+liftMaybe _ (Just x) = return x
+liftMaybe s _ = lift $ Left s
+
+-- | Copy anns attached to a into b with modification, then delete anns of a
+transferAnn :: (Data a, Data b) => Located a -> Located b -> (Annotation -> Annotation) -> TransformT (Either String) ()
+transferAnn la lb f = do
+  anns <- getAnnsT
+  let oldKey = mkAnnKey la
+      newKey = mkAnnKey lb
+  oldValue <- liftMaybe "Unable to find ann" $ Map.lookup oldKey anns
+  putAnnsT $ Map.delete oldKey $ Map.insert newKey (f oldValue) anns
+
+------------------------------------------------------------------------------
+extendImport :: Maybe String -> String -> LImportDecl GhcPs -> Rewrite
+extendImport mparent identifier lDecl@(L l _) =
+  Rewrite l $ \df -> do
+    case mparent of
+      Just parent -> extendImportViaParent df parent identifier lDecl
+      _ -> extendImportTopLevel df identifier lDecl
+
+-- | Add an identifier to import list
+--
+-- extendImportTopLevel "foo" AST:
+--
+-- import A --> Error
+-- import A (bar) --> import A (bar, foo)
+extendImportTopLevel :: DynFlags -> String -> LImportDecl GhcPs -> TransformT (Either String) (LImportDecl GhcPs)
+extendImportTopLevel df idnetifier (L l it@ImportDecl {..})
+  | Just (hide, L l' lies) <- ideclHiding,
+    hasSibling <- not $ null lies = do
+    src <- uniqueSrcSpanT
+    top <- uniqueSrcSpanT
+    rdr <- liftParseAST df idnetifier
+    let lie = L src $ IEName rdr
+        x = L top $ IEVar noExtField lie
+    when hasSibling $
+      addTrailingCommaT (last lies)
+    addSimpleAnnT x (DP (0, if hasSibling then 1 else 0)) []
+    addSimpleAnnT rdr dp00 $ unqalDP $ hasParen idnetifier
+    -- Parens are attachted to `lies`, so if `lies` was empty previously,
+    -- we need change the ann key from `[]` to `:` to keep parens and other anns.
+    unless hasSibling $
+      transferAnn (L l' lies) (L l' [x]) id
+    return $ L l it {ideclHiding = Just (hide, L l' $ lies ++ [x])}
+extendImportTopLevel _ _ _ = lift $ Left "Unable to extend the import list"
+
+-- | Add an identifier with its parent to import list
+--
+-- extendImportViaParent "Bar" "Cons" AST:
+--
+-- import A --> Error
+-- import A () --> import A (Bar(Cons))
+-- import A (Foo, Bar) --> import A (Foo, Bar(Cons))
+-- import A (Foo, Bar()) --> import A (Foo, Bar(Cons))
+extendImportViaParent :: DynFlags -> String -> String -> LImportDecl GhcPs -> TransformT (Either String) (LImportDecl GhcPs)
+extendImportViaParent df parent child (L l it@ImportDecl {..})
+  | Just (hide, L l' lies) <- ideclHiding = go hide l' [] lies
+  where
+    go :: Bool -> SrcSpan -> [LIE GhcPs] -> [LIE GhcPs] -> TransformT (Either String) (LImportDecl GhcPs)
+    go hide l' pre (lAbs@(L ll' (IEThingAbs _ absIE@(L _ ie))) : xs)
+      -- ThingAbs ie => ThingWith ie child
+      | parent == unIEWrappedName ie = do
+        srcChild <- uniqueSrcSpanT
+        childRdr <- liftParseAST df child
+        let childLIE = L srcChild $ IEName childRdr
+            x :: LIE GhcPs = L ll' $ IEThingWith noExtField absIE NoIEWildcard [childLIE] []
+        -- take anns from ThingAbs, and attatch parens to it
+        transferAnn lAbs x $ \old -> old {annsDP = annsDP old ++ [(G AnnOpenP, DP (0, 1)), (G AnnCloseP, dp00)]}
+        addSimpleAnnT childRdr dp00 [(G AnnVal, dp00)]
+        return $ L l it {ideclHiding = Just (hide, L l' $ reverse pre ++ [x] ++ xs)}
+    go hide l' pre ((L l'' (IEThingWith _ twIE@(L _ ie) _ lies' _)) : xs)
+      -- ThingWith ie lies' => ThingWith ie (lies' ++ [child])
+      | parent == unIEWrappedName ie,
+        hasSibling <- not $ null lies' =
+        do
+          srcChild <- uniqueSrcSpanT
+          childRdr <- liftParseAST df child
+          when hasSibling $
+            addTrailingCommaT (last lies')
+          let childLIE = L srcChild $ IEName childRdr
+          addSimpleAnnT childRdr (DP (0, if hasSibling then 1 else 0)) $ unqalDP $ hasParen child
+          return $ L l it {ideclHiding = Just (hide, L l' $ reverse pre ++ [L l'' (IEThingWith noExtField twIE NoIEWildcard (lies' ++ [childLIE]) [])] ++ xs)}
+    go hide l' pre (x : xs) = go hide l' (x : pre) xs
+    go hide l' pre []
+      | hasSibling <- not $ null pre = do
+        -- [] => ThingWith parent [child]
+        l'' <- uniqueSrcSpanT
+        srcParent <- uniqueSrcSpanT
+        srcChild <- uniqueSrcSpanT
+        parentRdr <- liftParseAST df parent
+        childRdr <- liftParseAST df child
+        when hasSibling $
+          addTrailingCommaT (head pre)
+        let parentLIE = L srcParent $ IEName parentRdr
+            childLIE = L srcChild $ IEName childRdr
+            x :: LIE GhcPs = L l'' $ IEThingWith noExtField parentLIE NoIEWildcard [childLIE] []
+        addSimpleAnnT parentRdr (DP (0, if hasSibling then 1 else 0)) $ unqalDP $ hasParen parent
+        addSimpleAnnT childRdr (DP (0, 0)) $ unqalDP $ hasParen child
+        addSimpleAnnT x (DP (0, 0)) [(G AnnOpenP, DP (0, 1)), (G AnnCloseP, DP (0, 0))]
+        -- Parens are attachted to `pre`, so if `pre` was empty previously,
+        -- we need change the ann key from `[]` to `:` to keep parens and other anns.
+        unless hasSibling $
+          transferAnn (L l' $ reverse pre) (L l' [x]) id
+        return $ L l it {ideclHiding = Just (hide, L l' $ reverse pre ++ [x])}
+extendImportViaParent _ _ _ _ = lift $ Left "Unable to extend the import list via parent"
+
+unIEWrappedName :: IEWrappedName (IdP GhcPs) -> String
+unIEWrappedName (occName -> occ) = showSDocUnsafe $ parenSymOcc occ (ppr occ)
+
+hasParen :: String -> Bool
+hasParen ('(' : _) = True
+hasParen _ = False
+
+unqalDP :: Bool -> [(KeywordId, DeltaPos)]
+unqalDP paren =
+  ( if paren
+      then \x -> (G AnnOpenP, dp00) : x : [(G AnnCloseP, dp00)]
+      else pure
+  )
+    (G AnnVal, dp00)

ghcide/test/exe/Main.hs

@@ -1111,7 +1111,7 @@ extendImportTests = testGroup "extend import actions"
             ["Add stuffA to the import list of ModuleA"]
             (T.unlines
                     [ "module ModuleB where"
-                    , "import ModuleA as A (stuffA, stuffB)"
+                    , "import ModuleA as A (stuffB, stuffA)"
                     , "main = print (stuffA, stuffB)"
                     ])
         , testSession "extend single line import with operator" $ template
@@ -1131,7 +1131,7 @@ extendImportTests = testGroup "extend import actions"
             ["Add (.*) to the import list of ModuleA"]
             (T.unlines
                     [ "module ModuleB where"
-                    , "import ModuleA as A ((.*), stuffB)"
+                    , "import ModuleA as A (stuffB, (.*))"
                     , "main = print (stuffB .* stuffB)"
                     ])
         , testSession "extend single line import with type" $ template
@@ -1168,7 +1168,26 @@ extendImportTests = testGroup "extend import actions"
             ["Add A(Constructor) to the import list of ModuleA"]
             (T.unlines
                     [ "module ModuleB where"
-                    , "import ModuleA (A(Constructor))"
+                    , "import ModuleA (A (Constructor))"
+                    , "b :: A"
+                    , "b = Constructor"
+                    ])        
+        , testSession "extend single line import with constructor (with comments)" $ template
+            [("ModuleA.hs", T.unlines
+                    [ "module ModuleA where"
+                    , "data A = Constructor"
+                    ])]
+            ("ModuleB.hs", T.unlines
+                    [ "module ModuleB where"
+                    , "import ModuleA (A ({-Constructor-}))"
+                    , "b :: A"
+                    , "b = Constructor"
+                    ])
+            (Range (Position 2 5) (Position 2 5))
+            ["Add A(Constructor) to the import list of ModuleA"]
+            (T.unlines
+                    [ "module ModuleB where"
+                    , "import ModuleA (A (Constructor{-Constructor-}))"
                     , "b :: A"
                     , "b = Constructor"
                     ])
@@ -1180,15 +1199,15 @@ extendImportTests = testGroup "extend import actions"
                     ])]
             ("ModuleB.hs", T.unlines
                     [ "module ModuleB where"
-                    , "import ModuleA (A(ConstructorBar), a)"
+                    , "import ModuleA (A (ConstructorBar), a)"
                     , "b :: A"
                     , "b = ConstructorFoo"
                     ])
             (Range (Position 2 5) (Position 2 5))
             ["Add A(ConstructorFoo) to the import list of ModuleA"]
             (T.unlines
                     [ "module ModuleB where"
-                    , "import ModuleA (A(ConstructorFoo, ConstructorBar), a)"
+                    , "import ModuleA (A (ConstructorBar, ConstructorFoo), a)"
                     , "b :: A"
                     , "b = ConstructorFoo"
                     ])
@@ -1209,7 +1228,7 @@ extendImportTests = testGroup "extend import actions"
             ["Add stuffA to the import list of ModuleA"]
             (T.unlines
                     [ "module ModuleB where"
-                    , "import qualified ModuleA as A (stuffA, stuffB)"
+                    , "import qualified ModuleA as A (stuffB, stuffA)"
                     , "main = print (A.stuffA, A.stuffB)"
                     ])
         , testSession "extend multi line import with value" $ template
@@ -1230,7 +1249,7 @@ extendImportTests = testGroup "extend import actions"
             ["Add stuffA to the import list of ModuleA"]
             (T.unlines
                     [ "module ModuleB where"
-                    , "import ModuleA (stuffA, stuffB"
+                    , "import ModuleA (stuffB, stuffA"
                     , "               )"
                     , "main = print (stuffA, stuffB)"
                     ])
@@ -1251,7 +1270,7 @@ extendImportTests = testGroup "extend import actions"
              "Add m2 to the import list of ModuleA"]
             (T.unlines
                     [ "module ModuleB where"
-                    , "import ModuleA (C(m2, m1))"
+                    , "import ModuleA (C(m1, m2))"
                     , "b = m2"
                     ])
         , testSession "extend single line import with method without class" $ template
@@ -1271,7 +1290,7 @@ extendImportTests = testGroup "extend import actions"
              "Add C(m2) to the import list of ModuleA"]
             (T.unlines
                     [ "module ModuleB where"
-                    , "import ModuleA (m2, C(m1))"
+                    , "import ModuleA (C(m1), m2)"
                     , "b = m2"
                     ])
         , testSession "extend import list with multiple choices" $ template
@@ -1312,7 +1331,7 @@ extendImportTests = testGroup "extend import actions"
             ["Add (:~:)(Refl) to the import list of Data.Type.Equality"]
             (T.unlines
                     [ "module ModuleA where"
-                    , "import Data.Type.Equality ((:~:)(Refl))"
+                    , "import Data.Type.Equality ((:~:) (Refl))"
                     , "x :: (:~:) [] []"
                     , "x = Refl"
                     ])