{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

-- | Preprocessing for input source code.
module Ormolu.Processing.Preprocess
  ( preprocess,
  )
where

import Control.Monad
import Data.Array as A
import Data.Bifunctor (bimap)
import Data.Char (isSpace)
import Data.Function ((&))
import Data.IntMap (IntMap)
import Data.IntMap.Strict qualified as IntMap
import Data.IntSet (IntSet)
import Data.IntSet qualified as IntSet
import Data.List qualified as L
import Data.Maybe (isJust)
import Data.Text (Text)
import Data.Text qualified as T
import Ormolu.Config (RegionDeltas (..))
import Ormolu.Processing.Common
import Ormolu.Processing.Cpp

-- | Preprocess the specified region of the input into raw snippets
-- and subregions to be formatted.
preprocess ::
  -- | Whether CPP is enabled
  Bool ->
  RegionDeltas ->
  Text ->
  [Either Text RegionDeltas]
preprocess :: Bool -> RegionDeltas -> Text -> [Either Text RegionDeltas]
preprocess Bool
cppEnabled RegionDeltas
region Text
rawInput = [Either Text RegionDeltas]
rawSnippetsAndRegionsToFormat
  where
    (IntSet
linesNotToFormat', IntMap Text
replacementLines) = Bool -> RegionDeltas -> Text -> (IntSet, IntMap Text)
linesNotToFormat Bool
cppEnabled RegionDeltas
region Text
rawInput
    regionsToFormat :: [RegionDeltas]
regionsToFormat =
      Key -> IntSet -> [RegionDeltas]
intSetToRegions Key
rawLineLength (IntSet -> [RegionDeltas]) -> IntSet -> [RegionDeltas]
forall a b. (a -> b) -> a -> b
$
        [Key] -> IntSet
IntSet.fromAscList [Key
1 .. Key
rawLineLength] IntSet -> IntSet -> IntSet
IntSet.\\ IntSet
linesNotToFormat'
    regionsNotToFormat :: [RegionDeltas]
regionsNotToFormat = Key -> IntSet -> [RegionDeltas]
intSetToRegions Key
rawLineLength IntSet
linesNotToFormat'
    -- We want to interleave the regionsToFormat and regionsNotToFormat.
    -- If the first non-formattable region starts at the first line, it is
    -- the first interleaved region, otherwise, we start with the first
    -- region to format.
    interleave' :: [a] -> [a] -> [a]
interleave' = case [RegionDeltas]
regionsNotToFormat of
      RegionDeltas
r : [RegionDeltas]
_ | RegionDeltas -> Key
regionPrefixLength RegionDeltas
r Key -> Key -> Bool
forall a. Eq a => a -> a -> Bool
== Key
0 -> [a] -> [a] -> [a]
forall {a}. [a] -> [a] -> [a]
interleave
      [RegionDeltas]
_ -> ([a] -> [a] -> [a]) -> [a] -> [a] -> [a]
forall a b c. (a -> b -> c) -> b -> a -> c
flip [a] -> [a] -> [a]
forall {a}. [a] -> [a] -> [a]
interleave
    rawSnippets :: [Text]
rawSnippets = (RegionDeltas -> Text -> Text) -> Text -> RegionDeltas -> Text
forall a b c. (a -> b -> c) -> b -> a -> c
flip RegionDeltas -> Text -> Text
linesInRegion Text
updatedInput (RegionDeltas -> Text) -> [RegionDeltas] -> [Text]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [RegionDeltas]
regionsNotToFormat
      where
        updatedInput :: Text
updatedInput = [Text] -> Text
T.unlines ([Text] -> Text) -> (Text -> [Text]) -> Text -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Key, Text) -> Text) -> [(Key, Text)] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Key, Text) -> Text
updateLine ([(Key, Text)] -> [Text])
-> (Text -> [(Key, Text)]) -> Text -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Key] -> [Text] -> [(Key, Text)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Key
1 ..] ([Text] -> [(Key, Text)])
-> (Text -> [Text]) -> Text -> [(Key, Text)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> [Text]
T.lines (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ Text
rawInput
        updateLine :: (Key, Text) -> Text
updateLine (Key
i, Text
line) = Text -> Key -> IntMap Text -> Text
forall a. a -> Key -> IntMap a -> a
IntMap.findWithDefault Text
line Key
i IntMap Text
replacementLines
    rawSnippetsAndRegionsToFormat :: [Either Text RegionDeltas]
rawSnippetsAndRegionsToFormat =
      [Either Text RegionDeltas]
-> [Either Text RegionDeltas] -> [Either Text RegionDeltas]
forall {a}. [a] -> [a] -> [a]
interleave' (Text -> Either Text RegionDeltas
forall a b. a -> Either a b
Left (Text -> Either Text RegionDeltas)
-> [Text] -> [Either Text RegionDeltas]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Text]
rawSnippets) (RegionDeltas -> Either Text RegionDeltas
forall a b. b -> Either a b
Right (RegionDeltas -> Either Text RegionDeltas)
-> [RegionDeltas] -> [Either Text RegionDeltas]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [RegionDeltas]
regionsToFormat)
        [Either Text RegionDeltas]
