{-# OPTIONS_GHC -fno-warn-orphans #-}

module Agda.Compiler.Treeless.Pretty () where

import Prelude hiding ((!!)) -- don't use partial functions!

import Control.Arrow (first)
import Control.Monad.Reader
import Data.Maybe
import qualified Data.Map as Map

import Agda.Syntax.Treeless
import Agda.Compiler.Treeless.Subst
import Agda.Utils.Pretty
import Agda.Utils.List

import Agda.Utils.Impossible

data PEnv = PEnv { PEnv -> Int
pPrec :: Int
                 , PEnv -> [[Char]]
pFresh :: [String]
                 , PEnv -> [[Char]]
pBound :: [String] }

type P = Reader PEnv

--UNUSED Liang-Ting Chen 2019-07-16
--withName :: (String -> P a) -> P a
--withName k = withNames 1 $ \[x] -> k x

withNames :: Int -> ([String] -> P a) -> P a
withNames :: forall a. Int -> ([[Char]] -> P a) -> P a
withNames Int
n [[Char]] -> P a
k = do
  ([[Char]]
xs, [[Char]]
ys) <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a -> b) -> a -> b
$ forall a. Int -> [a] -> ([a], [a])
splitAt Int
n forall b c a. (b -> c) -> (a -> b) -> a -> c
. PEnv -> [[Char]]
pFresh
  forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (\ PEnv
e -> PEnv
e { pFresh :: [[Char]]
pFresh = [[Char]]
ys }) ([[Char]] -> P a
k [[Char]]
xs)

-- | Don't generate fresh names for unused variables.
withNames' :: HasFree a => Int -> a -> ([String] -> P b) -> P b
withNames' :: forall a b. HasFree a => Int -> a -> ([[Char]] -> P b) -> P b
withNames' Int
n a
tm [[Char]] -> P b
k = forall a. Int -> ([[Char]] -> P a) -> P a
withNames Int
n' forall a b. (a -> b) -> a -> b
$ [[Char]] -> P b
k forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[Char]] -> [[Char]]
insBlanks
  where
    fv :: Map Int Occurs
fv = forall a. HasFree a => a -> Map Int Occurs
freeVars a
tm
    n' :: Int
n'  = forall (t :: * -> *) a. Foldable t => t a -> Int
length forall a b. (a -> b) -> a -> b
$ forall a. (a -> Bool) -> [a] -> [a]
filter (forall a. Ord a => a -> a -> Bool
< Int
n) forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [k]
Map.keys Map Int Occurs
fv
    insBlanks :: [[Char]] -> [[Char]]
insBlanks = Int -> [[Char]] -> [[Char]]
go Int
n
      where
        go :: Int -> [[Char]] -> [[Char]]
go Int
0 [[Char]]
_ = []
        go Int
i xs0 :: [[Char]]
xs0@(~([Char]
x : [[Char]]
xs))
          | forall k a. Ord k => k -> Map k a -> Bool
Map.member (Int
i forall a. Num a => a -> a -> a
- Int
1) Map Int Occurs
fv = [Char]
x   forall a. a -> [a] -> [a]
: Int -> [[Char]] -> [[Char]]
go (Int
i forall a. Num a => a -> a -> a
- Int
1) [[Char]]
xs
          | Bool
otherwise             = [Char]
"_" forall a. a -> [a] -> [a]
: Int -> [[Char]] -> [[Char]]
go (Int
i forall a. Num a => a -> a -> a
- Int
1) [[Char]]
xs0

bindName :: String -> P a -> P a
bindName :: forall a. [Char] -> P a -> P a
bindName [Char]
x = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall a b. (a -> b) -> a -> b
$ \ PEnv
e -> PEnv
e { pBound :: [[Char]]
pBound = [Char]
x forall a. a -> [a] -> [a]
: PEnv -> [[Char]]
pBound PEnv
e }

bindNames :: [String] -> P a -> P a
bindNames :: forall a. [[Char]] -> P a -> P a
bindNames [[Char]]
xs P a
p = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall a. [Char] -> P a -> P a
bindName P a
p [[Char]]
xs

