{-# LANGUAGE OverloadedStrings #-}

module JS.Pretty where

import qualified Data.Map as M
import qualified Data.Text as T
import JS.Ast
import Prettyprinter
import Prettyprinter.Render.Text

pp :: (a -> Doc ann) -> a -> T.Text
pp :: (a -> Doc ann) -> a -> Text
pp a -> Doc ann
f = Doc ann -> Text
forall a. Doc a -> Text
render (Doc ann -> Text) -> (a -> Doc ann) -> a -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Doc ann
f

render :: Doc a -> T.Text
render :: Doc a -> Text
render = SimpleDocStream a -> Text
forall ann. SimpleDocStream ann -> Text
renderStrict (SimpleDocStream a -> Text)
-> (Doc a -> SimpleDocStream a) -> Doc a -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LayoutOptions -> Doc a -> SimpleDocStream a
forall ann. LayoutOptions -> Doc ann -> SimpleDocStream ann
layoutPretty LayoutOptions
defaultLayoutOptions

ppFile :: File -> Doc a
ppFile :: File -> Doc a
ppFile (File [Statement]
file) = [Doc a] -> Doc a
forall ann. [Doc ann] -> Doc ann
vsep ([Doc a] -> Doc a) -> [Doc a] -> Doc a
forall a b. (a -> b) -> a -> b
$ Statement -> Doc a
forall a. Statement -> Doc a
ppStatement (Statement -> Doc a) -> [Statement] -> [Doc a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Statement]
file

ppRecord :: (a -> Doc ann) -> Record a -> Doc ann
ppRecord :: (a -> Doc ann) -> Record a -> Doc ann
ppRecord a -> Doc ann
p Record a
r =
  Doc ann -> Doc ann -> Doc ann -> [Doc ann] -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann -> [Doc ann] -> Doc ann
encloseSep Doc ann
"{" Doc ann
"}" Doc ann
", " ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$
    ((Text, Doc ann) -> Doc ann) -> [(Text, Doc ann)] -> [Doc ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
      (\(Text
k, Doc ann
v) -> Doc ann -> Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann -> Doc ann
surround Doc ann
": " (String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (String -> Doc ann) -> String -> Doc ann
forall a b. (a -> b) -> a -> b
$ Text -> String
forall a. Show a => a -> String
show Text
k) Doc ann
v)
      (Map Text (Doc ann) -> [(Text, Doc ann)]
forall k a. Map k a -> [(k, a)]
M.toList ((a -> Doc ann) -> Record a -> Map Text (Doc ann)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Doc ann
p Record a
r))

ppLit :: Lit -> Doc a
ppLit :: Lit -> Doc a
ppLit = \case
  LInt Int
int -> Int -> Doc a
forall a ann. Pretty a => a -> Doc ann
pretty Int
int
  LFloat Float
float -> Float -> Doc a
forall a ann. Pretty a => a -> Doc ann
pretty Float
float
  LString Text
str -> Doc a
"\"" Doc a -> Doc a -> Doc a
forall a. Semigroup a => a -> a -> a
<> Text -> Doc a
forall a ann. Pretty a => a -> Doc ann
pretty Text
str Doc a -> Doc a -> Doc a
forall a. Semigroup a => a -> a -> a
<> Doc a
"\""
  LBool Bool
True -> Doc a
"true"
  LBool Bool
False -> Doc a
"false"

ppSub :: Sub -> Doc a
ppSub :: [Statement] -> Doc a
ppSub [Statement]
sub = [Doc a] -> Doc a
forall ann. [Doc ann] -> Doc ann
vsep (Statement -> Doc a
forall a. Statement -> Doc a
ppStatement (Statement -> Doc a) -> [Statement] -> [Doc a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Statement]
sub)

ppDef :: Definition -> Doc a
ppDef :: Definition -> Doc a
ppDef = \case
  Variable Text
name Expr
expr -> Doc a
"const" Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Text -> Doc a
forall a ann. Pretty a => a -> Doc ann
pretty Text
name Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc a
"=" Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Expr -> Doc a
forall a. Expr -> Doc a
ppExpr Expr
expr Doc a -> Doc a -> Doc a
forall a. Semigroup a => a -> a -> a
<> Doc a
";"
  Function Text
name [Text]
args [Statement]
body ->
    let arguments :: Doc ann
arguments = (Doc ann -> Doc ann -> Doc ann) -> [Doc ann] -> Doc ann
forall (t :: * -> *) ann.
Foldable t =>
(Doc ann -> Doc ann -> Doc ann) -> t (Doc ann) -> Doc ann
concatWith (Doc ann -> Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann -> Doc ann
surround Doc ann
", ") (Text -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Text -> Doc ann) -> [Text] -> [Doc ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Text]
args)
     in Doc a
"const" Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Text -> Doc a
forall a ann. Pretty a => a -> Doc ann
pretty Text
name Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc a
"="
          Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> [Doc a] -> Doc a
forall ann. [Doc ann] -> Doc ann
vsep
            [ Doc a
"function" Doc a -> Doc a -> Doc a
forall a. Semigroup a => a -> a -> a
<> Doc a
"(" Doc a -> Doc a -> Doc a
forall a. Semigroup a => a -> a -> a
<> Doc a
forall ann. Doc ann
arguments Doc a -> Doc a -> Doc a
forall a. Semigroup a => a -> a -> a
<> Doc a
")" Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc a
"{",
              Int -> Doc a -> Doc a
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (Doc a -> Doc a) -> Doc a -> Doc a
forall a b. (a -> b) -> a -> b
$ [Statement] -> Doc a
forall a. [Statement] -> Doc a
ppSub [Statement]
body Doc a -> Doc a -> Doc a
forall a. Semigroup a => a -> a -> a
<> Doc a
";",
              Doc a
"}"
            ]

ppStatement :: Statement -> Doc a
ppStatement :: Statement -> Doc a
ppStatement = \case
  SExpr Expr
expr -> Expr -> Doc a
forall a. Expr -> Doc a
ppExpr Expr
expr
  SRet Expr
expr -> Doc a
"return" Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Expr -> Doc a
forall a. Expr -> Doc a
ppExpr Expr
expr
  SDef Definition
def -> Definition -> Doc a
forall a. Definition -> Doc a
ppDef Definition
def
  SIf Expr
cond [Statement]
sub ->
    [Doc a] -> Doc a
forall ann. [Doc ann] -> Doc ann
vsep
      [ Doc a
"if" Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc a -> Doc a
forall ann. Doc ann -> Doc ann
parens (Expr -> Doc a
forall a. Expr -> Doc a
ppExpr Expr
cond) Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc a
"{",
        Int -> Doc a -> Doc a
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (Doc a -> Doc a) -> Doc a -> Doc a
forall a b. (a -> b) -> a -> b
$ [Statement] -> Doc a
forall a. [Statement] -> Doc a
ppSub [Statement]
sub,
        Doc a
"}"
      ]

ppExpr :: Expr -> Doc a
ppExpr :: Expr -> Doc a
ppExpr = \case
  ELit Lit
lit -> Lit -> Doc a
forall a. Lit -> Doc a
ppLit Lit
lit
  EVar Text
var -> Text -> Doc a
forall a ann. Pretty a => a -> Doc ann
pretty Text
var
  EFun [Text]
args [Statement]
body ->
    let arguments :: Doc ann
arguments = (Doc ann -> Doc ann -> Doc ann) -> [Doc ann] -> Doc ann
forall (t :: * -> *) ann.
Foldable t =>
(Doc ann -> Doc ann -> Doc ann) -> t (Doc ann) -> Doc ann
concatWith (Doc ann -> Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann -> Doc ann
surround Doc ann
", ") (Text -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Text -> Doc ann) -> [Text] -> [Doc ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Text]
args)
     in [Doc a] -> Doc a
