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

-- | Random utilities used by the code.
module Ormolu.Utils
  ( RelativePos (..),
    attachRelativePos,
    combineSrcSpans',
    notImplemented,
    showOutputable,
    splitDocString,
    incSpanLine,
    separatedByBlank,
    separatedByBlankNE,
    onTheSameLine,
    HasSrcSpan (..),
    getLoc',
    matchAddEpAnn,
    textToStringBuffer,
    ghcModuleNameToCabal,
  )
where

import Data.List (dropWhileEnd)
import Data.List.NonEmpty (NonEmpty (..))
import Data.List.NonEmpty qualified as NE
import Data.Maybe (fromMaybe)
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.Foreign qualified as TFFI
import Distribution.ModuleName (ModuleName)
import Distribution.ModuleName qualified as ModuleName
import Foreign (pokeElemOff, withForeignPtr)
import GHC.Data.Strict qualified as Strict
import GHC.Data.StringBuffer (StringBuffer (..))
import GHC.Driver.Ppr
import GHC.DynFlags (baseDynFlags)
import GHC.ForeignPtr (mallocPlainForeignPtrBytes)
import GHC.Hs hiding (ModuleName)
import GHC.IO.Unsafe (unsafePerformIO)
import GHC.Types.SrcLoc
import GHC.Utils.Outputable (Outputable (..))
import Language.Haskell.Syntax.Module.Name qualified as GHC

-- | Relative positions in a list.
data RelativePos
  = SinglePos
  | FirstPos
  | MiddlePos
  | LastPos
  deriving (RelativePos -> RelativePos -> Bool
(RelativePos -> RelativePos -> Bool)
-> (RelativePos -> RelativePos -> Bool) -> Eq RelativePos
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: RelativePos -> RelativePos -> Bool
== :: RelativePos -> RelativePos -> Bool
$c/= :: RelativePos -> RelativePos -> Bool
/= :: RelativePos -> RelativePos -> Bool
Eq, Int -> RelativePos -> ShowS
[RelativePos] -> ShowS
RelativePos -> String
(Int -> RelativePos -> ShowS)
-> (RelativePos -> String)
-> ([RelativePos] -> ShowS)
-> Show RelativePos
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> RelativePos -> ShowS
showsPrec :: Int -> RelativePos -> ShowS
$cshow :: RelativePos -> String
show :: RelativePos -> String
$cshowList :: [RelativePos] -> ShowS
showList :: [RelativePos] -> ShowS
Show)

-- | Attach 'RelativePos'es to elements of a given list.
attachRelativePos :: [a] -> [(RelativePos, a)]
attachRelativePos :: forall a. [a] -> [(RelativePos, a)]
attachRelativePos = \case
  [] -> []
  [a
x] -> [(RelativePos
SinglePos, a
x)]
  (a
x : [a]
xs) -> (RelativePos
FirstPos, a
x) (RelativePos, a) -> [(RelativePos, a)] -> [(RelativePos, a)]
forall a. a -> [a] -> [a]
: [a] -> [(RelativePos, a)]
forall a. [a] -> [(RelativePos, a)]
markLast [a]
xs
  where
    markLast :: [b] -> [(RelativePos, b)]
markLast [] = []
    markLast [b
x] = [(RelativePos
LastPos, b
x)]
    markLast (b
x : [b]
xs) = (RelativePos
MiddlePos, b
x) (RelativePos, b) -> [(RelativePos, b)] -> [(RelativePos, b)]
forall a. a -> [a] -> [a]
: [b] -> [(RelativePos, b)]
markLast [b]
xs

-- | Combine all source spans from the given list.
combineSrcSpans' :: NonEmpty SrcSpan -> SrcSpan
combineSrcSpans' :: NonEmpty SrcSpan -> SrcSpan
combineSrcSpans' (SrcSpan
x :| [SrcSpan]
xs) = (SrcSpan -> SrcSpan -> SrcSpan) -> SrcSpan -> [SrcSpan] -> SrcSpan
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr SrcSpan -> SrcSpan -> SrcSpan
combineSrcSpans SrcSpan
x [SrcSpan]
xs

