1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE LambdaCase #-}
3{-# LANGUAGE OverloadedStrings #-}
4{-# LANGUAGE RecordWildCards #-}
5
6-- | Preprocessing for input source code.
7module Ormolu.Processing.Preprocess
8  ( preprocess,
9  )
10where
11
12import Control.Monad
13import Data.Array as A
14import Data.Bifunctor (bimap)
15import Data.Char (isSpace)
16import Data.Function ((&))
17import Data.IntMap (IntMap)
18import qualified Data.IntMap.Strict as IntMap
19import Data.IntSet (IntSet)
20import qualified Data.IntSet as IntSet
21import qualified Data.List as L
22import Data.Maybe (isJust)
23import Data.Text (Text)
24import qualified Data.Text as T
25import Ormolu.Config (RegionDeltas (..))
26import Ormolu.Processing.Common
27import Ormolu.Processing.Cpp
28
29-- | Preprocess the specified region of the input into raw snippets
30-- and subregions to be formatted.
31preprocess ::
32  -- | Whether CPP is enabled
33  Bool ->
34  RegionDeltas ->
35  String ->
36  [Either Text RegionDeltas]
37preprocess cppEnabled region rawInput = rawSnippetsAndRegionsToFormat
38  where
39    (linesNotToFormat', replacementLines) = linesNotToFormat cppEnabled region rawInput
40    regionsToFormat =
41      intSetToRegions rawLineLength $
42        IntSet.fromAscList [1 .. rawLineLength] IntSet.\\ linesNotToFormat'
43    regionsNotToFormat = intSetToRegions rawLineLength linesNotToFormat'
44    -- We want to interleave the regionsToFormat and regionsNotToFormat.
45    -- If the first non-formattable region starts at the first line, it is
46    -- the first interleaved region, otherwise, we start with the first
47    -- region to format.
48    interleave' = case regionsNotToFormat of
49      r : _ | regionPrefixLength r == 0 -> interleave
50      _ -> flip interleave
51    rawSnippets = T.pack . flip linesInRegion updatedInput <$> regionsNotToFormat
52      where
53        updatedInput = unlines . fmap updateLine . zip [1 ..] . lines $ rawInput
54        updateLine (i, line) = IntMap.findWithDefault line i replacementLines
55    rawSnippetsAndRegionsToFormat =
56      interleave' (Left <$> rawSnippets) (Right <$> regionsToFormat)
57        >>= patchSeparatingBlankLines
58        & dropWhile isBlankRawSnippet
59        & L.dropWhileEnd isBlankRawSnippet
60    -- For every formattable region, we want to ensure that it is separated by
61    -- a blank line from preceding/succeeding raw snippets if it starts/ends
62    -- with a blank line.
63    -- Empty formattable regions are replaced by a blank line instead.
64    -- Extraneous raw snippets at the start/end are dropped afterwards.
65    patchSeparatingBlankLines = \case
66      Right r@RegionDeltas {..} ->
67        if all isSpace (linesInRegion r rawInput)
68          then [blankRawSnippet]
69          else
70            [blankRawSnippet | isBlankLine regionPrefixLength] <> [Right r]
71              <> [blankRawSnippet | isBlankLine (rawLineLength - regionSuffixLength - 1)]
72      Left r -> [Left r]
73      where
74        blankRawSnippet = Left "\n"
75        isBlankLine i = isJust . mfilter (all isSpace) $ rawLines !!? i
76    isBlankRawSnippet = \case
77      Left r | T.all isSpace r -> True
78      _ -> False
79
80    rawLines = A.listArray (0, length rawLines' - 1) rawLines'
81      where
82        rawLines' = lines rawInput
83    rawLineLength = length rawLines
84
85    interleave [] bs = bs
86    interleave (a : as) bs = a : interleave bs as
87
88    xs !!? i = if A.bounds rawLines `A.inRange` i then Just $ xs A.! i else Nothing
89
90-- | All lines we are not supposed to format, and a set of replacements
91-- for specific lines.
92linesNotToFormat ::
93  -- | Whether CPP is enabled
94  Bool ->
95  RegionDeltas ->
96  String ->
97  (IntSet, IntMap String)
98linesNotToFormat cppEnabled region@RegionDeltas {..} input =
99  (unconsidered <> magicDisabled <> otherDisabled, lineUpdates)
100  where
101    unconsidered =
102      IntSet.fromAscList $
103        [1 .. regionPrefixLength] <> [totalLines - regionSuffixLength + 1 .. totalLines]
104    totalLines = length (lines input)
105    regionLines = linesInRegion region input
106    (magicDisabled, lineUpdates) = magicDisabledLines regionLines
107    otherDisabled = (mconcat allLines) regionLines
108      where
109        allLines = [shebangLines, linePragmaLines] <> [cppLines | cppEnabled]
110
111-- | Ormolu state.
112data OrmoluState
113  = -- | Enabled
114    OrmoluEnabled
115  | -- | Disabled
116    OrmoluDisabled
117  deriving (Eq, Show)
118
119-- | All lines which are disabled by Ormolu's magic comments,
120-- as well as normalizing replacements.
121magicDisabledLines :: String -> (IntSet, IntMap String)
122magicDisabledLines input =
123  bimap IntSet.fromAscList IntMap.fromAscList . mconcat $
124    go OrmoluEnabled (lines input `zip` [1 ..])
125  where
126    go _ [] = []
127    go state ((line, i) : ls)
128      | isMagicComment ormoluDisable line,
129        state == OrmoluEnabled =
130          ([i], [(i, magicComment ormoluDisable)]) : go OrmoluDisabled ls
131      | isMagicComment ormoluEnable line,
132        state == OrmoluDisabled =
133          ([i], [(i, magicComment ormoluEnable)]) : go OrmoluEnabled ls
134      | otherwise = iIfDisabled : go state ls
135      where
136        iIfDisabled = case state of
137          OrmoluDisabled -> ([i], [])
138          OrmoluEnabled -> ([], [])
139
140-- | All lines which satisfy a predicate.
141linesFiltered :: (String -> Bool) -> String -> IntSet
142linesFiltered p =
143  IntSet.fromAscList . fmap snd . filter (p . fst) . (`zip` [1 ..]) . lines
144
145-- | Lines which contain a shebang.
146shebangLines :: String -> IntSet
147shebangLines = linesFiltered ("#!" `L.isPrefixOf`)
148
149-- | Lines which contain a LINE pragma.
150linePragmaLines :: String -> IntSet
151linePragmaLines = linesFiltered ("{-# LINE" `L.isPrefixOf`)
152
153-- | Inner text of a magic enabling marker.
154ormoluEnable :: String
155ormoluEnable = "ORMOLU_ENABLE"
156
157-- | Inner text of a magic disabling marker.
158ormoluDisable :: String
159ormoluDisable = "ORMOLU_DISABLE"
160
161-- | Creates a magic comment with the given inner text.
162magicComment :: String -> String
163magicComment t = "{- " <> t <> " -}"
164
165-- | Construct a function for whitespace-insensitive matching of string.
166isMagicComment ::
167  -- | What to expect
168  String ->
169  -- | String to test
170  String ->
171  -- | Whether or not the two strings watch
172  Bool
173isMagicComment expected s0 = isJust $ do
174  let trim = dropWhile isSpace
175  s1 <- trim <$> L.stripPrefix "{-" (trim s0)
176  s2 <- trim <$> L.stripPrefix expected s1
177  s3 <- L.stripPrefix "-}" s2
178  guard (all isSpace s3)
179