-- 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
_ = String -> String
forall a. a -> a
id

addColor :: ColoriseFun
addColor :: ColoriseFun
addColor ColorIntensity
intensity Color
color String
x = [String] -> String
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
_ -> String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate (ColoriseFun
colorise ColorIntensity
Dull Color
Cyan String
":")
#endif
    [ ColoriseFun
colorise ColorIntensity
Dull Color
Magenta (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ FastString -> String
unpackFS (FastString -> String) -> FastString -> String
forall a b. (a -> b) -> a -> b
$ RealSrcLoc -> FastString
srcLocFile RealSrcLoc
loc
    , ColoriseFun
colorise ColorIntensity
Dull Color
Green (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ Int -> String
forall a. Show a => a -> String
show (Int -> String) -> Int -> String
forall a b. (a -> b) -> a -> b
$ RealSrcLoc -> Int
srcLocLine RealSrcLoc
loc
    , ColoriseFun
colorise ColorIntensity
Dull Color
Green (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ Int -> String
forall a. Show a => a -> String
show (Int -> String) -> Int -> String
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 = [String] -> Maybe [String] -> [String]
forall a. a -> Maybe a -> a
fromMaybe [String
replacement] (Maybe [String] -> [String]) -> Maybe [String] -> [String]
forall a b. (a -> b) -> a -> b
$ do
  RealSrcLoc
startPos <- SrcLoc -> Maybe RealSrcLoc
getRealLoc (SrcLoc -> Maybe RealSrcLoc) -> SrcLoc -> Maybe RealSrcLoc
forall a b. (a -> b) -> a -> b
$ SrcSpan -> SrcLoc
srcSpanStart SrcSpan
spn
  RealSrcLoc
endPos <- SrcLoc -> Maybe RealSrcLoc
getRealLoc (SrcLoc -> Maybe RealSrcLoc) -> SrcLoc -> Maybe RealSrcLoc
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
  [String] -> Maybe [String]
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return ([String] -> Maybe [String]) -> [String] -> Maybe [String]
forall a b. (a -> b) -> a -> b
$ String -> [String]
lines (String -> [String]) -> String -> [String]
forall a b. (a -> b) -> a -> b
$ [String] -> String
forall a. Monoid a => [a] -> a
mconcat
    [ Int -> String -> String
forall a. Int -> [a] -> [a]
take (RealSrcLoc -> Int
srcLocCol RealSrcLoc
startPos Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) String
startLine
    , (Char -> Bool) -> String -> String
forall a. (a -> Bool) -> [a] -> [a]
dropWhile Char -> Bool
isSpace String
replacement
    , Int -> String -> String
forall a. Int -> [a] -> [a]
drop (RealSrcLoc -> Int
srcLocCol RealSrcLoc
endPos Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) String
endLine
    ]
  where
    getLine' :: RealSrcLoc -> Maybe String
getLine' RealSrcLoc
pos = Int -> HashMap Int String -> Maybe String
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 = [(Int, String)] -> HashMap Int String
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
HashMap.fromList ([(Int, String)] -> HashMap Int String)
-> (String -> [(Int, String)]) -> String -> HashMap Int String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [String] -> [(Int, String)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
1..] ([String] -> [(Int, String)])
-> (String -> [String]) -> String -> [(Int, String)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> [String]
lines (String -> HashMap Int String)
-> IO String -> IO (HashMap Int String)
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 = (Char -> Bool) -> String -> String
forall a. (a -> Bool) -> [a] -> [a]
dropWhileEnd Char -> Bool
isSpace (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> String -> String
forall a. (a -> Bool) -> [a] -> [a]
dropWhile Char -> Bool
isSpace