paren :: Int -> P Doc -> P Doc
paren :: Int -> P Doc -> P Doc
paren Int
p P Doc
doc = do
  Int
n <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks PEnv -> Int
pPrec
  (if Int
p forall a. Ord a => a -> a -> Bool
< Int
n then Doc -> Doc
parens else forall a. a -> a
id) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> P Doc
doc

prec :: Int -> P a -> P a
prec :: forall a. Int -> P a -> P a
prec Int
p = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall a b. (a -> b) -> a -> b
$ \ PEnv
e -> PEnv
e { pPrec :: Int
pPrec = Int
p }

name :: Int -> P String
name :: Int -> P [Char]
name Int
x = forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks
  forall a b. (a -> b) -> a -> b
$ (\ [[Char]]
xs -> forall a. a -> [a] -> Int -> a
indexWithDefault forall a. HasCallStack => a
__IMPOSSIBLE__ [[Char]]
xs Int
x)
  forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map (([Char]
"^" forall a. [a] -> [a] -> [a]
++) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> [Char]
show) [Integer
1..])
  forall b c a. (b -> c) -> (a -> b) -> a -> c
. PEnv -> [[Char]]
pBound

runP :: P a -> a
runP :: forall a. P a -> a
runP P a
p = forall r a. Reader r a -> r -> a
runReader P a
p PEnv{ pPrec :: Int
pPrec = Int
0, pFresh :: [[Char]]
pFresh = [[Char]]
names, pBound :: [[Char]]
pBound = [] }
  where
    names :: [[Char]]
names = [ [Char]
x forall a. [a] -> [a] -> [a]
++ [Char]
i | [Char]
i <- [Char]
"" forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map forall a. Show a => a -> [Char]
show [Integer
1..], [Char]
x <- forall a b. (a -> b) -> [a] -> [b]
map (forall a. a -> [a] -> [a]
:[]) [Char
'a'..Char
'z'] ]

instance Pretty TTerm where
  prettyPrec :: Int -> TTerm -> Doc
prettyPrec Int
p TTerm
t = forall a. P a -> a
runP forall a b. (a -> b) -> a -> b
$ forall a. Int -> P a -> P a
prec Int
p (TTerm -> P Doc
pTerm TTerm
t)

opName :: TPrim -> String
opName :: TPrim -> [Char]
opName TPrim
PAdd = [Char]
"+"
opName TPrim
PSub = [Char]
"-"
opName TPrim
PMul = [Char]
"*"
opName TPrim
PQuot = [Char]
"quot"
opName TPrim
PRem = [Char]
"rem"
opName TPrim
PGeq = [Char]
">="
opName TPrim
PLt  = [Char]
"<"
opName TPrim
PEqI = [Char]
"==I"
opName TPrim
PAdd64 = [Char]
"+64"
opName TPrim
PSub64 = [Char]
"-64"
opName TPrim
PMul64 = [Char]
"*64"
opName TPrim
PQuot64 = [Char]
"quot64"
opName TPrim
PRem64 = [Char]
"rem64"
opName TPrim
PLt64  = [Char]
"<64"
opName TPrim
PEq64 = [Char]
"==64"
opName TPrim
PEqF = [Char]
"==F"
opName TPrim
PEqS = [Char]
"==S"
opName TPrim
PEqC = [Char]
"==C"
opName TPrim
PEqQ = [Char]
"==Q"
opName TPrim
PIf  = [Char]
"if_then_else_"
opName TPrim
PSeq = [Char]
"seq"
opName TPrim
PITo64 = [Char]
"toWord64"
opName TPrim
P64ToI = [Char]
"fromWord64"


isInfix :: TPrim -> Maybe (Int, Int, Int)
isInfix :: TPrim -> Maybe (Int, Int, Int)
isInfix TPrim
op =
  case TPrim
op of
    TPrim
PMul -> forall {c}. Num c => c -> Maybe (c, c, c)
l Int
7
    TPrim
PAdd -> forall {c}. Num c => c -> Maybe (c, c, c)
l Int
6
    TPrim
