{-# LANGUAGE OverloadedStrings #-}

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

import Data.Char (isSpace)
import Data.IntSet (IntSet)
import qualified Data.IntSet as IntSet
import qualified Data.List as L
import Data.Maybe (isJust)

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

-- | Return an 'IntSet' containing all lines which are affected by CPP.
cppLines :: String -> IntSet
cppLines :: String -> IntSet
cppLines String
input = [Int] -> IntSet
IntSet.fromAscList ([Int] -> IntSet) -> [Int] -> IntSet
forall a b. (a -> b) -> a -> b
$ State -> [(String, Int)] -> [Int]
forall a. State -> [(String, a)] -> [a]
go State
Outside (String -> [String]
lines String
input [String] -> [Int] -> [(String, Int)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [Int
1 ..])
  where
    go :: State -> [(String, a)] -> [a]
go State
_ [] = []
    go State
state ((String
line, a
i) : [(String, a)]
ls)
      | (String -> Bool) -> [String] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any String -> Bool
for [String
"define ", String
"include ", String
"undef "] =
        a
i a -> [a] -> [a]
forall a. a -> [a] -> [a]
: State -> [(String, a)] -> [a]
go State
contState [(String, a)]
ls
      | (String -> Bool) -> [String] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any String -> Bool
for [String
"ifdef ", String
"ifndef ", String
"if "] =
        let state' :: State
state' = case State
state of
              InConditional nc -> Int -> State
InConditional (Int
nc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
              State
_ -> Int -> State
InConditional Int
1
         in a
i a -> [a] -> [a]
forall a. a -> [a] -> [a]
: State -> [(String, a)] -> [a]
go State
state' [(String, a)]
ls
      | String -> Bool
for String
"endif" =
        let state' :: State
state' = case State
state of
              InConditional nc | Int
nc Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 -> Int -> State
InConditional (Int
nc Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
              State
_ -> State
Outside
         in a
i a -> [a] -> [a]
forall a. a -> [a] -> [a]
: State -> [(String, a)] -> [a]
go State
state' [(String, 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 -> [(String, a)] -> [a]
go State
state' [(String, a)]
ls
      where
        for :: String -> Bool
for String
directive = Maybe String -> Bool
forall a. Maybe a -> Bool
isJust (Maybe String -> Bool) -> Maybe String -> Bool
forall a b. (a -> b) -> a -> b
$ do
          String
s <- (Char -> Bool) -> ShowS
forall a. (a -> Bool) -> [a] -> [a]
dropWhile Char -> Bool
isSpace ShowS -> Maybe String -> Maybe String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> String -> Maybe String
forall a. Eq a => [a] -> [a] -> Maybe [a]
L.stripPrefix String
"#" String
line
          String -> String -> Maybe String
forall a. Eq a => [a] -> [a] -> Maybe [a]
L.stripPrefix String
directive String
s
        contState :: State
contState =
          if String
"\\" String -> String -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`L.isSuffixOf` String
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