From 95e54c710544d47ac25e69229bad462ee0b236c3 Mon Sep 17 00:00:00 2001
From: Potato Hatsue <1793913507@qq.com>
Date: Fri, 22 Jan 2021 19:39:46 +0800
Subject: [PATCH 1/8] Implement extendImport using exact print

---
 ghcide/src/Development/IDE/GHC/ExactPrint.hs  |   4 +
 .../IDE/Plugin/CodeAction/ExactPrint.hs       | 105 +++++++++++++++++-
 2 files changed, 108 insertions(+), 1 deletion(-)

diff --git a/ghcide/src/Development/IDE/GHC/ExactPrint.hs b/ghcide/src/Development/IDE/GHC/ExactPrint.hs
index 46a1654933..4ac6eb69ab 100644
--- a/ghcide/src/Development/IDE/GHC/ExactPrint.hs
+++ b/ghcide/src/Development/IDE/GHC/ExactPrint.hs
@@ -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
diff --git a/ghcide/src/Development/IDE/Plugin/CodeAction/ExactPrint.hs b/ghcide/src/Development/IDE/Plugin/CodeAction/ExactPrint.hs
index 18caf7fa71..c9ae6f3d95 100644
--- a/ghcide/src/Development/IDE/Plugin/CodeAction/ExactPrint.hs
+++ b/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 RdrName
 
 ------------------------------------------------------------------------------
 
@@ -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,103 @@ headMaybe (a : _) = Just a
 lastMaybe :: [a] -> Maybe a
 lastMaybe [] = Nothing
 lastMaybe other = Just $ last other