PSub -> forall {c}. Num c => c -> Maybe (c, c, c)
l Int
6
    TPrim
PGeq -> forall {c}. Num c => c -> Maybe (c, c, c)
non Int
4
    TPrim
PLt  -> forall {c}. Num c => c -> Maybe (c, c, c)
non Int
4
    TPrim
PMul64 -> forall {c}. Num c => c -> Maybe (c, c, c)
l Int
7
    TPrim
PAdd64 -> forall {c}. Num c => c -> Maybe (c, c, c)
l Int
6
    TPrim
PSub64 -> forall {c}. Num c => c -> Maybe (c, c, c)
l Int
6
    TPrim
PLt64  -> forall {c}. Num c => c -> Maybe (c, c, c)
non Int
4
    TPrim
p | TPrim -> Bool
isPrimEq TPrim
p -> forall {c}. Num c => c -> Maybe (c, c, c)
non Int
4
    TPrim
_    -> forall a. Maybe a
Nothing
  where
    l :: c -> Maybe (c, c, c)
l c
n   = forall a. a -> Maybe a
Just (c
n, c
n, c
n forall a. Num a => a -> a -> a
+ c
1)
    r :: c -> Maybe (c, c, c)
r c
n   = forall a. a -> Maybe a
Just (c
n, c
n forall a. Num a => a -> a -> a
+ c
1, c
n) -- NB:: Defined but not used
    non :: c -> Maybe (c, c, c)
non c
n = forall a. a -> Maybe a
Just (c
n, c
n forall a. Num a => a -> a -> a
+ c
1, c
n forall a. Num a => a -> a -> a
+ c
1)

pTerm' :: Int -> TTerm -> P Doc
pTerm' :: Int -> TTerm -> P Doc
pTerm' Int
p = forall a. Int -> P a -> P a
prec Int
p forall b c a. (b -> c) -> (a -> b) -> a -> c
. TTerm -> P Doc
pTerm

pTerm :: TTerm -> P Doc
pTerm :: TTerm -> P Doc
pTerm = \case
  TVar Int
x -> [Char] -> Doc
text forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> P [Char]
name Int
x
  TApp (TPrim TPrim
op) [TTerm
a, TTerm
b] | Just (Int
c, Int
l, Int
r) <- TPrim -> Maybe (Int, Int, Int)
isInfix TPrim
op ->
    Int -> P Doc -> P Doc
paren Int
c forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *). Foldable t => t Doc -> Doc
sep forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [ Int -> TTerm -> P Doc
pTerm' Int
l TTerm
a
                               , forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [Char] -> Doc
text forall a b. (a -> b) -> a -> b
$ TPrim -> [Char]
opName TPrim
op
                               , Int -> TTerm -> P Doc
pTerm' Int
r TTerm
b ]
  TApp (TPrim TPrim
PIf) [TTerm
a, TTerm
b, TTerm
c] ->
    Int -> P Doc -> P Doc
paren Int
0 forall a b. (a -> b) -> a -> b
$ (\ Doc
a Doc
b Doc
c -> forall (t :: * -> *). Foldable t => t Doc -> Doc
sep [ Doc
"if" Doc -> Doc -> Doc
<+> Doc
a
                              , Int -> Doc -> Doc
nest Int
2 forall a b. (a -> b) -> a -> b
$ Doc
"then" Doc -> Doc -> Doc
<+> Doc
b
                              , Int -> Doc -> Doc
nest Int
2 forall a b. (a -> b) -> a -> b
$ Doc
"else" Doc -> Doc -> Doc
<+> Doc
c ])
              forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> TTerm -> P Doc
pTerm' Int
0 TTerm
a
              forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> TTerm -> P Doc
pTerm' Int
0 TTerm
b
              forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TTerm -> P Doc
pTerm TTerm
c
  TDef QName
f -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> Doc
pretty QName
f
  TCon QName
c -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> Doc
pretty QName
c
  TLit Literal
l -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> Doc
pretty Literal
l
  TPrim TPrim
