{-# LANGUAGE DerivingVia               #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE OverloadedStrings         #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Disco.Pretty
-- Copyright   :  disco team and contributors
-- Maintainer  :  byorgey@gmail.com
--
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Various pretty-printing facilities for disco.
--
-----------------------------------------------------------------------------

-- TODO: the calls to 'error' should be replaced with logging/error capabilities.

module Disco.Pretty
  ( module Disco.Pretty.DSL
  , module Disco.Pretty
  , module Disco.Pretty.Prec
  , Doc
  )
  where

import           Prelude                          hiding ((<>))

import           Data.Bifunctor
import           Data.Char                        (isAlpha)
import           Data.Map                         (Map)
import qualified Data.Map                         as M
import           Data.Ratio
import           Data.Set                         (Set)
import qualified Data.Set                         as S

import           Disco.Effects.LFresh
import           Polysemy

import           Polysemy.Reader

import           Text.PrettyPrint                 (Doc)
import           Unbound.Generics.LocallyNameless (Name)

import           Disco.Pretty.DSL
import           Disco.Pretty.Prec
import           Disco.Syntax.Operators

------------------------------------------------------------
-- Utilities for handling precedence and associativity

-- | Convenience function combining 'setPA' and 'mparens', since we
--   often want to simultaneously indicate what the precedence and
--   associativity of a term is, and optionally surround it with
--   parentheses depending on the precedence and associativity of its
--   parent.
withPA :: Member (Reader PA) r => PA -> Sem r Doc -> Sem r Doc
withPA :: PA -> Sem r Doc -> Sem r Doc
withPA PA
pa = PA -> Sem r Doc -> Sem r Doc
forall (r :: EffectRow).
Member (Reader PA) r =>
PA -> Sem r Doc -> Sem r Doc
mparens PA
pa (Sem r Doc -> Sem r Doc)
-> (Sem r Doc -> Sem r Doc) -> Sem r Doc -> Sem r Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PA -> Sem r Doc -> Sem r Doc
forall (r :: EffectRow) a.
Member (Reader PA) r =>
PA -> Sem r a -> Sem r a
setPA PA
pa

-- | Locally set the precedence and associativity within a
--   subcomputation.
setPA :: Member (Reader PA) r => PA -> Sem r a -> Sem r a
setPA :: PA -> Sem r a -> Sem r a
setPA = (PA -> PA) -> Sem r a -> Sem r a
forall i (r :: EffectRow) a.
Member (Reader i) r =>
(i -> i) -> Sem r a -> Sem r a
local ((PA -> PA) -> Sem r a -> Sem r a)
-> (PA -> PA -> PA) -> PA -> Sem r a -> Sem r a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PA -> PA -> PA
forall a b. a -> b -> a
const

-- | Mark a subcomputation as pretty-printing a term on the left of an
--   operator (so parentheses can be inserted appropriately, depending
--   on the associativity).
lt :: Member (Reader PA) r => Sem r Doc -> Sem r Doc
lt :: Sem r Doc -> Sem r Doc
lt = (PA -> PA) -> Sem r Doc -> Sem r Doc
forall i (r :: EffectRow) a.
Member (Reader i) r =>
(i -> i) -> Sem r a -> Sem r a
local (\(PA Prec
p BFixity
_) -> Prec -> BFixity -> PA
PA Prec
p BFixity
InL)

-- | Mark a subcomputation as pretty-printing a term on the right of
--   an operator (so parentheses can be inserted appropriately,
--   depending on the associativity).
rt :: Member (Reader PA) r => Sem r Doc -> Sem r Doc
rt :: Sem r Doc -> Sem r Doc
rt = (PA -> PA) -> Sem r Doc -> Sem r Doc
forall i (r :: EffectRow) a.
Member (Reader i) r =>
(i -> i) -> Sem r a -> Sem r a
local (\(PA Prec
p BFixity
_) -> Prec -> BFixity -> PA
PA Prec
p BFixity
InR)

-- | Optionally surround a pretty-printed term with parentheses,
--   depending on its precedence and associativity (given as the 'PA'
--   argument) and that of its context (given by the ambient 'Reader
--   PA' effect).
mparens :: Member (Reader PA) r => PA -> Sem r Doc -> Sem r Doc
mparens :: PA -> Sem r Doc -> Sem r Doc
mparens PA
pa Sem r Doc
doc = do
  PA
parentPA <- Sem r PA
forall i (r :: EffectRow). Member (Reader i) r => Sem r i
ask
  (if PA
pa PA -> PA -> Bool
forall a. Ord a => a -> a -> Bool
< PA
parentPA then Sem r Doc -> Sem r Doc
forall (f :: * -> *). Functor f => f Doc -> f Doc
parens else Sem r Doc -> Sem r Doc
forall a. a -> a
id) Sem r Doc
doc

------------------------------------------------------------
-- Pretty type class

class Pretty t where
  pretty :: Members '[Reader PA, LFresh] r => t -> Sem r Doc

prettyStr :: Pretty t => t -> Sem r String
prettyStr :: t -> Sem r String
prettyStr = Sem (Reader PA : r) Doc -> Sem r String
forall (r :: EffectRow). Sem (Reader PA : r) Doc -> Sem r String
renderDoc (Sem (Reader PA : r) Doc -> Sem r String)
-> (t -> Sem (Reader PA : r) Doc) -> t -> Sem r String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sem (LFresh : Reader PA : r) Doc -> Sem (Reader PA : r) Doc
forall (r :: EffectRow) a. Sem (LFresh : r) a -> Sem r a
runLFresh (Sem (LFresh : Reader PA : r) Doc -> Sem (Reader PA : r) Doc)
-> (t -> Sem (LFresh : Reader PA : r) Doc)
-> t
-> Sem (Reader PA : r) Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t -> Sem (LFresh : Reader PA : r) Doc
forall t (r :: EffectRow).
(Pretty t, Members '[Reader PA, LFresh] r) =>
t -> Sem r Doc
pretty

pretty' :: Pretty t => t -> Sem r Doc
pretty' :: t -> Sem r Doc
pretty' = PA -> Sem (Reader PA : r) Doc -> Sem r Doc
forall i (r :: EffectRow) a. i -> Sem (Reader i : r) a -> Sem r a
runReader PA
initPA (Sem (Reader PA : r) Doc -> Sem r Doc)
-> (t -> Sem (Reader PA : r) Doc) -> t -> Sem r Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sem (LFresh : Reader PA : r) Doc -> Sem (Reader PA : r) Doc
forall (r :: EffectRow) a. Sem (LFresh : r) a -> Sem r a
runLFresh (Sem (LFresh : Reader PA : r) Doc -> Sem (Reader PA : r) Doc)
-> (t -> Sem (LFresh : Reader PA : r) Doc)
-> t
-> Sem (Reader PA : r) Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t -> Sem (LFresh : Reader PA : r) Doc
forall t (r :: EffectRow).
(Pretty t, Members '[Reader PA, LFresh] r) =>
t -> Sem r Doc
pretty

------------------------------------------------------------
-- Some standard instances

instance Pretty a => Pretty [a] where
  pretty :: [a] -> Sem r Doc
pretty = Sem r Doc -> Sem r Doc
forall (f :: * -> *). Functor f => f Doc -> f Doc
brackets (Sem r Doc -> Sem r Doc) -> ([a] -> Sem r Doc) -> [a] -> Sem r Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sem r Doc -> [Sem r Doc] -> Sem r Doc
forall (f :: * -> *). Monad f => f Doc -> [f Doc] -> f Doc
intercalate Sem r Doc
"," ([Sem r Doc] -> Sem r Doc)
-> ([a] -> [Sem r Doc]) -> [a] -> Sem r Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Sem r Doc) -> [a] -> [Sem r Doc]
forall a b. (a -> b) -> [a] -> [b]
map a -> Sem r Doc
forall t (r :: EffectRow).
(Pretty t, Members '[Reader PA, LFresh] r) =>
t -> Sem r Doc
pretty

instance (Pretty k, Pretty v) => Pretty (Map k v) where
  pretty :: Map k v -> Sem r Doc
pretty Map k v
m = do
    let es :: [Sem r Doc]
es = ((k, v) -> Sem r Doc) -> [(k, v)] -> [Sem r Doc]
forall a b. (a -> b) -> [a] -> [b]
map (\(k
k,v
v) -> k -> Sem r Doc
forall t (r :: EffectRow).
(Pretty t, Members '[Reader PA, LFresh] r) =>
t -> Sem r Doc
pretty k
k Sem r Doc -> Sem r Doc -> Sem r Doc
forall (f :: * -> *). Applicative f => f Doc -> f Doc -> f Doc
<+> Sem r Doc
"->" Sem r Doc -> Sem r Doc -> Sem r Doc
forall (f :: * -> *). Applicative f => f Doc -> f Doc -> f Doc
<+> v -> Sem r Doc
forall t (r :: EffectRow).
(Pretty t, Members '[Reader PA, LFresh] r) =>
t -> Sem r Doc
pretty v
v) (Map k v -> [(k, v)]
forall k a. Map k a -> [(k, a)]
M.assocs Map k v
m)
    [Sem r Doc]
ds <- PA -> Sem r [Sem r Doc] -> Sem r [Sem r Doc]
forall (r :: EffectRow) a.
Member (Reader PA) r =>
PA -> Sem r a -> Sem r a
setPA PA
initPA (Sem r [Sem r Doc] -> Sem r [Sem r Doc])
-> Sem r [Sem r Doc] -> Sem r [Sem r Doc]
forall a b. (a -> b) -> a -> b
$ Sem r Doc -> [Sem r Doc] -> Sem r [Sem r Doc]
forall (f :: * -> *).
Applicative f =>
f Doc -> [f Doc] -> f [f Doc]
punctuate Sem r Doc
"," [Sem r Doc]
es
    Sem r Doc -> Sem r Doc
forall (f :: * -> *). Functor f => f Doc -> f Doc
braces ([Sem r Doc] -> Sem r Doc
forall (f :: * -> *). Applicative f => [f Doc] -> f Doc
hsep [Sem r Doc]
ds)

instance Pretty a => Pretty (Set a) where
  pretty :: Set a -> Sem r Doc
pretty = Sem r Doc -> Sem r Doc
forall (f :: * -> *). Functor f => f Doc -> f Doc
braces (Sem r Doc -> Sem r Doc)
-> (Set a -> Sem r Doc) -> Set a -> Sem r Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sem r Doc -> [Sem r Doc] -> Sem r Doc
forall (f :: * -> *). Monad f => f Doc -> [f Doc] -> f Doc
intercalate Sem r Doc
"," ([Sem r Doc] -> Sem r Doc)
-> (Set a -> [Sem r Doc]) -> Set a -> Sem r Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Sem r Doc) -> [a] -> [Sem r Doc]
forall a b. (a -> b) -> [a] -> [b]
map a -> Sem r Doc
forall t (r :: EffectRow).
(Pretty t, Members '[Reader PA, LFresh] r) =>
t -> Sem r Doc
pretty ([a] -> [Sem r Doc]) -> (Set a -> [a]) -> Set a -> [Sem r Doc]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set a -> [a]
forall a. Set a -> [a]
S.toList

------------------------------------------------------------
-- Some Disco instances

instance Pretty (Name a) where
  pretty :: Name a -> Sem r Doc
pretty = String -> Sem r Doc
forall (m :: * -> *). Applicative m => String -> m Doc
text (String -> Sem r Doc) -> (Name a -> String) -> Name a -> Sem r Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name a -> String
forall a. Show a => a -> String
show

instance Pretty TyOp where
  pretty :: TyOp -> Sem r Doc
pretty = \case
    TyOp
Enumerate -> String -> Sem r Doc
forall (m :: * -> *). Applicative m => String -> m Doc
text String
"enumerate"
    TyOp
Count     -> String -> Sem r Doc
forall (m :: * -> *). Applicative m => String -> m Doc
text String
"count"

-- | Pretty-print a unary operator, by looking up its concrete syntax
--   in the 'uopMap'.
instance Pretty UOp where
  pretty :: UOp -> Sem r Doc
pretty UOp
op = case UOp -> Map UOp OpInfo -> Maybe OpInfo
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup UOp
op Map UOp OpInfo
uopMap of
    Just (OpInfo OpFixity
_ (String
syn:[String]
_) Prec
_) ->
      String -> Sem r Doc
forall (m :: * -> *). Applicative m => String -> m Doc
text (String -> Sem r Doc) -> String -> Sem r Doc
forall a b. (a -> b) -> a -> b
$ String
syn String -> String -> String
forall a. [a] -> [a] -> [a]
++ (if (Char -> Bool) -> String -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Char -> Bool
isAlpha String
syn then String
" " else String
"")
    Maybe OpInfo
_ -> String -> Sem r Doc
forall a. HasCallStack => String -> a
error (String -> Sem r Doc) -> String -> Sem r Doc
forall a b. (a -> b) -> a -> b
$ String
"UOp " String -> String -> String
forall a. [a] -> [a] -> [a]
++ UOp -> String
forall a. Show a => a -> String
show UOp
op String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" not in uopMap!"

-- | Pretty-print a binary operator, by looking up its concrete syntax
--   in the 'bopMap'.
instance Pretty BOp where
  pretty :: BOp -> Sem r Doc
pretty BOp
op = case BOp -> Map BOp OpInfo -> Maybe OpInfo
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup BOp
op Map BOp OpInfo
bopMap of
    Just (OpInfo OpFixity
_ (String
syn:[String]
_) Prec
_) -> String -> Sem r Doc
forall (m :: * -> *). Applicative m => String -> m Doc
text String
syn
    Maybe OpInfo
_                         -> String -> Sem r Doc
forall a. HasCallStack => String -> a
error (String -> Sem r Doc) -> String -> Sem r Doc
forall a b. (a -> b) -> a -> b
$ String
"BOp " String -> String -> String
forall a. [a] -> [a] -> [a]
++ BOp -> String
forall a. Show a => a -> String
show BOp
op String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" not in bopMap!"

--------------------------------------------------
-- Pretty-printing decimals

-- | Pretty-print a rational number using its decimal expansion, in
--   the format @nnn.prefix[rep]...@, with any repeating digits enclosed
--   in square brackets.
prettyDecimal :: Rational -> String
prettyDecimal :: Rational -> String
prettyDecimal Rational
r = String
printedDecimal
   where
     (Integer
n,Rational
d) = Rational -> (Integer, Rational)
forall a b. (RealFrac a, Integral b) => a -> (b, a)
properFraction Rational
r :: (Integer, Rational)
     ([Integer]
expan, Prec
len) = Integer -> Integer -> Integer -> ([Integer], Prec)
digitalExpansion Integer
10 (Rational -> Integer
forall a. Ratio a -> a
numerator Rational
d) (Rational -> Integer
forall a. Ratio a -> a
denominator Rational
d)
     printedDecimal :: String
printedDecimal
       | [Integer] -> Prec
forall (t :: * -> *) a. Foldable t => t a -> Prec
length [Integer]
first102 Prec -> Prec -> Bool
forall a. Ord a => a -> a -> Bool
> Prec
101 Bool -> Bool -> Bool
|| [Integer] -> Prec
forall (t :: * -> *) a. Foldable t => t a -> Prec
length [Integer]
first102 Prec -> Prec -> Bool
forall a. Eq a => a -> a -> Bool
== Prec
101 Bool -> Bool -> Bool
&& [Integer] -> Integer
forall a. [a] -> a
last [Integer]
first102 Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= Integer
0
         = Integer -> String
forall a. Show a => a -> String
show Integer
n String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"." String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Integer -> String) -> [Integer] -> String
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Integer -> String
forall a. Show a => a -> String
show (Prec -> [Integer] -> [Integer]
forall a. Prec -> [a] -> [a]
take Prec
100 [Integer]
expan) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"..."
       | [Integer]
rep [Integer] -> [Integer] -> Bool
forall a. Eq a => a -> a -> Bool
== [Integer
0]
         = Integer -> String
forall a. Show a => a -> String
show Integer
n String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"." String -> String -> String
forall a. [a] -> [a] -> [a]
++ (if [Integer] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Integer]
pre then String
"0" else (Integer -> String) -> [Integer] -> String
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Integer -> String
forall a. Show a => a -> String
show [Integer]
pre)
       | Bool
otherwise
         = Integer -> String
forall a. Show a => a -> String
show Integer
n String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"." String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Integer -> String) -> [Integer] -> String
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Integer -> String
forall a. Show a => a -> String
show [Integer]
pre String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"[" String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Integer -> String) -> [Integer] -> String
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Integer -> String
forall a. Show a => a -> String
show [Integer]
rep String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"]"
       where
         ([Integer]
pre, [Integer]
rep) = Prec -> [Integer] -> ([Integer], [Integer])
forall a. Prec -> [a] -> ([a], [a])
splitAt Prec
len [Integer]
expan
         first102 :: [Integer]
first102   = Prec -> [Integer] -> [Integer]
forall a. Prec -> [a] -> [a]
take Prec
102 [Integer]
expan

-- Given a list, find the indices of the list giving the first and
-- second occurrence of the first element to repeat, or Nothing if
-- there are no repeats.
findRep :: Ord a => [a] -> ([a], Int)
findRep :: [a] -> ([a], Prec)
findRep = Map a Prec -> Prec -> [a] -> ([a], Prec)
forall a. Ord a => Map a Prec -> Prec -> [a] -> ([a], Prec)
findRep' Map a Prec
forall k a. Map k a
M.empty Prec
0

findRep' :: Ord a => M.Map a Int -> Int -> [a] -> ([a], Int)
findRep' :: Map a Prec -> Prec -> [a] -> ([a], Prec)
findRep' Map a Prec
_ Prec
_ [] = String -> ([a], Prec)
forall a. HasCallStack => String -> a
error String
"Impossible. Empty list in findRep'"
findRep' Map a Prec
prevs Prec
ix (a
x:[a]
xs)
  | a
x a -> Map a Prec -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map a Prec
prevs = ([], Map a Prec
prevs Map a Prec -> a -> Prec
forall k a. Ord k => Map k a -> k -> a
M.! a
x)
  | Bool
otherwise          = ([a] -> [a]) -> ([a], Prec) -> ([a], Prec)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:) (([a], Prec) -> ([a], Prec)) -> ([a], Prec) -> ([a], Prec)
forall a b. (a -> b) -> a -> b
$ Map a Prec -> Prec -> [a] -> ([a], Prec)
forall a. Ord a => Map a Prec -> Prec -> [a] -> ([a], Prec)
findRep' (a -> Prec -> Map a Prec -> Map a Prec
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
x Prec
ix Map a Prec
prevs) (Prec
ixPrec -> Prec -> Prec
forall a. Num a => a -> a -> a
+Prec
1) [a]
xs

