{-# LANGUAGE AllowAmbiguousTypes #-} -- for reifying

module TypeLevelShow.Doc
  (
  -- * Term level
    Doc(..)
  , renderDoc

  -- * Type level
  , PDoc
  , RenderPDoc
  , ErrorPDoc

  -- * Singleton
  , SDoc(..)
  , SingDoc(singDoc)
  , demoteDoc
  , reifyDoc
  , errorPDoc
  ) where

import GHC.TypeLits qualified as TE -- TE = TypeError
import GHC.TypeLits hiding ( ErrorMessage(..) )
import Singleraeh.Demote
import Singleraeh.SingI

-- | Simple pretty document ADT.
--
-- Designed to work on both type level (as a limited 'TE.ErrorMessage') and term
-- level (as a boring ADT).
--
-- Note that 'TE.ShowType' is magical (see
-- @compiler/GHC/Core/Type.hs#L1309@), so we need to remove it for term level.
--
-- singletons-base defines a version of this, but retains the 'TE.ShowType'
-- constructor and is in the singletons ecosystem.
data Doc s
  = Text s
  -- ^ plain ol' text
  | Doc s :<>: Doc s
  -- ^ append docs next to each other
  | Doc s :$$: Doc s
  -- ^ stack docs on top of each other (newline)
    deriving stock Int -> Doc s -> ShowS
[Doc s] -> ShowS
Doc s -> String
(Int -> Doc s -> ShowS)
-> (Doc s -> String) -> ([Doc s] -> ShowS) -> Show (Doc s)
forall s. Show s => Int -> Doc s -> ShowS
forall s. Show s => [Doc s] -> ShowS
forall s. Show s => Doc s -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall s. Show s => Int -> Doc s -> ShowS
showsPrec :: Int -> Doc s -> ShowS
$cshow :: forall s. Show s => Doc s -> String
show :: Doc s -> String
$cshowList :: forall s. Show s => [Doc s] -> ShowS
showList :: [Doc s] -> ShowS
Show

-- | Render a 'Doc' as a 'String', formatted how 'ErrorMessage's get displayed.
renderDoc :: Doc String -> String
renderDoc :: Doc String -> String
renderDoc Doc String
doc = String
"    • " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Doc String -> String
go Doc String
doc
  where
    go :: Doc String -> String
go = \case
      Text String
s   -> String
s
      Doc String
l :<>: Doc String
r -> Doc String -> String
go Doc String
l String -> ShowS
forall a. Semigroup a => a -> a -> a
<>               Doc String -> String
go Doc String
r
      Doc String
l :$$: Doc String
r -> Doc String -> String
go Doc String
l String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
"\n      " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Doc String -> String
go Doc String
r

-- | Promoted 'Doc'.
type PDoc = Doc Symbol

-- | Render a 'PDoc' as an 'ErrorMessage', for type-level error messages.
--
-- 'PDoc' is a subset of 'ErrorMessage', so this is very boring.
type RenderPDoc :: PDoc -> TE.ErrorMessage
type family RenderPDoc doc where
    RenderPDoc (Text s)   = TE.Text s
    RenderPDoc (l :<>: r) = RenderPDoc l TE.:<>: RenderPDoc r
    RenderPDoc (l :$$: r) = RenderPDoc l TE.:$$: RenderPDoc r

-- | Render a 'PDoc' as an 'ErrorMessage', and wrap it in a 'TypeError'.
--
-- Note that this must be a type family, or the 'PDoc' won't actually get
-- rendered.
type ErrorPDoc :: PDoc -> a
type family ErrorPDoc doc where ErrorPDoc doc = TypeError (RenderPDoc doc)

-- | Singleton 'Doc'.
data SDoc (doc :: PDoc) where
    SText   :: SSymbol s -> SDoc (Text s)
    (:$<>:) :: SDoc docl -> SDoc docr -> SDoc (docl :<>: docr)
    (:$$$:) :: SDoc docl -> SDoc docr -> SDoc (docl :$$: docr)

-- | Demote an 'SDoc' singleton to the corresponding base 'Doc'.
demoteDoc :: SDoc doc -> Doc String
demoteDoc :: forall (doc :: PDoc). SDoc doc -> Doc String
demoteDoc = \case
  SText SSymbol s
s   -> String -> Doc String
forall s. s -> Doc s
Text (String -> Doc String) -> String -> Doc String
forall a b. (a -> b) -> a -> b
$ SSymbol s -> String
forall (s :: Symbol). SSymbol s -> String
fromSSymbol SSymbol s
s
  SDoc docl
l :$<>: SDoc docr
r -> SDoc docl -> Doc String
forall (doc :: PDoc). SDoc doc -> Doc String
demoteDoc SDoc docl
l Doc String -> Doc String -> Doc String
forall s. Doc s -> Doc s -> Doc s
:<>: SDoc docr -> Doc String
forall (doc :: PDoc). SDoc doc -> Doc String
demoteDoc SDoc docr
r
  SDoc docl
l :$$$: SDoc docr
r -> SDoc docl -> Doc String
forall (doc :: PDoc). SDoc doc -> Doc String
demoteDoc SDoc docl
l Doc String -> Doc String -> Doc String
forall s. Doc s -> Doc s -> Doc s
:$$: SDoc docr -> Doc String
forall (doc :: PDoc). SDoc doc -> Doc String
demoteDoc SDoc docr
r

instance Demotable SDoc where
    type Demote SDoc = Doc String
    demote :: forall (k1 :: PDoc). SDoc k1 -> Demote SDoc
demote = SDoc k1 -> Demote SDoc
SDoc k1 -> Doc String
forall (doc :: PDoc). SDoc doc -> Doc String
demoteDoc

-- | Produce the singleton for the given promoted 'Doc'.
class SingDoc (doc :: PDoc) where
    singDoc :: SDoc doc

instance KnownSymbol s => SingDoc (Text s) where
    singDoc :: SDoc ('Text s)
singDoc = SSymbol s -> SDoc ('Text s)
forall (docl :: Symbol). SSymbol docl -> SDoc ('Text docl)
SText (forall (s :: Symbol). KnownSymbol s => SSymbol s
SSymbol @s)

instance (SingDoc l, SingDoc r) => SingDoc (l :<>: r) where
    singDoc :: SDoc (l ':<>: r)
singDoc = forall (doc :: PDoc). SingDoc doc => SDoc doc
singDoc @l SDoc l -> SDoc r -> SDoc (l ':<>: r)
forall (docl :: PDoc) (docr :: PDoc).
SDoc docl -> SDoc docr -> SDoc (docl ':<>: docr)
:$<>: forall (doc :: PDoc). SingDoc doc => SDoc doc
singDoc @r

instance (SingDoc l, SingDoc r) => SingDoc (l :$$: r) where
    singDoc :: SDoc (l ':$$: r)
singDoc = forall (doc :: PDoc). SingDoc doc => SDoc doc
singDoc @l SDoc l -> SDoc r -> SDoc (l ':$$: r)
forall (docl :: PDoc) (docr :: PDoc).
SDoc docl -> SDoc docr -> SDoc (docl ':$$: docr)
:$$$: forall (doc :: PDoc). SingDoc doc => SDoc doc
singDoc @r

instance SingDoc doc => SingI doc where
    type Sing = SDoc
    sing' :: Sing doc
sing' = Sing doc
SDoc doc
forall (doc :: PDoc). SingDoc doc => SDoc doc
singDoc

-- | Reify a promoted 'Doc' to the corresponding term-level one.
reifyDoc :: forall (doc :: PDoc). SingDoc doc => Doc String
reifyDoc :: forall (doc :: PDoc). SingDoc doc => Doc String
reifyDoc = SDoc doc -> Doc String
forall (doc :: PDoc). SDoc doc -> Doc String
demoteDoc (forall (doc :: PDoc). SingDoc doc => SDoc doc
singDoc @doc)

-- | Render a 'PDoc' as a 'String', and call 'error' on it.
--
-- This enables using the same code for type- and term- "runtime" errors.
-- This can't be typechecked, naturally, but it's still nice.
errorPDoc :: forall (doc :: PDoc) a. SingDoc doc => a
-- add extra newline because runtime errors print like @*** Exception: _@
errorPDoc :: forall (doc :: PDoc) a. SingDoc doc => a
errorPDoc = String -> a
forall a. HasCallStack => String -> a
error (String -> a) -> String -> a
forall a b. (a -> b) -> a -> b
$ (String
"\n" <>) ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$ Doc String -> String
renderDoc (Doc String -> String) -> Doc String -> String
forall a b. (a -> b) -> a -> b
$ forall (doc :: PDoc). SingDoc doc => Doc String
reifyDoc @doc