Insert pragmas after shebang or to existing pragma list

plugins/default/src/Ide/Plugin/Pragmas.hs

@@ -11,6 +11,7 @@ import           Control.Lens               hiding (List)
 import           Control.Monad              (join)
 import           Control.Monad.IO.Class
 import qualified Data.HashMap.Strict        as H
+import           Data.List
 import           Data.List.Extra            (nubOrdOn)
 import           Data.Maybe                 (catMaybes, listToMaybe)
 import qualified Data.Text                  as T
@@ -45,8 +46,9 @@ codeActionProvider state _plId (CodeActionParams _ _ docId _ (J.CodeActionContex
   let mFile = docId ^. J.uri & uriToFilePath <&> toNormalizedFilePath'
       uri = docId ^. J.uri
   pm <- liftIO $ fmap join $ runAction "Pragmas.GetParsedModule" state $ getParsedModule `traverse` mFile
+  mbContents <- liftIO $ fmap (join . fmap snd) $ runAction "Pragmas.GetFileContents" state $ getFileContents `traverse` mFile
   let dflags = ms_hspp_opts . pm_mod_summary <$> pm
-      insertRange = maybe (Range (Position 0 0) (Position 0 0)) endOfModuleHeader pm
+      insertRange = maybe (Range (Position 0 0) (Position 0 0)) endOfModuleHeader mbContents
       pedits = nubOrdOn snd . concat $ suggest dflags <$> diags
   return $ Right $ List $ pragmaEditToAction uri insertRange <$> pedits
 
@@ -178,14 +180,11 @@ completion _ide _ complParams = do
 
 -- ---------------------------------------------------------------------
 
--- | Find the first non-blank line before the first of (module name / imports / declarations).
+-- | Find first line after (last pragma / last shebang / beginning of file).
 -- Useful for inserting pragmas.
-endOfModuleHeader :: ParsedModule -> Range
-endOfModuleHeader pm =
-  let mod = unLoc $ pm_parsed_source pm
-      modNameLoc = getLoc <$> hsmodName mod
-      firstImportLoc = getLoc <$> listToMaybe (hsmodImports mod)
-      firstDeclLoc = getLoc <$> listToMaybe (hsmodDecls mod)
-      line = maybe 0 (_line . _start) (modNameLoc <|> firstImportLoc <|> firstDeclLoc >>= srcSpanToRange)
-      loc = Position line 0
-   in Range loc loc
+endOfModuleHeader :: T.Text -> Range
+endOfModuleHeader contents = Range loc loc
+    where
+        loc = Position line 0
+        line = maybe 0 succ (lastLineWithPrefix "{-#" <|> lastLineWithPrefix "#!")
+        lastLineWithPrefix pre = listToMaybe $ reverse $ findIndices (T.isPrefixOf pre) $ T.lines contents

test/functional/FunctionalCodeAction.hs

@@ -511,14 +511,8 @@ missingPragmaTests = testGroup "missing pragma warning code actions" [
             contents <- documentContents doc
 
             let expected =
--- TODO: Why CPP???
-#if __GLASGOW_HASKELL__ < 810
                     [ "{-# LANGUAGE ScopedTypeVariables #-}"
                     , "{-# LANGUAGE TypeApplications #-}"
-#else
-                    [ "{-# LANGUAGE TypeApplications #-}"
-                    , "{-# LANGUAGE ScopedTypeVariables #-}"
-#endif
                     , "module TypeApplications where"
                     , ""
                     , "foo :: forall a. a -> a"
@@ -555,7 +549,7 @@ missingPragmaTests = testGroup "missing pragma warning code actions" [
                     , "f Record{a, b} = a"
                     ]
             liftIO $ T.lines contents @?= expected
-    , testCase "After Shebang" $ do
+    , testCase "After shebang" $ do
         runSession hlsCommand fullCaps "test/testdata/addPragmas" $ do
             doc <- openDoc "AfterShebang.hs" "haskell"
 
@@ -571,8 +565,8 @@ missingPragmaTests = testGroup "missing pragma warning code actions" [
             let expected =
                     [ "#! /usr/bin/env nix-shell"
                     , "#! nix-shell --pure -i runghc -p \"haskellPackages.ghcWithPackages (hp: with hp; [ turtle ])\""
-                    , ""
                     , "{-# LANGUAGE NamedFieldPuns #-}"
+                    , ""
                     , "module AfterShebang where"
                     , ""
                     , "data Record = Record"
@@ -584,6 +578,67 @@ missingPragmaTests = testGroup "missing pragma warning code actions" [
                     , "f Record{a, b} = a"
                     ]
 
+            liftIO $ T.lines contents @?= expected
+    , testCase "Append to existing pragmas" $ do
+        runSession hlsCommand fullCaps "test/testdata/addPragmas" $ do
+            doc <- openDoc "AppendToExisting.hs" "haskell"
+
+            _ <- waitForDiagnosticsFrom doc
+            cas <- map fromAction <$> getAllCodeActions doc
+
+            liftIO $ "Add \"NamedFieldPuns\"" `elem` map (^. L.title) cas @? "Contains NamedFieldPuns code action"
+
+            executeCodeAction $ head cas
+
+            contents <- documentContents doc
+
+            let expected =
+                    [ "-- | Doc before pragma"
+                    , "{-# OPTIONS_GHC -Wno-dodgy-imports #-}"
+                    , "{-# LANGUAGE NamedFieldPuns #-}"
+                    , "module AppendToExisting where"
+                    , ""
+                    , "data Record = Record"
+                    , "  { a :: Int,"
+                    , "    b :: Double,"
+                    , "    c :: String"
+                    , "  }"
+                    , ""
+                    , "f Record{a, b} = a"
+                    ]
+
+            liftIO $ T.lines contents @?= expected
+    , testCase "Before Doc Comments" $ do
+        runSession hlsCommand fullCaps "test/testdata/addPragmas" $ do
+            doc <- openDoc "BeforeDocComment.hs" "haskell"
+
+            _ <- waitForDiagnosticsFrom doc
+            cas <- map fromAction <$> getAllCodeActions doc
+
+            liftIO $ "Add \"NamedFieldPuns\"" `elem` map (^. L.title) cas @? "Contains NamedFieldPuns code action"
+
+            executeCodeAction $ head cas
+
+            contents <- documentContents doc
+
+            let expected =
+                    [ "#! /usr/bin/env nix-shell"
+                    , "#! nix-shell --pure -i runghc -p \"haskellPackages.ghcWithPackages (hp: with hp; [ turtle ])\""
+                    , "{-# LANGUAGE NamedFieldPuns #-}"
+                    , "-- | Doc Comment"
+                    , "{- Block -}"
+                    , ""
+                    , "module BeforeDocComment where"
+                    , ""
+                    , "data Record = Record"
+                    , "  { a :: Int,"
+                    , "    b :: Double,"
+                    , "    c :: String"
+                    , "  }"
+                    , ""
+                    , "f Record{a, b} = a"
+                    ]
+
             liftIO $ T.lines contents @?= expected
     ]
 
@@ -614,9 +669,9 @@ disableWarningTests =
           ]
       , T.unlines
           [ "{-# OPTIONS_GHC -Wall #-}"
+          , "{-# OPTIONS_GHC -Wno-unused-imports #-}"
           , ""
           , ""
-          , "{-# OPTIONS_GHC -Wno-unused-imports #-}"
           , "module M where"
           , ""
           , "import Data.Functor"

test/testdata/addPragmas/AppendToExisting.hs

@@ -0,0 +1,11 @@
+-- | Doc before pragma
+{-# OPTIONS_GHC -Wno-dodgy-imports #-}
+module AppendToExisting where
+
+data Record = Record
+  { a :: Int,
+    b :: Double,
+    c :: String
+  }
+
+f Record{a, b} = a

test/testdata/addPragmas/BeforeDocComment.hs

@@ -0,0 +1,14 @@
+#! /usr/bin/env nix-shell
+#! nix-shell --pure -i runghc -p "haskellPackages.ghcWithPackages (hp: with hp; [ turtle ])"
+-- | Doc Comment
+{- Block -}
+
+module BeforeDocComment where
+
+data Record = Record
+  { a :: Int,
+    b :: Double,
+    c :: String
+  }
+
+f Record{a, b} = a