{-# LANGUAGE ConstraintKinds            #-}
{-# LANGUAGE DerivingStrategies         #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE FunctionalDependencies     #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE UndecidableInstances       #-}

-- | Configurable precedence-aware pretty-printing.
--
-- Look into @test/Expr.hs@ for an extended example.

module Text.PrettyBy.Fixity
    ( module Export
    , module Text.PrettyBy.Fixity
    ) where

import Text.Fixity as Export
import Text.Pretty
import Text.PrettyBy.Internal
import Text.PrettyBy.Internal.Utils
import Text.PrettyBy.Monad as Export

import Control.Monad.Reader
import Data.String
import Lens.Micro

-- | A constraint for \"'RenderContext' is a part of @config@\".
class HasRenderContext config where
    renderContext :: Lens' config RenderContext

instance HasRenderContext RenderContext where
    renderContext :: Lens' RenderContext RenderContext
renderContext = forall a. a -> a
id

-- | A constraint for \"@m@ is a 'Monad' supporting configurable precedence-aware pretty-printing\".
type MonadPrettyContext config env m = (MonadPretty config env m, HasRenderContext config)

-- | A @newtype@ wrapper around @a@ introduced for its 'HasPrettyConfig' instance.
newtype Sole a = Sole
    { forall a. Sole a -> a
unSole :: a
    }

-- | It's not possible to have @HasPrettyConfig config config@, because that would mean that every
-- environment is a pretty-printing config on its own, which doesn't make sense. We could have an
-- OVERLAPPABLE instance, but I'd rather not.
instance HasPrettyConfig (Sole config) config where
    prettyConfig :: Lens' (Sole config) config
prettyConfig config -> f config
f (Sole config
x) = forall a. a -> Sole a
Sole forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> config -> f config
f config
x

-- | A monad for precedence-aware pretty-printing.
newtype InContextM config a = InContextM
    { forall config a. InContextM config a -> Reader (Sole config) a
unInContextM :: Reader (Sole config) a
    } deriving newtype (forall a b. a -> InContextM config b -> InContextM config a
forall a b. (a -> b) -> InContextM config a -> InContextM config b
forall config a b. a -> InContextM config b -> InContextM config a
forall config a b.
(a -> b) -> InContextM config a -> InContextM config b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> InContextM config b -> InContextM config a
$c<$ :: forall config a b. a -> InContextM config b -> InContextM config a
fmap :: forall a b. (a -> b) -> InContextM config a -> InContextM config b
$cfmap :: forall config a b.
(a -> b) -> InContextM config a -> InContextM config b
Functor, forall config. Functor (InContextM config)
forall a. a -> InContextM config a
forall config a. a -> InContextM config a
forall a b.
InContextM config a -> InContextM config b -> InContextM config a
forall a b.
InContextM config a -> InContextM config b -> InContextM config b
forall a b.
InContextM config (a -> b)
-> InContextM config a -> InContextM config b
forall config a b.
InContextM config a -> InContextM config b -> InContextM config a
forall config a b.
InContextM config a -> InContextM config b -> InContextM config b
forall config a b.
InContextM config (a -> b)
-> InContextM config a -> InContextM config b
forall a b c.
(a -> b -> c)
-> InContextM config a
-> InContextM config b
-> InContextM config c
forall config a b c.
(a -> b -> c)
-> InContextM config a
-> InContextM config b
-> InContextM config c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: forall a b.
InContextM config a -> InContextM config b -> InContextM config a
$c<* :: forall config a b.
InContextM config a -> InContextM config b -> InContextM config a
*> :: forall a b.
InContextM config a -> InContextM config b -> InContextM config b
$c*> :: forall config a b.
InContextM config a -> InContextM config b -> InContextM config b
liftA2 :: forall a b c.
(a -> b -> c)
-> InContextM config a
-> InContextM config b
-> InContextM config c
$cliftA2 :: forall config a b c.
(a -> b -> c)
-> InContextM config a
-> InContextM config b
-> InContextM config c
<*> :: forall a b.
InContextM config (a -> b)
-> InContextM config a -> InContextM config b
$c<*> :: forall config a b.
InContextM config (a -> b)
-> InContextM config a -> InContextM config b
pure :: forall a. a -> InContextM config a
$cpure :: forall config a. a -> InContextM config a
Applicative, forall config. Applicative (InContextM config)
forall a. a -> InContextM config a
forall config a. a -> InContextM config a
forall a b.
InContextM config a -> InContextM config b -> InContextM config b
forall a b.
InContextM config a
-> (a -> InContextM config b) -> InContextM config b
forall config a b.
InContextM config a -> InContextM config b -> InContextM config b
forall config a b.
InContextM config a
-> (a -> InContextM config b) -> InContextM config b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: forall a. a -> InContextM config a
$creturn :: forall config a. a -> InContextM config a
>> :: forall a b.
InContextM config a -> InContextM config b -> InContextM config b
$c>> :: forall config a b.
InContextM config a -> InContextM config b -> InContextM config b
>>= :: forall a b.
InContextM config a
-> (a -> InContextM config b) -> InContextM config b
$c>>= :: forall config a b.
InContextM config a
-> (a -> InContextM config b) -> InContextM config b
Monad, MonadReader (Sole config))

-- | Run 'InContextM' by supplying a @config@.
runInContextM :: config -> InContextM config a -> a
runInContextM :: forall config a. config -> InContextM config a -> a
runInContextM config
config (InContextM Reader (Sole config) a
a) = forall r a. Reader r a -> r -> a
runReader Reader (Sole config) a
a forall a b. (a -> b) -> a -> b
$ forall a. a -> Sole a
Sole config
config

-- | Takes a monadic pretty-printer and turns it into one that receives a @config@ explicitly.
-- Useful for defining instances of 'PrettyBy' monadically when writing precedence-aware
-- pretty-printing code (and since all functions below are monadic, it's currenty the only option).
inContextM :: (a -> InContextM config (Doc ann)) -> config -> a -> Doc ann
inContextM :: forall a config ann.
(a -> InContextM config (Doc ann)) -> config -> a -> Doc ann
inContextM a -> InContextM config (Doc ann)
pM config
config = forall config a. config -> InContextM config a -> a
runInContextM config
config forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> InContextM config (Doc ann)
pM

-- | A string written in the 'InContextM' monad gets enclosed with 'unitDocM' automatically.
instance (HasRenderContext config, doc ~ Doc ann) => IsString (InContextM config doc) where
    fromString :: String -> InContextM config doc
fromString = forall config env (m :: * -> *) ann.
MonadPrettyContext config env m =>
Doc ann -> m (Doc ann)
unitDocM forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. IsString a => String -> a
fromString

-- TODO: when writing a precedence-aware pretty-printer we basically always want to specify a
-- fixity in each clause. Would be nice to enforce that in types.
-- | Enclose a 'Doc' in parentheses if required or leave it as is. The need for enclosing is
-- determined from an outer 'RenderContext' (stored in the environment of the monad) and the inner
-- fixity provided as an argument.
encloseM :: MonadPrettyContext config env m => Fixity -> Doc ann -> m (Doc ann)
encloseM :: forall config env (m :: * -> *) ann.
MonadPrettyContext config env m =>
Fixity -> Doc ann -> m (Doc ann)
encloseM Fixity
fixity Doc ann
doc =
    forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view (forall env config. HasPrettyConfig env config => Lens' env config
prettyConfig forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall config.
HasRenderContext config =>
Lens' config RenderContext
renderContext) forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \RenderContext
context ->
        forall prec a.
Ord prec =>
(a -> a) -> RenderContextOver prec -> FixityOver prec -> a -> a
encloseIn forall ann. Doc ann -> Doc ann
parens RenderContext
context Fixity
fixity Doc ann
doc

-- | The type of a general @config@-based pretty-printer.
type AnyToDoc config ann = forall a. PrettyBy config a => a -> Doc ann

-- | Instantiate a supplied continuation with a precedence-aware pretty-printer.
withPrettyIn
    :: MonadPrettyContext config env m
    => ((forall a. PrettyBy config a => Direction -> Fixity -> a -> Doc ann) -> m r) -> m r
withPrettyIn :: forall config env (m :: * -> *) ann r.
MonadPrettyContext config env m =>
((forall a.
  PrettyBy config a =>
  Direction -> Fixity -> a -> Doc ann)
 -> m r)
-> m r
withPrettyIn (forall a.
 PrettyBy config a =>
 Direction -> Fixity -> a -> Doc ann)
-> m r
cont = do
    config
config <- forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view forall env config. HasPrettyConfig env config => Lens' env config
prettyConfig
    (forall a.
 PrettyBy config a =>
 Direction -> Fixity -> a -> Doc ann)
-> m r
cont forall a b. (a -> b) -> a -> b
$ \Direction
dir Fixity
fixity -> forall config a ann. PrettyBy config a => config -> a -> Doc ann
prettyBy forall a b. (a -> b) -> a -> b
$ config
config forall a b. a -> (a -> b) -> b
& forall config.
HasRenderContext config =>
Lens' config RenderContext
renderContext forall s t a b. ASetter s t a b -> b -> s -> t
.~ forall prec. Direction -> FixityOver prec -> RenderContextOver prec
RenderContext Direction
dir Fixity
fixity

-- | Instantiate a supplied continuation with a pretty-printer specialized to supplied
-- 'Fixity' and 'Direction'.
withPrettyAt
    :: MonadPrettyContext config env m
    => Direction -> Fixity -> (AnyToDoc config ann -> m r) -> m r
withPrettyAt :: forall config env (m :: * -> *) ann r.
MonadPrettyContext config env m =>
Direction -> Fixity -> (AnyToDoc config ann -> m r) -> m r
withPrettyAt Direction
dir Fixity
fixity AnyToDoc config ann -> m r
cont = forall config env (m :: * -> *) ann r.
MonadPrettyContext config env m =>
((forall a.
  PrettyBy config a =>
  Direction -> Fixity -> a -> Doc ann)
 -> m r)
-> m r
withPrettyIn forall a b. (a -> b) -> a -> b
$ \forall a. PrettyBy config a => Direction -> Fixity -> a -> Doc ann
prettyIn -> AnyToDoc config ann -> m r
cont forall a b. (a -> b) -> a -> b
$ forall a. PrettyBy config a => Direction -> Fixity -> a -> Doc ann
prettyIn Direction
dir Fixity
fixity

-- | Call 'encloseM' on 'unitFixity'.
unitDocM :: MonadPrettyContext config env m => Doc ann -> m (Doc ann)
unitDocM :: forall config env (m :: * -> *) ann.
MonadPrettyContext config env m =>
Doc ann -> m (Doc ann)
unitDocM = forall config env (m :: * -> *) ann.
MonadPrettyContext config env m =>
Fixity -> Doc ann -> m (Doc ann)
encloseM Fixity
unitFixity

-- | Instantiate a supplied continuation with a pretty-printer and apply 'encloseM',
-- specialized to supplied 'Fixity', to the result.
compoundDocM
    :: MonadPrettyContext config env m
    => Fixity
    -> ((forall a. PrettyBy config a => Direction -> Fixity -> a -> Doc ann) -> Doc ann)
    -> m (Doc ann)
compoundDocM :: forall config env (m :: * -> *) ann.
MonadPrettyContext config env m =>
Fixity
-> ((forall a.
     PrettyBy config a =>
     Direction -> Fixity -> a -> Doc ann)
    -> Doc ann)
-> m (Doc ann)
compoundDocM Fixity
fixity (forall a.
 PrettyBy config a =>
 Direction -> Fixity -> a -> Doc ann)
-> Doc ann
k = forall config env (m :: * -> *) ann r.
MonadPrettyContext config env m =>
((forall a.
  PrettyBy config a =>
  Direction -> Fixity -> a -> Doc ann)
 -> m r)
-> m r
withPrettyIn forall a b. (a -> b) -> a -> b
$ \forall a. PrettyBy config a => Direction -> Fixity -> a -> Doc ann
prettyIn -> forall config env (m :: * -> *) ann.
MonadPrettyContext config env m =>
Fixity -> Doc ann -> m (Doc ann)
encloseM Fixity
fixity forall a b. (a -> b) -> a -> b
$ (forall a.
 PrettyBy config a =>
 Direction -> Fixity -> a -> Doc ann)
-> Doc ann
k forall a. PrettyBy config a => Direction -> Fixity -> a -> Doc ann
prettyIn

-- | Instantiate a supplied continuation with a pretty-printer specialized to supplied
-- 'Fixity' and 'Direction' and apply 'encloseM' specialized to the provided fixity to the result.
-- This can be useful for pretty-printing a sequence of values (possibly consisting of a single
-- value).
sequenceDocM
    :: MonadPrettyContext config env m
    => Direction -> Fixity -> (AnyToDoc config ann -> Doc ann) -> m (Doc ann)
sequenceDocM :: forall config env (m :: * -> *) ann.
MonadPrettyContext config env m =>
Direction
-> Fixity -> (AnyToDoc config ann -> Doc ann) -> m (Doc ann)
sequenceDocM Direction
dir Fixity
fixity AnyToDoc config ann -> Doc ann
k = forall config env (m :: * -> *) ann.
MonadPrettyContext config env m =>
Fixity
-> ((forall a.
     PrettyBy config a =>
     Direction -> Fixity -> a -> Doc ann)
    -> Doc ann)
-> m (Doc ann)
compoundDocM Fixity
fixity forall a b. (a -> b) -> a -> b
$ \forall a. PrettyBy config a => Direction -> Fixity -> a -> Doc ann
prettyIn -> AnyToDoc config ann -> Doc ann
k forall a b. (a -> b) -> a -> b
$ forall a. PrettyBy config a => Direction -> Fixity -> a -> Doc ann
prettyIn Direction
dir Fixity
fixity

-- | Instantiate a supplied continuation with two pretty-printers (one is going in the 'ToTheLeft'
-- direction, the other is in the 'ToTheRight' direction) specialized to supplied 'Fixity'
-- and apply 'encloseM', specialized to the same fixity, to the result.
-- The idea is that to the outside an infix operator has the same inner fixity as
-- it has the outer fixity to inner subexpressions.
infixDocM
    :: MonadPrettyContext config env m
    => Fixity
    -> (AnyToDoc config ann -> AnyToDoc config ann -> Doc ann)
    -> m (Doc ann)
infixDocM :: forall config env (m :: * -> *) ann.
MonadPrettyContext config env m =>
Fixity
-> (AnyToDoc config ann -> AnyToDoc config ann -> Doc ann)
-> m (Doc ann)
infixDocM Fixity
fixity AnyToDoc config ann -> AnyToDoc config ann -> Doc ann
k =
    forall config env (m :: * -> *) ann.
MonadPrettyContext config env m =>
Fixity
-> ((forall a.
     PrettyBy config a =>
     Direction -> Fixity -> a -> Doc ann)
    -> Doc ann)
-> m (Doc ann)
compoundDocM Fixity
fixity forall a b. (a -> b) -> a -> b
$ \forall a. PrettyBy config a => Direction -> Fixity -> a -> Doc ann
prettyIn ->
        AnyToDoc config ann -> AnyToDoc config ann -> Doc ann
k (forall a. PrettyBy config a => Direction -> Fixity -> a -> Doc ann
prettyIn Direction
ToTheLeft Fixity
fixity) (forall a. PrettyBy config a => Direction -> Fixity -> a -> Doc ann
prettyIn Direction
ToTheRight Fixity
fixity)

-- | Pretty-print two things with a space between them. The fixity of the context in which the
-- arguments get pretty-printed is set to 'juxtFixity'.
juxtPrettyM
    :: (MonadPrettyContext config env m, PrettyBy config a, PrettyBy config b)
    => a -> b -> m (Doc ann)
juxtPrettyM :: forall config env (m :: * -> *) a b ann.
(MonadPrettyContext config env m, PrettyBy config a,
 PrettyBy config b) =>
a -> b -> m (Doc ann)
juxtPrettyM a
fun b
arg =
    forall config env (m :: * -> *) ann.
MonadPrettyContext config env m =>
Fixity
-> (AnyToDoc config ann -> AnyToDoc config ann -> Doc ann)
-> m (Doc ann)
infixDocM Fixity
juxtFixity forall a b. (a -> b) -> a -> b
$ \AnyToDoc config ann
prettyL AnyToDoc config ann
prettyR -> AnyToDoc config ann
prettyL a
fun forall ann. Doc ann -> Doc ann -> Doc ann
<+> AnyToDoc config ann
prettyR b
arg