module Language.Haskell.Tools.PrettyPrint (prettyPrint, toRoseTree) where
import FastString (fsLit)
import SrcLoc
import Language.Haskell.Tools.AST (SrcTemplateStage, SourceInfoTraversal(..))
import Language.Haskell.Tools.PrettyPrint.RoseTree (RoseSourceInfo(..), RoseTree(..), toRoseTree)
import Language.Haskell.Tools.Transform.SourceTemplate
import Control.Monad.State
import Control.Reference
import Data.Foldable (Foldable(..), concat)
import Data.List as List
import Data.List.Split (splitOn)
import Data.Sequence hiding (null, replicate)
import Debug.Trace
prettyPrint :: (SourceInfoTraversal node) => node dom SrcTemplateStage -> String
prettyPrint = toList . printRose . toRoseTree
printRose :: RoseTree SrcTemplateStage -> Seq Char
printRose rt = evalState (printRose' startLoc rt) startLoc
where startLoc = mkRealSrcLoc (fsLit "") 1 1
type PPState = State RealSrcLoc
printRose' :: RealSrcLoc -> RoseTree SrcTemplateStage -> PPState (Seq Char)
printRose' parent (RoseTree (RoseSpan (SourceTemplateNode rng elems minInd relInd)) children)
= do slide <- calculateSlide rng
let printTemplateElems :: [SourceTemplateElem] -> [RoseTree SrcTemplateStage] -> PPState (Seq Char)
printTemplateElems (TextElem txtElems _ : rest) children = putString slide min txt >+< printTemplateElems rest children
where txt = concatMap (^. sourceTemplateText) txtElems
printTemplateElems (ChildElem : rest) (child : children) = printRose' parent child >+< printTemplateElems rest children
printTemplateElems [] [] = return empty
printTemplateElems _ [] = error $ "More child elem in template than actual children (elems: " ++ show elems ++ ", children: " ++ show children ++ ")"
printTemplateElems [] _ = error $ "Not all children are used to pretty printing. (elems: " ++ show elems ++ ", children: " ++ show children ++ ")"
min = minInd `max` getPosByRelative parent relInd
printTemplateElems elems children
printRose' _ (RoseTree (RoseList (SourceTemplateList {})) []) = return empty
printRose' parent (RoseTree (RoseList (SourceTemplateList rng bef aft defSep indented seps minInd relInd)) children)
= do slide <- calculateSlide rng
actRng <- get
let min = minInd `max` getPosByRelative parent relInd
putString slide min bef
>+< (maybe printListWithSeps printListWithSepsIndented indented) actRng slide min actualSeps children
>+< putString slide min aft
where stringSeps :: [String]
stringSeps = map (concatMap (^. sourceTemplateText)) (map fst seps)
actualSeps = case stringSeps of [] -> repeat defSep
_ -> stringSeps ++ repeat (last stringSeps)
printRose' _ (RoseTree (RoseOptional (SourceTemplateOpt {})) []) = return empty
printRose' parent (RoseTree (RoseOptional (SourceTemplateOpt rng bef aft minInd relInd)) [child])
= do slide <- calculateSlide rng
actRng <- get
let min = minInd `max` getPosByRelative parent relInd
putString slide min bef >+< printRose' actRng child >+< putString slide min aft
printRose' _ (RoseTree (RoseOptional _) _) = error "More than one child element in an optional node."
getPosByRelative :: RealSrcLoc -> Maybe Int -> Int
getPosByRelative sp (Just i) = srcLocCol sp + i 1
getPosByRelative _ _ = 0
calculateSlide :: SrcSpan -> PPState Int
calculateSlide (RealSrcSpan originalSpan) = do
actualSpan <- get
return $ srcLocCol actualSpan srcLocCol (realSrcSpanStart originalSpan)
calculateSlide _ = return 0
putString :: Int -> Int -> String -> PPState (Seq Char)
putString slide minInd s
= do modify $ advanceStr newStr
return (fromList newStr)
where start:rest = splitOn "\n" s
newStr = concat $ intersperse ("\n" ++ replicate slide ' ') (start : map (extendToNSpaces minInd) rest)
extendToNSpaces n str = replicate n ' ' ++ (List.dropWhile (== ' ') $ List.take n str) ++ List.drop n str
advanceStr :: String -> RealSrcLoc -> RealSrcLoc
advanceStr s loc = foldl advanceSrcLoc loc s
untilReaches :: String -> RealSrcLoc -> RealSrcLoc -> (String, Int)
untilReaches s start end
= let ls = splitOn "\n" s
in case ls of _:_:_ -> (unlines (init ls) ++)
`mapFst` untilReaches' (last ls) (advanceSrcLoc start '\n') end
_ -> (s, srcLocCol $ foldl advanceSrcLoc start s)
where
untilReaches' [] curr _ = ([], srcLocCol curr)
untilReaches' (c:rest) curr until | srcLocCol advancedLoc <= srcLocCol until
= (c:) `mapFst` untilReaches' rest advancedLoc until
where advancedLoc = advanceSrcLoc curr c
untilReaches' _ curr _ = ([], srcLocCol curr)
mapFst :: (a -> b) -> (a, x) -> (b, x)
mapFst f (a, x) = (f a, x)
(>+<) :: PPState (Seq Char) -> PPState (Seq Char) -> PPState (Seq Char)
(>+<) = liftM2 (><)
printListWithSeps :: RealSrcLoc -> Int -> Int -> [String] -> [RoseTree SrcTemplateStage] -> PPState (Seq Char)
printListWithSeps = printListWithSeps' (const putString) 0
printListWithSepsIndented :: [Bool] -> RealSrcLoc -> Int -> Int -> [String] -> [RoseTree SrcTemplateStage] -> PPState (Seq Char)
printListWithSepsIndented indentedChildren parent slide minInd seps children
= do base <- get
let putCorrectSep i _ min s | isIndented i
= do curr <- get
let (shortened, currCol) = untilReaches s curr base
putString 0 min $ shortened ++ replicate (srcLocCol base currCol) ' '
putCorrectSep _ slide minInd s = putString slide minInd s
printListWithSeps' putCorrectSep 0 parent slide minInd seps children
where
isIndented i = case List.drop i indentedChildren of False:_ -> False; _ -> True
printListWithSeps' :: (Int -> Int -> Int -> String -> PPState (Seq Char)) -> Int -> RealSrcLoc -> Int -> Int -> [String] -> [RoseTree SrcTemplateStage] -> PPState (Seq Char)
printListWithSeps' _ _ _ _ _ _ [] = return empty
printListWithSeps' _ _ parent _ _ _ [child] = printRose' parent child
printListWithSeps' putCorrectSep i parent slide minInd (sep:seps) (child:children)
= printRose' parent child >+< putCorrectSep i slide minInd sep >+< printListWithSeps' putCorrectSep (i+1) parent slide minInd seps children
printListWithSeps' _ _ _ _ _ [] _ = error "printListWithSeps': the number of elements and separators does not match"