+
+------------------------------------------------------------------------------
+extendImport :: Maybe String -> String -> LImportDecl GhcPs -> Rewrite
+extendImport mparent identifier lDecl@(L l _) = Rewrite l $ \_ -> do
+  go lDecl
+  where
+    go decl
+      | Just parent <- mparent =
+        extendImportViaParent parent identifier decl
+      | otherwise = extendImportTopLevel identifier decl
+
+-- | Add an identifier to import list
+--
+-- extendImportTopLevel "foo" AST:
+--
+-- import A --> Error
+-- import A (bar) --> import A (bar, foo)
+extendImportTopLevel :: String -> LImportDecl GhcPs -> TransformT (Either String) (LImportDecl GhcPs)
+extendImportTopLevel idnetifier (L l it@ImportDecl {..})
+  | Just (hide, L l' lies) <- ideclHiding,
+    hasSibling <- not $ null lies = do
+    src <- uniqueSrcSpanT
+    top <- uniqueSrcSpanT
+    let rdr = L src $ mkVarName idnetifier
+        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 [(G AnnVal, dp00)]
+    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 :: String -> String -> LImportDecl GhcPs -> TransformT (Either String) (LImportDecl GhcPs)
+extendImportViaParent 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 => ThingWith ie child
+      | parent == unIEWrappedName ie = do
+        srcChild <- uniqueSrcSpanT
+        let childRdr = L srcChild $ mkVarName child
+            childLIE = L srcChild $ IEName childRdr
+            x :: LIE GhcPs = L ll' $ IEThingWith NoExtField absIE NoIEWildcard [childLIE] []
+        modifyAnnsT $ \anns ->
+          let oldKey = mkAnnKey lAbs
+              oldValue = anns Map.! oldKey
+              newKey = mkAnnKey x
+           in Map.insert newKey oldValue {annsDP = annsDP oldValue ++ [(G AnnOpenP, DP (0, 1)), (G AnnCloseP, dp00)]} $ Map.delete oldKey anns
+        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 => ThingWith ie (lies' ++ [child])
+      | parent == unIEWrappedName ie,
+        hasSibling <- not $ null lies' =
+        do
+          srcChild <- uniqueSrcSpanT
+          when hasSibling $
+            addTrailingCommaT (last lies')
+          let childRdr = L srcChild $ mkVarName child
+              childLIE = L srcChild $ IEName childRdr
+          addSimpleAnnT childRdr (DP (0, if hasSibling then 1 else 0)) [(G AnnVal, dp00)]
+          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
+        when hasSibling $
+          addTrailingCommaT (head pre)
+        let parentRdr = L srcParent $ mkTcClsName parent
+            parentLIE = L srcParent $ IEName parentRdr
+            childRdr = L srcChild $ mkVarName child
+            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)) [(G AnnVal, DP (0, 0))]
+        addSimpleAnnT childRdr (DP (0, 0)) [(G AnnVal, DP (0, 0))]
+        addSimpleAnnT x (DP (0, 0)) [(G AnnOpenP, DP (0, 1)), (G AnnCloseP, DP (0, 0))]
+        return $ L l it {ideclHiding = Just (hide, L l' $ reverse pre ++ [x])}
+extendImportViaParent _ _ _ = lift $ Left "Unable to extend the import list via parent"
+
+mkTcClsName :: String -> RdrName
+mkTcClsName = mkRdrUnqual . mkOccName tcClsName
+
+mkVarName :: String -> RdrName
+mkVarName = mkRdrUnqual . mkOccName varName
+
+unIEWrappedName :: IEWrappedName (IdP GhcPs) -> String
+unIEWrappedName = occNameString . occName

From aab4875a0586d9f9956582d70b1d707a49af7fea Mon Sep 17 00:00:00 2001
From: Potato Hatsue <1793913507@qq.com>
Date: Fri, 22 Jan 2021 20:19:28 +0800
Subject: [PATCH 2/8] Avoid using RdrName directly

---
 .../IDE/Plugin/CodeAction/ExactPrint.hs       | 48 ++++++++-----------
 1 file changed, 19 insertions(+), 29 deletions(-)

diff --git a/ghcide/src/Development/IDE/Plugin/CodeAction/ExactPrint.hs b/ghcide/src/Development/IDE/Plugin/CodeAction/ExactPrint.hs
index c9ae6f3d95..8649c4fa50 100644
--- a/ghcide/src/Development/IDE/Plugin/CodeAction/ExactPrint.hs
+++ b/ghcide/src/Development/IDE/Plugin/CodeAction/ExactPrint.hs
@@ -30,7 +30,6 @@ import Language.Haskell.GHC.ExactPrint
 import Language.Haskell.GHC.ExactPrint.Types (DeltaPos (DP), KeywordId (G), mkAnnKey)
 import Language.Haskell.LSP.Types
 import OccName
-import RdrName
 
 ------------------------------------------------------------------------------
 
@@ -179,13 +178,10 @@ lastMaybe other = Just $ last other
 
 ------------------------------------------------------------------------------
 extendImport :: Maybe String -> String -> LImportDecl GhcPs -> Rewrite
-extendImport mparent identifier lDecl@(L l _) = Rewrite l $ \_ -> do
-  go lDecl
-  where
-    go decl
-      | Just parent <- mparent =
-        extendImportViaParent parent identifier decl
-      | otherwise = extendImportTopLevel identifier decl
+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
 --
@@ -193,21 +189,21 @@ extendImport mparent identifier lDecl@(L l _) = Rewrite l $ \_ -> do
 --
 -- import A --> Error
 -- import A (bar) --> import A (bar, foo)
-extendImportTopLevel :: String -> LImportDecl GhcPs -> TransformT (Either String) (LImportDecl GhcPs)
-extendImportTopLevel idnetifier (L l it@ImportDecl {..})
+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
-    let rdr = L src $ mkVarName idnetifier
-        lie = L src $ IEName rdr
+    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 [(G AnnVal, dp00)]
     return $ L l it {ideclHiding = Just (hide, L l' $ lies ++ [x])}
-extendImportTopLevel _ _ = lift $ Left "Unable to extend the import list"
+extendImportTopLevel _ _ _ = lift $ Left "Unable to extend the import list"
 
 -- | Add an identifier with its parent to import list
 --
@@ -217,8 +213,8 @@ extendImportTopLevel _ _ = lift $ Left "Unable to extend the import list"
 -- 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 :: String -> String -> LImportDecl GhcPs -> TransformT (Either String) (LImportDecl GhcPs)
-extendImportViaParent parent child (L l it@ImportDecl {..})
+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)
@@ -226,8 +222,8 @@ extendImportViaParent parent child (L l it@ImportDecl {..})
       -- ThingAbs => ThingWith ie child
       | parent == unIEWrappedName ie = do
         srcChild <- uniqueSrcSpanT
-        let childRdr = L srcChild $ mkVarName child
-            childLIE = L srcChild $ IEName childRdr
+        childRdr <- liftParseAST df child
+        let childLIE = L srcChild $ IEName childRdr
             x :: LIE GhcPs = L ll' $ IEThingWith NoExtField absIE NoIEWildcard [childLIE] []
         modifyAnnsT $ \anns ->
           let oldKey = mkAnnKey lAbs
@@ -242,10 +238,10 @@ extendImportViaParent parent child (L l it@ImportDecl {..})
         hasSibling <- not $ null lies' =
         do
           srcChild <- uniqueSrcSpanT
+          childRdr <- liftParseAST df child
           when hasSibling $
             addTrailingCommaT (last lies')
-          let childRdr = L srcChild $ mkVarName child
-              childLIE = L srcChild $ IEName childRdr
+          let childLIE = L srcChild $ IEName childRdr
           addSimpleAnnT childRdr (DP (0, if hasSibling then 1 else 0)) [(G AnnVal, dp00)]
           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
@@ -255,24 +251,18 @@ extendImportViaParent parent child (L l it@ImportDecl {..})
         l'' <- uniqueSrcSpanT
         srcParent <- uniqueSrcSpanT
         srcChild <- uniqueSrcSpanT
+        parentRdr <- liftParseAST df parent
+        childRdr <- liftParseAST df child
         when hasSibling $
           addTrailingCommaT (head pre)
-        let parentRdr = L srcParent $ mkTcClsName parent
-            parentLIE = L srcParent $ IEName parentRdr
-            childRdr = L srcChild $ mkVarName child
+        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)) [(G AnnVal, DP (0, 0))]
         addSimpleAnnT childRdr (DP (0, 0)) [(G AnnVal, DP (0, 0))]
         addSimpleAnnT x (DP (0, 0)) [(G AnnOpenP, DP (0, 1)), (G AnnCloseP, DP (0, 0))]
         return $ L l it {ideclHiding = Just (hide, L l' $ reverse pre ++ [x])}
-extendImportViaParent _ _ _ = lift $ Left "Unable to extend the import list via parent"
-
-mkTcClsName :: String -> RdrName
-mkTcClsName = mkRdrUnqual . mkOccName tcClsName
-
-mkVarName :: String -> RdrName
-mkVarName = mkRdrUnqual . mkOccName varName
+extendImportViaParent _ _ _ _ = lift $ Left "Unable to extend the import list via parent"
 
 unIEWrappedName :: IEWrappedName (IdP GhcPs) -> String
 unIEWrappedName = occNameString . occName

From 66f1094102ec4dd5607f59411823c717040d58cf Mon Sep 17 00:00:00 2001
From: Potato Hatsue <1793913507@qq.com>
Date: Fri, 22 Jan 2021 20:56:55 +0800
Subject: [PATCH 3/8] Replace NoExtField with noExtField

---
 .../src/Development/IDE/Plugin/CodeAction/ExactPrint.hs   | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/ghcide/src/Development/IDE/Plugin/CodeAction/ExactPrint.hs b/ghcide/src/Development/IDE/Plugin/CodeAction/ExactPrint.hs
index 8649c4fa50..915eb5a4c2 100644
--- a/ghcide/src/Development/IDE/Plugin/CodeAction/ExactPrint.hs
+++ b/ghcide/src/Development/IDE/Plugin/CodeAction/ExactPrint.hs
@@ -197,7 +197,7 @@ extendImportTopLevel df idnetifier (L l it@ImportDecl {..})
     top <- uniqueSrcSpanT
     rdr <- liftParseAST df idnetifier
     let lie = L src $ IEName rdr
-        x = L top $ IEVar NoExtField lie
+        x = L top $ IEVar noExtField lie
     when hasSibling $
       addTrailingCommaT (last lies)
     addSimpleAnnT x (DP (0, if hasSibling then 1 else 0)) []
@@ -224,7 +224,7 @@ extendImportViaParent df parent child (L l it@ImportDecl {..})
         srcChild <- uniqueSrcSpanT
         childRdr <- liftParseAST df child
         let childLIE = L srcChild $ IEName childRdr
-            x :: LIE GhcPs = L ll' $ IEThingWith NoExtField absIE NoIEWildcard [childLIE] []
+            x :: LIE GhcPs = L ll' $ IEThingWith noExtField absIE NoIEWildcard [childLIE] []
         modifyAnnsT $ \anns ->
           let oldKey = mkAnnKey lAbs
               oldValue = anns Map.! oldKey
@@ -243,7 +243,7 @@ extendImportViaParent df parent child (L l it@ImportDecl {..})
             addTrailingCommaT (last lies')
           let childLIE = L srcChild $ IEName childRdr
           addSimpleAnnT childRdr (DP (0, if hasSibling then 1 else 0)) [(G AnnVal, dp00)]
-          return $ L l it {ideclHiding = Just (hide, L l' $ reverse pre ++ [L l'' (IEThingWith NoExtField twIE NoIEWildcard (lies' ++ [childLIE]) [])] ++ xs)}
+          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
@@ -257,7 +257,7 @@ extendImportViaParent df parent child (L l it@ImportDecl {..})
           addTrailingCommaT (head pre)
         let parentLIE = L srcParent $ IEName parentRdr
             childLIE = L srcChild $ IEName childRdr
-            x :: LIE GhcPs = L l'' $ IEThingWith NoExtField parentLIE NoIEWildcard [childLIE] []
+            x :: LIE GhcPs = L l'' $ IEThingWith noExtField parentLIE NoIEWildcard [childLIE] []
         addSimpleAnnT parentRdr (DP (0, if hasSibling then 1 else 0)) [(G AnnVal, DP (0, 0))]
         addSimpleAnnT childRdr (DP (0, 0)) [(G AnnVal, DP (0, 0))]
         addSimpleAnnT x (DP (0, 0)) [(G AnnOpenP, DP (0, 1)), (G AnnCloseP, DP (0, 0))]

From 6eafe90c0cf3bdbbd5f53e67cae38014519fa2a2 Mon Sep 17 00:00:00 2001
From: Potato Hatsue <1793913507@qq.com>
Date: Sat, 23 Jan 2021 14:14:08 +0800
Subject: [PATCH 4/8] Fix parens

---
 .../IDE/Plugin/CodeAction/ExactPrint.hs       | 64 ++++++++++++++-----
 1 file changed, 48 insertions(+), 16 deletions(-)

diff --git a/ghcide/src/Development/IDE/Plugin/CodeAction/ExactPrint.hs b/ghcide/src/Development/IDE/Plugin/CodeAction/ExactPrint.hs
index 915eb5a4c2..a35c793c16 100644
--- a/ghcide/src/Development/IDE/Plugin/CodeAction/ExactPrint.hs
+++ b/ghcide/src/Development/IDE/Plugin/CodeAction/ExactPrint.hs
@@ -30,6 +30,7 @@ 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)
 
 ------------------------------------------------------------------------------
 
@@ -176,12 +177,26 @@ 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
+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
 --
@@ -201,7 +216,11 @@ extendImportTopLevel df idnetifier (L l it@ImportDecl {..})
     when hasSibling $
       addTrailingCommaT (last lies)
     addSimpleAnnT x (DP (0, if hasSibling then 1 else 0)) []
-    addSimpleAnnT rdr dp00 [(G AnnVal, dp00)]
+    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"
 
@@ -219,21 +238,18 @@ extendImportViaParent df parent child (L l it@ImportDecl {..})
   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 => ThingWith ie child
+      -- 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] []
-        modifyAnnsT $ \anns ->
-          let oldKey = mkAnnKey lAbs
-              oldValue = anns Map.! oldKey
-              newKey = mkAnnKey x
-           in Map.insert newKey oldValue {annsDP = annsDP oldValue ++ [(G AnnOpenP, DP (0, 1)), (G AnnCloseP, dp00)]} $ Map.delete oldKey anns
+        -- 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 => ThingWith ie (lies' ++ [child])
+      -- ThingWith ie lies' => ThingWith ie (lies' ++ [child])
       | parent == unIEWrappedName ie,
         hasSibling <- not $ null lies' =
         do
@@ -242,7 +258,7 @@ extendImportViaParent df parent child (L l it@ImportDecl {..})
           when hasSibling $
             addTrailingCommaT (last lies')
           let childLIE = L srcChild $ IEName childRdr
-          addSimpleAnnT childRdr (DP (0, if hasSibling then 1 else 0)) [(G AnnVal, dp00)]
+          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 []
@@ -258,11 +274,27 @@ extendImportViaParent df parent child (L l it@ImportDecl {..})
         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)) [(G AnnVal, DP (0, 0))]
