Add fix for correct placement of import (#2100) (#2116)

* Remove HasSrcSpan import

* Amend extendImportHandler' function

* Amend old import tests

* Adapt tests for windows newline issue

Co-authored-by: Javier Neira <atreyu.bbb@gmail.com>

Author: nini-faroux

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

@@ -234,8 +234,8 @@ findDeclContainingLoc loc = find (\(L l _) -> loc `isInsideSrcSpan` l)
 --  imported from ‘Data.ByteString’ at B.hs:6:1-22
 --  imported from ‘Data.ByteString.Lazy’ at B.hs:8:1-27
 --  imported from ‘Data.Text’ at B.hs:7:1-16
-suggestHideShadow :: ParsedSource -> Maybe TcModuleResult -> Maybe HieAstResult -> Diagnostic -> [(T.Text, [Either TextEdit Rewrite])]
-suggestHideShadow ps@(L _ HsModule {hsmodImports}) mTcM mHar Diagnostic {_message, _range}
+suggestHideShadow :: ParsedSource -> T.Text -> Maybe TcModuleResult -> Maybe HieAstResult -> Diagnostic -> [(T.Text, [Either TextEdit Rewrite])]
+suggestHideShadow ps@(L _ HsModule {hsmodImports}) fileContents mTcM mHar Diagnostic {_message, _range}
   | Just [identifier, modName, s] <-
       matchRegexUnifySpaces
         _message
@@ -260,7 +260,7 @@ suggestHideShadow ps@(L _ HsModule {hsmodImports}) mTcM mHar Diagnostic {_messag
         mDecl <- findImportDeclByModuleName hsmodImports $ T.unpack modName,
         title <- "Hide " <> identifier <> " from " <> modName =
         if modName == "Prelude" && null mDecl
-          then maybeToList $ (\(_, te) -> (title, [Left te])) <$> newImportToEdit (hideImplicitPreludeSymbol identifier) ps
+          then maybeToList $ (\(_, te) -> (title, [Left te])) <$> newImportToEdit (hideImplicitPreludeSymbol identifier) ps fileContents
           else maybeToList $ (title,) . pure . pure . hideSymbol (T.unpack identifier) <$> mDecl
       | otherwise = []
 
@@ -887,9 +887,10 @@ suggestImportDisambiguation ::
     DynFlags ->
     Maybe T.Text ->
     ParsedSource ->
+    T.Text ->
     Diagnostic ->
     [(T.Text, [Either TextEdit Rewrite])]
-suggestImportDisambiguation df (Just txt) ps@(L _ HsModule {hsmodImports}) diag@Diagnostic {..}
+suggestImportDisambiguation df (Just txt) ps@(L _ HsModule {hsmodImports}) fileContents diag@Diagnostic {..}
     | Just [ambiguous] <-
         matchRegexUnifySpaces
             _message
@@ -930,7 +931,7 @@ suggestImportDisambiguation df (Just txt) ps@(L _ HsModule {hsmodImports}) diag@
         suggestionsImpl symbol targetsWithRestImports =
             sortOn fst
             [ ( renderUniquify mode modNameText symbol
-              , disambiguateSymbol ps diag symbol mode
+              , disambiguateSymbol ps fileContents diag symbol mode
               )
             | (modTarget, restImports) <- targetsWithRestImports
             , let modName = targetModuleName modTarget
@@ -964,7 +965,7 @@ suggestImportDisambiguation df (Just txt) ps@(L _ HsModule {hsmodImports}) diag@
                 <> T.pack (moduleNameString qual)
                 <> "."
                 <> symbol
-suggestImportDisambiguation _ _ _ _ = []
+suggestImportDisambiguation _ _ _ _ _ = []
 
 occursUnqualified :: T.Text -> ImportDecl GhcPs -> Bool
 occursUnqualified symbol ImportDecl{..}
@@ -989,19 +990,20 @@ targetModuleName (ExistingImp _) =
 
 disambiguateSymbol ::
     ParsedSource ->
+    T.Text ->
     Diagnostic ->
     T.Text ->
     HidingMode ->
     [Either TextEdit Rewrite]
-disambiguateSymbol pm Diagnostic {..} (T.unpack -> symbol) = \case
+disambiguateSymbol pm fileContents Diagnostic {..} (T.unpack -> symbol) = \case
     (HideOthers hiddens0) ->
         [ Right $ hideSymbol symbol idecl
         | ExistingImp idecls <- hiddens0
         , idecl <- NE.toList idecls
         ]
             ++ mconcat
                 [ if null imps
-                    then maybeToList $ Left . snd <$> newImportToEdit (hideImplicitPreludeSymbol $ T.pack symbol) pm
+                    then maybeToList $ Left . snd <$> newImportToEdit (hideImplicitPreludeSymbol $ T.pack symbol) pm fileContents
                     else Right . hideSymbol symbol <$> imps
                 | ImplicitPrelude imps <- hiddens0
                 ]
@@ -1203,8 +1205,8 @@ removeRedundantConstraints df (L _ HsModule {hsmodDecls}) Diagnostic{..}
 
 -------------------------------------------------------------------------------------------------
 
-suggestNewOrExtendImportForClassMethod :: ExportsMap -> ParsedSource -> Diagnostic -> [(T.Text, CodeActionKind, [Either TextEdit Rewrite])]
-suggestNewOrExtendImportForClassMethod packageExportsMap ps Diagnostic {_message}
+suggestNewOrExtendImportForClassMethod :: ExportsMap -> ParsedSource -> T.Text -> Diagnostic -> [(T.Text, CodeActionKind, [Either TextEdit Rewrite])]
+suggestNewOrExtendImportForClassMethod packageExportsMap ps fileContents Diagnostic {_message}
   | Just [methodName, className] <-
       matchRegexUnifySpaces
         _message
@@ -1229,7 +1231,7 @@ suggestNewOrExtendImportForClassMethod packageExportsMap ps Diagnostic {_message
             ]
           -- new
           _
-            | Just (range, indent) <- newImportInsertRange ps
+            | Just (range, indent) <- newImportInsertRange ps fileContents
             ->
              (\(kind, unNewImport -> x) -> (x, kind, [Left $ TextEdit range (x <> "\n" <> T.replicate indent " ")])) <$>
             [ (quickFixImportKind' "new" style, newUnqualImport moduleNameText rendered False)
@@ -1239,8 +1241,8 @@ suggestNewOrExtendImportForClassMethod packageExportsMap ps Diagnostic {_message
               <> [(quickFixImportKind "new.all", newImportAll moduleNameText)]
             | otherwise -> []
 
-suggestNewImport :: ExportsMap -> ParsedSource -> Diagnostic -> [(T.Text, CodeActionKind, TextEdit)]
-suggestNewImport packageExportsMap ps@(L _ HsModule {..}) Diagnostic{_message}
+suggestNewImport :: ExportsMap -> ParsedSource -> T.Text -> Diagnostic -> [(T.Text, CodeActionKind, TextEdit)]
+suggestNewImport packageExportsMap ps@(L _ HsModule {..}) fileContents Diagnostic{_message}
   | msg <- unifySpaces _message
   , Just thingMissing <- extractNotInScopeName msg
   , qual <- extractQualifiedModuleName msg
@@ -1249,13 +1251,13 @@ suggestNewImport packageExportsMap ps@(L _ HsModule {..}) Diagnostic{_message}
         >>= (findImportDeclByModuleName hsmodImports . T.unpack)
         >>= ideclAs . unLoc
         <&> T.pack . moduleNameString . unLoc
-  , Just (range, indent) <- newImportInsertRange ps
+  , Just (range, indent) <- newImportInsertRange ps fileContents
   , extendImportSuggestions <- matchRegexUnifySpaces msg
     "Perhaps you want to add ‘[^’]*’ to the import list in the import of ‘([^’]*)’"
   = sortOn fst3 [(imp, kind, TextEdit range (imp <> "\n" <> T.replicate indent " "))
     | (kind, unNewImport -> imp) <- constructNewImportSuggestions packageExportsMap (qual <|> qual', thingMissing) extendImportSuggestions
     ]
-suggestNewImport _ _ _ = []
+suggestNewImport _ _ _ _ = []
 
 constructNewImportSuggestions
   :: ExportsMap -> (Maybe T.Text, NotInScope) -> Maybe [T.Text] -> [(CodeActionKind, NewImport)]
@@ -1282,26 +1284,70 @@ constructNewImportSuggestions exportsMap (qual, thingMissing) notTheseModules =
 newtype NewImport = NewImport {unNewImport :: T.Text}
   deriving (Show, Eq, Ord)
 
-newImportToEdit :: NewImport -> ParsedSource -> Maybe (T.Text, TextEdit)
-newImportToEdit (unNewImport -> imp) ps
-  | Just (range, indent) <- newImportInsertRange ps
+newImportToEdit :: NewImport -> ParsedSource -> T.Text -> Maybe (T.Text, TextEdit)
+newImportToEdit (unNewImport -> imp) ps fileContents
+  | Just (range, indent) <- newImportInsertRange ps fileContents
   = Just (imp, TextEdit range (imp <> "\n" <> T.replicate indent " "))
   | otherwise = Nothing
 
-newImportInsertRange :: ParsedSource -> Maybe (Range, Int)
-newImportInsertRange (L _ HsModule {..})
+-- | Finds the next valid position for inserting a new import declaration
+-- * If the file already has existing imports it will be inserted under the last of these,
+-- it is assumed that the existing last import declaration is in a valid position
+-- * If the file does not have existing imports, but has a (module ... where) declaration,
+-- the new import will be inserted directly under this declaration (accounting for explicit exports)
+-- * If the file has neither existing imports nor a module declaration,
+-- the import will be inserted at line zero if there are no pragmas,
+-- * otherwise inserted one line after the last file-header pragma
+newImportInsertRange :: ParsedSource -> T.Text -> Maybe (Range, Int)
+newImportInsertRange (L _ HsModule {..}) fileContents
   |  Just (uncurry Position -> insertPos, col) <- case hsmodImports of
-      [] -> case getLoc (head hsmodDecls) of
-        OldRealSrcSpan s -> let col = srcLocCol (realSrcSpanStart s) - 1
-              in Just ((srcLocLine (realSrcSpanStart s) - 1, col), col)
-        _            -> Nothing
-      _ -> case  getLoc (last hsmodImports) of
-        OldRealSrcSpan s -> let col = srcLocCol (realSrcSpanStart s) - 1
-            in Just ((srcLocLine $ realSrcSpanEnd s,col), col)
-        _            -> Nothing
+      [] -> findPositionNoImports hsmodName hsmodExports fileContents
+      _  -> findPositionFromImportsOrModuleDecl hsmodImports last True
     = Just (Range insertPos insertPos, col)
   | otherwise = Nothing
 
+-- | Insert the import under the Module declaration exports if they exist, otherwise just under the module declaration.
+-- If no module declaration exists, then no exports will exist either, in that case
+-- insert the import after any file-header pragmas or at position zero if there are no pragmas
+findPositionNoImports :: Maybe (Located ModuleName) -> Maybe (Located [LIE name]) -> T.Text -> Maybe ((Int, Int), Int)
+findPositionNoImports Nothing _ fileContents = findNextPragmaPosition fileContents
+findPositionNoImports _ (Just hsmodExports) _ = findPositionFromImportsOrModuleDecl hsmodExports id False
+findPositionNoImports (Just hsmodName) _ _ = findPositionFromImportsOrModuleDecl hsmodName id False
+
+findPositionFromImportsOrModuleDecl :: HasSrcSpan a => t -> (t -> a) -> Bool -> Maybe ((Int, Int), Int)
+findPositionFromImportsOrModuleDecl hsField f hasImports = case getLoc (f hsField) of
+  OldRealSrcSpan s ->
+    let col = calcCol s
+     in Just ((srcLocLine (realSrcSpanEnd s), col), col)
+  _ -> Nothing
+  where calcCol s = if hasImports then srcLocCol (realSrcSpanStart s) - 1 else 0
+
+-- | Find the position one after the last file-header pragma
+-- Defaults to zero if there are no pragmas in file
+findNextPragmaPosition :: T.Text -> Maybe ((Int, Int), Int)
+findNextPragmaPosition contents = Just ((lineNumber, 0), 0)
+  where
+    lineNumber = afterLangPragma . afterOptsGhc $ afterShebang
+    afterLangPragma = afterPragma "LANGUAGE" contents'
+    afterOptsGhc = afterPragma "OPTIONS_GHC" contents'
+    afterShebang = lastLineWithPrefix (T.isPrefixOf "#!") contents' 0
+    contents' = T.lines contents
+
+afterPragma :: T.Text -> [T.Text] -> Int -> Int
+afterPragma name contents lineNum = lastLineWithPrefix (checkPragma name) contents lineNum
+
+lastLineWithPrefix :: (T.Text -> Bool) -> [T.Text] -> Int -> Int
+lastLineWithPrefix p contents lineNum = max lineNum next
+  where
+    next = maybe lineNum succ $ listToMaybe . reverse $ findIndices p contents
+
+checkPragma :: T.Text -> T.Text -> Bool
+checkPragma name = check
+  where
+    check l = isPragma l && getName l == name
+    getName l = T.take (T.length name) $ T.dropWhile isSpace $ T.drop 3 l
+    isPragma = T.isPrefixOf "{-#"
+
 -- | Construct an import declaration with at most one symbol
 newImport
   :: T.Text -- ^ module name

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

@@ -187,13 +187,14 @@ extendImportHandler' ideState ExtendImport {..}
   | Just fp <- uriToFilePath doc,
     nfp <- toNormalizedFilePath' fp =
     do
-      (ModSummaryResult {..}, ps) <- MaybeT $ liftIO $
+      (ModSummaryResult {..}, ps, contents) <- MaybeT $ liftIO $
         runAction "extend import" ideState $
           runMaybeT $ do
             -- We want accurate edits, so do not use stale data here
             msr <- MaybeT $ use GetModSummaryWithoutTimestamps nfp
             ps <- MaybeT $ use GetAnnotatedParsedSource nfp
-            return (msr, ps)
+            (_, contents) <- MaybeT $ use GetFileContents nfp
+            return (msr, ps, contents)
       let df = ms_hspp_opts msrModSummary
           wantedModule = mkModuleName (T.unpack importName)
           wantedQual = mkModuleName . T.unpack <$> importQual
@@ -209,7 +210,7 @@ extendImportHandler' ideState ExtendImport {..}
                 it = case thingParent of
                   Nothing -> newThing
                   Just p  -> p <> "(" <> newThing <> ")"
-            t <- liftMaybe $ snd <$> newImportToEdit n (astA ps)
+            t <- liftMaybe $ snd <$> newImportToEdit n (astA ps) (fromMaybe "" contents)
             return (nfp, WorkspaceEdit {_changes=Just (fromList [(doc,List [t])]), _documentChanges=Nothing, _changeAnnotations=Nothing})
   | otherwise =
     mzero

ghcide/test/data/import-placement/CommentAtTop.expected.hs

@@ -0,0 +1,9 @@
+module Test
+( SomeData(..)
+) where
+import Data.Monoid
+
+-- | Some comment
+class Semigroup a => SomeData a
+
+instance SomeData All

ghcide/test/data/import-placement/CommentAtTop.hs

@@ -0,0 +1,8 @@
+module Test
+( SomeData(..)
+) where
+
+-- | Some comment
+class Semigroup a => SomeData a
+
+instance SomeData All

ghcide/test/data/import-placement/CommentAtTopMultipleComments.expected.hs

@@ -0,0 +1,12 @@
+module Test
+( SomeData(..)
+) where
+import Data.Monoid
+
+-- | Another comment
+data SomethingElse = SomethingElse
+
+-- | Some comment
+class Semigroup a => SomeData a
+
+instance SomeData All

ghcide/test/data/import-placement/CommentAtTopMultipleComments.hs

@@ -0,0 +1,11 @@
+module Test
+( SomeData(..)
+) where
+
+-- | Another comment
+data SomethingElse = SomethingElse
+
+-- | Some comment
+class Semigroup a => SomeData a
+
+instance SomeData All

ghcide/test/data/import-placement/CommentCurlyBraceAtTop.expected.hs

@@ -0,0 +1,9 @@
+module Test
+( SomeData(..)
+) where
+import Data.Monoid
+
+{- Some comment -}
+class Semigroup a => SomeData a
+
+instance SomeData All

ghcide/test/data/import-placement/CommentCurlyBraceAtTop.hs

@@ -0,0 +1,8 @@
+module Test
+( SomeData(..)
+) where
+
+{- Some comment -}
+class Semigroup a => SomeData a
+
+instance SomeData All

ghcide/test/data/import-placement/DataAtTop.expected.hs

@@ -0,0 +1,11 @@
+module Test
+( SomeData(..)
+) where
+import Data.Monoid
+
+data Something = Something
+
+-- | some comment
+class Semigroup a => SomeData a
+
+instance SomeData All

ghcide/test/data/import-placement/DataAtTop.hs

@@ -0,0 +1,10 @@
+module Test
+( SomeData(..)
+) where
+
+data Something = Something
+
+-- | some comment
+class Semigroup a => SomeData a
+
+instance SomeData All

ghcide/test/data/import-placement/ImportAtTop.expected.hs

@@ -0,0 +1,14 @@
+module Test
+( SomeData(..)
+) where
+
+import Data.Char
+import Data.Monoid
+
+{- Some multi 
+   line comment 
+-}
+class Semigroup a => SomeData a
+
+-- | a comment
+instance SomeData All

ghcide/test/data/import-placement/ImportAtTop.hs

@@ -0,0 +1,13 @@
+module Test
+( SomeData(..)
+) where
+
+import Data.Char
+
+{- Some multi 
+   line comment 
+-}
+class Semigroup a => SomeData a
+
+-- | a comment
+instance SomeData All

ghcide/test/data/import-placement/LangPragmaModuleAtTop.expected.hs

@@ -0,0 +1,8 @@
+{-# LANGUAGE OverloadedStrings #-}
+
+module Test where
+import Data.Monoid
+
+class Semigroup a => SomeData a
+
+instance SomeData All

ghcide/test/data/import-placement/LangPragmaModuleAtTop.hs

@@ -0,0 +1,7 @@
+{-# LANGUAGE OverloadedStrings #-}
+
+module Test where
+
+class Semigroup a => SomeData a
+
+instance SomeData All

ghcide/test/data/import-placement/LangPragmaModuleExplicitExports.expected.hs

@@ -0,0 +1,10 @@
+{-# LANGUAGE OverloadedStrings #-}
+
+module Test
+( SomeData(..)
+) where
+import Data.Monoid
+
+class Semigroup a => SomeData a
+
+instance SomeData All

ghcide/test/data/import-placement/LangPragmaModuleExplicitExports.hs

@@ -0,0 +1,9 @@
+{-# LANGUAGE OverloadedStrings #-}
+
+module Test
+( SomeData(..)
+) where
+
+class Semigroup a => SomeData a
+
+instance SomeData All

ghcide/test/data/import-placement/LangPragmaModuleWithComment.expected.hs

@@ -0,0 +1,11 @@
+{-# LANGUAGE OverloadedStrings #-}
+
+module Test
+( SomeData(..)
+) where
+import Data.Monoid
+
+-- comment
+class Semigroup a => SomeData a
+
+instance SomeData All