    {-# LANGUAGE OverloadedStrings #-}

    -- | Support for CPP.
    module Ormolu.Processing.Cpp
      ( cppLines,
        eraseCppLines,
      )
    where

    import Data.IntSet (IntSet)
    import Data.IntSet qualified as IntSet
    import Data.Maybe (isJust)
    import Data.Text (Text)
    import Data.Text qualified as T

    -- | State of the CPP processor.
    data State
      = -- | Outside of CPP directives
        Outside
      | -- | In a conditional expression, with a positive nesting count
        InConditional Int
      | -- | In a continuation (after @\\@), but not in a conditional expression
        InContinuation
      deriving (State -> State -> Bool
    (State -> State -> Bool) -> (State -> State -> Bool) -> Eq State
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: State -> State -> Bool
    == :: State -> State -> Bool
    $c/= :: State -> State -> Bool
    /= :: State -> State -> Bool
    Eq, Key -> State -> ShowS
    [State] -> ShowS
    State -> String
    (Key -> State -> ShowS)
    -> (State -> String) -> ([State] -> ShowS) -> Show State
    forall a.
    (Key -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Key -> State -> ShowS
    showsPrec :: Key -> State -> ShowS
    $cshow :: State -> String
    show :: State -> String
    $cshowList :: [State] -> ShowS
    showList :: [State] -> ShowS
    Show)

    -- | Return an 'IntSet' containing all lines which are affected by CPP.
    cppLines :: Text -> IntSet
    cppLines :: Text -> IntSet
    cppLines Text
    input = [Key] -> IntSet
    IntSet.fromAscList ([Key] -> IntSet) -> [Key] -> IntSet
    forall a b. (a -> b) -> a -> b
    $ State -> [(Text, Key)] -> [Key]
    forall {a}. State -> [(Text, a)] -> [a]
    go State
    Outside (Text -> [Text]
    T.lines Text
    input [Text] -> [Key] -> [(Text, Key)]
    forall a b. [a] -> [b] -> [(a, b)]
    `zip` [Key
    1 ..])
      where
        go :: State -> [(Text, a)] -> [a]
    go State
    _ [] = []
        go State
    state ((Text
    line, a
    i) : [(Text, a)]
    ls)
          | (Text -> Bool) -> [Text] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any Text -> Bool
    for [Text
    "define ", Text
    "include ", Text
    "undef "] =
              a
    i a -> [a] -> [a]
    forall a. a -> [a] -> [a]
    : State -> [(Text, a)] -> [a]
    go State
    contState [(Text, a)]
    ls
          | (Text -> Bool) -> [Text] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any Text -> Bool
    for [Text
    "ifdef ", Text
    "ifndef ", Text
    "if "] =
              let state' :: State
    state' = case State
    state of
                    InConditional Key
    nc -> Key -> State
    InConditional (Key
    nc Key -> Key -> Key
    forall a. Num a => a -> a -> a
    + Key
    1)
                    State
    _ -> Key -> State
    InConditional Key
    1
               in a
    i a -> [a] -> [a]
    forall a. a -> [a] -> [a]
    : State -> [(Text, a)] -> [a]
    go State
    state' [(Text, a)]
    ls
          | Text -> Bool
    for Text
    "endif" =
              let state' :: State
    state' = case State
    state of
                    InConditional Key
    nc | Key
    nc Key -> Key -> Bool
    forall a. Ord a => a -> a -> Bool
    > Key
    1 -> Key -> State
    InConditional (Key
    nc Key -> Key -> Key
    forall a. Num a => a -> a -> a
    - Key
    1)
                    State
    _ -> State
    Outside
               in a
    i a -> [a] -> [a]
    forall a. a -> [a] -> [a]
    : State -> [(Text, a)] -> [a]
    go State
    state' [(Text, a)]
    ls
          | Bool
    otherwise =
              let is :: [a]
    is = case State
    state of
                    State
    Outside -> []
                    State
    _ -> [a
    i]
                  state' :: State
    state' = case State
    state of
                    State
    InContinuation -> State
    contState
                    State
    _ -> State
    state
               in [a]
    is [a] -> [a] -> [a]
    forall a. Semigroup a => a -> a -> a
    <> State -> [(Text, a)] -> [a]
    go State
    state' [(Text, a)]
    ls
          where
            for :: Text -> Bool
    for Text
    directive = Maybe Text -> Bool
    forall a. Maybe a -> Bool
    isJust (Maybe Text -> Bool) -> Maybe Text -> Bool
    forall a b. (a -> b) -> a -> b
    $ do
              s <- Text -> Text
    T.stripStart (Text -> Text) -> Maybe Text -> Maybe Text
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Text -> Text -> Maybe Text
    T.stripPrefix Text
    "#" Text
    line
              T.stripPrefix directive s
            contState :: State
    contState =
              if Text
    "\\" Text -> Text -> Bool
    `T.isSuffixOf` Text
    line Bool -> Bool -> Bool
    && Bool -> Bool
    not Bool
    inConditional
                then State
    InContinuation
                else State
    Outside
              where
                inConditional :: Bool
    inConditional = case State
    state of
                  InConditional {} -> Bool
    True
                  State
    _ -> Bool
    False

    -- | Replace all lines affected by CPP with blank lines.
    eraseCppLines :: Text -> Text
    eraseCppLines :: Text -> Text
    eraseCppLines Text
    input =
      [Text] -> Text
    T.unlines ([Text] -> Text)
    -> ([(Text, Key)] -> [Text]) -> [(Text, Key)] -> Text
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ((Text, Key) -> Text) -> [(Text, Key)] -> [Text]
    forall a b. (a -> b) -> [a] -> [b]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap (Text, Key) -> Text
    forall {a}. IsString a => (a, Key) -> a
    eraseCpp ([(Text, Key)] -> Text) -> [(Text, Key)] -> Text
    forall a b. (a -> b) -> a -> b
    $ Text -> [Text]
    T.lines Text
    input [Text] -> [Key] -> [(Text, Key)]
    forall a b. [a] -> [b] -> [(a, b)]
    `zip` [Key
    1 ..]
      where
        linesToErase :: IntSet
    linesToErase = Text -> IntSet
    cppLines Text
    input
        eraseCpp :: (a, Key) -> a
    eraseCpp (a
    x, Key
    i) =
          if Key
    i Key -> IntSet -> Bool
    `IntSet.member` IntSet
    linesToErase
            then a
    "\n"
            else a
    x