-- | Placeholder for things that are not yet implemented.
notImplemented :: String -> a
notImplemented :: forall a. String -> a
notImplemented String
msg = String -> a
forall a. HasCallStack => String -> a
error (String -> a) -> String -> a
forall a b. (a -> b) -> a -> b
$ String
"not implemented yet: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
msg

-- | Pretty-print an 'GHC.Outputable' thing.
showOutputable :: (Outputable o) => o -> String
showOutputable :: forall o. Outputable o => o -> String
showOutputable = DynFlags -> SDoc -> String
showSDoc DynFlags
baseDynFlags (SDoc -> String) -> (o -> SDoc) -> o -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. o -> SDoc
forall a. Outputable a => a -> SDoc
ppr

-- | Split and normalize a doc string. The result is a list of lines that
-- make up the comment.
splitDocString :: HsDocString -> [Text]
splitDocString :: HsDocString -> [Text]
splitDocString HsDocString
docStr =
  case [Text]
r of
    [] -> [Text
""]
    [Text]
_ -> [Text]
r
  where
    r :: [Text]
r =
      (Text -> Text) -> [Text] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Text -> Text
escapeLeadingDollar
        ([Text] -> [Text]) -> (String -> [Text]) -> String -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Text] -> [Text]
dropPaddingSpace
        ([Text] -> [Text]) -> (String -> [Text]) -> String -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text -> Bool) -> [Text] -> [Text]
forall a. (a -> Bool) -> [a] -> [a]
dropWhileEnd Text -> Bool
T.null
        ([Text] -> [Text]) -> (String -> [Text]) -> String -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String -> Text) -> [String] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Text -> Text
T.stripEnd (Text -> Text) -> (String -> Text) -> String -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Text
T.pack)
        ([String] -> [Text]) -> (String -> [String]) -> String -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> [String]
lines
        (String -> [Text]) -> String -> [Text]
forall a b. (a -> b) -> a -> b
$ HsDocString -> String
renderHsDocString HsDocString
docStr
    -- We cannot have the first character to be a dollar because in that
    -- case it'll be a parse error (apparently collides with named docs
    -- syntax @-- $name@ somehow).
    escapeLeadingDollar :: Text -> Text
escapeLeadingDollar Text
txt =
      case Text -> Maybe (Char, Text)
T.uncons Text
txt of
        Just (Char
'$', Text
_) -> Char -> Text -> Text
T.cons Char
'\\' Text
txt
        Maybe (Char, Text)
_ -> Text
txt
    dropPaddingSpace :: [Text] -> [Text]
dropPaddingSpace [Text]
xs =
      case (Text -> Bool) -> [Text] -> [Text]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile Text -> Bool
T.null [Text]
xs of
        [] -> []
        (Text
x : [Text]
_) ->
          let leadingSpace :: Text -> Bool
leadingSpace Text
txt = case Text -> Maybe (Char, Text)
T.uncons Text
txt of
                Just (Char
' ', Text
_) -> Bool
True
                Maybe (Char, Text)
_ -> Bool
False
              dropSpace :: Text -> Text
dropSpace Text
txt =
                if Text -> Bool
leadingSpace Text
txt
                  then Int -> Text -> Text
T.drop Int
1 Text
txt
                  else Text
txt
           in if Text -> Bool
leadingSpace Text
x
                then Text -> Text
dropSpace (Text -> Text) -> [Text] -> [Text]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Text]
xs
                else [Text]
xs

-- | Increment line number in a 'SrcSpan'.
incSpanLine :: Int -> SrcSpan -> SrcSpan
incSpanLine :: Int -> SrcSpan -> SrcSpan
incSpanLine Int
i = \case
  RealSrcSpan RealSrcSpan
s Maybe BufSpan
_ ->
    let start :: RealSrcLoc
start = RealSrcSpan -> RealSrcLoc
realSrcSpanStart RealSrcSpan
s
        end :: RealSrcLoc
end = RealSrcSpan -> RealSrcLoc
realSrcSpanEnd RealSrcSpan
s
        incLine :: RealSrcLoc -> RealSrcLoc
