{-# LANGUAGE LambdaCase #-}

module Transformations.Anf.PrettyPrinter (prettyPrint) where

import Control.Monad.State (MonadState (get), State, evalState, modify)
import Data.Text (unpack)
import Transformations.Anf.Anf
import Trees.Common

prettyPrint :: Program -> String
prettyPrint :: Program -> String
prettyPrint (Program [GlobalDeclaration]
decls) = [String] -> String
unlines ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ GlobalDeclaration -> String
prettyDecl (GlobalDeclaration -> String) -> [GlobalDeclaration] -> [String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [GlobalDeclaration]
decls

type IndentState = State IndentLevel

type IndentLevel = Int

prettyDecl :: GlobalDeclaration -> String
prettyDecl :: GlobalDeclaration -> String
prettyDecl GlobalDeclaration
decl = State IndentLevel String -> IndentLevel -> String
forall s a. State s a -> s -> a
evalState (GlobalDeclaration -> State IndentLevel String
prettyDecl' GlobalDeclaration
decl) IndentLevel
0 String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
";;"
  where
    prettyDecl' :: GlobalDeclaration -> IndentState String
    prettyDecl' :: GlobalDeclaration -> State IndentLevel String
prettyDecl' (GlobVarDecl Identifier'
name Expression
value) = do
      String
val' <- Expression -> State IndentLevel String
prettyExpr Expression
value
      String -> State IndentLevel String
forall a. a -> StateT IndentLevel Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> State IndentLevel String)
-> String -> State IndentLevel String
forall a b. (a -> b) -> a -> b
$ [String] -> String
unwords [String
"let", Identifier' -> String
prettyId Identifier'
name, String
"=", String
val']
    prettyDecl' (GlobFunDecl Identifier'
name [Identifier']
params Expression
body) = do
      String
val' <- Expression -> State IndentLevel String
prettyExpr Expression
body
      String -> State IndentLevel String
forall a. a -> StateT IndentLevel Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> State IndentLevel String)
-> String -> State IndentLevel String
forall a b. (a -> b) -> a -> b
$ [String] -> String
unwords [String
"let", Identifier' -> String
prettyId Identifier'
name, [String] -> String
unwords (Identifier' -> String
prettyId (Identifier' -> String) -> [Identifier'] -> [String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Identifier']
params), String
"=", String
val']

prettyExpr :: Expression -> IndentState String
prettyExpr :: Expression -> State IndentLevel String
prettyExpr (ExprAtom AtomicExpression
aexpr) = String -> State IndentLevel String
forall a. a -> StateT IndentLevel Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> State IndentLevel String)
-> String -> State IndentLevel String
forall a b. (a -> b) -> a -> b
$ AtomicExpression -> String
prettyAtomic AtomicExpression
aexpr
prettyExpr (ExprComp ComplexExpression
cexpr) = ComplexExpression -> State IndentLevel String
prettyComplex ComplexExpression
cexpr
prettyExpr (ExprLetIn (Identifier'
ident, Expression
val) Expression
expr) = do
  (IndentLevel -> IndentLevel) -> StateT IndentLevel Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (IndentLevel -> IndentLevel -> IndentLevel
forall a. Num a => a -> a -> a
+ IndentLevel
2)
  IndentLevel
indent <- StateT IndentLevel Identity IndentLevel
forall s (m :: * -> *). MonadState s m => m s
get
  String
val' <- Expression -> State IndentLevel String
prettyExpr Expression
val
  String
expr' <- Expression -> State IndentLevel String
prettyExpr Expression
expr
  let declText :: String
declText = IndentLevel -> String
createIndent IndentLevel
indent String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"let " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Identifier' -> String
prettyId Identifier'
ident String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" = " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
val'
  let exprText :: String
exprText = IndentLevel -> String
createIndent IndentLevel
indent String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"in " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
expr'
  (IndentLevel -> IndentLevel) -> StateT IndentLevel Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((IndentLevel -> IndentLevel) -> StateT IndentLevel Identity ())
-> (IndentLevel -> IndentLevel) -> StateT IndentLevel Identity ()
forall a b. (a -> b) -> a -> b
$ \IndentLevel
x -> IndentLevel
x IndentLevel -> IndentLevel -> IndentLevel
forall a. Num a => a -> a -> a
- IndentLevel
2
  String -> State IndentLevel String
forall a. a -> StateT IndentLevel Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> State IndentLevel String)
-> String -> State IndentLevel String
forall a b. (a -> b) -> a -> b
$ String
declText String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
exprText

prettyComplex :: ComplexExpression -> IndentState String
prettyComplex :: ComplexExpression -> State IndentLevel String
prettyComplex = \case
  CompApp Identifier'