-        addSimpleAnnT childRdr (DP (0, 0)) [(G AnnVal, DP (0, 0))]
+        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 = occNameString . occName
+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)

From eb557b3efc9ea3ca7516a19d761dd5cb9d2ad882 Mon Sep 17 00:00:00 2001
From: Potato Hatsue <1793913507@qq.com>
Date: Sat, 23 Jan 2021 14:14:39 +0800
Subject: [PATCH 5/8] Use exact print to suggestExtendImport

---
 .../src/Development/IDE/Plugin/CodeAction.hs  | 79 +++++--------------
 ghcide/test/exe/Main.hs                       | 20 ++---
 2 files changed, 30 insertions(+), 69 deletions(-)

diff --git a/ghcide/src/Development/IDE/Plugin/CodeAction.hs b/ghcide/src/Development/IDE/Plugin/CodeAction.hs
index f6ac664aa5..d9cf8d9698 100644
--- a/ghcide/src/Development/IDE/Plugin/CodeAction.hs
+++ b/ghcide/src/Development/IDE/Plugin/CodeAction.hs
@@ -108,13 +108,13 @@ codeAction lsp state (TextDocumentIdentifier uri) _range CodeActionContext{_diag
         | x <- xs, (title, tedit) <- suggestAction exportsMap ideOptions parsedModule text x
         , let edit = WorkspaceEdit (Just $ Map.singleton uri $ List tedit) Nothing
         ] <> caRemoveRedundantImports parsedModule text diag xs uri
