{-# LANGUAGE PatternGuards #-}
module Plugin.Pl.PrettyPrinter (
  prettyDecl,
  prettyExpr,
  prettyTopLevel,
 ) where

import Plugin.Pl.Common

import Data.Char
import Data.List (intercalate)

prettyDecl :: Decl -> String
prettyDecl :: Decl -> String
prettyDecl (Define String
f Expr
e) = String
f forall a. [a] -> [a] -> [a]
++ String
" = " forall a. [a] -> [a] -> [a]
++ Expr -> String
prettyExpr Expr
e

prettyDecls :: [Decl] -> String
prettyDecls :: [Decl] -> String
prettyDecls = forall a. [a] -> [[a]] -> [a]
intercalate String
"; " forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map Decl -> String
prettyDecl

prettyExpr :: Expr -> String
prettyExpr :: Expr -> String
prettyExpr = forall a. Show a => a -> String
show forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr -> SExpr
toSExpr

prettyTopLevel :: TopLevel -> String
prettyTopLevel :: TopLevel -> String
prettyTopLevel (TLE Expr
e) = Expr -> String
prettyExpr Expr
e
prettyTopLevel (TLD Bool
_ Decl
d) = Decl -> String
prettyDecl Decl
d

data SExpr
  = SVar !String
  | SLambda ![Pattern] !SExpr
  | SLet ![Decl] !SExpr
  | SApp !SExpr !SExpr
  | SInfix !String !SExpr !SExpr
  | LeftSection !String !SExpr  -- (x +)
  | RightSection !String !SExpr -- (+ x)
  | List ![SExpr]
  | Tuple ![SExpr]
  | Enum !Expr !(Maybe Expr) !(Maybe Expr)

{-# INLINE toSExprHead #-}
toSExprHead :: String -> [Expr] -> Maybe SExpr
toSExprHead :: String -> [Expr] -> Maybe SExpr
toSExprHead String
hd [Expr]
tl
  | forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (forall a. Eq a => a -> a -> Bool
==Char
',') String
hd, forall (t :: * -> *) a. Foldable t => t a -> Int
length String
hdforall a. Num a => a -> a -> a
+Int
1 forall a. Eq a => a -> a -> Bool
== forall (t :: * -> *) a. Foldable t => t a -> Int
length [Expr]
tl 
  = forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. [SExpr] -> SExpr
Tuple forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [a] -> [a]
reverse forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map Expr -> SExpr
toSExpr [Expr]
tl
  | Bool
otherwise = case (String
hd,forall a. [a] -> [a]
reverse [Expr]
tl) of
      (String
"enumFrom", [Expr
e])              -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Expr -> Maybe Expr -> Maybe Expr -> SExpr
Enum Expr
e forall a. Maybe a
Nothing   forall a. Maybe a
Nothing
      (String
"enumFromThen", [Expr
e,Expr
e'])       -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Expr -> Maybe Expr -> Maybe Expr -> SExpr
Enum Expr
e (forall a. a -> Maybe a
Just Expr
e') forall a. Maybe a
Nothing
      (String
"enumFromTo", [Expr
e,Expr
e'])         -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Expr -> Maybe Expr -> Maybe Expr -> SExpr
Enum Expr
e forall a. Maybe a
Nothing   (forall a. a -> Maybe a
Just Expr
e')
      (String
"enumFromThenTo", [Expr
e,Expr
e',Expr
e'']) -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Expr -> Maybe Expr -> Maybe Expr -> SExpr
Enum Expr
e (forall a. a -> Maybe a
Just Expr
e') (forall a. a -> Maybe a
Just Expr
e'')
      (String, [Expr])
_                              -> forall a. Maybe a
Nothing

toSExpr :: Expr -> SExpr
toSExpr :: Expr -> SExpr
toSExpr (Var Fixity
_ String
v) = String -> SExpr
SVar String
v
toSExpr (Lambda Pattern
v Expr
e) = case Expr -> SExpr
toSExpr Expr
e of
  (SLambda [Pattern]
vs SExpr
e') -> [Pattern] -> SExpr -> SExpr
SLambda (Pattern
vforall a. a -> [a] -> [a]
:[Pattern]
vs) SExpr
e'
  SExpr
e'              -> [Pattern] -> SExpr -> SExpr
SLambda [Pattern
v] SExpr
e'
toSExpr (Let [Decl]
ds Expr
e) = [Decl] -> SExpr -> SExpr
SLet [Decl]
ds forall a b. (a -> b) -> a -> b
$ Expr -> SExpr
toSExpr Expr
e
toSExpr Expr
e | Just (String
hd,[Expr]
tl) <- Expr -> Maybe (String, [Expr])
getHead Expr
e, Just SExpr
se <- String -> [Expr] -> Maybe SExpr
toSExprHead String
hd [Expr]
tl = SExpr
se
toSExpr Expr
e | ([Expr]
ls, Expr
tl) <- Expr -> ([Expr], Expr)
getList Expr
e, Expr
tl forall a. Eq a => a -> a -> Bool
== Expr
nil
  = [SExpr] -> SExpr
List forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map Expr -> SExpr
toSExpr [Expr]
ls
toSExpr (App Expr
e1 Expr
e2) = case Expr
e1 of
  App (Var Fixity
Inf String
v) Expr
e0 
    -> String -> SExpr -> SExpr -> SExpr
SInfix String
v (Expr -> SExpr
toSExpr Expr
e0) (Expr -> SExpr
toSExpr Expr
e2)
  Var Fixity
Inf String
v | String
v forall a. Eq a => a -> a -> Bool
/= String
"-"
    -> String -> SExpr -> SExpr
LeftSection String
v (Expr -> SExpr
toSExpr Expr
e2)

  Var Fixity
_ String
"flip" | Var Fixity
Inf String
v <- Expr
e2, String
v forall a. Eq a => a -> a -> Bool
== String
"-" -> Expr -> SExpr
toSExpr forall a b. (a -> b) -> a -> b
$ Fixity -> String -> Expr
Var Fixity
Pref String
"subtract"
    
  App (Var Fixity
_ String
"flip") (Var Fixity
pr String
v)
    | String
v forall a. Eq a => a -> a -> Bool
== String
"-"  -> Expr -> SExpr
toSExpr forall a b. (a -> b) -> a -> b
$ Fixity -> String -> Expr
Var Fixity
Pref String
"subtract" Expr -> Expr -> Expr
`App` Expr
e2
    | String
v forall a. Eq a => a -> a -> Bool
== String
"id" -> String -> SExpr -> SExpr
RightSection String
"$" (Expr -> SExpr
toSExpr Expr
e2)
    | Fixity
Inf <- Fixity
pr, forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (forall a. Eq a => a -> a -> Bool
/= Char
',') String
v -> String -> SExpr -> SExpr
RightSection String
v (Expr -> SExpr
toSExpr Expr
e2)
  Expr
_ -> SExpr -> SExpr -> SExpr
SApp (Expr -> SExpr
toSExpr Expr
e1) (Expr -> SExpr
toSExpr Expr
e2)

getHead :: Expr -> Maybe (String, [Expr])
getHead :: Expr -> Maybe (String, [Expr])
getHead (Var Fixity
_ String
v) = forall a. a -> Maybe a
Just (String
v, [])
getHead (App Expr
e1 Expr
e2) = forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (Expr
e2forall a. a -> [a] -> [a]
:) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` Expr -> Maybe (String, [Expr])
getHead Expr
e1
getHead Expr
_ = forall a. Maybe a
Nothing

instance Show SExpr where
  showsPrec :: Int -> SExpr -> ShowS
showsPrec Int
_ (SVar String
v) = (ShowS
getPrefName String
v forall a. [a] -> [a] -> [a]
++)
  showsPrec Int
p (SLambda [Pattern]
vs SExpr
e) = Bool -> ShowS -> ShowS
showParen (Int
p forall a. Ord a => a -> a -> Bool
> Int
minPrec) forall a b. (a -> b) -> a -> b
$ (Char
'\\'forall a. a -> [a] -> [a]
:) forall b c a. (b -> c) -> (a -> b) -> a -> c
. 
    forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall b c a. (b -> c) -> (a -> b) -> a -> c
(.) forall a. a -> a
id (forall a. a -> [a] -> [a]
intersperse (Char
' 'forall a. a -> [a] -> [a]
:) (forall a b. (a -> b) -> [a] -> [b]
map (Int -> Pattern -> ShowS
prettyPrecPattern forall a b. (a -> b) -> a -> b
$ Int
maxPrecforall a. Num a => a -> a -> a
+Int
1) [Pattern]
vs)) forall b c a. (b -> c) -> (a -> b) -> a -> c
.
    (String
" -> "forall a. [a] -> [a] -> [a]
++) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => Int -> a -> ShowS
showsPrec Int
minPrec SExpr
e
  showsPrec Int
p (SApp SExpr
e1 SExpr
e2) = Bool -> ShowS -> ShowS
showParen (Int
p forall a. Ord a => a -> a -> Bool
> Int
maxPrec) forall a b. (a -> b) -> a -> b
$
    forall a. Show a => Int -> a -> ShowS
showsPrec Int
maxPrec SExpr
e1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char
' 'forall a. a -> [a] -> [a]
:) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => Int -> a -> ShowS
showsPrec (Int
maxPrecforall a. Num a => a -> a -> a
+Int
1) SExpr
e2
  showsPrec Int
_ (LeftSection String
fx SExpr
e) = Bool -> ShowS -> ShowS
showParen Bool
True forall a b. (a -> b) -> a -> b
$ 
    forall a. Show a => Int -> a -> ShowS
showsPrec (forall a b. (a, b) -> b
snd (String -> (Assoc (), Int)
lookupFix String
fx) forall a. Num a => a -> a -> a
+ Int
1) SExpr
e forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char
' 'forall a. a -> [a] -> [a]
:) forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ShowS
getInfName String
fxforall a. [a] -> [a] -> [a]
++)
  showsPrec Int
_ (RightSection String
fx SExpr
e) = Bool -> ShowS -> ShowS
showParen Bool
True forall a b. (a -> b) -> a -> b
$ 
    (ShowS
getInfName String
fxforall a. [a] -> [a] -> [a]
++) forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char
' 'forall a. a -> [a] -> [a]
:) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => Int -> a -> ShowS
showsPrec (forall a b. (a, b) -> b
snd (String -> (Assoc (), Int)
lookupFix String
fx) forall a. Num a => a -> a -> a
+ Int
1) SExpr
e
  showsPrec Int
_ (Tuple [SExpr]
es) = Bool -> ShowS -> ShowS
showParen Bool
True forall a b. (a -> b) -> a -> b
$
    (forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall a. a -> a
`id` forall a. a -> [a] -> [a]
intersperse String
", " (forall a b. (a -> b) -> [a] -> [b]
map forall a. Show a => a -> String
show [SExpr]
es) forall a. [a] -> [a] -> [a]
++)
  
  showsPrec Int
_ (List [SExpr]
es) 
    | Just String
cs <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
(=<<) forall a. Read a => String -> Maybe a
readM forall b c a. (b -> c) -> (a -> b) -> a -> c
. SExpr -> Maybe String
fromSVar) [SExpr]
es = forall a. Show a => a -> ShowS
shows (String
cs::String)
    | Bool
otherwise = (Char
'['forall a. a -> [a] -> [a]
:) forall b c a. (b -> c) -> (a -> b) -> a -> c
. 
      (forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall a. a -> a
`id` forall a. a -> [a] -> [a]
intersperse String
", " (forall a b. (a -> b) -> [a] -> [b]
map forall a. Show a => a -> String
show [SExpr]
es) forall a. [a] -> [a] -> [a]
++) forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char
']'forall a. a -> [a] -> [a]
:)
    where fromSVar :: SExpr -> Maybe String
fromSVar (SVar String
str) = forall a. a -> Maybe a
Just String
str
          fromSVar SExpr
_          = forall a. Maybe a
Nothing
  showsPrec Int
_ (Enum Expr
fr Maybe Expr
tn Maybe Expr
to) = (Char
'['forall a. a -> [a] -> [a]
:) forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ShowS
showString (Expr -> String
prettyExpr Expr
fr) forall b c a. (b -> c) -> (a -> b) -> a -> c
. 
    forall {a}. Maybe [a] -> [a] -> [a]
showsMaybe (((Char
','forall a. a -> [a] -> [a]
:) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr -> String
prettyExpr) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` Maybe Expr
tn) forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
".."forall a. [a] -> [a] -> [a]
++) forall b c a. (b -> c) -> (a -> b) -> a -> c
. 
    forall {a}. Maybe [a] -> [a] -> [a]
showsMaybe (Expr -> String
prettyExpr forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` Maybe Expr
to) forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char
']'forall a. a -> [a] -> [a]
:)
      where showsMaybe :: Maybe [a] -> [a] -> [a]
showsMaybe = forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. a -> a
id forall a. [a] -> [a] -> [a]
(++)
  showsPrec Int
_ (SLet [Decl]
ds SExpr
e) = (String
"let "forall a. [a] -> [a] -> [a]
++) forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ShowS
showString ([Decl] -> String
prettyDecls [Decl]
ds forall a. [a] -> [a] -> [a]
++ String
" in ") forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> ShowS
shows SExpr
e


  showsPrec Int
p (SInfix String
fx SExpr
e1 SExpr
e2) = Bool -> ShowS -> ShowS
showParen (Int
p forall a. Ord a => a -> a -> Bool
> Int
fixity) forall a b. (a -> b) -> a -> b
$
    forall a. Show a => Int -> a -> ShowS
showsPrec Int
f1 SExpr
e1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char
' 'forall a. a -> [a] -> [a]
:) forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ShowS
getInfName String
fxforall a. [a] -> [a] -> [a]
++) forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char
' 'forall a. a -> [a] -> [a]
:) forall b c a. (b -> c) -> (a -> b) -> a -> c
. 
    forall a. Show a => Int -> a -> ShowS
showsPrec Int
f2 SExpr
e2 where
      fixity :: Int
fixity = forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ String -> (Assoc (), Int)
lookupFix String
fx
      (Int
f1, Int
f2) = case forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ String -> (Assoc (), Int)
lookupFix String
fx of
        AssocRight ()
_ -> (Int
fixityforall a. Num a => a -> a -> a
+Int
1, Int
fixity forall a. Num a => a -> a -> a
+ SExpr -> Assoc () -> Int -> Int
infixSafe SExpr
e2 (forall l. l -> Assoc l
AssocLeft ()) Int
fixity)
        AssocLeft ()
_ -> (Int
fixity forall a. Num a => a -> a -> a
+ SExpr -> Assoc () -> Int -> Int
infixSafe SExpr
e1 (forall l. l -> Assoc l
AssocRight ()) Int
fixity, Int
fixityforall a. Num a => a -> a -> a
+Int
1)
        AssocNone ()
_ -> (Int
fixityforall a. Num a => a -> a -> a
+Int
1, Int
fixityforall a. Num a => a -> a -> a
+Int
1)

      -- This is a little bit awkward, but at least seems to produce no false
      -- results anymore
      infixSafe :: SExpr -> Assoc () -> Int -> Int
      infixSafe :: SExpr -> Assoc () -> Int -> Int
infixSafe (SInfix String
fx'' SExpr
_ SExpr
_) Assoc ()
assoc Int
fx'
        | String -> (Assoc (), Int)
lookupFix String
fx'' forall a. Eq a => a -> a -> Bool
== (Assoc ()
assoc, Int
fx') = Int
1
        | Bool
otherwise = Int
0
      infixSafe SExpr
_ Assoc ()
_ Int
_ = Int
0 -- doesn't matter

prettyPrecPattern :: Int -> Pattern -> ShowS
prettyPrecPattern :: Int -> Pattern -> ShowS
prettyPrecPattern Int
_ (PVar String
v) = String -> ShowS
showString String
v
prettyPrecPattern Int
_ (PTuple Pattern
p1 Pattern
p2) = Bool -> ShowS -> ShowS
showParen Bool
True forall a b. (a -> b) -> a -> b
$
  Int -> Pattern -> ShowS
prettyPrecPattern Int
0 Pattern
p1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
", "forall a. [a] -> [a] -> [a]
++) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Pattern -> ShowS
prettyPrecPattern Int
0 Pattern
p2
prettyPrecPattern Int
p (PCons Pattern
p1 Pattern
p2) = Bool -> ShowS -> ShowS
showParen (Int
pforall a. Ord a => a -> a -> Bool
>Int
5) forall a b. (a -> b) -> a -> b
$
  Int -> Pattern -> ShowS
prettyPrecPattern Int
6 Pattern
p1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char
':'forall a. a -> [a] -> [a]
:) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Pattern -> ShowS
prettyPrecPattern Int
5 Pattern
p2
  
isOperator :: String -> Bool
isOperator :: String -> Bool
isOperator String
s =
  case forall a. (a -> Bool) -> [a] -> ([a], [a])
break (forall a. Eq a => a -> a -> Bool
== Char
'.') String
s of
    (String
_, String
"") -> String -> Bool
isUnqualOp String
s
    (String
before, Char
_dot : String
rest)
      | String -> Bool
isUnqualOp String
before -> String -> Bool
isUnqualOp String
rest
      | String -> Bool
isModule String
before -> String -> Bool
isOperator String
rest
      | Bool
otherwise -> Bool
False
  where
    isModule :: String -> Bool
isModule String
"" = Bool
False
    isModule (Char
c : String
cs) = Char -> Bool
isUpper Char
c Bool -> Bool -> Bool
&& forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Char
c -> Char -> Bool
isAlphaNum Char
c Bool -> Bool -> Bool
|| Char
c forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Char
'\'', Char
'_']) String
cs
    isUnqualOp :: String -> Bool
isUnqualOp String
s = String
s forall a. Eq a => a -> a -> Bool
/= String
"()" Bool -> Bool -> Bool
&& forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Char
c -> Char -> Bool
isSymbol Char
c Bool -> Bool -> Bool
|| Char -> Bool
isPunctuation Char
c) String
s

getInfName :: String -> String
getInfName :: ShowS
getInfName String
str = if String -> Bool
isOperator String
str then String
str else String
"`"forall a. [a] -> [a] -> [a]
++String
strforall a. [a] -> [a] -> [a]
++String
"`"

getPrefName :: String -> String
getPrefName :: ShowS
getPrefName String
str = if String -> Bool
isOperator String
str Bool -> Bool -> Bool
|| Char
',' forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` String
str then String
"("forall a. [a] -> [a] -> [a]
++String
strforall a. [a] -> [a] -> [a]
++String
")" else String
str

{-
instance Show Assoc where
  show AssocLeft  = "AssocLeft"
  show AssocRight = "AssocRight"
  show AssocNone  = "AssocNone"

instance Ord Assoc where
  AssocNone <= _ = True
  _ <= AssocNone = False
  AssocLeft <= _ = True
  _ <= AssocLeft = False
  _ <= _ = True
-}