From 3bfd8be6e1c6a5e90b1cdf9ee5723e93de26d409 Mon Sep 17 00:00:00 2001
From: Vaibhav Sagar <vaibhavsagar@gmail.com>
Date: Wed, 23 Aug 2017 22:19:51 +0800
Subject: Remove shebang lines at the beginning of file

---
 lib/Language/Haskell/Stylish/Parse.hs         |  8 ++++----
 tests/Language/Haskell/Stylish/Parse/Tests.hs | 11 +++++++++++
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/lib/Language/Haskell/Stylish/Parse.hs b/lib/Language/Haskell/Stylish/Parse.hs
index 724ebe2..01def63 100644
--- a/lib/Language/Haskell/Stylish/Parse.hs
+++ b/lib/Language/Haskell/Stylish/Parse.hs
@@ -44,11 +44,11 @@ unCpp = unlines . go False . lines
 
 
 --------------------------------------------------------------------------------
--- | Remove shebang from the first line
+-- | Remove shebang lines
 unShebang :: String -> String
-unShebang str
-    | "#!" `isPrefixOf` str = unlines $ ("" :) $ drop 1 $ lines str
-    | otherwise             = str
+unShebang str =
+    let (shebangs, other) = break (not . ("#!" `isPrefixOf`)) (lines str) in
+    unlines $ map (const "") shebangs ++ other
 
 
 --------------------------------------------------------------------------------
diff --git a/tests/Language/Haskell/Stylish/Parse/Tests.hs b/tests/Language/Haskell/Stylish/Parse/Tests.hs
index 9806be2..a8ebf39 100644
--- a/tests/Language/Haskell/Stylish/Parse/Tests.hs
+++ b/tests/Language/Haskell/Stylish/Parse/Tests.hs
@@ -22,6 +22,7 @@ tests = testGroup "Language.Haskell.Stylish.Parse"
     , testCase "Haskell2010 extension"       testHaskell2010
     , testCase "Shebang"                     testShebang
     , testCase "ShebangExt"                  testShebangExt
+    , testCase "ShebangDouble"               testShebangDouble
     , testCase "GADTs extension"             testGADTs
     , testCase "KindSignatures extension"    testKindSignatures
     , testCase "StandalonDeriving extension" testStandaloneDeriving
@@ -85,6 +86,16 @@ testShebang = assert $ isRight $ parseModule [] Nothing $ unlines
 
 --------------------------------------------------------------------------------
 
+testShebangDouble :: Assertion
+testShebangDouble = assert $ isRight $ parseModule [] Nothing $ unlines
+    [ "#!nix-shell"
+    , "#!nix-shell -i runhaskell -p haskellPackages.ghc"
+    , "module Main where"
+    , "main = return ()"
+    ]
+
+--------------------------------------------------------------------------------
+
 -- | These tests are for syntactic language extensions that should always be
 -- enabled for parsing, even when the pragma is absent.
 
-- 
cgit v1.2.3