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

{-# 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 :: forall a ann. Data a => a -> Doc ann
ppData = forall ann. MetaDoc ann -> Doc ann
mdPayload forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a ann. Data a => a -> MetaDoc ann
gpretty

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

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

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

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

        isTuple :: Bool
        isTuple :: Bool
isTuple = forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (forall a. Eq a => a -> a -> Bool
== Char
',') (forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. (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 a c.
Data a =>
(forall b. Data b => b -> c) -> a -> Maybe (c, c)
isPair forall b. Data b => b -> c
f a
x
  | String
constructorName forall a. Eq a => a -> a -> Bool
== String
"(,)" = forall a. a -> Maybe a
Just (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, 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                = forall a. Maybe a
Nothing
  where
    constructorName :: String
    constructorName :: String
constructorName = Constr -> String
showConstr forall a b. (a -> b) -> a -> b
$ 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 a c.
Data a =>
(forall b. Data b => b -> c) -> a -> Maybe [c]
listElements forall b. Data b => b -> c
f = forall d. Data d => d -> Maybe [c]
go
  where
    go :: Data d => d -> Maybe [c]
    go :: forall d. Data d => d -> Maybe [c]
go d
x
      | Bool
isNull    = forall a. a -> Maybe a
Just []
      | Bool
isCons    = (:) (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) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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 = forall a. Maybe a
Nothing
      where
        constructorName :: String
        constructorName :: String
constructorName = Constr -> String
showConstr forall a b. (a -> b) -> a -> b
$ forall a. Data a => a -> Constr
toConstr d
x
        isCons :: Bool
isCons = String
constructorName forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [String
"(:)", String
":|"]
        isNull :: Bool
isNull = String
constructorName forall a. Eq a => a -> a -> Bool
== String
"[]"