incLine RealSrcLoc
x =
          let file :: FastString
file = RealSrcLoc -> FastString
srcLocFile RealSrcLoc
x
              line :: Int
line = RealSrcLoc -> Int
srcLocLine RealSrcLoc
x
              col :: Int
col = RealSrcLoc -> Int
srcLocCol RealSrcLoc
x
           in FastString -> Int -> Int -> RealSrcLoc
mkRealSrcLoc FastString
file (Int
line Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i) Int
col
     in RealSrcSpan -> Maybe BufSpan -> SrcSpan
RealSrcSpan (RealSrcLoc -> RealSrcLoc -> RealSrcSpan
mkRealSrcSpan (RealSrcLoc -> RealSrcLoc
incLine RealSrcLoc
start) (RealSrcLoc -> RealSrcLoc
incLine RealSrcLoc
end)) Maybe BufSpan
forall a. Maybe a
Strict.Nothing
  UnhelpfulSpan UnhelpfulSpanReason
x -> UnhelpfulSpanReason -> SrcSpan
UnhelpfulSpan UnhelpfulSpanReason
x

-- | Do two declarations have a blank between them?
separatedByBlank :: (a -> SrcSpan) -> a -> a -> Bool
separatedByBlank :: forall a. (a -> SrcSpan) -> a -> a -> Bool
separatedByBlank a -> SrcSpan
loc a
a a
b =
  Bool -> Maybe Bool -> Bool
forall a. a -> Maybe a -> a
fromMaybe Bool
False (Maybe Bool -> Bool) -> Maybe Bool -> Bool
forall a b. (a -> b) -> a -> b
$ do
    Int
endA <- RealSrcSpan -> Int
srcSpanEndLine (RealSrcSpan -> Int) -> Maybe RealSrcSpan -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SrcSpan -> Maybe RealSrcSpan
srcSpanToRealSrcSpan (a -> SrcSpan
loc a
a)
    Int
startB <- RealSrcSpan -> Int
srcSpanStartLine (RealSrcSpan -> Int) -> Maybe RealSrcSpan -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SrcSpan -> Maybe RealSrcSpan
srcSpanToRealSrcSpan (a -> SrcSpan
loc a
b)
    Bool -> Maybe Bool
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
startB Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
endA Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2)

-- | Do two declaration groups have a blank between them?
separatedByBlankNE :: (a -> SrcSpan) -> NonEmpty a -> NonEmpty a -> Bool
separatedByBlankNE :: forall a. (a -> SrcSpan) -> NonEmpty a -> NonEmpty a -> Bool
separatedByBlankNE a -> SrcSpan
loc NonEmpty a
a NonEmpty a
b = (a -> SrcSpan) -> a -> a -> Bool
forall a. (a -> SrcSpan) -> a -> a -> Bool
separatedByBlank a -> SrcSpan
loc (NonEmpty a -> a
forall a. NonEmpty a -> a
NE.last NonEmpty a
a) (NonEmpty a -> a
forall a. NonEmpty a -> a
NE.head NonEmpty a
b)

-- | Return 'True' if one span ends on the same line the second one starts.
onTheSameLine :: SrcSpan -> SrcSpan -> Bool
onTheSameLine :: SrcSpan -> SrcSpan -> Bool
onTheSameLine SrcSpan
a SrcSpan
b =
  SrcSpan -> Bool
isOneLineSpan (SrcLoc -> SrcLoc -> SrcSpan
mkSrcSpan (SrcSpan -> SrcLoc
srcSpanEnd SrcSpan
a) (SrcSpan -> SrcLoc
srcSpanStart SrcSpan
b))

class HasSrcSpan l where
  loc' :: l -> SrcSpan

instance HasSrcSpan SrcSpan where
  loc' :: SrcSpan -> SrcSpan
loc' = SrcSpan -> SrcSpan
forall a. a -> a
id

instance HasSrcSpan RealSrcSpan where
  loc' :: RealSrcSpan -> SrcSpan
