{-# LANGUAGE OverloadedStrings #-}

module HIndent.CodeBlock
  ( CodeBlock(..)
  , cppSplitBlocks
  ) where

import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as S8
import Data.Monoid

-- | A block of code.
data CodeBlock
    = Shebang ByteString
    | HaskellSource Int ByteString
    -- ^ Includes the starting line (indexed from 0) for error reporting
    | CPPDirectives ByteString
     deriving (Int -> CodeBlock -> ShowS
[CodeBlock] -> ShowS
CodeBlock -> String
(Int -> CodeBlock -> ShowS)
-> (CodeBlock -> String)
-> ([CodeBlock] -> ShowS)
-> Show CodeBlock
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [CodeBlock] -> ShowS
$cshowList :: [CodeBlock] -> ShowS
show :: CodeBlock -> String
$cshow :: CodeBlock -> String
showsPrec :: Int -> CodeBlock -> ShowS
$cshowsPrec :: Int -> CodeBlock -> ShowS
Show, CodeBlock -> CodeBlock -> Bool
(CodeBlock -> CodeBlock -> Bool)
-> (CodeBlock -> CodeBlock -> Bool) -> Eq CodeBlock
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: CodeBlock -> CodeBlock -> Bool
$c/= :: CodeBlock -> CodeBlock -> Bool
== :: CodeBlock -> CodeBlock -> Bool
$c== :: CodeBlock -> CodeBlock -> Bool
Eq)

-- | Break a Haskell code string into chunks, using CPP as a delimiter.
-- Lines that start with '#if', '#end', or '#else' are their own chunks, and
-- also act as chunk separators. For example, the code
--
-- > #ifdef X
-- > x = X
-- > y = Y
-- > #else
-- > x = Y
-- > y = X
-- > #endif
--
-- will become five blocks, one for each CPP line and one for each pair of declarations.
cppSplitBlocks :: ByteString -> [CodeBlock]
cppSplitBlocks :: ByteString -> [CodeBlock]
cppSplitBlocks ByteString
inp =
  (CodeBlock -> CodeBlock) -> [CodeBlock] -> [CodeBlock]
forall a. (a -> a) -> [a] -> [a]
modifyLast ((ByteString -> ByteString) -> CodeBlock -> CodeBlock
inBlock (ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
trailing)) ([CodeBlock] -> [CodeBlock])
-> (ByteString -> [CodeBlock]) -> ByteString -> [CodeBlock]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
  [CodeBlock] -> [CodeBlock]
groupLines ([CodeBlock] -> [CodeBlock])
-> (ByteString -> [CodeBlock]) -> ByteString -> [CodeBlock]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Int, ByteString)] -> [CodeBlock]
classifyLines ([(Int, ByteString)] -> [CodeBlock])
-> (ByteString -> [(Int, ByteString)]) -> ByteString -> [CodeBlock]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [ByteString] -> [(Int, ByteString)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] ([ByteString] -> [(Int, ByteString)])
-> (ByteString -> [ByteString])
-> ByteString
-> [(Int, ByteString)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> [ByteString]
S8.lines (ByteString -> [CodeBlock]) -> ByteString -> [CodeBlock]
forall a b. (a -> b) -> a -> b
$
  ByteString
inp
  where
    groupLines :: [CodeBlock] -> [CodeBlock]
    groupLines :: [CodeBlock] -> [CodeBlock]
groupLines (CodeBlock
line1:CodeBlock
line2:[CodeBlock]
remainingLines) =
      case CodeBlock -> CodeBlock -> Maybe CodeBlock
mergeLines CodeBlock
line1 CodeBlock
line2 of
        Just CodeBlock
line1And2 -> [CodeBlock] -> [CodeBlock]
groupLines (CodeBlock
line1And2 CodeBlock -> [CodeBlock] -> [CodeBlock]
forall a. a -> [a] -> [a]
: [CodeBlock]
remainingLines)
        Maybe CodeBlock
Nothing -> CodeBlock
line1 CodeBlock -> [CodeBlock] -> [CodeBlock]
forall a. a -> [a] -> [a]
: [CodeBlock] -> [CodeBlock]
groupLines (CodeBlock
line2 CodeBlock -> [CodeBlock] -> [CodeBlock]
forall a. a -> [a] -> [a]
: [CodeBlock]
remainingLines)
    groupLines xs :: [CodeBlock]
xs@[CodeBlock
_] = [CodeBlock]
xs
    groupLines xs :: [CodeBlock]
xs@[] = [CodeBlock]
xs
    mergeLines :: CodeBlock -> CodeBlock -> Maybe CodeBlock
    mergeLines :: CodeBlock -> CodeBlock -> Maybe CodeBlock
mergeLines (CPPDirectives ByteString
src1) (CPPDirectives ByteString
src2) =
      CodeBlock -> Maybe CodeBlock
forall a. a -> Maybe a
Just (CodeBlock -> Maybe CodeBlock) -> CodeBlock -> Maybe CodeBlock
forall a b. (a -> b) -> a -> b
$ ByteString -> CodeBlock
CPPDirectives (ByteString
src1 ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
"\n" ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
src2)
    mergeLines (Shebang ByteString
src1) (Shebang ByteString
src2) =
      CodeBlock -> Maybe CodeBlock
forall a. a -> Maybe a
Just (CodeBlock -> Maybe CodeBlock) -> CodeBlock -> Maybe CodeBlock
forall a b. (a -> b) -> a -> b
$ ByteString -> CodeBlock
Shebang (ByteString
src1 ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
"\n" ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
src2)
    mergeLines (HaskellSource Int
lineNumber1 ByteString
src1) (HaskellSource Int
_lineNumber2 ByteString
src2) =
      CodeBlock -> Maybe CodeBlock
forall a. a -> Maybe a
Just (CodeBlock -> Maybe CodeBlock) -> CodeBlock -> Maybe CodeBlock
forall a b. (a -> b) -> a -> b
$ Int -> ByteString -> CodeBlock
HaskellSource Int
lineNumber1 (ByteString
src1 ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
"\n" ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
src2)
    mergeLines CodeBlock
_ CodeBlock
_ = Maybe CodeBlock
forall a. Maybe a
Nothing
    shebangLine :: ByteString -> Bool
    shebangLine :: ByteString -> Bool
shebangLine = ByteString -> ByteString -> Bool
S8.isPrefixOf ByteString
"#!"
    cppLine :: ByteString -> Bool
    cppLine :: ByteString -> Bool
cppLine ByteString
src =
      (ByteString -> Bool) -> [ByteString] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any
        (ByteString -> ByteString -> Bool
`S8.isPrefixOf` ByteString
src)
        [ByteString
"#if", ByteString
"#end", ByteString
"#else", ByteString
"#define", ByteString
"#undef", ByteString
"#elif", ByteString
"#include", ByteString
"#error", ByteString
"#warning"]
        -- Note: #ifdef and #ifndef are handled by #if
    hasEscapedTrailingNewline :: ByteString -> Bool
    hasEscapedTrailingNewline :: ByteString -> Bool
hasEscapedTrailingNewline ByteString
src = ByteString
"\\" ByteString -> ByteString -> Bool
`S8.isSuffixOf` ByteString
src
    classifyLines :: [(Int, ByteString)] -> [CodeBlock]
    classifyLines :: [(Int, ByteString)] -> [CodeBlock]
classifyLines allLines :: [(Int, ByteString)]
allLines@((Int
lineIndex, ByteString
src):[(Int, ByteString)]
nextLines)
      | ByteString -> Bool
cppLine ByteString
src =
        let ([(Int, ByteString)]
cppLines, [(Int, ByteString)]
nextLines') = [(Int, ByteString)] -> ([(Int, ByteString)], [(Int, ByteString)])
spanCPPLines [(Int, ByteString)]
allLines
         in ByteString -> CodeBlock
CPPDirectives (ByteString -> [ByteString] -> ByteString
S8.intercalate ByteString
"\n" (((Int, ByteString) -> ByteString)
-> [(Int, ByteString)] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map (Int, ByteString) -> ByteString
forall a b. (a, b) -> b
snd [(Int, ByteString)]
cppLines)) CodeBlock -> [CodeBlock] -> [CodeBlock]
forall a. a -> [a] -> [a]
:
            [(Int, ByteString)] -> [CodeBlock]
classifyLines [(Int, ByteString)]
nextLines'
      | ByteString -> Bool
shebangLine ByteString
src = ByteString -> CodeBlock
Shebang ByteString
src CodeBlock -> [CodeBlock] -> [CodeBlock]
forall a. a -> [a] -> [a]
: [(Int, ByteString)] -> [CodeBlock]
classifyLines [(Int, ByteString)]
nextLines
      | Bool
otherwise = Int -> ByteString -> CodeBlock
HaskellSource Int
lineIndex ByteString
src CodeBlock -> [CodeBlock] -> [CodeBlock]
forall a. a -> [a] -> [a]
: [(Int, ByteString)] -> [CodeBlock]
classifyLines [(Int, ByteString)]
nextLines
    classifyLines [] = []
    spanCPPLines ::
         [(Int, ByteString)] -> ([(Int, ByteString)], [(Int, ByteString)])
    spanCPPLines :: [(Int, ByteString)] -> ([(Int, ByteString)], [(Int, ByteString)])
spanCPPLines (line :: (Int, ByteString)
line@(Int
_, ByteString
src):[(Int, ByteString)]
nextLines)
      | ByteString -> Bool
hasEscapedTrailingNewline ByteString
src =
        let ([(Int, ByteString)]
cppLines, [(Int, ByteString)]
nextLines') = [(Int, ByteString)] -> ([(Int, ByteString)], [(Int, ByteString)])
spanCPPLines [(Int, ByteString)]
nextLines
         in ((Int, ByteString)
line (Int, ByteString) -> [(Int, ByteString)] -> [(Int, ByteString)]
forall a. a -> [a] -> [a]
: [(Int, ByteString)]
cppLines, [(Int, ByteString)]
nextLines')
      | Bool
otherwise = ([(Int, ByteString)
line], [(Int, ByteString)]
nextLines)
    spanCPPLines [] = ([], [])
    -- Hack to work around some parser issues in haskell-src-exts: Some pragmas
    -- need to have a newline following them in order to parse properly, so we include
    -- the trailing newline in the code block if it existed.
    trailing :: ByteString
    trailing :: ByteString
trailing =
      if ByteString -> ByteString -> Bool
S8.isSuffixOf ByteString
"\n" ByteString
inp
        then ByteString
"\n"
        else ByteString
""
    modifyLast :: (a -> a) -> [a] -> [a]
    modifyLast :: (a -> a) -> [a] -> [a]
modifyLast a -> a
_ [] = []
    modifyLast a -> a
f [a
x] = [a -> a
f a
x]
    modifyLast a -> a
f (a
x:[a]
xs) = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: (a -> a) -> [a] -> [a]
forall a. (a -> a) -> [a] -> [a]
modifyLast a -> a
f [a]
xs
    inBlock :: (ByteString -> ByteString) -> CodeBlock -> CodeBlock
    inBlock :: (ByteString -> ByteString) -> CodeBlock -> CodeBlock
inBlock ByteString -> ByteString
f (HaskellSource Int
line ByteString
txt) = Int -> ByteString -> CodeBlock
HaskellSource Int
line (ByteString -> ByteString
f ByteString
txt)
    inBlock ByteString -> ByteString
_ CodeBlock
dir = CodeBlock
dir