forall ann. [Doc ann] -> Doc ann
vsep
          [ Doc a
"function" Doc a -> Doc a -> Doc a
forall a. Semigroup a => a -> a -> a
<> Doc a
"(" Doc a -> Doc a -> Doc a
forall a. Semigroup a => a -> a -> a
<> Doc a
forall ann. Doc ann
arguments Doc a -> Doc a -> Doc a
forall a. Semigroup a => a -> a -> a
<> Doc a
")" Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc a
"{",
            Int -> Doc a -> Doc a
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (Doc a -> Doc a) -> Doc a -> Doc a
forall a b. (a -> b) -> a -> b
$ [Statement] -> Doc a
forall a. [Statement] -> Doc a
ppSub [Statement]
body Doc a -> Doc a -> Doc a
forall a. Semigroup a => a -> a -> a
<> Doc a
";",
            Doc a
"}"
          ]
  EFunCall Expr
fun [Expr]
args ->
    let encloseIfNotSimple :: Doc ann -> Doc ann
encloseIfNotSimple = if Expr -> Bool
isSimpleExpr Expr
fun then Doc ann -> Doc ann
forall a. a -> a
id else Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens
        arguments :: Doc ann
arguments = (Doc ann -> Doc ann -> Doc ann) -> [Doc ann] -> Doc ann
forall (t :: * -> *) ann.
Foldable t =>
(Doc ann -> Doc ann -> Doc ann) -> t (Doc ann) -> Doc ann
concatWith (Doc ann -> Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann -> Doc ann
surround Doc ann
", ") (Expr -> Doc ann
forall a. Expr -> Doc a
ppExpr (Expr -> Doc ann) -> [Expr] -> [Doc ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Expr]
args)
     in Doc a -> Doc a
