{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE RankNTypes #-}
module DataTreePrint
( simplePrintTree
, simplePrintTreeWithCustom
, printTree
, printTreeWithCustom
, showTree
, showTreeWithCustom
, DataToLayouter(..)
, LayouterF
, NodeLayouter(..)
, defaultLayouterF
)
where
import Data.Data
import Text.PrettyPrint
import Data.Generics.Aliases
import Data.Function (fix)
import Data.Functor ((<$>))
-- | The "simple" printer does not try to fit more than one node into the
-- same line, even if it would fit.
simplePrintTree :: Data a => a -> Doc
simplePrintTree = runDataToDoc (fix defaultToDocF)
-- | Allows to specialize the transformation for specific types. Use `syb`'s
-- `extQ` function(s). See the source of `defaultLayouterF` for an
-- example of how to do this.
simplePrintTreeWithCustom :: Data a => ToDocF -> a -> Doc
simplePrintTreeWithCustom toDocF = runDataToDoc (fix toDocF)
-------
-------
-- | Somewhat more intelligent printer that tries to fit multiple nodes
-- into the same line there is space given the specified number of total
-- columns.
-- For example, `(1,2,3)` will be printed as "(,,) (1) (2) (3)" instead
-- of "(,,)\n 1\n 2\n 3". Parentheses are added in these cases to prevent
-- syntactic ambiguities.
printTree :: forall a . Data a => Int -> a -> Doc
printTree startIndent node =
_lay_func (runDataToLayouter (fix defaultLayouterF) node) (Right startIndent)
printTreeWithCustom :: Data a => Int -> LayouterF -> a -> Doc
printTreeWithCustom startIndent layoutF node =
_lay_func (runDataToLayouter (fix layoutF) node) (Right startIndent)
showTree :: Data a => a -> String
showTree = render . printTree 100
showTreeWithCustom :: Data a => LayouterF -> a -> String
showTreeWithCustom layoutF node = render $ printTreeWithCustom 100 layoutF node
-- | This newtype is necessary so `fix` can be used in combination with
-- the constrained forall-quantification.
newtype DataToDoc = DataToDoc
{ runDataToDoc :: forall a . Data a => a -> Doc }
type ToDocF = DataToDoc -> DataToDoc
data NodeLayouter = NodeLayouter
{ _lay_llength :: Int -- ^ the length of this node, if printed
-- on a single line
, _lay_needsParens :: Bool
, _lay_func :: Either Bool Int -> Doc
-- ^ Left: one-line output, the boolean
-- indicates if parentheses are advisable
-- given the context. (They can be omitted
-- in cases like when there is only one
-- constructor).
-- Right: The Int is the remaining vertical
-- space left for this node.
}
-- | This newtype is necessary so `fix` can be used in combination with
-- the constrained forall-quantification.
newtype DataToLayouter = DataToLayouter
{ runDataToLayouter :: forall a . Data a => a -> NodeLayouter }
type LayouterF = DataToLayouter -> DataToLayouter
defaultToDocF :: ToDocF
defaultToDocF (DataToDoc lf) = DataToDoc $ genLayouter `extQ` string
`ext1Q` listLayouter
where
genLayouter n =
let cStr = showConstr $ toConstr n
childrenDoc = gmapQ lf n
in text cStr $$ nest 2 (vcat childrenDoc)
listLayouter :: forall b . Data b => [b] -> Doc
listLayouter [] = text "[]"
listLayouter (x1:xr) = text "[" $$ nest 2 d1
$$ vcat [text "," $$ nest 2 d | d <- dr]
$$ text "]"
where
d1 = lf x1
dr = lf <$> xr
string :: String -> Doc
string s = text $ show s
defaultLayouterF :: LayouterF
defaultLayouterF (DataToLayouter lf) = DataToLayouter
$ genLayouter `ext1Q` listLayouter
`extQ` string
where
genLayouter :: forall b . Data b => b -> NodeLayouter
genLayouter n = NodeLayouter llen needParens func
where
cs = show $ toConstr n
subs = gmapQ lf n
llen = length cs
+ length subs
+ sum [ if _lay_needsParens s
then _lay_llength s + 2
else _lay_llength s
| s <- subs
]
needParens = not $ null subs
func (Right i)
| llen<=i = text cs <+> hsep [_lay_func s (Left True) | s <- subs]
| otherwise = text cs
$$ nest 2 (vcat [_lay_func s (Right $ i-2) | s <- subs])
func (Left True)
= (if null subs then id else parens)
$ text cs <+> hsep [_lay_func s (Left True) | s <- subs]
func (Left False)
= text cs <+> hsep [_lay_func s (Left True) | s <- subs]
listLayouter :: forall b . Data b => [b] -> NodeLayouter
listLayouter [] = NodeLayouter 2 False $ \_ -> text "[]"
listLayouter xs@(_:_) = NodeLayouter llen False func
where
subs@(s1:sr) = lf <$> xs
llen = 1
+ length subs
+ sum (_lay_llength <$> subs)
func (Right i)
| llen<=i = text "["
<> hcat (punctuate (text ",") [_lay_func s (Left False) | s <- subs])
<> text "]"
| otherwise = text "[" $$ nest 2 (_lay_func s1 (Right $ i-2))
$$ vcat [text "," $$ nest 2 (_lay_func s (Right $ i-2)) | s <- sr]
$$ text "]"
func (Left _)
= func (Right 99999999)
string :: String -> NodeLayouter
string s = NodeLayouter (length s') False $ \_ -> text $ s'
where s' = show s