-> (Either Text RegionDeltas -> [Either Text RegionDeltas])
-> [Either Text RegionDeltas]
forall a b. [a] -> (a -> [b]) -> [b]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Either Text RegionDeltas -> [Either Text RegionDeltas]
patchSeparatingBlankLines
        [Either Text RegionDeltas]
-> ([Either Text RegionDeltas] -> [Either Text RegionDeltas])
-> [Either Text RegionDeltas]
forall a b. a -> (a -> b) -> b
& (Either Text RegionDeltas -> Bool)
-> [Either Text RegionDeltas] -> [Either Text RegionDeltas]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile Either Text RegionDeltas -> Bool
forall {b}. Either Text b -> Bool
isBlankRawSnippet
        [Either Text RegionDeltas]
-> ([Either Text RegionDeltas] -> [Either Text RegionDeltas])
-> [Either Text RegionDeltas]
forall a b. a -> (a -> b) -> b
& (Either Text RegionDeltas -> Bool)
-> [Either Text RegionDeltas] -> [Either Text RegionDeltas]
forall a. (a -> Bool) -> [a] -> [a]
L.dropWhileEnd Either Text RegionDeltas -> Bool
forall {b}. Either Text b -> Bool
isBlankRawSnippet
    -- For every formattable region, we want to ensure that it is separated by
    -- a blank line from preceding/succeeding raw snippets if it starts/ends
    -- with a blank line.
    -- Empty formattable regions are replaced by a blank line instead.
    -- Extraneous raw snippets at the start/end are dropped afterwards.
    patchSeparatingBlankLines :: Either Text RegionDeltas -> [Either Text RegionDeltas]
patchSeparatingBlankLines = \case
      Right r :: RegionDeltas
r@RegionDeltas {Key
regionPrefixLength :: RegionDeltas -> Key
regionPrefixLength :: Key
regionSuffixLength :: Key
regionSuffixLength :: RegionDeltas -> Key
..} ->
        if (Char -> Bool) -> Text -> Bool
T.all Char -> Bool
isSpace (RegionDeltas -> Text -> Text
linesInRegion RegionDeltas
r Text
rawInput)
          then [Either Text RegionDeltas
forall {b}. Either Text b
blankRawSnippet]
          else
            [Either Text RegionDeltas
forall {b}. Either Text b
blankRawSnippet | Key -> Bool
isBlankLine Key
regionPrefixLength]
              [Either Text RegionDeltas]
-> [Either Text RegionDeltas] -> [Either Text RegionDeltas]
forall a. Semigroup a => a -> a -> a
<> [RegionDeltas -> Either Text RegionDeltas
forall a b. b -> Either a b
Right RegionDeltas
r]
              [Either Text RegionDeltas]
-> [Either Text RegionDeltas] -> [Either Text RegionDeltas]
forall a. Semigroup a => a -> a -> a
<> [Either Text RegionDeltas
forall {b}. Either Text b
blankRawSnippet | Key -> Bool
isBlankLine (Key
rawLineLength Key -> Key -> Key
forall a. Num a => a -> a -> a
- Key
regionSuffixLength Key -> Key -> Key
forall a. Num a => a -> a -> a
- Key
1)]
      Left Text
r -> [Text -> Either Text RegionDeltas
forall a b. a -> Either a b
Left Text
r]
      where
        blankRawSnippet :: Either Text b
blankRawSnippet = Text -> Either Text b
forall a b. a -> Either a b
Left Text
"\n"
        isBlankLine :: Key -> Bool