op | forall a. Maybe a -> Bool
isJust (TPrim -> Maybe (Int, Int, Int)
isInfix TPrim
op) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [Char] -> Doc
text ([Char]
"_" forall a. [a] -> [a] -> [a]
++ TPrim -> [Char]
opName TPrim
op forall a. [a] -> [a] -> [a]
++ [Char]
"_")
           | Bool
otherwise -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [Char] -> Doc
text (TPrim -> [Char]
opName TPrim
op)
  TApp TTerm
f [TTerm]
es ->
    Int -> P Doc -> P Doc
paren Int
9 forall a b. (a -> b) -> a -> b
$ (\Doc
a [Doc]
bs -> forall (t :: * -> *). Foldable t => t Doc -> Doc
sep [Doc
a, Int -> Doc -> Doc
nest Int
2 forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *). Foldable t => t Doc -> Doc
fsep [Doc]
bs])
              forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> TTerm -> P Doc
pTerm' Int
9 TTerm
f
              forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Int -> TTerm -> P Doc
pTerm' Int
10) [TTerm]
es
  t :: TTerm
t@TLam{} -> Int -> P Doc -> P Doc
paren Int
0 forall a b. (a -> b) -> a -> b
$ forall a b. HasFree a => Int -> a -> ([[Char]] -> P b) -> P b
withNames' Int
n TTerm
b forall a b. (a -> b) -> a -> b
$ \ [[Char]]
xs -> forall a. [[Char]] -> P a -> P a
bindNames [[Char]]
xs forall a b. (a -> b) -> a -> b
$
    (\Doc
b -> forall (t :: * -> *). Foldable t => t Doc -> Doc
sep [ [Char] -> Doc
text ([Char]
"λ " forall a. [a] -> [a] -> [a]
++ [[Char]] -> [Char]
unwords [[Char]]
xs forall a. [a] -> [a] -> [a]
++ [Char]
" →")
               , Int -> Doc -> Doc
nest Int
2 Doc
b ]) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> TTerm -> P Doc
pTerm' Int
0 TTerm
b
    where
      (Int
n, TTerm
b) = TTerm -> (Int, TTerm)
tLamView TTerm
t
  t :: TTerm
t@TLet{} -> Int -> P Doc -> P Doc
paren Int
0 forall a b. (a -> b) -> a -> b
$ forall a. Int -> ([[Char]] -> P a) -> P a
withNames (forall (t :: * -> *) a. Foldable t => t a -> Int
length [TTerm]
es) forall a b. (a -> b) -> a -> b
$ \ [[Char]]
xs ->
    (\ ([([Char], Doc)]
binds, Doc
b) -> forall (t :: * -> *). Foldable t => t Doc -> Doc
sep [ Doc
"let" Doc -> Doc -> Doc
<+> forall (t :: * -> *). Foldable t => t Doc -> Doc
vcat [ forall (t :: * -> *). Foldable t => t Doc -> Doc
sep [ [Char] -> Doc
text [Char]
x Doc -> Doc -> Doc
<+> Doc
"="
                                                , Int -> Doc -> Doc
nest Int
2 Doc
e ] | ([Char]
x, Doc
e) <- [([Char], Doc)]
binds ]
                              Doc -> Doc -> Doc
<+> Doc
"in", Doc
b ])
      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [([Char], TTerm)] -> TTerm -> P ([([Char], Doc)], Doc)
pLets (forall a b. [a] -> [b] -> [(a, b)]
zip [[Char]]
xs [TTerm]
es) TTerm
b
    where
      ([TTerm]
es, TTerm
b) = TTerm -> ([TTerm], TTerm)
tLetView TTerm
t

      pLets :: [([Char], TTerm)] -> TTerm -> P ([([Char], Doc)], Doc)
pLets [] TTerm
b = ([],) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> TTerm -> P Doc
pTerm' Int
0 TTerm
b
      pLets (([Char]
x, TTerm
e) : [([Char], TTerm)]
bs) TTerm
b = do
        Doc
