----------------------------------------------------------------------------
-- |
-- Module      :  Prettyprinter.Data
-- Copyright   :  (c) Sergey Vinokurov 2018
-- License     :  Apache-2.0 (see LICENSE)
-- Maintainer  :  serg.foo@gmail.com
----------------------------------------------------------------------------

{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE PatternGuards       #-}
{-# LANGUAGE Rank2Types          #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Prettyprinter.Data
  ( ppData
  , ppDataSimple
  , Data
  ) where

import Data.Data
import Data.Generics qualified
import Prettyprinter
import Prettyprinter qualified as PP
import Prettyprinter.Combinators
import Prettyprinter.MetaDoc

-- $setup
-- >>> :set -XDeriveDataTypeable
-- >>> :set -XImportQualifiedPost
-- >>> import Data.Data
-- >>> import Data.List.NonEmpty (NonEmpty(..))
-- >>> import Data.List.NonEmpty qualified as NonEmpty
-- >>> import Data.Map.Strict (Map)
-- >>> import Data.Map.Strict qualified as Map
--
-- >>> :{
-- data Test =
--     Foo Int [Int] Double (Maybe Test)
--   | Bar (String, Int, Int) (Map String Int) (Map String Int) (Maybe Test) (NonEmpty Int)
--   deriving (Data)
-- :}

-- | Prettyprint using 'Data.Data' instance.
--
-- >>> :{
-- test =
--   Bar
--     ("foo", 10, 20)
--     (Map.fromList (zip ["foo", "bar", "baz"] [1..]))
--     (Map.fromList (zip ["foo", "bar", "baz", "quux", "fizz", "buzz", "frob", "wat"] [1..]))
--     (Just
--       (Foo
--          1
--          []
--          3.14159265358979323846264338327950288
--          (Just
--             (Foo
--                1
--                [2]
--                2.71828182
--                (Just (Bar ("x", 1, 2) mempty mempty Nothing (NonEmpty.fromList [42])))))))
--     (NonEmpty.fromList [1..42])
-- :}
--
-- >>> ppData test
-- Bar
--   (foo, 10, 20)
--   {bar -> 2, baz -> 3, foo -> 1}
--   { bar -> 2
--   , baz -> 3
--   , buzz -> 6
--   , fizz -> 5
--   , foo -> 1
--   , frob -> 7
--   , quux -> 4
--   , wat -> 8
--   }
--   Just Foo
--          1
--          {}
--          3.141592653589793
--          Just (Foo 1 [2] 2.71828182 (Just (Bar (x, 1, 2) {} {} Nothing [42])))
--   [ 1
--   , 2
--   , 3
--   , 4
--   , 5
--   , 6
--   , 7
--   , 8
--   , 9
--   , 10
--   , 11
--   , 12
--   , 13
--   , 14
--   , 15
--   , 16
--   , 17
--   , 18
--   , 19
--   , 20
--   , 21
--   , 22
--   , 23
--   , 24
--   , 25
--   , 26
--   , 27
--   , 28
--   , 29
--   , 30
--   , 31
--   , 32
--   , 33
--   , 34
--   , 35
--   , 36
--   , 37
--   , 38
--   , 39
--   , 40
--   , 41
--   , 42
--   ]
ppData :: Data a => a -> Doc ann
ppData :: a -> Doc ann
ppData = MetaDoc ann -> Doc ann
forall ann. MetaDoc ann -> Doc ann
mdPayload (MetaDoc ann -> Doc ann) -> (a -> MetaDoc ann) -> a -> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> MetaDoc ann
forall a ann. Data a => a -> MetaDoc ann
gpretty

ppDataSimple :: Data a => a -> Doc ann
ppDataSimple :: a -> Doc ann
ppDataSimple = String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (String -> Doc ann) -> (a -> String) -> a -> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> String
forall a. Data a => a -> String
Data.Generics.gshow

gpretty :: forall a ann. Data a => a -> MetaDoc ann
gpretty :: a -> MetaDoc ann
gpretty =
  a -> MetaDoc ann
forall b. Data b => b -> MetaDoc ann
go
    (a -> MetaDoc ann) -> (String -> MetaDoc ann) -> a -> MetaDoc ann
forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`Data.Generics.extQ` String -> MetaDoc ann
forall ann. String -> MetaDoc ann
stringMetaDoc
    (a -> MetaDoc ann) -> (Text -> MetaDoc ann) -> a -> MetaDoc ann
forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`Data.Generics.extQ` Text -> MetaDoc ann
forall ann. Text -> MetaDoc ann
strictTextMetaDoc
    (a -> MetaDoc ann) -> (Text -> MetaDoc ann) -> a -> MetaDoc ann
forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`Data.Generics.extQ` Text -> MetaDoc ann
forall ann. Text -> MetaDoc ann
lazyTextMetaDoc
    (a -> MetaDoc ann) -> (Int -> MetaDoc ann) -> a -> MetaDoc ann
forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`Data.Generics.extQ` Int -> MetaDoc ann
forall ann. Int -> MetaDoc ann
metaDocInt
    (a -> MetaDoc ann) -> (Float -> MetaDoc ann) -> a -> MetaDoc ann
forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`Data.Generics.extQ` Float -> MetaDoc ann
forall ann. Float -> MetaDoc ann
metaDocFloat
    (a -> MetaDoc ann) -> (Double -> MetaDoc ann) -> a -> MetaDoc ann
forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`Data.Generics.extQ` Double -> MetaDoc ann
forall ann. Double -> MetaDoc ann
metaDocDouble
    (a -> MetaDoc ann) -> (Integer -> MetaDoc ann) -> a -> MetaDoc ann
forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`Data.Generics.extQ` Integer -> MetaDoc ann
forall ann. Integer -> MetaDoc ann
metaDocInteger
    (a -> MetaDoc ann) -> (Word -> MetaDoc ann) -> a -> MetaDoc ann
forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`Data.Generics.extQ` Word -> MetaDoc ann
forall ann. Word -> MetaDoc ann
metaDocWord
    (a -> MetaDoc ann) -> (Word8 -> MetaDoc ann) -> a -> MetaDoc ann
forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`Data.Generics.extQ` Word8 -> MetaDoc ann
forall ann. Word8 -> MetaDoc ann
metaDocWord8
    (a -> MetaDoc ann) -> (Word16 -> MetaDoc ann) -> a -> MetaDoc ann
forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`Data.Generics.extQ` Word16 -> MetaDoc ann
forall ann. Word16 -> MetaDoc ann
metaDocWord16
    (a -> MetaDoc ann) -> (Word32 -> MetaDoc ann) -> a -> MetaDoc ann
forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`Data.Generics.extQ` Word32 -> MetaDoc ann
forall ann. Word32 -> MetaDoc ann
metaDocWord32
    (a -> MetaDoc ann) -> (Word64 -> MetaDoc ann) -> a -> MetaDoc ann
forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`Data.Generics.extQ` Word64 -> MetaDoc ann
forall ann. Word64 -> MetaDoc ann
metaDocWord64
    (a -> MetaDoc ann) -> (Int8 -> MetaDoc ann) -> a -> MetaDoc ann
forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`Data.Generics.extQ` Int8 -> MetaDoc ann
forall ann. Int8 -> MetaDoc ann
metaDocInt8
    (a -> MetaDoc ann) -> (Int16 -> MetaDoc ann) -> a -> MetaDoc ann
forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`Data.Generics.extQ` Int16 -> MetaDoc ann
forall ann. Int16 -> MetaDoc ann
metaDocInt16
    (a -> MetaDoc ann) -> (Int32 -> MetaDoc ann) -> a -> MetaDoc ann
forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`Data.Generics.extQ` Int32 -> MetaDoc ann
forall ann. Int32 -> MetaDoc ann
metaDocInt32
    (a -> MetaDoc ann) -> (Int64 -> MetaDoc ann) -> a -> MetaDoc ann
forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`Data.Generics.extQ` Int64 -> MetaDoc ann
forall ann. Int64 -> MetaDoc ann
metaDocInt64
    (a -> MetaDoc ann) -> (() -> MetaDoc ann) -> a -> MetaDoc ann
forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`Data.Generics.extQ` () -> MetaDoc ann
forall ann. () -> MetaDoc ann
metaDocUnit
    (a -> MetaDoc ann) -> (Bool -> MetaDoc ann) -> a -> MetaDoc ann
forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`Data.Generics.extQ` Bool -> MetaDoc ann
forall ann. Bool -> MetaDoc ann
metaDocBool
    (a -> MetaDoc ann) -> (Char -> MetaDoc ann) -> a -> MetaDoc ann
forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`Data.Generics.extQ` Char -> MetaDoc ann
forall ann. Char -> MetaDoc ann
metaDocChar
    -- Probably requires qualified constrtains...
    -- `Data.Generics.extQ`
    --   ((atomicMetaDoc . ppMapWith (mdPayload . gpretty) (mdPayload . gpretty)) ::
    --     forall k v. (Data k, Data v) => Map k v -> MetaDoc ann)
  where
    go :: Data b => b -> MetaDoc ann
    go :: b -> MetaDoc ann
go b
t
      | String
constructorName String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
"fromList"
      , Just [Maybe (MetaDoc ann, MetaDoc ann)]
mapItems  <- Int
-> (forall d.
    Data d =>
    d -> Maybe [Maybe (MetaDoc ann, MetaDoc ann)])
-> b
-> Maybe [Maybe (MetaDoc ann, MetaDoc ann)]
forall a u. Data a => Int -> (forall d. Data d => d -> u) -> a -> u
gmapQi Int
0 ((forall b. Data b => b -> Maybe (MetaDoc ann, MetaDoc ann))
-> d -> Maybe [Maybe (MetaDoc ann, MetaDoc ann)]
forall a c.
Data a =>
(forall b. Data b => b -> c) -> a -> Maybe [c]
listElements ((forall b. Data b => b -> MetaDoc ann)
-> b -> Maybe (MetaDoc ann, MetaDoc ann)
forall a c.
Data a =>
(forall b. Data b => b -> c) -> a -> Maybe (c, c)
isPair forall b. Data b => b -> MetaDoc ann
forall a ann. Data a => a -> MetaDoc ann
gpretty)) b
t
      , Just [(MetaDoc ann, MetaDoc ann)]
mapItems' <- [Maybe (MetaDoc ann, MetaDoc ann)]
-> Maybe [(MetaDoc ann, MetaDoc ann)]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [Maybe (MetaDoc ann, MetaDoc ann)]
mapItems
      = Doc ann -> MetaDoc ann
forall ann. Doc ann -> MetaDoc ann
atomicMetaDoc
      (Doc ann -> MetaDoc ann) -> Doc ann -> MetaDoc ann
forall a b. (a -> b) -> a -> b
$ (MetaDoc ann -> Doc ann)
-> (MetaDoc ann -> Doc ann)
-> [(MetaDoc ann, MetaDoc ann)]
-> Doc ann
forall k ann v.
(k -> Doc ann) -> (v -> Doc ann) -> [(k, v)] -> Doc ann
ppAssocListWith MetaDoc ann -> Doc ann
forall ann. MetaDoc ann -> Doc ann
mdPayload MetaDoc ann -> Doc ann
forall ann. MetaDoc ann -> Doc ann
mdPayload [(MetaDoc ann, MetaDoc ann)]
mapItems'
      | Just [Maybe (MetaDoc ann, MetaDoc ann)]
mapItems  <- (forall b. Data b => b -> Maybe (MetaDoc ann, MetaDoc ann))
-> b -> Maybe [Maybe (MetaDoc ann, MetaDoc ann)]
forall a c.
Data a =>
(forall b. Data b => b -> c) -> a -> Maybe [c]
listElements ((forall b. Data b => b -> MetaDoc ann)
-> b -> Maybe (MetaDoc ann, MetaDoc ann)
forall a c.
Data a =>
(forall b. Data b => b -> c) -> a -> Maybe (c, c)
isPair forall b. Data b => b -> MetaDoc ann
forall a ann. Data a => a -> MetaDoc ann
gpretty) b
t
      , Just [(MetaDoc ann, MetaDoc ann)]
mapItems' <- [Maybe (MetaDoc ann, MetaDoc ann)]
-> Maybe [(MetaDoc ann, MetaDoc ann)]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [Maybe (MetaDoc ann, MetaDoc ann)]
mapItems
      = Doc ann -> MetaDoc ann
forall ann. Doc ann -> MetaDoc ann
atomicMetaDoc
      (Doc ann -> MetaDoc ann) -> Doc ann -> MetaDoc ann
forall a b. (a -> b) -> a -> b
$ (MetaDoc ann -> Doc ann)
-> (MetaDoc ann -> Doc ann)
-> [(MetaDoc ann, MetaDoc ann)]
-> Doc ann
forall k ann v.
(k -> Doc ann) -> (v -> Doc ann) -> [(k, v)] -> Doc ann
ppAssocListWith MetaDoc ann -> Doc ann
forall ann. MetaDoc ann -> Doc ann
mdPayload MetaDoc ann -> Doc ann
forall ann. MetaDoc ann -> Doc ann
mdPayload [(MetaDoc ann, MetaDoc ann)]
mapItems'
      | Just [MetaDoc ann]
listItems <- (forall b. Data b => b -> MetaDoc ann) -> b -> Maybe [MetaDoc ann]
forall a c.
Data a =>
(forall b. Data b => b -> c) -> a -> Maybe [c]
listElements forall b. Data b => b -> MetaDoc ann
forall a ann. Data a => a -> MetaDoc ann
gpretty b
t
      = Doc ann -> MetaDoc ann
forall ann. Doc ann -> MetaDoc ann
atomicMetaDoc
      (Doc ann -> MetaDoc ann) -> Doc ann -> MetaDoc ann
forall a b. (a -> b) -> a -> b
$ (MetaDoc ann -> Doc ann) -> [MetaDoc ann] -> Doc ann
forall a ann. (a -> Doc ann) -> [a] -> Doc ann
ppListWith MetaDoc ann -> Doc ann
forall ann. MetaDoc ann -> Doc ann
mdPayload [MetaDoc ann]
listItems
      | Bool
isTuple
      = Doc ann -> MetaDoc ann
forall ann. Doc ann -> MetaDoc ann
atomicMetaDoc
      (Doc ann -> MetaDoc ann) -> Doc ann -> MetaDoc ann
forall a b. (a -> b) -> a -> b
$ Doc ann -> Doc ann -> [Doc ann] -> Doc ann
forall (f :: * -> *) ann.
Foldable f =>
Doc ann -> Doc ann -> f (Doc ann) -> Doc ann
ppListWithDelim Doc ann
forall ann. Doc ann
PP.lparen Doc ann
forall ann. Doc ann
PP.rparen
      ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$ (MetaDoc ann -> Doc ann) -> [MetaDoc ann] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map MetaDoc ann -> Doc ann
forall ann. MetaDoc ann -> Doc ann
mdPayload [MetaDoc ann]
fields
      | Bool
otherwise
      = MetaDoc ann -> [MetaDoc ann] -> MetaDoc ann
forall ann. MetaDoc ann -> [MetaDoc ann] -> MetaDoc ann
constructorAppMetaDoc MetaDoc ann
constructorDoc [MetaDoc ann]
fields
      where
        constructorDoc :: MetaDoc ann
        constructorDoc :: MetaDoc ann
constructorDoc = Doc ann -> MetaDoc ann
forall ann. Doc ann -> MetaDoc ann
atomicMetaDoc (Doc ann -> MetaDoc ann) -> Doc ann -> MetaDoc ann
forall a b. (a -> b) -> a -> b
$ String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty String
constructorName

        fields :: [MetaDoc ann]
        fields :: [MetaDoc ann]
fields = (forall b. Data b => b -> MetaDoc ann) -> b -> [MetaDoc ann]
forall a u. Data a => (forall d. Data d => d -> u) -> a -> [u]
gmapQ forall b. Data b => b -> MetaDoc ann
forall a ann. Data a => a -> MetaDoc ann
gpretty b
t

        constructorName :: String
        constructorName :: String
constructorName = Constr -> String
showConstr (Constr -> String) -> Constr -> String
forall a b. (a -> b) -> a -> b
$ b -> Constr
forall a. Data a => a -> Constr
toConstr b
t

        isTuple :: Bool
        isTuple :: Bool
isTuple = (Char -> Bool) -> String -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
',') ((Char -> Bool) -> String -> String
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Char -> Bool) -> Char -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> String -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` (String
"()" :: String))) String
constructorName)

isPair :: Data a => (forall b. Data b => b -> c) -> a -> Maybe (c, c)
isPair :: (forall b. Data b => b -> c) -> a -> Maybe (c, c)
isPair forall b. Data b => b -> c
f a
x
  | String
constructorName String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
"(,)" = (c, c) -> Maybe (c, c)
forall a. a -> Maybe a
Just (Int -> (forall b. Data b => b -> c) -> a -> c
forall a u. Data a => Int -> (forall d. Data d => d -> u) -> a -> u
gmapQi Int
0 forall b. Data b => b -> c
f a
x, Int -> (forall b. Data b => b -> c) -> a -> c
forall a u. Data a => Int -> (forall d. Data d => d -> u) -> a -> u
gmapQi Int
1 forall b. Data b => b -> c
f a
x)
  | Bool
otherwise                = Maybe (c, c)
forall a. Maybe a
Nothing
  where
    constructorName :: String
    constructorName :: String
constructorName = Constr -> String
showConstr (Constr -> String) -> Constr -> String
forall a b. (a -> b) -> a -> b
$ a -> Constr
forall a. Data a => a -> Constr
toConstr a
x

-- | Try to treat @a@ as a list and prettyprint its elements with @f@.
-- Returns Just on succes and Nothing if @a@ wasn't a list after all.
listElements :: forall a c. Data a => (forall b. Data b => b -> c) -> a -> Maybe [c]
listElements :: (forall b. Data b => b -> c) -> a -> Maybe [c]
listElements forall b. Data b => b -> c
f = a -> Maybe [c]
forall d. Data d => d -> Maybe [c]
go
  where
    go :: Data d => d -> Maybe [c]
    go :: d -> Maybe [c]
go d
x
      | Bool
isNull    = [c] -> Maybe [c]
forall a. a -> Maybe a
Just []
      | Bool
isCons    = (:) (Int -> (forall b. Data b => b -> c) -> d -> c
forall a u. Data a => Int -> (forall d. Data d => d -> u) -> a -> u
gmapQi Int
0 forall b. Data b => b -> c
f d
x) ([c] -> [c]) -> Maybe [c] -> Maybe [c]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> (forall d. Data d => d -> Maybe [c]) -> d -> Maybe [c]
forall a u. Data a => Int -> (forall d. Data d => d -> u) -> a -> u
gmapQi Int
1 forall d. Data d => d -> Maybe [c]
go d
x
      | Bool
otherwise = Maybe [c]
forall a. Maybe a
Nothing
      where
        constructorName :: String
        constructorName :: String
constructorName = Constr -> String
showConstr (Constr -> String) -> Constr -> String
forall a b. (a -> b) -> a -> b
$ d -> Constr
forall a. Data a => a -> Constr
toConstr d
x
        isCons :: Bool
isCons = String
constructorName String -> [String] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [String
"(:)", String
":|"]
        isNull :: Bool
isNull = String
constructorName String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
"[]"