-
+      
       actions' =
           [mkCA title [x] edit
           | 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
           ]
@@ -171,14 +171,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
@@ -191,7 +193,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
@@ -660,32 +661,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)
@@ -700,6 +700,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’
@@ -1109,49 +1112,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)
@@ -1243,6 +1203,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}
diff --git a/ghcide/test/exe/Main.hs b/ghcide/test/exe/Main.hs
index 6101e29e11..e08b5787c5 100644
--- a/ghcide/test/exe/Main.hs
+++ b/ghcide/test/exe/Main.hs
@@ -1110,7 +1110,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
@@ -1130,7 +1130,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
@@ -1167,7 +1167,7 @@ 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"
                     ])
@@ -1179,7 +1179,7 @@ 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"
                     ])
@@ -1187,7 +1187,7 @@ extendImportTests = testGroup "extend import actions"
             ["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"
                     ])
@@ -1208,7 +1208,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
@@ -1229,7 +1229,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)"
                     ])
@@ -1250,7 +1250,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
@@ -1270,7 +1270,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
@@ -1311,7 +1311,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"
                     ])

From 0555e09333731e51c46fe327a587f3bdd59aa481 Mon Sep 17 00:00:00 2001
From: Potato Hatsue <1793913507@qq.com>
Date: Sat, 23 Jan 2021 18:05:06 +0800
Subject: [PATCH 6/8] Preserve comments in GetAnnotatedParsedSource