e <- Int -> TTerm -> P Doc
pTerm' Int
0 TTerm
e
        forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first (([Char]
x, Doc
e) forall a. a -> [a] -> [a]
:) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. [Char] -> P a -> P a
bindName [Char]
x ([([Char], TTerm)] -> TTerm -> P ([([Char], Doc)], Doc)
pLets [([Char], TTerm)]
bs TTerm
b)

  TCase Int
x CaseInfo
_ TTerm
def [TAlt]
alts -> Int -> P Doc -> P Doc
paren Int
0 forall a b. (a -> b) -> a -> b
$
    (\ Doc
sc [Doc]
alts Doc
defd ->
      forall (t :: * -> *). Foldable t => t Doc -> Doc
sep [ Doc
"case" Doc -> Doc -> Doc
<+> Doc
sc Doc -> Doc -> Doc
<+> Doc
"of"
          , Int -> Doc -> Doc
nest Int
2 forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *). Foldable t => t Doc -> Doc
vcat ([Doc]
alts forall a. [a] -> [a] -> [a]
++ [ Doc
"_ →" Doc -> Doc -> Doc
<+> Doc
defd | forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Doc]
alts Bool -> Bool -> Bool
|| TTerm
def forall a. Eq a => a -> a -> Bool
/= TError -> TTerm
TError TError
TUnreachable ]) ]
    ) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> TTerm -> P Doc
pTerm' Int
0 (Int -> TTerm
TVar Int
x)
      forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TAlt -> P Doc
pAlt [TAlt]
alts
      forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> TTerm -> P Doc
pTerm' Int
0 TTerm
def
    where
      pAlt :: TAlt -> P Doc
pAlt (TALit Literal
l TTerm
b) = Doc -> Doc -> Doc
pAlt' forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> TTerm -> P Doc
pTerm' Int
0 (Literal -> TTerm
TLit Literal
l) forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> TTerm -> P Doc
pTerm' Int
0 TTerm
b
      pAlt (TAGuard TTerm
g TTerm
b) =
        Doc -> Doc -> Doc
pAlt' forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Doc
"_" Doc -> Doc -> Doc
<+> Doc
"|" Doc -> Doc -> Doc
<+>) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> TTerm -> P Doc
pTerm' Int
0 TTerm
g)
              forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Int -> TTerm -> P Doc
pTerm' Int
0 TTerm
b)
      pAlt (TACon QName
c Int
a TTerm
b) =
        forall a b. HasFree a => Int -> a -> ([[Char]] -> P b) -> P b
withNames' Int
a TTerm
b forall a b. (a -> b) -> a -> b
$ \ [[Char]]
xs -> forall a. [[Char]] -> P a -> P a
bindNames [[Char]]
xs forall a b. (a -> b) -> a -> b
$
        Doc -> Doc -> Doc
pAlt' forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> TTerm -> P Doc
pTerm' Int
0 (TTerm -> [TTerm] -> TTerm
TApp (QName -> TTerm
TCon QName
c) [Int -> TTerm
TVar Int
i | Int
i <- forall a. [a] -> [a]
reverse [Int
0..Int
a forall a. Num a => a -> a -> a
- Int
1]])
              forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> TTerm -> P Doc
pTerm' Int
0 TTerm
b
      pAlt' :: Doc -> Doc -> Doc
pAlt' Doc
p Doc
b = forall (t :: * -> *). Foldable t => t Doc -> Doc
sep [Doc
p Doc -> Doc -> Doc
<+> Doc
"→", Int -> Doc -> Doc
nest Int
2 Doc
b]

  TTerm
TUnit -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Doc
"()"
  TTerm
TSort -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Doc
"Set"
  TTerm
TErased -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Doc
"_"
  TError TError
err -> Int -> P Doc -> P Doc
paren Int
9 forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Doc
"error" Doc -> Doc -> Doc
<+> [Char] -> Doc
text (forall a. Show a => a -> [Char]
show (forall a. Show a => a -> [Char]
show TError
err))
  TCoerce TTerm
t -> Int -> P Doc -> P Doc
paren Int
9 forall a b. (a -> b) -> a -> b
$ (Doc
"coe" Doc -> Doc -> Doc
<+>) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> TTerm -> P Doc
pTerm' Int
10 TTerm
t