    {-# 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
    $cshowsPrec :: Int -> CodeBlock -> ShowS
    showsPrec :: Int -> CodeBlock -> ShowS
    $cshow :: CodeBlock -> String
    show :: CodeBlock -> String
    $cshowList :: [CodeBlock] -> ShowS
    showList :: [CodeBlock] -> ShowS
    Show, CodeBlock -> CodeBlock -> Bool
    (CodeBlock -> CodeBlock -> Bool)
    -> (CodeBlock -> CodeBlock -> Bool) -> Eq CodeBlock
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: CodeBlock -> CodeBlock -> Bool
    == :: CodeBlock -> CodeBlock -> Bool
    $c/= :: CodeBlock -> CodeBlock -> Bool
    /= :: 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 :: forall a. (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