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

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

-- |
-- Module      :  Disco.Pretty
-- Copyright   :  disco team and contributors
-- Maintainer  :  byorgey@gmail.com
--
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Various pretty-printing facilities for disco.
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 Disco.Pretty.DSL
import Disco.Pretty.Prec
import Disco.Syntax.Operators
import Polysemy
import Polysemy.Reader
import Prettyprinter (Doc)
import Unbound.Generics.LocallyNameless (Name)

------------------------------------------------------------
-- 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 ann) -> Sem r (Doc ann)
withPA :: forall (r :: EffectRow) ann.
Member (Reader PA) r =>
PA -> Sem r (Doc ann) -> Sem r (Doc ann)
withPA PA
pa = forall (r :: EffectRow) ann.
Member (Reader PA) r =>
PA -> Sem r (Doc ann) -> Sem r (Doc ann)
mparens PA
pa forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall (r :: EffectRow) a.
Member (Reader PA) r =>
PA -> Sem r a -> Sem r a
setPA = forall i (r :: EffectRow) a.
Member (Reader i) r =>
(i -> i) -> Sem r a -> Sem r a
local forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 ann) -> Sem r (Doc ann)
lt :: forall (r :: EffectRow) ann.
Member (Reader PA) r =>
Sem r (Doc ann) -> Sem r (Doc ann)
lt = forall i (r :: EffectRow) a.
Member (Reader i) r =>
(i -> i) -> Sem r a -> Sem r a
local (\(PA Int
p BFixity
_) -> Int -> BFixity -> PA
PA Int
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 ann) -> Sem r (Doc ann)
rt :: forall (r :: EffectRow) ann.
Member (Reader PA) r =>
Sem r (Doc ann) -> Sem r (Doc ann)
rt = forall i (r :: EffectRow) a.
Member (Reader i) r =>
(i -> i) -> Sem r a -> Sem r a
local (\(PA Int
p BFixity
_) -> Int -> BFixity -> PA
PA Int
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 ann) -> Sem r (Doc ann)
mparens :: forall (r :: EffectRow) ann.
Member (Reader PA) r =>
PA -> Sem r (Doc ann) -> Sem r (Doc ann)
mparens PA
pa Sem r (Doc ann)
doc = do
  PA
parentPA <- forall i (r :: EffectRow). Member (Reader i) r => Sem r i
ask
  (if PA
pa forall a. Ord a => a -> a -> Bool
< PA
parentPA then forall (f :: * -> *) ann. Functor f => f (Doc ann) -> f (Doc ann)
parens else forall a. a -> a
id) Sem r (Doc ann)
doc

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

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

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

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

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

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

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

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

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

instance Pretty (Name a) where
  pretty :: forall (r :: EffectRow) ann.
Members '[Reader PA, LFresh] r =>
Name a -> Sem r (Doc ann)
pretty = forall (m :: * -> *) ann. Applicative m => [Char] -> m (Doc ann)
text forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> [Char]
show

instance Pretty TyOp where
  pretty :: forall (r :: EffectRow) ann.
Members '[Reader PA, LFresh] r =>
TyOp -> Sem r (Doc ann)
pretty = \case
    TyOp
Enumerate -> forall (m :: * -> *) ann. Applicative m => [Char] -> m (Doc ann)
text [Char]
"enumerate"
    TyOp
Count -> forall (m :: * -> *) ann. Applicative m => [Char] -> m (Doc ann)
text [Char]
"count"

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

-- | Pretty-print a binary operator, by looking up its concrete syntax
--   in the 'bopMap'.
instance Pretty BOp where
  pretty :: forall (r :: EffectRow) ann.
Members '[Reader PA, LFresh] r =>
BOp -> Sem r (Doc ann)
pretty BOp
op = case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup BOp
op Map BOp OpInfo
bopMap of
    Just (OpInfo OpFixity
_ ([Char]
syn : [[Char]]
_) Int
_) -> forall (m :: * -> *) ann. Applicative m => [Char] -> m (Doc ann)
text [Char]
syn
    Maybe OpInfo
