{-# LANGUAGE OverloadedStrings #-}

{-| Custom width-aware pretty-printer for Template Haskell declarations.

Replaces @TH.pprint@ + @formatTHOutput@ for golden test output, using the
@prettyprinter@ library's Wadler\/Lindig algorithm to produce readable output
that respects a configurable page width (default 100 columns). -}
module PlutusTx.Test.THPretty (pprintDecs) where

import Prelude

import Data.Char (isAlphaNum)
import Language.Haskell.TH qualified as TH
import Prettyprinter
import Prettyprinter.Render.String (renderString)

-- | Pretty-print a list of TH declarations with width-aware layout.
pprintDecs :: [TH.Dec] -> String
pprintDecs :: [Dec] -> String
pprintDecs [Dec]
decs = SimpleDocStream Any -> String
forall ann. SimpleDocStream ann -> String
renderString (SimpleDocStream Any -> String) -> SimpleDocStream Any -> String
forall a b. (a -> b) -> a -> b
$ LayoutOptions -> Doc Any -> SimpleDocStream Any
forall ann. LayoutOptions -> Doc ann -> SimpleDocStream ann
layoutSmart LayoutOptions
opts (Doc Any -> SimpleDocStream Any) -> Doc Any -> SimpleDocStream Any
forall a b. (a -> b) -> a -> b
$ [Doc Any] -> Doc Any
forall ann. [Doc ann] -> Doc ann
vsep ((Dec -> Doc Any) -> [Dec] -> [Doc Any]
forall a b. (a -> b) -> [a] -> [b]
map Dec -> Doc Any
forall ann. Dec -> Doc ann
ppDec [Dec]
decs)
  where
    opts :: LayoutOptions
opts = PageWidth -> LayoutOptions
LayoutOptions (Int -> Double -> PageWidth
AvailablePerLine Int
100 Double
1.0)

------------------------------------------------------------
-- Declarations
------------------------------------------------------------

ppDec :: TH.Dec -> Doc ann
ppDec :: forall ann. Dec -> Doc ann
ppDec (TH.InstanceD Maybe Overlap
_ Cxt
cxt Type
typ [Dec]
decs) =
  Cxt -> Type -> Doc ann
forall ann. Cxt -> Type -> Doc ann
ppInstanceHead Cxt
cxt Type
typ Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
line Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 ([Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
vsep ((Dec -> Doc ann) -> [Dec] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Dec -> Doc ann
forall ann. Dec -> Doc ann
ppDec [Dec]
decs))
ppDec (TH.FunD Name
name [Clause]
clauses) =
  [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
vsep ((Clause -> Doc ann) -> [Clause] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map (Name -> Clause -> Doc ann
forall ann. Name -> Clause -> Doc ann
ppClause Name
name) [Clause]
clauses)
ppDec (TH.PragmaD Pragma
pragma) = Pragma -> Doc ann
forall ann. Pragma -> Doc ann
ppPragma Pragma
pragma
ppDec (TH.SigD Name
name Type
typ) = Name -> Doc ann
forall ann. Name -> Doc ann
ppNamePrefix Name
name Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"::" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Type -> Doc ann
forall ann. Type -> Doc ann
ppType Type
typ
ppDec Dec
d = String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Dec -> String
forall a. Ppr a => a -> String
TH.pprint Dec
d) -- fallback

ppInstanceHead :: TH.Cxt -> TH.Type -> Doc ann
ppInstanceHead :: forall ann. Cxt -> Type -> Doc ann
ppInstanceHead Cxt
cxt Type
typ =
  Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
group (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$ Doc ann
"instance" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Cxt -> Doc ann
forall ann. Cxt -> Doc ann
ppCxt Cxt
cxt Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Type -> Doc ann
forall ann. Type -> Doc ann
ppType Type
typ Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"where"

ppCxt :: TH.Cxt -> Doc ann
ppCxt :: forall ann. Cxt -> Doc ann
ppCxt [] = Doc ann
forall a. Monoid a => a
mempty
ppCxt [Type
t] = Type -> Doc ann
forall ann. Type -> Doc ann
ppType Type
t Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"=>" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
softline
ppCxt Cxt
ts =
  Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
group (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
", " ((Type -> Doc ann) -> Cxt -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Doc ann
forall ann. Type -> Doc ann
ppType Cxt
ts))
    Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"=>"
    Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
softline

------------------------------------------------------------
-- Pragmas
------------------------------------------------------------

ppPragma :: TH.Pragma -> Doc ann
ppPragma :: forall ann. Pragma -> Doc ann
ppPragma (TH.InlineP Name
name Inline
TH.Inlinable RuleMatch
TH.FunLike Phases
TH.AllPhases) =
  Doc ann
"{-#" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"INLINABLE" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Name -> Doc ann
forall ann. Name -> Doc ann
ppNamePrefix Name
name Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"#-}"
ppPragma (TH.InlineP Name
name Inline
TH.Inline RuleMatch
TH.FunLike Phases
TH.AllPhases) =
  Doc ann
"{-#" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"INLINE" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Name -> Doc ann
forall ann. Name -> Doc ann
ppNamePrefix Name
name Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"#-}"
ppPragma (TH.InlineP Name
name Inline
TH.NoInline RuleMatch
TH.FunLike Phases
TH.AllPhases) =
  Doc ann
"{-#" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"NOINLINE" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Name -> Doc ann
forall ann. Name -> Doc ann
ppNamePrefix Name
name Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"#-}"
ppPragma Pragma
p = String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Dec -> String
forall a. Ppr a => a -> String
TH.pprint (Pragma -> Dec
TH.PragmaD Pragma
p)) -- fallback

------------------------------------------------------------
-- Clauses and matches
------------------------------------------------------------

ppClause :: TH.Name -> TH.Clause -> Doc ann
ppClause :: forall ann. Name -> Clause -> Doc ann
ppClause Name
name (TH.Clause [Pat]
pats Body
body [Dec]
wheres) =
  Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
group (Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
hang Int
2 ([Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
sep (Name -> Doc ann
forall ann. Name -> Doc ann
ppNamePrefix Name
name Doc ann -> [Doc ann] -> [Doc ann]
forall a. a -> [a] -> [a]
: (Pat -> Doc ann) -> [Pat] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Pat -> Doc ann
forall ann. Pat -> Doc ann
ppPat [Pat]
pats)) Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Body -> Doc ann
forall ann. Body -> Doc ann
ppBody Body
body)
    Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> [Dec] -> Doc ann
forall ann. [Dec] -> Doc ann
ppWheres [Dec]
wheres

ppBody :: TH.Body -> Doc ann
ppBody :: forall ann. Body -> Doc ann
ppBody (TH.NormalB Exp
expr) = Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
nest Int
2 (Doc ann
forall ann. Doc ann
softline Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
"=" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExp Exp
expr)
ppBody (TH.GuardedB [(Guard, Exp)]
guards) = Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
nest Int
2 ([Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
vsep (((Guard, Exp) -> Doc ann) -> [(Guard, Exp)] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map (Guard, Exp) -> Doc ann
forall {ann}. (Guard, Exp) -> Doc ann
ppGuard [(Guard, Exp)]
guards))
  where
    ppGuard :: (Guard, Exp) -> Doc ann
ppGuard (Guard
guard', Exp
expr) = Doc ann
forall ann. Doc ann
softline Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
"|" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Guard -> Doc ann
forall {ann}. Guard -> Doc ann
ppGuardExp Guard
guard' Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"=" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExp Exp
expr
    ppGuardExp :: Guard -> Doc ann
ppGuardExp (TH.NormalG Exp
e) = Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExp Exp
e
    ppGuardExp (TH.PatG [Stmt]
stmts) = [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
hsep ((Stmt -> Doc ann) -> [Stmt] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map (String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (String -> Doc ann) -> (Stmt -> String) -> Stmt -> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stmt -> String
forall a. Ppr a => a -> String
TH.pprint) [Stmt]
stmts)

ppWheres :: [TH.Dec] -> Doc ann
ppWheres :: forall ann. [Dec] -> Doc ann
ppWheres [] = Doc ann
forall a. Monoid a => a
mempty
ppWheres [Dec]
ds = Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
nest Int
2 (Doc ann
forall ann. Doc ann
line Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
"where" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
line Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
vsep ((Dec -> Doc ann) -> [Dec] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Dec -> Doc ann
forall ann. Dec -> Doc ann
ppDec [Dec]
ds))

ppMatch :: TH.Match -> Doc ann
ppMatch :: forall ann. Match -> Doc ann
ppMatch (TH.Match Pat
pat Body
body [Dec]
wheres) =
  Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
group (Pat -> Doc ann
forall ann. Pat -> Doc ann
ppPat Pat
pat Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Body -> Doc ann
forall ann. Body -> Doc ann
ppBody Body
body) Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> [Dec] -> Doc ann
forall ann. [Dec] -> Doc ann
ppWheres [Dec]
wheres

------------------------------------------------------------
-- Types
------------------------------------------------------------

ppType :: TH.Type -> Doc ann
ppType :: forall ann. Type -> Doc ann
ppType (TH.ForallT [TyVarBndr Specificity]
bndrs Cxt
cxt Type
typ) =
  Doc ann
"forall" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
hsep ((TyVarBndr Specificity -> Doc ann)
-> [TyVarBndr Specificity] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map TyVarBndr Specificity -> Doc ann
forall a ann. TyVarBndr a -> Doc ann
ppTyVarBndr [TyVarBndr Specificity]
bndrs) Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
"." Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Cxt -> Doc ann
forall ann. Cxt -> Doc ann
ppCxt Cxt
cxt Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Type -> Doc ann
forall ann. Type -> Doc ann
ppType Type
typ
ppType (TH.AppT (TH.AppT Type
TH.ArrowT Type
a) Type
b) =
  Int -> Type -> Doc ann
forall ann. Int -> Type -> Doc ann
ppTypePrec Int
1 Type
a Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"->" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Type -> Doc ann
forall ann. Type -> Doc ann
ppType Type
b
ppType (TH.AppT Type
TH.ListT Type
a) =
  Doc ann
"[" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Type -> Doc ann
forall ann. Type -> Doc ann
ppType Type
a Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
"]"
ppType t :: Type
t@(TH.AppT Type
_ Type
_) = Type -> Doc ann
forall ann. Type -> Doc ann
ppTypeApp Type
t
ppType (TH.ConT Name
name) = Name -> Doc ann
forall ann. Name -> Doc ann
ppName Name
name
ppType (TH.VarT Name
name) = Name -> Doc ann
forall ann. Name -> Doc ann
ppName Name
name
ppType Type
TH.ArrowT = Doc ann
"(->)"
ppType (TH.TupleT Int
0) = Doc ann
"()"
ppType (TH.TupleT Int
n) = Doc ann
"(" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Int -> Char -> String
forall a. Int -> a -> [a]
replicate (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Char
',') Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
")"
ppType Type
TH.ListT = Doc ann
"[]"
ppType (TH.SigT Type
t Type
k) = Type -> Doc ann
forall ann. Type -> Doc ann
ppType Type
t Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"::" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Type -> Doc ann
forall ann. Type -> Doc ann
ppType Type
k
ppType (TH.ParensT Type
t) = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Type -> Doc ann
forall ann. Type -> Doc ann
ppType Type
t)
ppType Type
t = String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Type -> String
forall a. Ppr a => a -> String
TH.pprint Type
t) -- fallback

ppTypePrec :: Int -> TH.Type -> Doc ann
ppTypePrec :: forall ann. Int -> Type -> Doc ann
ppTypePrec Int
p t :: Type
t@(TH.AppT (TH.AppT Type
TH.ArrowT Type
_) Type
_)
  | Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Type -> Doc ann
forall ann. Type -> Doc ann
ppType Type
t)
ppTypePrec Int
p t :: Type
t@(TH.AppT Type
_ Type
_)
  | Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2, Bool -> Bool
not (Type -> Bool
isSingleArgApp Type
t) = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Type -> Doc ann
forall ann. Type -> Doc ann
ppType Type
t)
ppTypePrec Int
_ Type
t = Type -> Doc ann
forall ann. Type -> Doc ann
ppType Type
t

{-| Is this a type application with exactly one argument (e.g. @Maybe a@)?
Single-arg applications at precedence 2 don't need parens. -}
isSingleArgApp :: TH.Type -> Bool
isSingleArgApp :: Type -> Bool
isSingleArgApp (TH.AppT Type
f Type
_) = case Type
f of
  TH.ConT {} -> Bool
True
  TH.VarT {} -> Bool
True
  Type
_ -> Bool
False
isSingleArgApp Type
_ = Bool
False

ppTypeApp :: TH.Type -> Doc ann
ppTypeApp :: forall ann. Type -> Doc ann
ppTypeApp Type
t = case Type -> (Type, Cxt)
collectTypeApps Type
t of
  (TH.TupleT Int
_, Cxt
args) -> [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
tupled ((Type -> Doc ann) -> Cxt -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Doc ann
forall ann. Type -> Doc ann
ppType Cxt
args)
  (Type
fun, Cxt
args) -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
group (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$ Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
nest Int
2 (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$ Type -> Doc ann
forall ann. Type -> Doc ann
ppType Type
fun Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
softline Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
sep ((Type -> Doc ann) -> Cxt -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Type -> Doc ann
forall ann. Int -> Type -> Doc ann
ppTypePrec Int
2) Cxt
args)

collectTypeApps :: TH.Type -> (TH.Type, [TH.Type])
collectTypeApps :: Type -> (Type, Cxt)
collectTypeApps = Cxt -> Type -> (Type, Cxt)
go []
  where
    go :: Cxt -> Type -> (Type, Cxt)
go Cxt
acc (TH.AppT Type
f Type
a) = Cxt -> Type -> (Type, Cxt)
go (Type
a Type -> Cxt -> Cxt
forall a. a -> [a] -> [a]
: Cxt
acc) Type
f
    go Cxt
acc Type
t = (Type
t, Cxt
acc)

ppTyVarBndr :: TH.TyVarBndr a -> Doc ann
ppTyVarBndr :: forall a ann. TyVarBndr a -> Doc ann
ppTyVarBndr (TH.PlainTV Name
name a
_) = Name -> Doc ann
forall ann. Name -> Doc ann
ppName Name
name
ppTyVarBndr (TH.KindedTV Name
name a
_ Type
k) = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Name -> Doc ann
forall ann. Name -> Doc ann
ppName Name
name Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"::" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Type -> Doc ann
forall ann. Type -> Doc ann
ppType Type
k)

------------------------------------------------------------
-- Patterns
------------------------------------------------------------

ppPat :: TH.Pat -> Doc ann
ppPat :: forall ann. Pat -> Doc ann
ppPat (TH.VarP Name
name) = Name -> Doc ann
forall ann. Name -> Doc ann
ppName Name
name
ppPat Pat
TH.WildP = Doc ann
"_"
ppPat (TH.ConP Name
name Cxt
_ [Pat]
pats)
  | [Pat] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Pat]
pats = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Name -> Doc ann
forall ann. Name -> Doc ann
ppNamePrefix Name
name)
  | Bool
otherwise =
      Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$
        Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
group (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$
          Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
nest Int
2 (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$
            Name -> Doc ann
forall ann. Name -> Doc ann
ppNamePrefix Name
name Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
softline Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
sep ((Pat -> Doc ann) -> [Pat] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Pat -> Doc ann
forall ann. Pat -> Doc ann
ppPat [Pat]
pats)
ppPat (TH.LitP Lit
lit) = Lit -> Doc ann
forall ann. Lit -> Doc ann
ppLit Lit
lit
ppPat (TH.TupP [Pat]
pats) = [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
tupled ((Pat -> Doc ann) -> [Pat] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Pat -> Doc ann
forall ann. Pat -> Doc ann
ppPat [Pat]
pats)
ppPat (TH.BangP Pat
pat) = Doc ann
"!" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Pat -> Doc ann
forall ann. Pat -> Doc ann
ppPat Pat
pat
ppPat (TH.ViewP Exp
expr Pat
pat) = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExp Exp
expr Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"->" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Pat -> Doc ann
forall ann. Pat -> Doc ann
ppPat Pat
pat)
ppPat (TH.ParensP Pat
pat) = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Pat -> Doc ann
forall ann. Pat -> Doc ann
ppPat Pat
pat)
ppPat (TH.SigP Pat
pat Type
typ) = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Pat -> Doc ann
forall ann. Pat -> Doc ann
ppPat Pat
pat Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"::" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Type -> Doc ann
forall ann. Type -> Doc ann
ppType Type
typ)
ppPat (TH.AsP Name
name Pat
pat) = Name -> Doc ann
forall ann. Name -> Doc ann
ppName Name
name Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
"@" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Pat -> Doc ann
forall ann. Pat -> Doc ann
ppPat Pat
pat
ppPat (TH.InfixP Pat
p1 Name
name Pat
p2) = Pat -> Doc ann
forall ann. Pat -> Doc ann
ppPat Pat
p1 Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Name -> Doc ann
forall ann. Name -> Doc ann
ppName Name
name Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Pat -> Doc ann
forall ann. Pat -> Doc ann
ppPat Pat
p2
ppPat (TH.ListP [Pat]
pats) = [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
list ((Pat -> Doc ann) -> [Pat] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Pat -> Doc ann
forall ann. Pat -> Doc ann
ppPat [Pat]
pats)
ppPat Pat
p = String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Pat -> String
forall a. Ppr a => a -> String
TH.pprint Pat
p) -- fallback

------------------------------------------------------------
-- Expressions
------------------------------------------------------------

ppExp :: TH.Exp -> Doc ann
ppExp :: forall ann. Exp -> Doc ann
ppExp (TH.VarE Name
name) = Name -> Doc ann
forall ann. Name -> Doc ann
ppName Name
name
ppExp (TH.ConE Name
name) = Name -> Doc ann
forall ann. Name -> Doc ann
ppName Name
name
ppExp (TH.LitE Lit
lit) = Lit -> Doc ann
forall ann. Lit -> Doc ann
ppLit Lit
lit
ppExp (TH.AppE Exp
f Exp
x) = Exp -> Exp -> Doc ann
forall ann. Exp -> Exp -> Doc ann
ppAppExp Exp
f Exp
x
ppExp e :: Exp
e@TH.InfixE {} = Exp -> Doc ann
forall ann. Exp -> Doc ann
ppInfixExp Exp
e
ppExp (TH.LamE [Pat]
pats Exp
body) =
  Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
group (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$ Doc ann
"\\" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
hsep ((Pat -> Doc ann) -> [Pat] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Pat -> Doc ann
forall ann. Pat -> Doc ann
ppPat [Pat]
pats) Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"->" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
nest Int
2 (Doc ann
forall ann. Doc ann
softline Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExp Exp
body)
ppExp (TH.LamCaseE [Match]
matches) =
  Doc ann
"\\case" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
hardline Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 ([Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
bracesSemiList ((Match -> Doc ann) -> [Match] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Match -> Doc ann
forall ann. Match -> Doc ann
ppMatch [Match]
matches))
ppExp (TH.CaseE Exp
scrut [Match]
matches) =
  Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
group (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$
    Doc ann
"case"
      Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExp Exp
scrut
      Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"of"
      Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
hardline
      Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 ([Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
bracesSemiList ((Match -> Doc ann) -> [Match] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Match -> Doc ann
forall ann. Match -> Doc ann
ppMatch [Match]
matches))
ppExp (TH.LetE [Dec]
decs Exp
body) =
  Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
group (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$ Doc ann
"let" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
align ([Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
vsep ((Dec -> Doc ann) -> [Dec] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Dec -> Doc ann
forall ann. Dec -> Doc ann
ppDec [Dec]
decs)) Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
softline Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
"in" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExp Exp
body
ppExp (TH.ListE [Exp]
exprs) = [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
list ((Exp -> Doc ann) -> [Exp] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExp [Exp]
exprs)
ppExp (TH.TupE [Maybe Exp]
mexprs) = [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
tupled ((Maybe Exp -> Doc ann) -> [Maybe Exp] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map (Doc ann -> (Exp -> Doc ann) -> Maybe Exp -> Doc ann
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Doc ann
forall a. Monoid a => a
mempty Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExp) [Maybe Exp]
mexprs)
ppExp (TH.SigE Exp
expr Type
typ) = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExp Exp
expr Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"::" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Type -> Doc ann
forall ann. Type -> Doc ann
ppType Type
typ)
ppExp (TH.ParensE Exp
expr) = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExp Exp
expr)
ppExp Exp
e = String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Exp -> String
forall a. Ppr a => a -> String
TH.pprint Exp
e) -- fallback

-- | Pretty-print function application, collecting the spine.
ppAppExp :: TH.Exp -> TH.Exp -> Doc ann
ppAppExp :: forall ann. Exp -> Exp -> Doc ann
ppAppExp Exp
f Exp
x =
  case Exp -> (Exp, [Exp])
collectApps (Exp -> Exp -> Exp
TH.AppE Exp
f Exp
x) of
    (Exp
fun, [Exp]
args) -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
group (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$ Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
nest Int
2 (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$ Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExpPrec Exp
fun Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
softline Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
sep ((Exp -> Doc ann) -> [Exp] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExpPrec [Exp]
args)

collectApps :: TH.Exp -> (TH.Exp, [TH.Exp])
collectApps :: Exp -> (Exp, [Exp])
collectApps = [Exp] -> Exp -> (Exp, [Exp])
go []
  where
    go :: [Exp] -> Exp -> (Exp, [Exp])
go [Exp]
acc (TH.AppE Exp
f Exp
a) = [Exp] -> Exp -> (Exp, [Exp])
go (Exp
a Exp -> [Exp] -> [Exp]
forall a. a -> [a] -> [a]
: [Exp]
acc) Exp
f
    go [Exp]
acc Exp
e = (Exp
e, [Exp]
acc)

-- | Wrap in parens if the expression is a compound form.
ppExpPrec :: TH.Exp -> Doc ann
ppExpPrec :: forall ann. Exp -> Doc ann
ppExpPrec e :: Exp
e@TH.AppE {} = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExp Exp
e)
ppExpPrec e :: Exp
e@TH.InfixE {} = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExp Exp
e)
ppExpPrec e :: Exp
e@TH.LamE {} = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExp Exp
e)
ppExpPrec e :: Exp
e@TH.LamCaseE {} = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExp Exp
e)
ppExpPrec e :: Exp
e@TH.CaseE {} = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExp Exp
e)
ppExpPrec e :: Exp
e@TH.LetE {} = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExp Exp
e)
ppExpPrec Exp
e = Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExp Exp
e

-- | Pretty-print an infix expression, handling chained operators.
ppInfixExp :: TH.Exp -> Doc ann
ppInfixExp :: forall ann. Exp -> Doc ann
ppInfixExp (TH.InfixE (Just Exp
l) Exp
op (Just Exp
r)) =
  Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
group (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$ Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
nest Int
2 (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$ Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExpPrec Exp
l Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
softline Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExp Exp
op Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Exp -> Doc ann
forall ann. Exp -> Doc ann
ppInfixRight Exp
r
ppInfixExp (TH.InfixE Maybe Exp
Nothing Exp
op (Just Exp
r)) =
  Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExp Exp
op Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExpPrec Exp
r)
ppInfixExp (TH.InfixE (Just Exp
l) Exp
op Maybe Exp
Nothing) =
  Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExpPrec Exp
l Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExp Exp
op)
ppInfixExp Exp
e = Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExp Exp
e

{-| For the right side of an infix, don't add extra parens if it's the
same infix chain (like nested &&). -}
ppInfixRight :: TH.Exp -> Doc ann
ppInfixRight :: forall ann. Exp -> Doc ann
ppInfixRight = Exp -> Doc ann
forall ann. Exp -> Doc ann
ppExpPrec

------------------------------------------------------------
-- Literals
------------------------------------------------------------

ppLit :: TH.Lit -> Doc ann
ppLit :: forall ann. Lit -> Doc ann
ppLit (TH.IntegerL Integer
n) = Integer -> Doc ann
forall ann. Integer -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty Integer
n
ppLit (TH.StringL String
s) = String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (String -> String
forall a. Show a => a -> String
show String
s)
ppLit (TH.CharL Char
c) = String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Char -> String
forall a. Show a => a -> String
show Char
c)
ppLit (TH.RationalL Rational
r) = String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Rational -> String
forall a. Show a => a -> String
show Rational
r)
ppLit Lit
l = String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Exp -> String
forall a. Ppr a => a -> String
TH.pprint (Lit -> Exp
TH.LitE Lit
l)) -- fallback

