--------------------------------------------------------------------------------
{-# LANGUAGE LambdaCase    #-}
{-# LANGUAGE PatternGuards #-}
module Language.Haskell.Stylish.Util
    ( indent
    , padRight
    , everything
    , infoPoints
    , trimLeft
    , trimRight
    , wrap
    , wrapRest
    , wrapMaybe
    , wrapRestMaybe

    -- * Extra list functions
    , withHead
    , withInit
    , withTail
    , withLast
    , flagEnds

    , toRealSrcSpan

    , traceOutputable
    , traceOutputableM

    , unguardedRhsBody
    , rhsBody

    , getGuards
    ) where


--------------------------------------------------------------------------------
import           Data.Char                     (isSpace)
import           Data.Data                     (Data)
import qualified Data.Generics                 as G
import           Data.Maybe                    (maybeToList)
import           Data.Typeable                 (cast)
import           Debug.Trace                   (trace)
import qualified GHC.Hs                        as Hs
import qualified Outputable
import qualified SrcLoc                        as S


--------------------------------------------------------------------------------
import           Language.Haskell.Stylish.Step


--------------------------------------------------------------------------------
indent :: Int -> String -> String
indent :: Int -> String -> String
indent Int
len = (Int -> String
indentPrefix Int
len String -> String -> String
forall a. [a] -> [a] -> [a]
++)


--------------------------------------------------------------------------------
indentPrefix :: Int -> String
indentPrefix :: Int -> String
indentPrefix = (Int -> Char -> String
forall a. Int -> a -> [a]
`replicate` Char
' ')


--------------------------------------------------------------------------------
padRight :: Int -> String -> String
padRight :: Int -> String -> String
padRight Int
len String
str = String
str String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> Char -> String
forall a. Int -> a -> [a]
replicate (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- String -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
str) Char
' '


--------------------------------------------------------------------------------
everything :: (Data a, Data b) => a -> [b]
everything :: a -> [b]
everything = ([b] -> [b] -> [b]) -> GenericQ [b] -> GenericQ [b]
forall r. (r -> r -> r) -> GenericQ r -> GenericQ r
G.everything [b] -> [b] -> [b]
forall a. [a] -> [a] -> [a]
(++) (Maybe b -> [b]
forall a. Maybe a -> [a]
maybeToList (Maybe b -> [b]) -> (a -> Maybe b) -> a -> [b]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Maybe b
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast)


--------------------------------------------------------------------------------
infoPoints :: [S.Located pass] -> [((Int, Int), (Int, Int))]
infoPoints :: [Located pass] -> [((Int, Int), (Int, Int))]
infoPoints = (Located pass -> ((Int, Int), (Int, Int)))
-> [Located pass] -> [((Int, Int), (Int, Int))]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (SrcSpan -> ((Int, Int), (Int, Int))
helper (SrcSpan -> ((Int, Int), (Int, Int)))
-> (Located pass -> SrcSpan)
-> Located pass
-> ((Int, Int), (Int, Int))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Located pass -> SrcSpan
forall a. HasSrcSpan a => a -> SrcSpan
S.getLoc)
  where
    helper :: S.SrcSpan -> ((Int, Int), (Int, Int))
    helper :: SrcSpan -> ((Int, Int), (Int, Int))
helper (S.RealSrcSpan RealSrcSpan
s) = do
               let
                start :: RealSrcLoc
start = RealSrcSpan -> RealSrcLoc
S.realSrcSpanStart RealSrcSpan
s
                end :: RealSrcLoc
end = RealSrcSpan -> RealSrcLoc
S.realSrcSpanEnd RealSrcSpan
s
               ((RealSrcLoc -> Int
S.srcLocLine RealSrcLoc
start, RealSrcLoc -> Int
S.srcLocCol RealSrcLoc
start), (RealSrcLoc -> Int
S.srcLocLine RealSrcLoc
end, RealSrcLoc -> Int
S.srcLocCol RealSrcLoc
end))
    helper SrcSpan
_                   = ((-Int
1,-Int
1), (-Int
1,-Int
1))


--------------------------------------------------------------------------------
trimLeft :: String -> String
trimLeft :: String -> String
trimLeft  = (Char -> Bool) -> String -> String
forall a. (a -> Bool) -> [a] -> [a]
dropWhile Char -> Bool
isSpace


--------------------------------------------------------------------------------
trimRight :: String -> String
trimRight :: String -> String
trimRight = String -> String
forall a. [a] -> [a]
reverse (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> String
trimLeft (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> String
forall a. [a] -> [a]
reverse


--------------------------------------------------------------------------------
wrap :: Int       -- ^ Maximum line width
     -> String    -- ^ Leading string
     -> Int       -- ^ Indentation
     -> [String]  -- ^ Strings to add/wrap
     -> Lines     -- ^ Resulting lines
wrap :: Int -> String -> Int -> [String] -> [String]
wrap Int
maxWidth String
leading Int
ind = String -> [String] -> [String]
wrap' String
leading
  where
    wrap' :: String -> [String] -> [String]
wrap' String
ss [] = [String
ss]
    wrap' String
ss (String
str:[String]
strs)
        | String -> String -> Bool
forall (t :: * -> *) (t :: * -> *) a a.
(Foldable t, Foldable t) =>
t a -> t a -> Bool
overflows String
ss String
str =
            String
ss String -> [String] -> [String]
forall a. a -> [a] -> [a]
: Int -> Int -> [String] -> [String]
wrapRest Int
maxWidth Int
ind (String
strString -> [String] -> [String]
forall a. a -> [a] -> [a]
:[String]
strs)
        | Bool
otherwise = String -> [String] -> [String]
wrap' (String
ss String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
str) [String]
strs

    overflows :: t a -> t a -> Bool
overflows t a
ss t a
str = t a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length t a
ss Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxWidth Bool -> Bool -> Bool
||
        ((t a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length t a
ss Int -> Int -> Int
forall a. Num a => a -> a -> a
+ t a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length t a
str) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
maxWidth Bool -> Bool -> Bool
&& Int
ind Int -> Int -> Int
forall a. Num a => a -> a -> a
+ t a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length t a
str  Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxWidth)


--------------------------------------------------------------------------------
wrapMaybe :: Maybe Int -- ^ Maximum line width (maybe)
          -> String    -- ^ Leading string
          -> Int       -- ^ Indentation
          -> [String]  -- ^ Strings to add/wrap
          -> Lines     -- ^ Resulting lines
wrapMaybe :: Maybe Int -> String -> Int -> [String] -> [String]
wrapMaybe (Just Int
maxWidth) = Int -> String -> Int -> [String] -> [String]
wrap Int
maxWidth
wrapMaybe Maybe Int
Nothing         = String -> Int -> [String] -> [String]
noWrap


--------------------------------------------------------------------------------
noWrap :: String    -- ^ Leading string
       -> Int       -- ^ Indentation
       -> [String]  -- ^ Strings to add
       -> Lines     -- ^ Resulting lines
noWrap :: String -> Int -> [String] -> [String]
noWrap String
leading Int
_ind = String -> [String] -> [String]
noWrap' String
leading
  where
    noWrap' :: String -> [String] -> [String]
noWrap' String
ss []         = [String
ss]
    noWrap' String
ss (String
str:[String]
strs) = String -> [String] -> [String]
noWrap' (String
ss String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
str) [String]
strs


--------------------------------------------------------------------------------
wrapRest :: Int
         -> Int
         -> [String]
         -> Lines
wrapRest :: Int -> Int -> [String] -> [String]
wrapRest Int
maxWidth Int
ind = [String] -> [String]
forall a. [a] -> [a]
reverse ([String] -> [String])
-> ([String] -> [String]) -> [String] -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [String] -> String -> [String] -> [String]
wrapRest' [] String
""
  where
    wrapRest' :: [String] -> String -> [String] -> [String]
wrapRest' [String]
ls String
ss []
        | String -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null String
ss = [String]
ls
        | Bool
otherwise = String
ssString -> [String] -> [String]
forall a. a -> [a] -> [a]
:[String]
ls
    wrapRest' [String]
ls String
ss (String
str:[String]
strs)
        | String -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null String
ss = [String] -> String -> [String] -> [String]
wrapRest' [String]
ls (Int -> String -> String
indent Int
ind String
str) [String]
strs
        | String -> String -> Bool
forall (t :: * -> *) (t :: * -> *) a a.
(Foldable t, Foldable t) =>
t a -> t a -> Bool
overflows String
ss String
str = [String] -> String -> [String] -> [String]
wrapRest' (String
ssString -> [String] -> [String]
forall a. a -> [a] -> [a]
:[String]
ls) String
"" (String
strString -> [String] -> [String]
forall a. a -> [a] -> [a]
:[String]
strs)
        | Bool
otherwise = [String] -> String -> [String] -> [String]
wrapRest' [String]
ls (String
ss String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
str) [String]
strs

    overflows :: t a -> t a -> Bool
overflows t a
ss t a
str = (t a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length t a
ss Int -> Int -> Int
forall a. Num a => a -> a -> a
+ t a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length t a
str Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
maxWidth


--------------------------------------------------------------------------------
wrapRestMaybe :: Maybe Int
              -> Int
              -> [String]
              -> Lines
wrapRestMaybe :: Maybe Int -> Int -> [String] -> [String]
wrapRestMaybe (Just Int
maxWidth) = Int -> Int -> [String] -> [String]
wrapRest Int
maxWidth
wrapRestMaybe Maybe Int
Nothing         = Int -> [String] -> [String]
noWrapRest


--------------------------------------------------------------------------------
noWrapRest :: Int
           -> [String]
           -> Lines
noWrapRest :: Int -> [String] -> [String]
noWrapRest Int
ind = [String] -> [String]
forall a. [a] -> [a]
reverse ([String] -> [String])
-> ([String] -> [String]) -> [String] -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [String] -> String -> [String] -> [String]
noWrapRest' [] String
""
  where
    noWrapRest' :: [String] -> String -> [String] -> [String]
noWrapRest' [String]
ls String
ss []
        | String -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null String
ss = [String]
ls
        | Bool
otherwise = String
ssString -> [String] -> [String]
forall a. a -> [a] -> [a]
:[String]
ls
    noWrapRest' [String]
ls String
ss (String
str:[String]
strs)
        | String -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null String
ss = [String] -> String -> [String] -> [String]
noWrapRest' [String]
ls (Int -> String -> String
indent Int
ind String
str) [String]
strs
        | Bool
otherwise = [String] -> String -> [String] -> [String]
noWrapRest' [String]
ls (String
ss String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
str) [String]
strs


--------------------------------------------------------------------------------
withHead :: (a -> a) -> [a] -> [a]
withHead :: (a -> a) -> [a] -> [a]
withHead a -> a
_ []       = []
withHead a -> a
f (a
x : [a]
xs) = a -> a
f a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
xs


--------------------------------------------------------------------------------
withLast :: (a -> a) -> [a] -> [a]
withLast :: (a -> a) -> [a] -> [a]
withLast a -> a
_ []       = []
withLast a -> a
f [a
x]      = [a -> a
f a
x]
withLast a -> a
f (a
x : [a]
xs) = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: (a -> a) -> [a] -> [a]
forall a. (a -> a) -> [a] -> [a]
withLast a -> a
f [a]
xs


--------------------------------------------------------------------------------
withInit :: (a -> a) -> [a] -> [a]
withInit :: (a -> a) -> [a] -> [a]
withInit a -> a
_ []       = []
withInit a -> a
_ [a
x]      = [a
x]
withInit a -> a
f (a
x : [a]
xs) = a -> a
f a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: (a -> a) -> [a] -> [a]
forall a. (a -> a) -> [a] -> [a]
withInit a -> a
f [a]
xs


--------------------------------------------------------------------------------
withTail :: (a -> a) -> [a] -> [a]
withTail :: (a -> a) -> [a] -> [a]
withTail a -> a
_ []       = []
withTail a -> a
f (a
x : [a]
xs) = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map a -> a
f [a]
xs



--------------------------------------------------------------------------------
-- | Utility for traversing through a list and knowing when you're at the
-- first and last element.
flagEnds :: [a] -> [(a, Bool, Bool)]
flagEnds :: [a] -> [(a, Bool, Bool)]
flagEnds = \case
    [] -> []
    [a
x] -> [(a
x, Bool
True, Bool
True)]
    a
x : a
y : [a]
zs -> (a
x, Bool
True, Bool
False) (a, Bool, Bool) -> [(a, Bool, Bool)] -> [(a, Bool, Bool)]
forall a. a -> [a] -> [a]
: [a] -> [(a, Bool, Bool)]
forall a. [a] -> [(a, Bool, Bool)]
go (a
y a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
zs)
  where
    go :: [a] -> [(a, Bool, Bool)]
go (a
x : a
y : [a]
zs) = (a
x, Bool
False, Bool
False) (a, Bool, Bool) -> [(a, Bool, Bool)] -> [(a, Bool, Bool)]
forall a. a -> [a] -> [a]
: [a] -> [(a, Bool, Bool)]
go (a
y a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
zs)
    go [a
x]          = [(a
x, Bool
False, Bool
True)]
    go []           = []


--------------------------------------------------------------------------------
traceOutputable :: Outputable.Outputable a => String -> a -> b -> b
traceOutputable :: String -> a -> b -> b
traceOutputable String
title a
x =
    String -> b -> b
forall a. String -> a -> a
trace (String
title String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
": " String -> String -> String
forall a. [a] -> [a] -> [a]
++ (SDoc -> String
Outputable.showSDocUnsafe (SDoc -> String) -> SDoc -> String
forall a b. (a -> b) -> a -> b
$ a -> SDoc
forall a. Outputable a => a -> SDoc
Outputable.ppr a
x))


--------------------------------------------------------------------------------
traceOutputableM :: (Outputable.Outputable a, Monad m) => String -> a -> m ()
traceOutputableM :: String -> a -> m ()
traceOutputableM String
title a
x = String -> a -> m () -> m ()
forall a b. Outputable a => String -> a -> b -> b
traceOutputable String
title a
x (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()


--------------------------------------------------------------------------------
-- take the (Maybe) RealSrcSpan out of the SrcSpan
toRealSrcSpan :: S.SrcSpan -> Maybe S.RealSrcSpan
toRealSrcSpan :: SrcSpan -> Maybe RealSrcSpan
toRealSrcSpan (S.RealSrcSpan RealSrcSpan
s) = RealSrcSpan -> Maybe RealSrcSpan
forall a. a -> Maybe a
Just RealSrcSpan
s
toRealSrcSpan SrcSpan
_                 = Maybe RealSrcSpan
forall a. Maybe a
Nothing


--------------------------------------------------------------------------------
-- Utility: grab the body out of guarded RHSs if it's a single unguarded one.
unguardedRhsBody :: Hs.GRHSs Hs.GhcPs a -> Maybe a
unguardedRhsBody :: GRHSs GhcPs a -> Maybe a
unguardedRhsBody (Hs.GRHSs XCGRHSs GhcPs a
_ [LGRHS GhcPs a
grhs] LHsLocalBinds GhcPs
_)
    | Hs.GRHS _ [] body <- LGRHS GhcPs a -> SrcSpanLess (LGRHS GhcPs a)
forall a. HasSrcSpan a => a -> SrcSpanLess a
S.unLoc LGRHS GhcPs a
grhs = a -> Maybe a
forall a. a -> Maybe a
Just a
body
unguardedRhsBody GRHSs GhcPs a
_ = Maybe a
forall a. Maybe a
Nothing


-- Utility: grab the body out of guarded RHSs
rhsBody :: Hs.GRHSs Hs.GhcPs a -> Maybe a
rhsBody :: GRHSs GhcPs a -> Maybe a
rhsBody (Hs.GRHSs XCGRHSs GhcPs a
_ [LGRHS GhcPs a
grhs] LHsLocalBinds GhcPs
_)
    | Hs.GRHS _ _ body <- LGRHS GhcPs a -> SrcSpanLess (LGRHS GhcPs a)
forall a. HasSrcSpan a => a -> SrcSpanLess a
S.unLoc LGRHS GhcPs a
grhs = a -> Maybe a
forall a. a -> Maybe a
Just a
body
rhsBody GRHSs GhcPs a
_ = Maybe a
forall a. Maybe a
Nothing


--------------------------------------------------------------------------------
-- get guards in a guarded rhs of a Match
getGuards :: Hs.Match Hs.GhcPs (Hs.LHsExpr Hs.GhcPs) -> [Hs.GuardLStmt Hs.GhcPs]
getGuards :: Match GhcPs (LHsExpr GhcPs) -> [GuardLStmt GhcPs]
getGuards (Hs.Match XCMatch GhcPs (LHsExpr GhcPs)
_ HsMatchContext (NameOrRdrName (IdP GhcPs))
_ [LPat GhcPs]
_ GRHSs GhcPs (LHsExpr GhcPs)
grhss) =
  let
    lgrhs :: [LGRHS GhcPs (LHsExpr GhcPs)]
lgrhs = GRHSs GhcPs (LHsExpr GhcPs) -> [LGRHS GhcPs (LHsExpr GhcPs)]
getLocGRHS GRHSs GhcPs (LHsExpr GhcPs)
grhss -- []
    grhs :: [GRHS GhcPs (LHsExpr GhcPs)]
grhs  = (LGRHS GhcPs (LHsExpr GhcPs) -> GRHS GhcPs (LHsExpr GhcPs))
-> [LGRHS GhcPs (LHsExpr GhcPs)] -> [GRHS GhcPs (LHsExpr GhcPs)]
forall a b. (a -> b) -> [a] -> [b]
map LGRHS GhcPs (LHsExpr GhcPs) -> GRHS GhcPs (LHsExpr GhcPs)
forall a. HasSrcSpan a => a -> SrcSpanLess a
S.unLoc [LGRHS GhcPs (LHsExpr GhcPs)]
lgrhs
  in
    (GRHS GhcPs (LHsExpr GhcPs) -> [GuardLStmt GhcPs])
-> [GRHS GhcPs (LHsExpr GhcPs)] -> [GuardLStmt GhcPs]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap GRHS GhcPs (LHsExpr GhcPs) -> [GuardLStmt GhcPs]
getGuardLStmts [GRHS GhcPs (LHsExpr GhcPs)]
grhs
getGuards (Hs.XMatch XXMatch GhcPs (LHsExpr GhcPs)
x) = NoExtCon -> [GuardLStmt GhcPs]
forall a. NoExtCon -> a
Hs.noExtCon NoExtCon
XXMatch GhcPs (LHsExpr GhcPs)
x


getLocGRHS :: Hs.GRHSs Hs.GhcPs (Hs.LHsExpr Hs.GhcPs) -> [Hs.LGRHS Hs.GhcPs (Hs.LHsExpr Hs.GhcPs)]
getLocGRHS :: GRHSs GhcPs (LHsExpr GhcPs) -> [LGRHS GhcPs (LHsExpr GhcPs)]
getLocGRHS (Hs.GRHSs XCGRHSs GhcPs (LHsExpr GhcPs)
_ [LGRHS GhcPs (LHsExpr GhcPs)]
guardeds LHsLocalBinds GhcPs
_) = [LGRHS GhcPs (LHsExpr GhcPs)]
guardeds
getLocGRHS (Hs.XGRHSs XXGRHSs GhcPs (LHsExpr GhcPs)
x)           = NoExtCon -> [LGRHS GhcPs (LHsExpr GhcPs)]
forall a. NoExtCon -> a
Hs.noExtCon NoExtCon
XXGRHSs GhcPs (LHsExpr GhcPs)
x


getGuardLStmts :: Hs.GRHS Hs.GhcPs (Hs.LHsExpr Hs.GhcPs) -> [Hs.GuardLStmt Hs.GhcPs]
getGuardLStmts :: GRHS GhcPs (LHsExpr GhcPs) -> [GuardLStmt GhcPs]
getGuardLStmts (Hs.GRHS XCGRHS GhcPs (LHsExpr GhcPs)
_ [GuardLStmt GhcPs]
guards LHsExpr GhcPs
_) = [GuardLStmt GhcPs]
guards
getGuardLStmts (Hs.XGRHS XXGRHS GhcPs (LHsExpr GhcPs)
x)         = NoExtCon -> [GuardLStmt GhcPs]
forall a. NoExtCon -> a
Hs.noExtCon NoExtCon
XXGRHS GhcPs (LHsExpr GhcPs)
x