{-# LANGUAGE ConstraintKinds            #-}
{-# LANGUAGE DerivingStrategies         #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE TypeOperators              #-}
{-# 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 = (RenderContext -> f RenderContext)
-> RenderContext -> f 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) = config -> Sole config
forall a. a -> Sole a
Sole (config -> Sole config) -> f config -> f (Sole config)
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 -> b) -> InContextM config a -> InContextM config b)
-> (forall a b. a -> InContextM config b -> InContextM config a)
-> Functor (InContextM config)
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
$cfmap :: forall config a b.
(a -> b) -> InContextM config a -> InContextM config b
fmap :: forall a b. (a -> b) -> InContextM config a -> InContextM config b
$c<$ :: forall config a b. a -> InContextM config b -> InContextM config a
<$ :: forall a b. a -> InContextM config b -> InContextM config a
Functor, Functor (InContextM config)
Functor (InContextM config) =>
(forall a. a -> InContextM config a)
-> (forall 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 a b.
    InContextM config a -> InContextM config b -> InContextM config b)
-> (forall a b.
    InContextM config a -> InContextM config b -> InContextM config a)
-> Applicative (InContextM config)
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
$cpure :: forall config a. a -> InContextM config a
pure :: forall a. a -> InContextM config a
$c<*> :: forall config a b.
InContextM config (a -> b)
-> InContextM config a -> InContextM config b
<*> :: forall a b.
InContextM config (a -> b)
-> InContextM config a -> InContextM config b
$cliftA2 :: forall config a b c.
(a -> b -> c)
-> InContextM config a
-> InContextM config b
-> InContextM config c
liftA2 :: forall a b c.
(a -> b -> c)
-> InContextM config a
-> InContextM config b
-> InContextM config c
$c*> :: forall config a b.
InContextM config a -> InContextM config b -> InContextM config b
*> :: 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 a
<* :: forall a b.
InContextM config a -> InContextM config b -> InContextM config a
Applicative, Applicative (InContextM config)
Applicative (InContextM config) =>
(forall a b.
 InContextM config a
 -> (a -> InContextM config b) -> InContextM config b)
-> (forall a b.
    InContextM config a -> InContextM config b -> InContextM config b)
-> (forall a. a -> InContextM config a)
-> Monad (InContextM config)
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
$c>>= :: forall config a b.
InContextM config a
-> (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 -> InContextM config b -> InContextM config b
>> :: forall a b.
InContextM config a -> InContextM config b -> InContextM config b
$creturn :: forall config a. a -> InContextM config a
return :: forall a. a -> InContextM config a
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) = Reader (Sole config) a -> Sole config -> a
forall r a. Reader r a -> r -> a
runReader Reader (Sole config) a
a (Sole config -> a) -> Sole config -> a
forall a b. (a -> b) -> a -> b
$ config -> Sole config
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 = config -> InContextM config (Doc ann) -> Doc ann
forall config a. config -> InContextM config a -> a
runInContextM config
config (InContextM config (Doc ann) -> Doc ann)
-> (a -> InContextM config (Doc ann)) -> a -> Doc ann
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 = Doc ann -> InContextM config doc
Doc ann -> InContextM config (Doc ann)
forall config env (m :: * -> *) ann.
MonadPrettyContext config env m =>
Doc ann -> m (Doc ann)
unitDocM (Doc ann -> InContextM config doc)
-> (String -> Doc ann) -> String -> InContextM config doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Doc ann
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 =
    Getting RenderContext env RenderContext -> m RenderContext
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view ((config -> Const RenderContext config)
-> env -> Const RenderContext env
forall env config. HasPrettyConfig env config => Lens' env config
Lens' env config
prettyConfig ((config -> Const RenderContext config)
 -> env -> Const RenderContext env)
-> ((RenderContext -> Const RenderContext RenderContext)
    -> config -> Const RenderContext config)
-> Getting RenderContext env RenderContext
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (RenderContext -> Const RenderContext RenderContext)
-> config -> Const RenderContext config
forall config.
HasRenderContext config =>
Lens' config RenderContext
Lens' config RenderContext
renderContext) m RenderContext -> (RenderContext -> Doc ann) -> m (Doc ann)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \RenderContext
context ->
        (Doc ann -> Doc ann)
-> RenderContext -> Fixity -> Doc ann -> Doc ann
forall prec a.
Ord prec =>
(a -> a) -> RenderContextOver prec -> FixityOver prec -> a -> a
encloseIn Doc ann -> Doc ann
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 <- Getting config env config -> m config
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting config env config
forall env config. HasPrettyConfig env config => Lens' env config
Lens' env config
prettyConfig
    (forall a.
 PrettyBy config a =>
 Direction -> Fixity -> a -> Doc ann)
-> m r
cont ((forall a.
  PrettyBy config a =>
  Direction -> Fixity -> a -> Doc ann)
 -> m r)
-> (forall a.
    PrettyBy config a =>
    Direction -> Fixity -> a -> Doc ann)
