-- Copyright (c) Facebook, Inc. and its affiliates.
--
-- This source code is licensed under the MIT license found in the
-- LICENSE file in the root directory of this source tree.
--
{-# LANGUAGE CPP #-}
module Retrie.Pretty
  ( noColor
  , addColor
  , ppSrcSpan
  , ColoriseFun
  , strip
  , ppRepl
  , linesMap
  ) where

import Data.Char
import Data.List
import Data.Maybe
import qualified Data.HashMap.Strict as HashMap
import System.Console.ANSI

import Retrie.GHC

type ColoriseFun = ColorIntensity -> Color -> String -> String

noColor :: ColoriseFun
noColor :: ColoriseFun
noColor ColorIntensity
_ Color
_ = forall a. a -> a
id

addColor :: ColoriseFun
addColor :: ColoriseFun
addColor ColorIntensity
intensity Color
color String
x = forall a. Monoid a => [a] -> a
mconcat
  [ [SGR] -> String
setSGRCode [ConsoleLayer -> ColorIntensity -> Color -> SGR
SetColor ConsoleLayer
Foreground ColorIntensity
intensity Color
color]
  , String
x
  , [SGR] -> String
setSGRCode [SGR
Reset]
  ]

-- | Pretty print location of the file.
ppSrcSpan :: ColoriseFun -> SrcSpan -> String
ppSrcSpan :: ColoriseFun -> SrcSpan -> String
ppSrcSpan ColoriseFun
colorise SrcSpan
spn = case SrcSpan -> SrcLoc
srcSpanStart SrcSpan
spn of
  UnhelpfulLoc FastString
x -> FastString -> String
unpackFS FastString
x
#if __GLASGOW_HASKELL__ < 900
  RealSrcLoc loc -> intercalate (colorise Dull Cyan ":")
#else
  RealSrcLoc RealSrcLoc
loc Maybe BufPos
_ -> forall a. [a] -> [[a]] -> [a]
intercalate (ColoriseFun
colorise ColorIntensity
Dull Color
Cyan String
":")
#endif
    [ ColoriseFun
colorise ColorIntensity
Dull Color
Magenta forall a b. (a -> b) -> a -> b
$ FastString -> String
unpackFS forall a b. (a -> b) -> a -> b
$ RealSrcLoc -> FastString
srcLocFile RealSrcLoc
loc
    , ColoriseFun
colorise ColorIntensity
Dull Color
Green forall a b. (a -> b) -> a -> b
$ forall a. Show a => a -> String
show forall a b. (a -> b) -> a -> b
$ RealSrcLoc -> Int
srcLocLine RealSrcLoc
loc
    , ColoriseFun
colorise ColorIntensity
Dull Color
Green forall a b. (a -> b) -> a -> b
$ forall a. Show a => a -> String
show forall a b. (a -> b) -> a -> b
$ RealSrcLoc -> Int
srcLocCol RealSrcLoc
loc
    , String
""
    ]

-- | Get lines covering span and replace span with replacement string.
ppRepl :: HashMap.HashMap Int String -> SrcSpan -> String -> [String]
ppRepl :: HashMap Int String -> SrcSpan -> String -> [String]
ppRepl HashMap Int String
lMap SrcSpan
spn String
replacement = forall a. a -> Maybe a -> a
fromMaybe [String
replacement] forall a b. (a -> b) -> a -> b
$ do
  RealSrcLoc
startPos <- SrcLoc -> Maybe RealSrcLoc
getRealLoc forall a b. (a -> b) -> a -> b
$ SrcSpan -> SrcLoc
srcSpanStart SrcSpan
spn
  RealSrcLoc
endPos <- SrcLoc -> Maybe RealSrcLoc
getRealLoc forall a b. (a -> b) -> a -> b
$ SrcSpan -> SrcLoc
srcSpanEnd SrcSpan
spn
  String
startLine <- RealSrcLoc -> Maybe String
getLine' RealSrcLoc
startPos
  String
endLine <- RealSrcLoc -> Maybe String
getLine' RealSrcLoc
endPos
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ String -> [String]
lines forall a b. (a -> b) -> a -> b
$ forall a. Monoid a => [a] -> a
mconcat
    [ forall a. Int -> [a] -> [a]
take (RealSrcLoc -> Int
srcLocCol RealSrcLoc
startPos forall a. Num a => a -> a -> a
- Int
1) String
startLine
    , forall a. (a -> Bool) -> [a] -> [a]
dropWhile Char -> Bool
isSpace String
replacement
    , forall a. Int -> [a] -> [a]
drop (RealSrcLoc -> Int
srcLocCol RealSrcLoc
endPos forall a. Num a => a -> a -> a
- Int
1) String
endLine
    ]
  where
    getLine' :: RealSrcLoc -> Maybe String
getLine' RealSrcLoc
pos = forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup (RealSrcLoc -> Int
srcLocLine RealSrcLoc
pos) HashMap Int String
lMap

-- | Return HashMap from line number to line of a file.
linesMap :: String -> IO (HashMap.HashMap Int String)
linesMap :: String -> IO (HashMap Int String)
linesMap String
fp = forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
HashMap.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. [a] -> [b] -> [(a, b)]
zip [Int
1..] forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> [String]
lines forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO String
readFile String
fp

strip :: String -> String
strip :: String -> String
strip = forall a. (a -> Bool) -> [a] -> [a]
dropWhileEnd Char -> Bool
isSpace forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> Bool) -> [a] -> [a]
dropWhile Char -> Bool
isSpace