Utils.hs

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

-- | Random utilities used by the code.
module Ormolu.Utils
  ( RelativePos (..),
    attachRelativePos,
    combineSrcSpans',
    notImplemented,
    showOutputable,
    splitDocString,
    typeArgToType,
    unSrcSpan,
    incSpanLine,
    separatedByBlank,
    separatedByBlankNE,
    onTheSameLine,
    removeIndentation,
    groupBy',
  )
where

import Data.Char (isSpace)
import Data.List (dropWhileEnd)
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty as NE
import Data.Maybe (fromMaybe)
import Data.Text (Text)
import qualified Data.Text as T
import GHC
import GHC.DynFlags (baseDynFlags)
import qualified Outputable as GHC

-- | Relative positions in a list.
data RelativePos
  = SinglePos
  | FirstPos
  | MiddlePos
  | LastPos
  deriving (Eq, Show)

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

-- | Combine all source spans from the given list.
combineSrcSpans' :: NonEmpty SrcSpan -> SrcSpan
combineSrcSpans' (x :| xs) = foldr combineSrcSpans x xs

-- | Placeholder for things that are not yet implemented.
notImplemented :: String -> a
notImplemented msg = error $ "not implemented yet: " ++ msg

-- | Pretty-print an 'GHC.Outputable' thing.
showOutputable :: GHC.Outputable o => o -> String
showOutputable = GHC.showSDoc baseDynFlags . GHC.ppr

-- | Split and normalize a doc string. The result is a list of lines that
-- make up the comment.
splitDocString :: HsDocString -> [Text]
splitDocString docStr =
  case r of
    [] -> [""]
    _ -> r
  where
    r =
      fmap (escapeLeadingDollar . escapeCommentBraces)
        . dropWhileEnd T.null
        . fmap (T.stripEnd . T.pack)
        . lines
        $ unpackHDS 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 txt =
      case T.uncons txt of
        Just ('$', _) -> T.cons '\\' txt
        _ -> txt
    escapeCommentBraces = T.replace "{-" "{\\-" . T.replace "-}" "-\\}"

-- | Get 'LHsType' out of 'LHsTypeArg'.
typeArgToType :: LHsTypeArg p -> LHsType p
typeArgToType = \case
  HsValArg tm -> tm
  HsTypeArg _ ty -> ty
  HsArgPar _ -> notImplemented "HsArgPar"

-- | Get 'RealSrcSpan' out of 'SrcSpan' if the span is “helpful”.
unSrcSpan :: SrcSpan -> Maybe RealSrcSpan
unSrcSpan = \case
  RealSrcSpan r -> Just r
  UnhelpfulSpan _ -> Nothing

-- | Increment line number in a 'SrcSpan'.
incSpanLine :: Int -> SrcSpan -> SrcSpan
incSpanLine i = \case
  RealSrcSpan s ->
    let start = realSrcSpanStart s
        end = realSrcSpanEnd s
        incLine x =
          let file = srcLocFile x
              line = srcLocLine x
              col = srcLocCol x
           in mkRealSrcLoc file (line + i) col
     in RealSrcSpan (mkRealSrcSpan (incLine start) (incLine end))
  UnhelpfulSpan x -> UnhelpfulSpan x

-- | Do two declarations have a blank between them?
separatedByBlank :: (a -> SrcSpan) -> a -> a -> Bool
separatedByBlank loc a b =
  fromMaybe False $ do
    endA <- srcSpanEndLine <$> unSrcSpan (loc a)
    startB <- srcSpanStartLine <$> unSrcSpan (loc b)
    pure (startB - endA >= 2)

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

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

-- | Remove indentation from a given 'String'. Return the input with
-- indentation removed and the detected indentation level.
removeIndentation :: String -> (String, Int)
removeIndentation (lines -> xs) = (unlines (drop n <$> xs), n)
  where
    n = minimum (getIndent <$> xs)
    getIndent y =
      if all isSpace y
        then 0
        else length (takeWhile isSpace y)

-- | A generalisation of 'groupBy' to functions which aren't equivalences - a group ends
-- when comparison fails with the previous element, rather than the first of the group.
groupBy' :: (a -> a -> Bool) -> [a] -> [NonEmpty a]
groupBy' eq = flip foldr [] $ \x -> \case
  [] -> [pure x]
  (y :| ys) : zs ->
    if x `eq` y
      then (x :| y : ys) : zs
      else pure x : (y :| ys) : zs