_ -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"BOp " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show BOp
op forall a. [a] -> [a] -> [a]
++ [Char]
" 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 -> [Char]
prettyDecimal Rational
r = [Char]
printedDecimal
 where
  (Integer
n, Rational
d) = forall a b. (RealFrac a, Integral b) => a -> (b, a)
properFraction Rational
r :: (Integer, Rational)
  ([Integer]
expan, Int
len) = Integer -> Integer -> Integer -> ([Integer], Int)
digitalExpansion Integer
10 (forall a. Ratio a -> a
numerator Rational
d) (forall a. Ratio a -> a
denominator Rational
d)
  printedDecimal :: [Char]
printedDecimal
    | forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
first102 forall a. Ord a => a -> a -> Bool
> Int
101 Bool -> Bool -> Bool
|| forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
first102 forall a. Eq a => a -> a -> Bool
== Int
101 Bool -> Bool -> Bool
&& forall a. [a] -> a
last [Integer]
first102 forall a. Eq a => a -> a -> Bool
/= Integer
0 =
        forall a. Show a => a -> [Char]
show Integer
n forall a. [a] -> [a] -> [a]
++ [Char]
"." forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a. Show a => a -> [Char]
show (forall a. Int -> [a] -> [a]
take Int
100 [Integer]
expan) forall a. [a] -> [a] -> [a]
++ [Char]
"..."
    | [Integer]
rep forall a. Eq a => a -> a -> Bool
== [Integer
0] =
        forall a. Show a => a -> [Char]
show Integer
n forall a. [a] -> [a] -> [a]
++ [Char]
"." forall a. [a] -> [a] -> [a]
++ (if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Integer]
pre then [Char]
"0" else forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a. Show a => a -> [Char]
show [Integer]
pre)
    | Bool
otherwise =
        forall a. Show a => a -> [Char]
show Integer
n forall a. [a] -> [a] -> [a]
++ [Char]
"." forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a. Show a => a -> [Char]
show [Integer]
pre forall a. [a] -> [a] -> [a]
++ [Char]
"[" forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a. Show a => a -> [Char]
show [Integer]
rep forall a. [a] -> [a] -> [a]
++ [Char]
"]"
   where
    ([Integer]
pre, [Integer]
rep) = forall a. Int -> [a] -> ([a], [a])
splitAt Int
len [Integer]
expan
    first102 :: [Integer]
first102 = forall a. Int -> [a] -> [a]
take Int
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 :: forall a. Ord a => [a] -> ([a], Int)
findRep = forall a. Ord a => Map a Int -> Int -> [a] -> ([a], Int)
findRep' forall k a. Map k a
M.empty Int
0

findRep' :: Ord a => M.Map a Int -> Int -> [a] -> ([a], Int)
findRep' :: forall a. Ord a => Map a Int -> Int -> [a] -> ([a], Int)
findRep' Map a Int
_ Int
_ [] = forall a. HasCallStack => [Char] -> a
error [Char]
"Impossible. Empty list in findRep'"
findRep' Map a Int
prevs Int
ix (a
x : [a]
xs)
  | a
x forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map a Int
prevs = ([], Map a Int
prevs forall k a. Ord k => Map k a -> k -> a
M.! a
x)
  | Bool
otherwise = forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (a
x forall a. a -> [a] -> [a]
:) forall a b. (a -> b) -> a -> b
$ forall a. Ord a => Map a Int -> Int -> [a] -> ([a], Int)
findRep' (forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
x Int
ix Map a Int
prevs) (Int
ix forall a. Num a => a -> a -> a
+ Int
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], Int)
digitalExpansion Integer
b Integer
n Integer
d = ([Integer], Int)
digits
 where
  longDivStep :: (Integer, Integer) -> (Integer, Integer)
longDivStep (Integer
_, Integer
r) = (Integer
b forall a. Num a => a -> a -> a
* Integer
r) forall a. Integral a => a -> a -> (a, a)
`divMod` Integer
d
  res :: [(Integer, Integer)]
res = forall a. [a] -> [a]
tail forall a b. (a -> b) -> a -> b
$ forall a. (a -> a) -> a -> [a]
iterate (Integer, Integer) -> (Integer, Integer)
longDivStep (Integer
0, Integer
n)
  digits :: ([Integer], Int)
digits = forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst) (forall a. Ord a => [a] -> ([a], Int)
findRep [(Integer, Integer)]
res)