module Language.Haskell.Tools.PrettyPrint (prettyPrint) where
import SrcLoc
import FastString
import Language.Haskell.Tools.AST
import Language.Haskell.Tools.PrettyPrint.RoseTree
import Language.Haskell.Tools.AnnTrf.SourceTemplate
import Control.Monad.State
import Control.Reference
import Data.Maybe
import Data.List.Split
import Data.Foldable
import Data.StructuralTraversal
import Data.Sequence hiding (null, replicate)
prettyPrint :: (StructuralTraversable node) => node (NodeInfo sema SourceTemplate) -> String
prettyPrint = toList . printRose . toRoseTree
printRose :: RoseTree (NodeInfo sema SourceTemplate) -> Seq Char
printRose rt = evalState (printRose' $ fmap (^. sourceInfo&sourceTemplateElems) rt) (mkRealSrcLoc (fsLit "") 1 1)
type PPState = State RealSrcLoc
printRose' :: RoseTree [SourceTemplateElem] -> PPState (Seq Char)
printRose' (RoseTree (TextElem txt : rest) children)
= putString txt >+< printRose' (RoseTree rest children)
printRose' (RoseTree (ChildElem : rest) (child : children))
= printRose' child >+< printRose' (RoseTree rest children)
printRose' (RoseTree [ChildListElem {}] []) = return empty
printRose' (RoseTree [ChildListElem bef aft defSep indented []] children)
= putString bef >+< (if indented then printListWithSepsIndented else printListWithSeps) (repeat defSep) children >+< putString aft
printRose' (RoseTree [ChildListElem bef aft _ indented seps] children)
= putString bef >+< (if indented then printListWithSepsIndented else printListWithSeps) (seps ++ repeat (last seps)) children >+< putString aft
printRose' (RoseTree [OptionalChildElem _ _] []) = return empty
printRose' (RoseTree [OptionalChildElem bef aft] [child]) = putString bef >+< printRose' child >+< putString aft
printRose' (RoseTree [OptionalChildElem _ _] _) = error "More than one child element in an optional node."
printRose' (RoseTree [] []) = return empty
printRose' r@(RoseTree (ChildElem : rest) []) = error ("More child elem in template than actual children. In: " ++ show r)
printRose' r@(RoseTree [] children) = error ("Not all children are used to pretty printing. In: " ++ show r)
printRose' r = error ("printRose': unexpected input: " ++ show r)
putString :: String -> PPState (Seq Char)
putString s = do modify $ advanceStr s
return (fromList s)
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 start)
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 :: [String] -> [RoseTree [SourceTemplateElem]] -> PPState (Seq Char)
printListWithSeps = printListWithSeps' putString
printListWithSepsIndented :: [String] -> [RoseTree [SourceTemplateElem]] -> PPState (Seq Char)
printListWithSepsIndented seps children
= do base <- get
let putCorrectSep s = do curr <- get
let (shortened, currCol) = untilReaches s curr base
putString $ shortened ++ replicate (srcLocCol base currCol) ' '
printListWithSeps' putCorrectSep seps children
printListWithSeps' :: (String -> PPState (Seq Char)) -> [String] -> [RoseTree [SourceTemplateElem]] -> PPState (Seq Char)
printListWithSeps' _ _ [] = return empty
printListWithSeps' putCorrectSep _ [child] = printRose' child
printListWithSeps' putCorrectSep (sep:seps) (child:children)
= printRose' child >+< putCorrectSep sep >+< printListWithSeps' putCorrectSep seps children