isBlankLine Key
i = Maybe Text -> Bool
forall a. Maybe a -> Bool
isJust (Maybe Text -> Bool)
-> (Maybe Text -> Maybe Text) -> Maybe Text -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text -> Bool) -> Maybe Text -> Maybe Text
forall (m :: * -> *) a. MonadPlus m => (a -> Bool) -> m a -> m a
mfilter ((Char -> Bool) -> Text -> Bool
T.all Char -> Bool
isSpace) (Maybe Text -> Bool) -> Maybe Text -> Bool
forall a b. (a -> b) -> a -> b
$ Array Key Text
rawLines Array Key Text -> Key -> Maybe Text
forall {a}. Array Key a -> Key -> Maybe a
!!? Key
i
    isBlankRawSnippet :: Either Text b -> Bool
isBlankRawSnippet = \case
      Left Text
r | (Char -> Bool) -> Text -> Bool
T.all Char -> Bool
isSpace Text
r -> Bool
True
      Either Text b
_ -> Bool
False

    rawLines :: Array Key Text
rawLines = (Key, Key) -> [Text] -> Array Key Text
forall i e. Ix i => (i, i) -> [e] -> Array i e
A.listArray (Key
0, [Text] -> Key
forall a. [a] -> Key
forall (t :: * -> *) a. Foldable t => t a -> Key
length [Text]
rawLines' Key -> Key -> Key
forall a. Num a => a -> a -> a
- Key
1) [Text]
rawLines'
      where
        rawLines' :: [Text]
rawLines' = Text -> [Text]
T.lines Text
rawInput
    rawLineLength :: Key
rawLineLength = Array Key Text -> Key
forall a. Array Key a -> Key
forall (t :: * -> *) a. Foldable t => t a -> Key
length Array Key Text
rawLines

    interleave :: [a] -> [a] -> [a]
interleave [] [a]
bs = [a]
bs
    interleave (a
a : [a]
as) [a]
bs = a
a a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
interleave [a]
bs [a]
as

    Array Key a
xs !!? :: Array Key a -> Key -> Maybe a
!!? Key
i = if Array Key Text -> (Key, Key)
forall i e. Array i e -> (i, i)
A.bounds Array Key Text
rawLines (Key, Key) -> Key -> Bool
forall a. Ix a => (a, a) -> a -> Bool
`A.inRange` Key
i then a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ Array Key a
xs Array Key a -> Key -> a
forall i e. Ix i => Array i e -> i -> e
A.! Key
i else Maybe a
forall a. Maybe a
Nothing

-- | All lines we are not supposed to format, and a set of replacements
-- for specific lines.
linesNotToFormat ::
  -- | Whether CPP is enabled
  Bool ->
  RegionDeltas ->
  Text ->
  (IntSet, IntMap Text)
linesNotToFormat :: Bool -> RegionDeltas -> Text -> (IntSet, IntMap Text)
linesNotToFormat Bool
cppEnabled region :: RegionDeltas
region@RegionDeltas {Key
regionPrefixLength :: RegionDeltas -> Key
regionSuffixLength :: RegionDeltas -> Key
regionPrefixLength :: Key
regionSuffixLength :: Key
..} Text
input =
  (IntSet
unconsidered IntSet -> IntSet -> IntSet
forall a. Semigroup a => a -> a -> a
<> IntSet
magicDisabled IntSet -> IntSet -> IntSet
forall a. Semigroup a => a -> a -> a
<> IntSet
otherDisabled, IntMap Text
lineUpdates)
  where
    unconsidered :: IntSet
unconsidered =
      [Key] -> IntSet
IntSet.fromAscList ([Key] -> IntSet) -> [Key] -> IntSet
forall a b. (a -> b) -> a -> b
$
        [Key
1 .. Key
regionPrefixLength] [Key] -> [Key] -> [Key]
forall a. Semigroup a => a -> a -> a
<> [Key
totalLines Key -> Key -> Key
forall a. Num a => a -> a -> a
- Key
regionSuffixLength Key -> Key -> Key
forall a. Num a => a -> a -> a
+ Key
1 .. Key
totalLines]
    totalLines :: Key
totalLines = [Text] -> Key
forall a. [a] -> Key
forall (t :: * -> *) a. Foldable t => t a -> Key
length (Text -> [Text]
T.lines Text
input)
    regionLines :: Text