------------------------------------------------------------
-- Names
------------------------------------------------------------

-- | Render a name. Operators are rendered without parens (for infix use).
ppName :: TH.Name -> Doc ann
ppName :: forall ann. Name -> Doc ann
ppName Name
name = String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Name -> String
forall a. Ppr a => a -> String
TH.pprint Name
name)

-- | Render a name in prefix position. Operators get wrapped in parens.
ppNamePrefix :: TH.Name -> Doc ann
ppNamePrefix :: forall ann. Name -> Doc ann
ppNamePrefix Name
name
  | Name -> Bool
isOperator Name
name = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Name -> Doc ann
forall ann. Name -> Doc ann
ppName Name
name)
  | Bool
otherwise = Name -> Doc ann
forall ann. Name -> Doc ann
ppName Name
name

-- | Check if a TH name is an operator (symbol-based, not alphanumeric).
isOperator :: TH.Name -> Bool
isOperator :: Name -> Bool
isOperator Name
name =
  case Name -> String
TH.nameBase Name
name of
    [] -> Bool
False
    (Char
c : String
_) -> Bool -> Bool
not (Char -> Bool
isAlphaNum Char
c) Bool -> Bool -> Bool
&& Char
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
'_'

------------------------------------------------------------
-- Helpers
------------------------------------------------------------

-- | Render items in braces with semicolons (used for case alternatives).
bracesSemiList :: [Doc ann] -> Doc ann
bracesSemiList :: forall ann. [Doc ann] -> Doc ann
bracesSemiList [] = Doc ann
"{}"
bracesSemiList [Doc ann
d] = Doc ann
"{" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
d Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
"}"
bracesSemiList [Doc ann]
ds = Doc ann
"{" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> (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
a Doc ann
b -> Doc ann
a Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
hardline Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
" " Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
b) [Doc ann]
ds Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
"}"