f AtomicExpression
arg -> String -> State IndentLevel String
forall a. a -> StateT IndentLevel Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> State IndentLevel String)
-> String -> State IndentLevel String
forall a b. (a -> b) -> a -> b
$ String -> String
parens (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ Identifier' -> String
prettyId Identifier'
f String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> AtomicExpression -> String
prettyAtomic AtomicExpression
arg
  CompIte AtomicExpression
c Expression
t Expression
e -> do
    (IndentLevel -> IndentLevel) -> StateT IndentLevel Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (IndentLevel -> IndentLevel -> IndentLevel
forall a. Num a => a -> a -> a
+ IndentLevel
4)
    IndentLevel
indent <- StateT IndentLevel Identity IndentLevel
forall s (m :: * -> *). MonadState s m => m s
get
    let cText :: String
cText = IndentLevel -> String
createIndent (IndentLevel
indent IndentLevel -> IndentLevel -> IndentLevel
forall a. Num a => a -> a -> a
- IndentLevel
2) String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"if " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> AtomicExpression -> String
prettyAtomic AtomicExpression
c
    String
t' <- Expression -> State IndentLevel String
prettyExpr Expression
t
    let tText :: String
tText = IndentLevel -> String
createIndent IndentLevel
indent String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"then " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
t'
    String
e' <- Expression -> State IndentLevel String
prettyExpr Expression
e
    let eText :: String
eText = IndentLevel -> String
createIndent IndentLevel
indent String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"else " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
e'
    (IndentLevel -> IndentLevel) -> StateT IndentLevel Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((IndentLevel -> IndentLevel) -> StateT IndentLevel Identity ())
-> (IndentLevel -> IndentLevel) -> StateT IndentLevel Identity ()
forall a b. (a -> b) -> a -> b
$ \IndentLevel
x -> IndentLevel
x IndentLevel -> IndentLevel -> IndentLevel
forall a. Num a => a -> a -> a
- IndentLevel
4
    String -> State IndentLevel String
forall a. a -> StateT IndentLevel Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> State IndentLevel String)
-> String -> State IndentLevel String
forall a b. (a -> b) -> a -> b
$ String
cText String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
tText String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
eText
  CompBinOp BinaryOperator
op AtomicExpression
lhs AtomicExpression
rhs -> String -> State IndentLevel String
forall a. a -> StateT IndentLevel Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> State IndentLevel String)
-> String -> State IndentLevel String
forall a b. (a -> b) -> a -> b
$ String -> String
parens ([String] -> String
unwords [AtomicExpression -> String
prettyAtomic AtomicExpression
lhs, BinaryOperator -> String
prettyBinOp BinaryOperator
op, AtomicExpression -> String
prettyAtomic AtomicExpression
rhs])
  CompUnOp UnaryOperator
op AtomicExpression
x -> String -> State IndentLevel String
forall a. a -> StateT IndentLevel Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> State IndentLevel String)
-> String -> State IndentLevel String
forall a b. (a -> b) -> a -> b
$ String -> String
parens (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ UnaryOperator -> String
prettyUnOp UnaryOperator
op String -> String -> String
forall a. Semigroup a => a -> a -> a
<> AtomicExpression -> String
prettyAtomic AtomicExpression
x

prettyAtomic :: AtomicExpression -> String
prettyAtomic :: AtomicExpression -> String
prettyAtomic = \case
  AtomId Identifier'
name -> Identifier' -> String
prettyId Identifier'
name
  AtomicExpression
AtomUnit -> String
"()"
  AtomBool Bool
value -> if Bool
value then String
"true" else String
"false"
  AtomInt Int64
value -> Int64 -> String
forall a. Show a => a -> String
show Int64
value

prettyId :: Identifier' -> String
prettyId :: Identifier' -> String
prettyId (Txt Identifier
n) = Identifier -> String
unpack Identifier
n
prettyId (Gen IdCnt
n Identifier
ident) = Identifier -> String
unpack Identifier
ident String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"'" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> IdCnt -> String
forall a. Show a => a -> String
show IdCnt
n

prettyBinOp :: BinaryOperator -> String
prettyBinOp :: BinaryOperator -> String
prettyBinOp = \case
  BoolOp BooleanOperator
AndOp -> String
"&&"
  BoolOp BooleanOperator
OrOp -> String
"||"
  ArithOp ArithmeticOperator
PlusOp -> String
"+"
  ArithOp ArithmeticOperator
MinusOp -> String
"-"
  ArithOp ArithmeticOperator
MulOp -> String
"*"
  ArithOp ArithmeticOperator
DivOp -> String
"/"
  CompOp ComparisonOperator
EqOp -> String
"="
  CompOp ComparisonOperator
NeOp -> String
"<>"
  CompOp ComparisonOperator
LtOp -> String
"<"
  CompOp ComparisonOperator
LeOp -> String
"<="
  CompOp ComparisonOperator
GtOp -> String
">"
  CompOp ComparisonOperator
GeOp -> String
">="

prettyUnOp :: UnaryOperator -> String
prettyUnOp :: UnaryOperator -> String
prettyUnOp UnaryOperator
UnMinusOp = String
"-"

createIndent :: Int -> String
createIndent :: IndentLevel -> String
createIndent IndentLevel
indent = String
"\n" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> IndentLevel -> Char -> String
forall a. IndentLevel -> a -> [a]
replicate IndentLevel
indent Char
' '

parens :: String -> String
parens :: String -> String
parens String
val = String
"(" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
val String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
")"