Fix parens

Author: berberman

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)