{-# 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 as PP

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 :: forall a. Data a => a -> Doc
simplePrintTree = DataToDoc -> forall a. Data a => a -> Doc
runDataToDoc ((DataToDoc -> DataToDoc) -> DataToDoc
forall a. (a -> a) -> a
fix DataToDoc -> DataToDoc
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 :: forall a. Data a => (DataToDoc -> DataToDoc) -> a -> Doc
simplePrintTreeWithCustom DataToDoc -> DataToDoc
toDocF = DataToDoc -> forall a. Data a => a -> Doc
runDataToDoc ((DataToDoc -> DataToDoc) -> DataToDoc
forall a. (a -> a) -> a
fix DataToDoc -> DataToDoc
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 :: forall a. Data a => Int -> a -> Doc
printTree Int
startIndent a
node =
  NodeLayouter -> Either Bool Int -> Doc
_lay_func (DataToLayouter -> forall a. Data a => a -> NodeLayouter
runDataToLayouter ((DataToLayouter -> DataToLayouter) -> DataToLayouter
forall a. (a -> a) -> a
fix DataToLayouter -> DataToLayouter
defaultLayouterF) a
node) (Int -> Either Bool Int
forall a b. b -> Either a b
Right Int
startIndent)

printTreeWithCustom :: Data a => Int -> LayouterF -> a -> Doc
printTreeWithCustom :: forall a.
Data a =>
Int -> (DataToLayouter -> DataToLayouter) -> a -> Doc
printTreeWithCustom Int
startIndent DataToLayouter -> DataToLayouter
layoutF a
node =
  NodeLayouter -> Either Bool Int -> Doc
_lay_func (DataToLayouter -> forall a. Data a => a -> NodeLayouter
runDataToLayouter ((DataToLayouter -> DataToLayouter) -> DataToLayouter
forall a. (a -> a) -> a
fix DataToLayouter -> DataToLayouter
layoutF) a
node) (Int -> Either Bool Int
forall a b. b -> Either a b
Right Int
startIndent)

showTree :: Data a => a -> String
showTree :: forall a. Data a => a -> String
showTree = Doc -> String
render (Doc -> String) -> (a -> Doc) -> a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> a -> Doc
forall a. Data a => Int -> a -> Doc
printTree Int
100

showTreeWithCustom :: Data a => LayouterF -> a -> String
showTreeWithCustom :: forall a.
Data a =>
(DataToLayouter -> DataToLayouter) -> a -> String
showTreeWithCustom DataToLayouter -> DataToLayouter
layoutF a
node = Doc -> String
render (Doc -> String) -> Doc -> String
forall a b. (a -> b) -> a -> b
$ Int -> (DataToLayouter -> DataToLayouter) -> a -> Doc
forall a.
Data a =>
Int -> (DataToLayouter -> DataToLayouter) -> a -> Doc
printTreeWithCustom Int
100 DataToLayouter -> DataToLayouter
layoutF a
node

-- | This newtype is necessary so `fix` can be used in combination with
--   the constrained forall-quantification.
newtype DataToDoc = DataToDoc
  { DataToDoc -> forall a. Data a => a -> Doc
runDataToDoc :: forall a . Data a => a -> Doc }

type ToDocF = DataToDoc -> DataToDoc

data NodeLayouter = NodeLayouter
  { NodeLayouter -> Int
_lay_llength :: Int -- ^ the length of this node, if printed
                        --   on a single line
  , NodeLayouter -> Bool
_lay_needsParens :: Bool
  , NodeLayouter -> Either Bool Int -> Doc
_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
  { DataToLayouter -> forall a. Data a => a -> NodeLayouter
runDataToLayouter :: forall a . Data a => a -> NodeLayouter }

type LayouterF = DataToLayouter -> DataToLayouter

defaultToDocF :: ToDocF
defaultToDocF :: DataToDoc -> DataToDoc
defaultToDocF (DataToDoc forall a. Data a => a -> Doc
lf) = (forall a. Data a => a -> Doc) -> DataToDoc
DataToDoc ((forall a. Data a => a -> Doc) -> DataToDoc)
-> (forall a. Data a => a -> Doc) -> DataToDoc
forall a b. (a -> b) -> a -> b
$ a -> Doc
forall a. Data a => a -> Doc
genLayouter (a -> Doc) -> (forall e. Data e => [e] -> Doc) -> a -> Doc
forall d (t :: * -> *) q.
(Data d, Typeable t) =>
(d -> q) -> (forall e. Data e => t e -> q) -> d -> q
`ext1Q` [e] -> Doc
forall e. Data e => [e] -> Doc
listLayouter
                                                       (a -> Doc) -> (String -> Doc) -> a -> Doc
forall a b r.
(Typeable a, Typeable b) =>
(a -> r) -> (b -> r) -> a -> r
`extQ` String -> Doc
string
  where
    genLayouter :: p -> Doc
genLayouter p
n =
      let cStr :: String
cStr = Constr -> String
showConstr (Constr -> String) -> Constr -> String
forall a b. (a -> b) -> a -> b
$ p -> Constr
forall a. Data a => a -> Constr
toConstr p
n
          childrenDoc :: [Doc]
childrenDoc = (forall a. Data a => a -> Doc) -> p -> [Doc]
forall a u. Data a => (forall d. Data d => d -> u) -> a -> [u]
forall u. (forall d. Data d => d -> u) -> p -> [u]
gmapQ d -> Doc
forall a. Data a => a -> Doc
lf p
n
      in  String -> Doc
text String
cStr Doc -> Doc -> Doc
$$ Int -> Doc -> Doc
nest Int
2 ([Doc] -> Doc
vcat [Doc]
childrenDoc)
    listLayouter :: forall b . Data b => [b] -> Doc
    listLayouter :: forall e. Data e => [e] -> Doc
listLayouter [] = String -> Doc
text String
"[]"
    listLayouter (b
x1:[b]
xr) = String -> Doc
text String
"[" Doc -> Doc -> Doc
$$ Int -> Doc -> Doc
nest Int
2 Doc
d1
                        Doc -> Doc -> Doc
$$ [Doc] -> Doc
vcat [String -> Doc
text String
"," Doc -> Doc -> Doc
$$ Int -> Doc -> Doc
nest Int
2 Doc
d | Doc
d <- [Doc]
dr]
                        Doc -> Doc -> Doc
$$ String -> Doc
text String
"]"
      where
        d1 :: Doc
d1 = b -> Doc
forall a. Data a => a -> Doc
lf b
x1
        dr :: [Doc]
dr = b -> Doc
forall a. Data a => a -> Doc
lf (b -> Doc) -> [b] -> [Doc]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [b]
xr
    string :: String -> Doc
    string :: String -> Doc
string String
s = String -> Doc
text (String -> Doc) -> String -> Doc
forall a b. (a -> b) -> a -> b
$ String -> String
forall a. Show a => a -> String
show String
s

defaultLayouterF :: LayouterF
defaultLayouterF :: DataToLayouter -> DataToLayouter
defaultLayouterF (DataToLayouter forall a. Data a => a -> NodeLayouter
lf) = (forall a. Data a => a -> NodeLayouter) -> DataToLayouter
DataToLayouter
                                     ((forall a. Data a => a -> NodeLayouter) -> DataToLayouter)
-> (forall a. Data a => a -> NodeLayouter) -> DataToLayouter
forall a b. (a -> b) -> a -> b
$ a -> NodeLayouter
forall a. Data a => a -> NodeLayouter
genLayouter (a -> NodeLayouter)
-> (forall e. Data e => [e] -> NodeLayouter) -> a -> NodeLayouter
forall d (t :: * -> *) q.
(Data d, Typeable t) =>
(d -> q) -> (forall e. Data e => t e -> q) -> d -> q
`ext1Q` [e] -> NodeLayouter
forall e. Data e => [e] -> NodeLayouter
listLayouter
                                                   (a -> NodeLayouter)
-> (String -> NodeLayouter) -> a -> NodeLayouter
forall a b r.
(Typeable a, Typeable b) =>
(a -> r) -> (b -> r) -> a -> r
`extQ` String -> NodeLayouter
string
  where
    genLayouter :: forall b . Data b => b -> NodeLayouter
    genLayouter :: forall a. Data a => a -> NodeLayouter
genLayouter b
n = Int -> Bool -> (Either Bool Int -> Doc) -> NodeLayouter
NodeLayouter Int
llen Bool
needParens Either Bool Int -> Doc
func
      where
        cs :: String
cs = Constr -> String
forall a. Show a => a -> String
show (Constr -> String) -> Constr -> String
forall a b. (a -> b) -> a -> b
$ b -> Constr
forall a. Data a => a -> Constr
toConstr b
n
        subs :: [NodeLayouter]
subs = (forall a. Data a => a -> NodeLayouter) -> b -> [NodeLayouter]
forall a u. Data a => (forall d. Data d => d -> u) -> a -> [u]
forall u. (forall d. Data d => d -> u) -> b -> [u]
gmapQ d -> NodeLayouter
forall a. Data a => a -> NodeLayouter
lf b
n
        llen :: Int
llen = String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
cs
             Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [NodeLayouter] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [NodeLayouter]
subs
             Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [ if NodeLayouter -> Bool
_lay_needsParens NodeLayouter
s
                       then NodeLayouter -> Int
_lay_llength NodeLayouter
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2
                       else NodeLayouter -> Int
_lay_llength NodeLayouter
s
                   | NodeLayouter
s <- [NodeLayouter]
subs
                   ]
        needParens :: Bool
needParens = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [NodeLayouter] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [NodeLayouter]
subs
        func :: Either Bool Int -> Doc
func (Right Int
i)
          | Int
llenInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<=Int
i = String -> Doc
text String
cs Doc -> Doc -> Doc
<+> [Doc] -> Doc
hsep [NodeLayouter -> Either Bool Int -> Doc
_lay_func NodeLayouter
s (Bool -> Either Bool Int
forall a b. a -> Either a b
Left Bool
True) | NodeLayouter
s <- [NodeLayouter]
subs]
          | Bool
otherwise =  String -> Doc
text String
cs
                      Doc -> Doc -> Doc
$$ Int -> Doc -> Doc
nest Int
2 ([Doc] -> Doc
vcat [NodeLayouter -> Either Bool Int -> Doc
_lay_func NodeLayouter
s (Int -> Either Bool Int
forall a b. b -> Either a b
Right (Int -> Either Bool Int) -> Int -> Either Bool Int
forall a b. (a -> b) -> a -> b
$ Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
2) | NodeLayouter
s <- [NodeLayouter]
subs])
        func (Left Bool
True)
          = (if [NodeLayouter] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [NodeLayouter]
subs then Doc -> Doc
forall a. a -> a
id else Doc -> Doc
parens)
          (Doc -> Doc) -> Doc -> Doc
forall a b. (a -> b) -> a -> b
$ String -> Doc
text String
cs Doc -> Doc -> Doc
<+> [Doc] -> Doc
hsep [NodeLayouter -> Either Bool Int -> Doc
_lay_func NodeLayouter
s (Bool -> Either Bool Int
forall a b. a -> Either a b
Left Bool
True) | NodeLayouter
s <- [NodeLayouter]
subs]
        func (Left Bool
False)
          = String -> Doc
text String
cs Doc -> Doc -> Doc
<+> [Doc] -> Doc
hsep [NodeLayouter -> Either Bool Int -> Doc
_lay_func NodeLayouter
s (Bool -> Either Bool Int
forall a b. a -> Either a b
Left Bool
True) | NodeLayouter
s <- [NodeLayouter]
subs]
    listLayouter :: forall b . Data b => [b] -> NodeLayouter
    listLayouter :: forall e. Data e => [e] -> NodeLayouter
listLayouter [] = Int -> Bool -> (Either Bool Int -> Doc) -> NodeLayouter
NodeLayouter Int
2 Bool
False ((Either Bool Int -> Doc) -> NodeLayouter)
-> (Either Bool Int -> Doc) -> NodeLayouter
forall a b. (a -> b) -> a -> b
$ \Either Bool Int
_ -> String -> Doc
text String
"[]"
    listLayouter xs :: [b]
xs@(b
_:[b]
_) = Int -> Bool -> (Either Bool Int -> Doc) -> NodeLayouter
NodeLayouter Int
llen Bool
False Either Bool Int -> Doc
forall {a}. Either a Int -> Doc
func
      where
        subs :: [NodeLayouter]
subs@(NodeLayouter
s1:[NodeLayouter]
sr) = b -> NodeLayouter
forall a. Data a => a -> NodeLayouter
lf (b -> NodeLayouter) -> [b] -> [NodeLayouter]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [b]
xs
        llen :: Int
llen = Int
1
             Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [NodeLayouter] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [NodeLayouter]
subs
             Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (NodeLayouter -> Int
_lay_llength (NodeLayouter -> Int) -> [NodeLayouter] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [NodeLayouter]
subs)
        func :: Either a Int -> Doc
func (Right Int
i)
          | Int
llenInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<=Int
i = String -> Doc
text String
"["
                   Doc -> Doc -> Doc
PP.<> [Doc] -> Doc
hcat (Doc -> [Doc] -> [Doc]
punctuate (String -> Doc
text String
",") [NodeLayouter -> Either Bool Int -> Doc
_lay_func NodeLayouter
s (Bool -> Either Bool Int
forall a b. a -> Either a b
Left Bool
False) | NodeLayouter
s <- [NodeLayouter]
subs])
                   Doc -> Doc -> Doc
PP.<> String -> Doc
text String
"]"
          | Bool
otherwise = String -> Doc
text String
"[" Doc -> Doc -> Doc
$$ Int -> Doc -> Doc
nest Int
2 (NodeLayouter -> Either Bool Int -> Doc
_lay_func NodeLayouter
s1 (Int -> Either Bool Int
forall a b. b -> Either a b
Right (Int -> Either Bool Int) -> Int -> Either Bool Int
forall a b. (a -> b) -> a -> b
$ Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
2))
                     Doc -> Doc -> Doc
$$ [Doc] -> Doc
vcat [String -> Doc
text String
"," Doc -> Doc -> Doc
$$ Int -> Doc -> Doc
nest Int
2 (NodeLayouter -> Either Bool Int -> Doc
_lay_func NodeLayouter
s (Int -> Either Bool Int
forall a b. b -> Either a b
Right (Int -> Either Bool Int) -> Int -> Either Bool Int
forall a b. (a -> b) -> a -> b
$ Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
2)) | NodeLayouter
s <- [NodeLayouter]
sr]
                     Doc -> Doc -> Doc
$$ String -> Doc
text String
"]"
        func (Left a
_)
          = Either a Int -> Doc
func (Int -> Either a Int
forall a b. b -> Either a b
Right Int
99999999)
    string :: String -> NodeLayouter
    string :: String -> NodeLayouter
string String
s = Int -> Bool -> (Either Bool Int -> Doc) -> NodeLayouter
NodeLayouter (String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
s') Bool
False ((Either Bool Int -> Doc) -> NodeLayouter)
-> (Either Bool Int -> Doc) -> NodeLayouter
forall a b. (a -> b) -> a -> b
$ \Either Bool Int
_ -> String -> Doc
text (String -> Doc) -> String -> Doc
forall a b. (a -> b) -> a -> b
$ String
s'
      where s' :: String
s' = String -> String
forall a. Show a => a -> String
show String
s