forall ann. Doc ann -> Doc ann
encloseIfNotSimple (Expr -> Doc a
forall a. Expr -> Doc a
ppExpr Expr
fun) Doc a -> Doc a -> Doc a
forall a. Semigroup a => a -> a -> a
<> Doc a
"(" Doc a -> Doc a -> Doc a
forall a. Semigroup a => a -> a -> a
<> Doc a
forall ann. Doc ann
arguments Doc a -> Doc a -> Doc a
forall a. Semigroup a => a -> a -> a
<> Doc a
")"
  ERecord Record Expr
record -> (Expr -> Doc a) -> Record Expr -> Doc a
forall a ann. (a -> Doc ann) -> Record a -> Doc ann
ppRecord Expr -> Doc a
forall a. Expr -> Doc a
ppExpr Record Expr
record
  EAnd [Expr]
exprs -> (Doc a -> Doc a -> Doc a) -> [Doc a] -> Doc a
forall (t :: * -> *) ann.
Foldable t =>
(Doc ann -> Doc ann -> Doc ann) -> t (Doc ann) -> Doc ann
concatWith (Doc a -> Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann -> Doc ann
surround Doc a
" && ") (Expr -> Doc a
forall a. Expr -> Doc a
ppExpr (Expr -> Doc a) -> [Expr] -> [Doc a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Expr]
exprs)
  EEqual Expr
a Expr
b -> Expr -> Doc a
forall a. Expr -> Doc a
ppExpr (Expr -> Doc a) -> Expr -> Doc a
forall a b. (a -> b) -> a -> b
$ Text -> Expr -> Expr -> Expr
EBinOp Text
"===" Expr
a Expr
b
  ERecordAccess Expr
expr Text
label ->
    (if Expr -> Bool
isSimpleExpr Expr
expr then Doc a -> Doc a
forall a. a -> a
id else Doc a -> Doc a
forall ann. Doc ann -> Doc ann
parens) (Expr -> Doc a
forall a. Expr -> Doc a
ppExpr Expr
expr)
      Doc a -> Doc a -> Doc a
forall a. Semigroup a => a -> a -> a
<> Doc a
"."
      Doc a -> Doc a -> Doc a
forall a. Semigroup a => a -> a -> a
<> Text -> Doc a
forall a ann. Pretty a => a -> Doc ann
pretty Text
label
  EBinOp Text
op Expr
a Expr
b -> Expr -> Doc a
forall a. Expr -> Doc a
ppExpr Expr
a Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Text -> Doc a
forall a ann. Pretty a => a -> Doc ann
pretty Text
op Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Expr -> Doc a
forall a. Expr -> Doc a
ppExpr Expr
b
  ERaw Text
raw -> Text -> Doc a
forall a ann. Pretty a => a -> Doc ann
pretty Text
raw

-- need surrounding parens?
isSimpleExpr :: Expr -> Bool
isSimpleExpr :: Expr -> Bool
isSimpleExpr = \case
  ERecord {} -> Bool
False
  EFun {} -> Bool
False
  EEqual {} -> Bool
False
  EBinOp {} -> Bool
False
  ERaw {} -> Bool
False
  Expr
_ -> Bool
True