regionLines = RegionDeltas -> Text -> Text
linesInRegion RegionDeltas
region Text
input
    (IntSet
magicDisabled, IntMap Text
lineUpdates) = Text -> (IntSet, IntMap Text)
magicDisabledLines Text
regionLines
    otherDisabled :: IntSet
otherDisabled = [Text -> IntSet] -> Text -> IntSet
forall a. Monoid a => [a] -> a
mconcat [Text -> IntSet]
allLines Text
regionLines
      where
        allLines :: [Text -> IntSet]
allLines = [Text -> IntSet
shebangLines, Text -> IntSet
linePragmaLines] [Text -> IntSet] -> [Text -> IntSet] -> [Text -> IntSet]
forall a. Semigroup a => a -> a -> a
<> [Text -> IntSet
cppLines | Bool
cppEnabled]

-- | Ormolu state.
data OrmoluState
  = -- | Enabled
    OrmoluEnabled
  | -- | Disabled
    OrmoluDisabled
  deriving (OrmoluState -> OrmoluState -> Bool
(OrmoluState -> OrmoluState -> Bool)
-> (OrmoluState -> OrmoluState -> Bool) -> Eq OrmoluState
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: OrmoluState -> OrmoluState -> Bool
== :: OrmoluState -> OrmoluState -> Bool
$c/= :: OrmoluState -> OrmoluState -> Bool
/= :: OrmoluState -> OrmoluState -> Bool
Eq, Key -> OrmoluState -> ShowS
[OrmoluState] -> ShowS
OrmoluState -> String
(Key -> OrmoluState -> ShowS)
-> (OrmoluState -> String)
-> ([OrmoluState] -> ShowS)
-> Show OrmoluState
forall a.
(Key -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Key -> OrmoluState -> ShowS
showsPrec :: Key -> OrmoluState -> ShowS
$cshow :: OrmoluState -> String
show :: OrmoluState -> String
$cshowList :: [OrmoluState] -> ShowS
showList :: [OrmoluState] -> ShowS
Show)

-- | All lines which are disabled by Ormolu's magic comments,
-- as well as normalizing replacements.
magicDisabledLines :: Text -> (IntSet, IntMap Text)
magicDisabledLines :: Text -> (IntSet, IntMap Text)
magicDisabledLines Text
input =
  ([Key] -> IntSet)
-> ([(Key, Text)] -> IntMap Text)
-> ([Key], [(Key, Text)])
-> (IntSet, IntMap Text)
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap [Key] -> IntSet
IntSet.fromAscList [(Key, Text)] -> IntMap Text
forall a. [(Key, a)] -> IntMap a
IntMap.fromAscList (([Key], [(Key, Text)]) -> (IntSet, IntMap Text))
-> ([([Key], [(Key, Text)])] -> ([Key], [(Key, Text)]))
-> [([Key], [(Key, Text)])]
-> (IntSet, IntMap Text)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [([Key], [(Key, Text)])] -> ([Key], [(Key, Text)])
forall a. Monoid a => [a] -> a
mconcat ([([Key], [(Key, Text)])] -> (IntSet, IntMap Text))
-> [([Key], [(Key, Text)])] -> (IntSet, IntMap Text)
forall a b. (a -> b) -> a -> b
$
    OrmoluState -> [(Text, Key)] -> [([Key], [(Key, Text)])]
forall {a}. OrmoluState -> [(Text, a)] -> [([a], [(a, Text)])]
go OrmoluState
OrmoluEnabled (Text -> [Text]
T.lines Text
input [Text] -> [Key] -> [(Text, Key)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [Key
1 ..])
  where
    go :: OrmoluState -> [(Text, a)] -> [([a], [(a, Text)])]
go OrmoluState
_ [] = []
    go OrmoluState
state ((Text
line, a
i) : [(Text, a)]
ls)
      | Just Text
rest <- Text -> Text -> Maybe Text
isMagicComment Text
ormoluDisable Text
line,
        OrmoluState
state OrmoluState -> OrmoluState -> Bool
forall a. Eq a => a -> a -> Bool
== OrmoluState
OrmoluEnabled =
          ([a
i], [(a
i, Text -> Text
magicComment Text
ormoluDisable Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
rest)]) ([a], [(a, Text)]) -> [([a], [(a, Text)])] -> [([a], [(a, Text)])]
forall a. a -> [a] -> [a]
: OrmoluState -> [(Text, a)] -> [([a], [(a, Text)])]
go OrmoluState
OrmoluDisabled [(Text, a)]
ls
      | Just Text
rest <- Text -> Text -> Maybe Text
isMagicComment Text
ormoluEnable Text
line,
        OrmoluState
state OrmoluState -> OrmoluState -> Bool
forall a. Eq a => a -> a -> Bool
== OrmoluState
OrmoluDisabled =
          ([a
i], [(a
i, Text -> Text
magicComment Text
ormoluEnable Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
rest)]) ([a], [(a, Text)]) -> [([a], [(a, Text)])] -> [([a], [(a, Text)])]
forall a. a -> [a] -> [a]
: OrmoluState -> [(Text, a)] -> [([a], [(a, Text)])]
go OrmoluState
OrmoluEnabled [(Text, a)]
ls
      | Bool
otherwise = ([a], [(a, Text)])
forall {a}. ([a], [a])
iIfDisabled ([a], [(a, Text)]) -> [([a], [(a, Text)])] -> [([a], [(a, Text)])]
forall a. a -> [a] -> [a]
: OrmoluState -> [(Text, a)] -> [([a], [(a, Text)])]
go OrmoluState
state [(Text, a)]
ls
      where
        iIfDisabled :: ([a], [a])
iIfDisabled = case OrmoluState
state of
          OrmoluState
OrmoluDisabled -> ([a
i], [])
          OrmoluState
OrmoluEnabled -> ([], [])

-- | All lines which satisfy a predicate.
linesFiltered :: (Text -> Bool) -> Text -> IntSet
linesFiltered :: (Text -> Bool) -> Text -> IntSet
linesFiltered Text -> Bool
p =
  [Key] -> IntSet
IntSet.fromAscList ([Key] -> IntSet) -> (Text -> [Key]) -> Text -> IntSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Text, Key) -> Key) -> [(Text, Key)] -> [Key]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Text, Key) -> Key
forall a b. (a, b) -> b
snd ([(Text, Key)] -> [Key])
-> (Text -> [(Text, Key)]) -> Text -> [Key]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Text, Key) -> Bool) -> [(Text, Key)] -> [(Text, Key)]
forall a. (a -> Bool) -> [a] -> [a]
filter (Text -> Bool
p (Text -> Bool) -> ((Text, Key) -> Text) -> (Text, Key) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text, Key) -> Text
forall a b. (a, b) -> a
fst) ([(Text, Key)] -> [(Text, Key)])
-> (Text -> [(Text, Key)]) -> Text -> [(Text, Key)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Text] -> [Key] -> [(Text, Key)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [Key
1 ..]) ([Text] -> [(Text, Key)])
-> (Text -> [Text]) -> Text -> [(Text, Key)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> [Text]
T.lines

-- | Lines which contain a shebang.
shebangLines :: Text -> IntSet
shebangLines :: Text -> IntSet
shebangLines = (Text -> Bool) -> Text -> IntSet
linesFiltered (Text
"#!" `T.isPrefixOf`)

-- | Lines which contain a LINE pragma.
linePragmaLines :: Text -> IntSet
linePragmaLines :: Text -> IntSet
linePragmaLines = (Text -> Bool) -> Text -> IntSet
linesFiltered (Text
"{-# LINE" `T.isPrefixOf`)

-- | Inner text of a magic enabling marker.
ormoluEnable :: Text
ormoluEnable :: Text
ormoluEnable = Text
"ORMOLU_ENABLE"

-- | Inner text of a magic disabling marker.
ormoluDisable :: Text
ormoluDisable :: Text
ormoluDisable = Text
"ORMOLU_DISABLE"

-- | Creates a magic comment with the given inner text.
magicComment :: Text -> Text
magicComment :: Text -> Text
magicComment Text
t = Text
"{- " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
t Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" -}"

-- | Construct a function for whitespace-insensitive matching of string.
isMagicComment ::
  -- | What to expect
  Text ->
  -- | String to test
  Text ->
  -- | If the two strings match, we return the rest of the line.
  Maybe Text
isMagicComment :: Text -> Text -> Maybe Text
isMagicComment Text
expected Text
s0 = do
  Text
s1 <- 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 -> Text
T.stripStart Text
s0)
  Text
s2 <- 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
expected Text
s1
  Text -> Text -> Maybe Text
T.stripPrefix Text
"-}" Text
s2