loc' RealSrcSpan
l = RealSrcSpan -> Maybe BufSpan -> SrcSpan
RealSrcSpan RealSrcSpan
l Maybe BufSpan
forall a. Maybe a
Strict.Nothing

instance HasSrcSpan (SrcSpanAnn' ann) where
  loc' :: SrcSpanAnn' ann -> SrcSpan
loc' = SrcSpanAnn' ann -> SrcSpan
forall ann. SrcSpanAnn' ann -> SrcSpan
locA

getLoc' :: (HasSrcSpan l) => GenLocated l a -> SrcSpan
getLoc' :: forall l a. HasSrcSpan l => GenLocated l a -> SrcSpan
getLoc' = l -> SrcSpan
forall l. HasSrcSpan l => l -> SrcSpan
loc' (l -> SrcSpan)
-> (GenLocated l a -> l) -> GenLocated l a -> SrcSpan
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GenLocated l a -> l
forall l e. GenLocated l e -> l
getLoc

-- | Check whether the given 'AnnKeywordId' or its Unicode variant is in an
-- 'AddEpAnn', and return the 'EpaLocation' if so.
matchAddEpAnn :: AnnKeywordId -> AddEpAnn -> Maybe EpaLocation
matchAddEpAnn :: AnnKeywordId -> AddEpAnn -> Maybe EpaLocation
matchAddEpAnn AnnKeywordId
annId (AddEpAnn AnnKeywordId
annId' EpaLocation
loc)
  | AnnKeywordId
annId AnnKeywordId -> AnnKeywordId -> Bool
forall a. Eq a => a -> a -> Bool
== AnnKeywordId
annId' Bool -> Bool -> Bool
|| AnnKeywordId -> AnnKeywordId
unicodeAnn AnnKeywordId
annId AnnKeywordId -> AnnKeywordId -> Bool
forall a. Eq a => a -> a -> Bool
== AnnKeywordId
annId' = EpaLocation -> Maybe EpaLocation
forall a. a -> Maybe a
Just EpaLocation
loc
  | Bool
otherwise = Maybe EpaLocation
forall a. Maybe a
Nothing

-- | Convert 'Text' to a 'StringBuffer' by making a copy.
textToStringBuffer :: Text -> StringBuffer
textToStringBuffer :: Text -> StringBuffer
textToStringBuffer Text
txt = IO StringBuffer -> StringBuffer
forall a. IO a -> a
unsafePerformIO (IO StringBuffer -> StringBuffer)
-> IO StringBuffer -> StringBuffer
forall a b. (a -> b) -> a -> b
$ do
  ForeignPtr Word8
buf <- Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
mallocPlainForeignPtrBytes (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3)
  ForeignPtr Word8 -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
buf ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr -> do
    Text -> Ptr Word8 -> IO ()
TFFI.unsafeCopyToPtr Text
txt Ptr Word8
ptr
    -- last three bytes have to be zero for easier decoding
    Ptr Word8 -> Int -> Word8 -> IO ()
forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff Ptr Word8
ptr Int
len Word8
0
    Ptr Word8 -> Int -> Word8 -> IO ()
forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff Ptr Word8
ptr (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Word8
0
    Ptr Word8 -> Int -> Word8 -> IO ()
forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff Ptr Word8
ptr (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2) Word8
0
  StringBuffer -> IO StringBuffer
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure StringBuffer {ForeignPtr Word8
buf :: ForeignPtr Word8
buf :: ForeignPtr Word8
buf, Int
len :: Int
len :: Int
len, cur :: Int
cur = Int
0}
  where
    len :: Int
len = Text -> Int
TFFI.lengthWord8 Text
txt

-- | Convert GHC's 'ModuleName' into the one used by Cabal.
ghcModuleNameToCabal :: GHC.ModuleName -> ModuleName
ghcModuleNameToCabal :: ModuleName -> ModuleName
ghcModuleNameToCabal = String -> ModuleName
forall a. IsString a => String -> a
ModuleName.fromString (String -> ModuleName)
-> (ModuleName -> String) -> ModuleName -> ModuleName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ModuleName -> String
GHC.moduleNameString