{-# LANGUAGE GADTs
, OverloadedStrings
, KindSignatures
, DataKinds
, FlexibleContexts
, UndecidableInstances
, LambdaCase
#-}
{-# OPTIONS_GHC -Wall -fwarn-tabs #-}
module Language.Hakaru.Pretty.Haskell
(
pretty
, prettyString
, prettyPrec
, prettyAssoc
, prettyPrecAssoc
, prettyType
, ppVariable
, ppVariables
, ppBinder
, ppCoerceTo
, ppUnsafeFrom
, ppRatio
, Associativity(..)
, ppBinop
, Pretty(..)
) where
import Data.Ratio
import Text.PrettyPrint (Doc, (<>), (<+>))
import qualified Text.PrettyPrint as PP
import qualified Data.Foldable as F
import qualified Data.List.NonEmpty as L
import qualified Data.Text as Text
import qualified Data.Sequence as Seq
import Prelude hiding ((<>))
import Data.Number.Nat (fromNat)
import Data.Number.Natural (fromNatural)
import Language.Hakaru.Syntax.IClasses (fmap11, foldMap11, List1(..)
,Foldable21(..))
import Language.Hakaru.Types.DataKind
import Language.Hakaru.Types.Coercion
import Language.Hakaru.Types.HClasses
import Language.Hakaru.Types.Sing
import Language.Hakaru.Syntax.AST
import Language.Hakaru.Syntax.Datum
import Language.Hakaru.Syntax.Reducer
import Language.Hakaru.Syntax.ABT
pretty :: (ABT Term abt) => abt '[] a -> Doc
pretty = prettyPrec 0
prettyString :: (ABT Term abt)
=> Sing a
-> abt '[] a
-> Doc
prettyString typ ast =
PP.text $ Text.unpack (Text.unlines $ header ++ [ Text.pack (prettyProg "prog" typ ast)])
prettyProg :: (ABT Term abt)
=> String
-> Sing a
-> abt '[] a
-> String
prettyProg name typ ast =
PP.renderStyle PP.style
( PP.sep [PP.text (name ++ " ::"), PP.nest 2 (prettyType typ)]
PP.$+$ PP.sep [PP.text (name ++ " =") , PP.nest 2 (pretty ast)] )
prettyPrec :: (ABT Term abt) => Int -> abt '[] a -> Doc
prettyPrec p = toDoc . prettyPrec_ p . LC_
prettyAssoc :: (ABT Term abt) => Assoc (abt '[]) -> Doc
prettyAssoc = prettyPrecAssoc 0
prettyPrecAssoc :: (ABT Term abt) => Int -> Assoc (abt '[]) -> Doc
prettyPrecAssoc p (Assoc x e) =
toDoc $ ppFun p "Assoc"
[ ppVariable x
, prettyPrec 11 e
]
prettyType :: Sing (a :: Hakaru) -> Doc
prettyType SInt = PP.text "Int"
prettyType SNat = PP.text "Int"
prettyType SReal = PP.text "Double"
prettyType SProb = PP.text "Prob"
prettyType (SArray t) =
let t' = PP.nest 2 (prettyType t) in
PP.parens (PP.sep [PP.text "MayBoxVec", t', t'])
prettyType (SMeasure t) =
PP.parens (PP.sep [PP.text "Measure", PP.nest 2 (prettyType t)])
prettyType (SFun t1 t2) =
PP.parens (PP.sep [prettyType t1 <+> PP.text "->", prettyType t2])
prettyType (SData _ (SDone `SPlus` SVoid)) =
PP.text "()"
prettyType (SData _ (SDone `SPlus` SDone `SPlus` SVoid)) =
PP.text "Bool"
prettyType (SData _ (SDone `SPlus` (SKonst t `SEt` SDone) `SPlus` SVoid)) =
PP.parens (PP.sep [PP.text "Maybe", PP.nest 2 (prettyType t)])
prettyType (SData _ ((SKonst t1 `SEt` SDone) `SPlus`
(SKonst t2 `SEt` SDone) `SPlus` SVoid)) =
PP.parens (PP.sep [PP.text "Either", PP.nest 2 (prettyType t1),
PP.nest 2 (prettyType t2)])
prettyType (SData _ ((SKonst t1 `SEt` SKonst t2 `SEt` SDone) `SPlus` SVoid)) =
PP.parens (PP.sep [prettyType t1 <> PP.comma, prettyType t2])
prettyType s = error ("TODO: prettyType: " ++ show s)
class Pretty (f :: Hakaru -> *) where
prettyPrec_ :: Int -> f a -> Docs
type Docs = [Doc]
toDoc :: Docs -> Doc
toDoc = PP.sep
ppVariable :: Variable (a :: Hakaru) -> Doc
ppVariable x = hint <> (PP.int . fromNat . varID) x
where
hint
| Text.null (varHint x) = PP.char 'x'
| otherwise = (PP.text . Text.unpack . varHint) x
ppVariables :: List1 Variable (xs :: [Hakaru]) -> Docs
ppVariables = ppList . go
where
go :: List1 Variable (xs :: [Hakaru]) -> Docs
go Nil1 = []
go (Cons1 x xs) = ppVariable x : go xs
ppBinder :: (ABT Term abt) => abt xs a -> Docs
ppBinder e =
case ppViewABT e of
([], body) -> body
(vars,body) -> PP.char '\\' <+> PP.sep vars <+> PP.text "->" : body
ppUncurryBinder :: (ABT Term abt) => abt xs a -> Docs
ppUncurryBinder e =
case ppViewABT e of
(vars,body) -> PP.char '\\' <+> unc vars <+> PP.text "->" : body
where
unc :: [Doc] -> Doc
unc [] = PP.text "()"
unc (x:xs) = PP.parens (x <> PP.comma <> unc xs)
ppViewABT :: (ABT Term abt) => abt xs a -> ([Doc], Docs)
ppViewABT e = go [] (viewABT e)
where
go :: (ABT Term abt) => [Doc] -> View (Term abt) xs a -> ([Doc],Docs)
go xs (Syn t) = (reverse xs, prettyPrec_ 0 (LC_ (syn t)))
go xs (Var x) = (reverse xs, [ppVariable x])
go xs (Bind x v) =
let x' = if True
then ppVariable x
else PP.char '_'
in go (x' : xs) v
instance (ABT Term abt) => Pretty (LC_ abt) where
prettyPrec_ p (LC_ e) =
caseVarSyn e ((:[]) . ppVariable) $ \t ->
case t of
o :$ es -> ppSCon p o es
NaryOp_ o es ->
let prettyNaryOp :: NaryOp a -> (String, Int, Maybe String)
prettyNaryOp And = ("&&", 3, Just "true")
prettyNaryOp Or = ("||", 2, Just "false")
prettyNaryOp Xor = ("`xor`", 0, Just "false")
prettyNaryOp Iff = ("`iff`", 0, Just "true")
prettyNaryOp (Min _) = ("`min`", 5, Nothing)
prettyNaryOp (Max _) = ("`max`", 5, Nothing)
prettyNaryOp (Sum _) = ("+", 6, Just "zero")
prettyNaryOp (Prod _) = ("*", 7, Just "one")
in
let (opName,opPrec,maybeIdentity) = prettyNaryOp o in
if Seq.null es
then
case maybeIdentity of
Just identity -> [PP.text identity]
Nothing ->
ppFun p "syn"
[ toDoc $ ppFun 11 "NaryOp_"
[ PP.text (showsPrec 11 o "")
, PP.text "(Seq.fromList [])"
]]
else
parens (p > opPrec)
. PP.punctuate (PP.space <> PP.text opName)
. map (prettyPrec opPrec)
$ F.toList es
Literal_ v -> prettyPrec_ p v
Empty_ _ -> [PP.text "empty"]
Array_ e1 e2 ->
ppFun 11 "array"
[ ppArg e1 <+> PP.char '$'
, toDoc $ ppBinder e2
]
ArrayLiteral_ es -> ppFun 11 "arrayLit" (ppList $ map (prettyPrec 0) es)
Datum_ d -> prettyPrec_ p (fmap11 LC_ d)
Case_ e1 bs ->
ppFun p "case_"
[ ppArg e1
, toDoc $ ppList (map (toDoc . prettyPrec_ 0) bs)
]
Bucket b ee r ->
ppFun p "bucket"
[ ppArg b
, ppArg ee
, toDoc $ parens True (prettyPrec_ p r)
]
Superpose_ pes ->
case pes of
(e1,e2) L.:| [] ->
ppFun 11 "pose"
[ ppArg e1 <+> PP.char '$'
, ppArg e2
]
_ ->
ppFun p "superpose"
[ toDoc
. ppList
. map (\(e1,e2) -> ppTuple [pretty e1, pretty e2])
$ L.toList pes
]
Reject_ _ -> [PP.text "reject"]
ppSCon :: (ABT Term abt) => Int -> SCon args a -> SArgs abt args -> Docs
ppSCon p Lam_ = \(e1 :* End) ->
parens (p > 0) $ adjustHead (PP.text "lam $" <+>) (ppBinder e1)
ppSCon p App_ = \(e1 :* e2 :* End) -> ppBinop "`app`" 9 LeftAssoc p e1 e2
ppSCon p Let_ = \(e1 :* e2 :* End) ->
parens (p > 0) $
adjustHead
(PP.text "let_" <+> ppArg e1 <+> PP.char '$' <+>)
(ppBinder e2)
ppSCon p (PrimOp_ o) = \es -> ppPrimOp p o es
ppSCon p (ArrayOp_ o) = \es -> ppArrayOp p o es
ppSCon p (CoerceTo_ c) = \(e1 :* End) -> ppCoerceTo p c e1
ppSCon p (UnsafeFrom_ c) = \(e1 :* End) -> ppUnsafeFrom p c e1
ppSCon p (MeasureOp_ o) = \es -> ppMeasureOp p o es
ppSCon p Dirac = \(e1 :* End) -> ppApply1 p "dirac" e1
ppSCon p MBind = \(e1 :* e2 :* End) ->
parens (p > 1) $
adjustHead
(prettyPrec 1 e1 <+> PP.text ">>=" <+>)
(ppBinder e2)
ppSCon p (Transform_ t) = ppTransform p t
ppSCon p Integrate = \(e1 :* e2 :* e3 :* End) ->
ppFun p "integrate"
[ ppArg e1
, ppArg e2
, toDoc $ parens True (ppBinder e3)
]
ppSCon p (Summate _ _) = \(e1 :* e2 :* e3 :* End) ->
ppFun p "summate"
[ ppArg e1
, ppArg e2
, toDoc $ parens True (ppBinder e3)
]
ppSCon p (Product _ _) = \(e1 :* e2 :* e3 :* End) ->
ppFun p "product"
[ ppArg e1
, ppArg e2
, toDoc $ parens True (ppBinder e3)
]
ppSCon _ Plate = \(e1 :* e2 :* End) ->
ppFun 11 "plate"
[ ppArg e1 <+> PP.char '$'
, toDoc $ ppBinder e2
]
ppSCon _ Chain = \(e1 :* e2 :* e3 :* End) ->
ppFun 11 "chain"
[ ppArg e1
, ppArg e2 <+> PP.char '$'
, toDoc $ ppBinder e3
]
ppTransform :: (ABT Term abt)
=> Int -> Transform args a -> SArgs abt args -> Docs
ppTransform p t es =
case t of
Expect ->
case es of
e1 :* e2 :* End ->
parens (p > 0) $
adjustHead
(PP.text "expect" <+> ppArg e1 <+> PP.char '$' <+>)
(ppBinder e2)
_ -> ppApply p (transformName t) es
ppCoerceTo :: ABT Term abt => Int -> Coercion a b -> abt '[] a -> Docs
ppCoerceTo =
\p c e -> ppFun p (prettyShow c) [ppArg e]
where
prettyShow (CCons (Signed HRing_Real) CNil) = "fromProb"
prettyShow (CCons (Signed HRing_Int) CNil) = "nat2int"
prettyShow (CCons (Continuous HContinuous_Real) CNil) = "fromInt"
prettyShow (CCons (Continuous HContinuous_Prob) CNil) = "nat2prob"
prettyShow (CCons (Continuous HContinuous_Prob)
(CCons (Signed HRing_Real) CNil)) = "nat2real"
prettyShow (CCons (Signed HRing_Int)
(CCons (Continuous HContinuous_Real) CNil)) = "nat2real"
prettyShow c = "coerceTo_ " ++ showsPrec 11 c ""
ppUnsafeFrom :: ABT Term abt => Int -> Coercion a b -> abt '[] b -> Docs
ppUnsafeFrom =
\p c e -> ppFun p (prettyShow c) [ppArg e]
where
prettyShow (CCons (Signed HRing_Real) CNil) = "unsafeProb"
prettyShow (CCons (Signed HRing_Int) CNil) = "unsafeNat"
prettyShow c = "unsafeFrom_ " ++ showsPrec 11 c ""
ppPrimOp
:: (ABT Term abt, typs ~ UnLCs args, args ~ LCs typs)
=> Int -> PrimOp typs a -> SArgs abt args -> Docs
ppPrimOp p Not = \(e1 :* End) -> ppApply1 p "not" e1
ppPrimOp p Impl = \(e1 :* e2 :* End) ->
ppFun p "syn"
[ toDoc $ ppFun 11 "Impl"
[ ppArg e1
, ppArg e2
]]
ppPrimOp p Diff = \(e1 :* e2 :* End) ->
ppFun p "syn"
[ toDoc $ ppFun 11 "Diff"
[ ppArg e1
, ppArg e2
]]
ppPrimOp p Nand = \(e1 :* e2 :* End) -> ppApply2 p "nand" e1 e2
ppPrimOp p Nor = \(e1 :* e2 :* End) -> ppApply2 p "nor" e1 e2
ppPrimOp _ Pi = \End -> [PP.text "pi"]
ppPrimOp p Sin = \(e1 :* End) -> ppApply1 p "sin" e1
ppPrimOp p Cos = \(e1 :* End) -> ppApply1 p "cos" e1
ppPrimOp p Tan = \(e1 :* End) -> ppApply1 p "tan" e1
ppPrimOp p Asin = \(e1 :* End) -> ppApply1 p "asin" e1
ppPrimOp p Acos = \(e1 :* End) -> ppApply1 p "acos" e1
ppPrimOp p Atan = \(e1 :* End) -> ppApply1 p "atan" e1
ppPrimOp p Sinh = \(e1 :* End) -> ppApply1 p "sinh" e1
ppPrimOp p Cosh = \(e1 :* End) -> ppApply1 p "cosh" e1
ppPrimOp p Tanh = \(e1 :* End) -> ppApply1 p "tanh" e1
ppPrimOp p Asinh = \(e1 :* End) -> ppApply1 p "asinh" e1
ppPrimOp p Acosh = \(e1 :* End) -> ppApply1 p "acosh" e1
ppPrimOp p Atanh = \(e1 :* End) -> ppApply1 p "atanh" e1
ppPrimOp p RealPow = \(e1 :* e2 :* End) -> ppBinop "**" 8 RightAssoc p e1 e2
ppPrimOp p Choose = \(e1 :* e2 :* End) -> ppApply2 p "choose" e1 e2
ppPrimOp p Exp = \(e1 :* End) -> ppApply1 p "exp" e1
ppPrimOp p Log = \(e1 :* End) -> ppApply1 p "log" e1
ppPrimOp _ (Infinity _) = \End -> [PP.text "infinity"]
ppPrimOp p GammaFunc = \(e1 :* End) -> ppApply1 p "gammaFunc" e1
ppPrimOp p BetaFunc = \(e1 :* e2 :* End) -> ppApply2 p "betaFunc" e1 e2
ppPrimOp p (Equal _) = \(e1 :* e2 :* End) -> ppBinop "==" 4 NonAssoc p e1 e2
ppPrimOp p (Less _) = \(e1 :* e2 :* End) -> ppBinop "<" 4 NonAssoc p e1 e2
ppPrimOp p (NatPow _) = \(e1 :* e2 :* End) -> ppBinop "^" 8 RightAssoc p e1 e2
ppPrimOp p (Negate _) = \(e1 :* End) -> ppApply1 p "negate" e1
ppPrimOp p (Abs _) = \(e1 :* End) -> ppApply1 p "abs_" e1
ppPrimOp p (Signum _) = \(e1 :* End) -> ppApply1 p "signum" e1
ppPrimOp p (Recip _) = \(e1 :* End) -> ppApply1 p "recip" e1
ppPrimOp p (NatRoot _) = \(e1 :* e2 :* End) ->
ppBinop "`thRootOf`" 9 LeftAssoc p e2 e1
ppPrimOp p (Erf _) = \(e1 :* End) -> ppApply1 p "erf" e1
ppPrimOp p Floor = \(e1 :* End) -> ppApply1 p "floor" e1
ppArrayOp
:: (ABT Term abt, typs ~ UnLCs args, args ~ LCs typs)
=> Int -> ArrayOp typs a -> SArgs abt args -> Docs
ppArrayOp p (Index _) = \(e1 :* e2 :* End) ->
ppBinop "!" 9 LeftAssoc p e1 e2
ppArrayOp p (Size _) = \(e1 :* End) ->
ppApply1 p "size" e1
ppArrayOp p (Reduce _) = \(e1 :* e2 :* e3 :* End) ->
ppFun p "reduce"
[ ppArg e1
, ppArg e2
, ppArg e3
]
ppMeasureOp
:: (ABT Term abt, typs ~ UnLCs args, args ~ LCs typs)
=> Int -> MeasureOp typs a -> SArgs abt args -> Docs
ppMeasureOp p Lebesgue = \(e1 :* e2 :* End) -> ppApply2 p "lebesgue" e1 e2
ppMeasureOp _ Counting = \End -> [PP.text "counting"]
ppMeasureOp p Categorical = \(e1 :* End) -> ppApply1 p "categorical" e1
ppMeasureOp p Uniform = \(e1 :* e2 :* End) -> ppApply2 p "uniform" e1 e2
ppMeasureOp p Normal = \(e1 :* e2 :* End) -> ppApply2 p "normal" e1 e2
ppMeasureOp p Poisson = \(e1 :* End) -> ppApply1 p "poisson" e1
ppMeasureOp p Gamma = \(e1 :* e2 :* End) -> ppApply2 p "gamma" e1 e2
ppMeasureOp p Beta = \(e1 :* e2 :* End) -> ppApply2 p "beta" e1 e2
instance Pretty Literal where
prettyPrec_ p (LNat n) = ppFun p "nat_" [PP.integer (fromNatural n)]
prettyPrec_ p (LInt i) = ppFun p "int_" [PP.integer i]
prettyPrec_ p (LProb l) = ppFun p "prob_" [ppRatio 11 l]
prettyPrec_ p (LReal r) = ppFun p "real_" [ppRatio 11 r]
instance Pretty f => Pretty (Datum f) where
prettyPrec_ p (Datum hint _typ d)
| Text.null hint =
ppFun p "datum_"
[error "TODO: prettyPrec_@Datum"]
| otherwise =
ppFun p "ann_"
[ PP.parens . PP.text . show $ _typ
, PP.parens . toDoc $ ppFun p (Text.unpack hint)
(foldMap11 ((:[]) . toDoc . prettyPrec_ 11) d)
]
ppPattern :: Int -> Pattern xs a -> Docs
ppPattern _ PWild = [PP.text "PWild"]
ppPattern _ PVar = [PP.text "PVar"]
ppPattern p (PDatum hint d0)
| Text.null hint = error "TODO: prettyPrec_@Pattern"
| otherwise = ppFun p ("p" ++ Text.unpack hint) (goCode d0)
where
goCode :: PDatumCode xss vars a -> Docs
goCode (PInr d) = goCode d
goCode (PInl d) = goStruct d
goStruct :: PDatumStruct xs vars a -> Docs
goStruct PDone = []
goStruct (PEt d1 d2) = goFun d1 ++ goStruct d2
goFun :: PDatumFun x vars a -> Docs
goFun (PKonst d) = [toDoc $ ppPattern 11 d]
goFun (PIdent d) = [toDoc $ ppPattern 11 d]
instance Pretty (Pattern xs) where
prettyPrec_ = ppPattern
instance (ABT Term abt) => Pretty (Branch a abt) where
prettyPrec_ p (Branch pat e) =
ppFun p "branch"
[ toDoc $ prettyPrec_ 11 pat
, PP.parens . toDoc $ ppBinder e
]
instance (ABT Term abt) => Pretty (Reducer abt xs) where
prettyPrec_ p (Red_Fanout r1 r2) =
ppFun p "r_fanout"
[ toDoc $ prettyPrec_ 11 r1
, toDoc $ prettyPrec_ 11 r2
]
prettyPrec_ p (Red_Index n o e) =
ppFun p "r_index"
[ toDoc $ parens True $ ppUncurryBinder n
, toDoc $ parens True $ ppUncurryBinder o
, toDoc $ prettyPrec_ 11 e
]
prettyPrec_ p (Red_Split b r1 r2) =
ppFun p "r_split"
[ toDoc $ parens True (ppUncurryBinder b)
, toDoc $ prettyPrec_ 11 r1
, toDoc $ prettyPrec_ 11 r2
]
prettyPrec_ _ Red_Nop =
[ PP.text "r_nop" ]
prettyPrec_ p (Red_Add _ e) =
ppFun p "r_add"
[ toDoc $ parens True (ppUncurryBinder e)]
adjustHead :: (Doc -> Doc) -> Docs -> Docs
adjustHead f [] = [f (toDoc [])]
adjustHead f (d:ds) = f d : ds
parens :: Bool -> Docs -> Docs
parens True ds = [PP.parens (PP.nest 1 (toDoc ds))]
parens False ds = ds
ppList :: [Doc] -> Docs
ppList = (:[]) . PP.brackets . PP.nest 1 . PP.fsep . PP.punctuate PP.comma
ppTuple :: [Doc] -> Doc
ppTuple = PP.parens . PP.sep . PP.punctuate PP.comma
ppFun :: Int -> String -> [Doc] -> Docs
ppFun _ f [] = [PP.text f]
ppFun p f ds =
parens (p > 9) [PP.text f <+> PP.nest (1 + length f) (PP.sep ds)]
ppArg :: (ABT Term abt) => abt '[] a -> Doc
ppArg = prettyPrec 11
ppApply1 :: (ABT Term abt) => Int -> String -> abt '[] a -> Docs
ppApply1 p f e1 = ppFun p f [ppArg e1]
ppApply2
:: (ABT Term abt) => Int -> String -> abt '[] a -> abt '[] b -> Docs
ppApply2 p f e1 e2 = ppFun p f [ppArg e1, ppArg e2]
ppApply
:: (ABT Term abt) => Int -> String -> SArgs abt as -> Docs
ppApply p f es = ppFun p f $ foldMap21 ppBinder es
ppRatio :: (Show a, Integral a) => Int -> Ratio a -> Doc
ppRatio p r
| d == 1 = ppShowS $ showsPrec p n
| n < 0 =
ppShowS
. showParen (p > 7)
$ showChar '-'
. showsPrec 8 (negate n)
. showChar '/'
. showsPrec 8 d
| otherwise =
ppShowS
. showParen (p > 7)
$ showsPrec 8 n
. showChar '/'
. showsPrec 8 d
where
d = denominator r
n = numerator r
ppShowS s = PP.text (s [])
data Associativity = LeftAssoc | RightAssoc | NonAssoc
ppBinop
:: (ABT Term abt)
=> String -> Int -> Associativity
-> Int -> abt '[] a -> abt '[] b -> Docs
ppBinop op p0 assoc =
let (p1,p2) =
case assoc of
LeftAssoc -> (p0, 1 + p0)
RightAssoc -> (1 + p0, p0)
NonAssoc -> (1 + p0, 1 + p0)
in \p e1 e2 ->
parens (p > p0)
[ prettyPrec p1 e1
, PP.text op
<+> prettyPrec p2 e2
]
header :: [Text.Text]
header =
[ "{-# LANGUAGE DataKinds, NegativeLiterals #-}"
, "module Prog where"
, ""
, "import Data.Number.LogFloat (LogFloat)"
, "import Prelude hiding (product, exp, log, (**), pi)"
, "import Language.Hakaru.Runtime.LogFloatPrelude"
, "import Language.Hakaru.Runtime.CmdLine"
, "import Language.Hakaru.Types.Sing"
, "import qualified System.Random.MWC as MWC"
, "import Control.Monad"
, "import System.Environment (getArgs)"
, "" ]