-> m r
forall a b. (a -> b) -> a -> b
$ \Direction
dir Fixity
fixity -> config -> a -> Doc ann
forall ann. config -> a -> Doc ann
forall config a ann. PrettyBy config a => config -> a -> Doc ann
prettyBy (config -> a -> Doc ann) -> config -> a -> Doc ann
forall a b. (a -> b) -> a -> b
$ config
config config -> (config -> config) -> config
forall a b. a -> (a -> b) -> b
& (RenderContext -> Identity RenderContext)
-> config -> Identity config
forall config.
HasRenderContext config =>
Lens' config RenderContext
Lens' config RenderContext
renderContext ((RenderContext -> Identity RenderContext)
 -> config -> Identity config)
-> RenderContext -> config -> config
forall s t a b. ASetter s t a b -> b -> s -> t
.~ Direction -> Fixity -> RenderContext
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 a.
  PrettyBy config a =>
  Direction -> Fixity -> a -> Doc ann)
 -> m r)
-> m r
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)
 -> m r)
-> ((forall a.
     PrettyBy config a =>
     Direction -> Fixity -> a -> Doc ann)
    -> m r)
-> m r
forall a b. (a -> b) -> a -> b
$ \forall a. PrettyBy config a => Direction -> Fixity -> a -> Doc ann
prettyIn -> AnyToDoc config ann -> m r
cont (AnyToDoc config ann -> m r) -> AnyToDoc config ann -> m r
forall a b. (a -> b) -> a -> b
$ Direction -> Fixity -> a -> Doc ann
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 = Fixity -> Doc ann -> m (Doc ann)
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 a.
  PrettyBy config a =>
  Direction -> Fixity -> a -> Doc ann)
 -> m (Doc ann))
-> m (Doc ann)
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 (Doc ann))
 -> m (Doc ann))
-> ((forall a.
     PrettyBy config a =>
     Direction -> Fixity -> a -> Doc ann)
    -> m (Doc ann))
-> m (Doc ann)
forall a b. (a -> b) -> a -> b
$ \forall a. PrettyBy config a => Direction -> Fixity -> a -> Doc ann
prettyIn -> Fixity -> Doc ann -> m (Doc ann)
forall config env (m :: * -> *) ann.
MonadPrettyContext config env m =>
Fixity -> Doc ann -> m (Doc ann)
encloseM Fixity
fixity (Doc ann -> m (Doc ann)) -> Doc ann -> m (Doc ann)
forall a b. (a -> b) -> a -> b
$ (forall a.
 PrettyBy config a =>
 Direction -> Fixity -> a -> Doc ann)
-> Doc ann
k Direction -> Fixity -> a -> Doc ann
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 = Fixity
-> ((forall {a}.
     PrettyBy config a =>
     Direction -> Fixity -> a -> Doc ann)
    -> Doc ann)
-> m (Doc ann)
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)
 -> m (Doc ann))
-> ((forall {a}.
     PrettyBy config a =>
     Direction -> Fixity -> a -> Doc ann)
    -> Doc ann)
-> m (Doc ann)
forall a b. (a -> b) -> a -> b
$ \forall {a}.
PrettyBy config a =>
Direction -> Fixity -> a -> Doc ann
prettyIn -> AnyToDoc config ann -> Doc ann
k (AnyToDoc config ann -> Doc ann) -> AnyToDoc config ann -> Doc ann
forall a b. (a -> b) -> a -> b
$ Direction -> Fixity -> a -> Doc ann
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 =
    Fixity
-> ((forall {a}.
     PrettyBy config a =>
     Direction -> Fixity -> a -> Doc ann)
    -> Doc ann)
-> m (Doc ann)
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)
 -> m (Doc ann))
-> ((forall {a}.
     PrettyBy config a =>
     Direction -> Fixity -> a -> Doc ann)
    -> Doc ann)
-> m (Doc ann)
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 (Direction -> Fixity -> a -> Doc ann
forall {a}.
PrettyBy config a =>
Direction -> Fixity -> a -> Doc ann
prettyIn Direction
ToTheLeft Fixity
fixity) (Direction -> Fixity -> a -> Doc ann
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 =
    Fixity
-> ((forall {a}. PrettyBy config a => a -> Doc ann)
    -> (forall {a}. PrettyBy config a => a -> Doc ann) -> Doc ann)
-> m (Doc ann)
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}. PrettyBy config a => a -> Doc ann)
  -> (forall {a}. PrettyBy config a => a -> Doc ann) -> Doc ann)
 -> m (Doc ann))
-> ((forall {a}. PrettyBy config a => a -> Doc ann)
    -> (forall {a}. PrettyBy config a => a -> Doc ann) -> Doc ann)
-> m (Doc ann)
forall a b. (a -> b) -> a -> b
$ \forall {a}. PrettyBy config a => a -> Doc ann
prettyL forall {a}. PrettyBy config a => a -> Doc ann
prettyR -> a -> Doc ann
forall {a}. PrettyBy config a => a -> Doc ann
prettyL a
fun Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> b -> Doc ann
forall {a}. PrettyBy config a => a -> Doc ann
prettyR b
arg