---
 ghcide/src/Development/IDE/GHC/ExactPrint.hs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ghcide/src/Development/IDE/GHC/ExactPrint.hs b/ghcide/src/Development/IDE/GHC/ExactPrint.hs
index 4ac6eb69ab..b0636174a1 100644
--- a/ghcide/src/Development/IDE/GHC/ExactPrint.hs
+++ b/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

From ad84d45f2789f6893f4ac337b81b669df8bcadc2 Mon Sep 17 00:00:00 2001
From: Potato Hatsue <1793913507@qq.com>
Date: Sat, 23 Jan 2021 18:07:55 +0800
Subject: [PATCH 7/8] Add a new test case

---
 ghcide/test/exe/Main.hs | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/ghcide/test/exe/Main.hs b/ghcide/test/exe/Main.hs
index e08b5787c5..c24128a70e 100644
--- a/ghcide/test/exe/Main.hs
+++ b/ghcide/test/exe/Main.hs
@@ -1170,6 +1170,25 @@ extendImportTests = testGroup "extend import actions"
                     , "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"
                     ])
         , testSession "extend single line import with mixed constructors" $ template
             [("ModuleA.hs", T.unlines

From 36573ec34a866a09e3f3d28236206d1dca95f737 Mon Sep 17 00:00:00 2001
From: Potato Hatsue <1793913507@qq.com>
Date: Sat, 23 Jan 2021 19:15:22 +0800
Subject: [PATCH 8/8] Resolve merge conflict

---
 ghcide/src/Development/IDE/Plugin/CodeAction.hs | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/ghcide/src/Development/IDE/Plugin/CodeAction.hs b/ghcide/src/Development/IDE/Plugin/CodeAction.hs
index d29e34753e..8b3a8512b9 100644
--- a/ghcide/src/Development/IDE/Plugin/CodeAction.hs
+++ b/ghcide/src/Development/IDE/Plugin/CodeAction.hs
@@ -860,10 +860,7 @@ suggestInstanceConstraint df (L _ HsModule {hsmodDecls}) Diagnostic {..} missing
         = Nothing
 
       readPositionNumber :: T.Text -> Int
-      readPositionNumber = T.unpack 
-      
-      
-      read
+      readPositionNumber = T.unpack >>> read
 
       actionTitle :: T.Text -> T.Text
       actionTitle constraint = "Add `" <> constraint