-- | @digitalExpansion b n d@ takes the numerator and denominator of a
--   fraction n/d between 0 and 1, and returns a pair of (1) a list of
--   digits @ds@, and (2) a nonnegative integer k such that @splitAt k
--   ds = (prefix, rep)@, where the infinite base-b expansion of
--   n/d is 0.@(prefix ++ cycle rep)@.  For example,
--
--   > digitalExpansion 10 1 4  = ([2,5,0], 2)
--   > digitalExpansion 10 1 7  = ([1,4,2,8,5,7], 0)
--   > digitalExpansion 10 3 28 = ([1,0,7,1,4,2,8,5], 2)
--   > digitalExpansion 2  1 5  = ([0,0,1,1], 0)
--
--   It works by performing the standard long division algorithm, and
--   looking for the first time that the remainder repeats.
digitalExpansion :: Integer -> Integer -> Integer -> ([Integer], Int)
digitalExpansion :: Integer -> Integer -> Integer -> ([Integer], Prec)
digitalExpansion Integer
b Integer
n Integer
d = ([Integer], Prec)
digits
  where
    longDivStep :: (Integer, Integer) -> (Integer, Integer)
longDivStep (Integer
_, Integer
r) = (Integer
bInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
*Integer
r) Integer -> Integer -> (Integer, Integer)
forall a. Integral a => a -> a -> (a, a)
`divMod` Integer
d
    res :: [(Integer, Integer)]
res       = [(Integer, Integer)] -> [(Integer, Integer)]
forall a. [a] -> [a]
tail ([(Integer, Integer)] -> [(Integer, Integer)])
-> [(Integer, Integer)] -> [(Integer, Integer)]
forall a b. (a -> b) -> a -> b
$ ((Integer, Integer) -> (Integer, Integer))
-> (Integer, Integer) -> [(Integer, Integer)]
forall a. (a -> a) -> a -> [a]
iterate (Integer, Integer) -> (Integer, Integer)
longDivStep (Integer
0,Integer
n)
    digits :: ([Integer], Prec)
digits    = ([(Integer, Integer)] -> [Integer])
-> ([(Integer, Integer)], Prec) -> ([Integer], Prec)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (((Integer, Integer) -> Integer)
-> [(Integer, Integer)] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Integer, Integer) -> Integer
forall a b. (a, b) -> a
fst) ([(Integer, Integer)] -> ([(Integer, Integer)], Prec)
forall a. Ord a => [a] -> ([a], Prec)
findRep [(Integer, Integer)]
res)