{-# LANGUAGE CPP , GADTs , KindSignatures , DataKinds , ScopedTypeVariables , PatternGuards , Rank2Types , TypeOperators , FlexibleContexts , UndecidableInstances #-} module Language.Hakaru.Pretty.SExpression where #if __GLASGOW_HASKELL__ < 710 import Data.Foldable (foldMap) import Control.Applicative ((<$>)) #endif import Data.Ratio import Data.Text (Text) import Data.Sequence (Seq) import qualified Data.Text as Text import Data.Number.Nat (fromNat) import Data.Number.Natural (fromNonNegativeRational) import qualified Data.List.NonEmpty as L import Data.Text.IO as IO import Language.Hakaru.Command (parseAndInfer) import Language.Hakaru.Syntax.IClasses (jmEq1, TypeEq(..)) import Language.Hakaru.Types.Coercion import Language.Hakaru.Types.DataKind import Language.Hakaru.Types.HClasses import Language.Hakaru.Types.Sing import Language.Hakaru.Summary import Language.Hakaru.Syntax.ABT import Language.Hakaru.Syntax.AST import Language.Hakaru.Syntax.AST.Transforms import Language.Hakaru.Syntax.Datum import Language.Hakaru.Syntax.Reducer import Language.Hakaru.Syntax.TypeCheck import Language.Hakaru.Syntax.TypeOf import Prelude hiding ((<>)) import Text.PrettyPrint (Doc, (<>), (<+>)) import Text.PrettyPrint as PP pretty :: (ABT Term abt) => abt '[] a -> Doc pretty a = PP.brackets (caseVarSyn a prettyVariable prettyTerm <+> PP.colon <+> prettyType (typeOf a)) prettyTerm :: (ABT Term abt) => Term abt a -> Doc prettyTerm (o :$ es) = PP.parens $ prettySCons o es prettyTerm (NaryOp_ op es) = PP.parens $ prettyNary op es prettyTerm (Literal_ v) = prettyLiteral v prettyTerm (Array_ e1 e2) = PP.parens $ (PP.text "array") <+> (caseBind e2 $ \x e2' -> PP.parens (prettyVariable x <+> pretty e1) <+> pretty e2') prettyTerm (Case_ e1 bs) = PP.parens $ PP.text "match" <+> pretty e1 <+> Prelude.foldl (<+>) PP.empty (prettyBranch <$> bs) prettyTerm (Bucket b e r) = PP.parens $ ( PP.text "bucket" <+> pretty b <+> pretty e <+> prettyReducer r) prettyTerm (Reject_ _) = PP.parens $ PP.text "reject" prettyTerm (Empty_ _) = PP.parens $ PP.text "empty" prettyTerm (ArrayLiteral_ es) = PP.parens $ (PP.text "array-literal" <+> foldMap pretty es) prettyTerm (Superpose_ pes) = case pes of (e1,e2) L.:| [] -> PP.parens $ (PP.text "pose" <+> pretty e1 <+> pretty e2) _ -> PP.parens $ (PP.text "superpose" <+> foldMap (\(e1,e2) -> PP.parens (pretty e1 <+> pretty e2)) (L.toList pes)) -- prettyTerm (Datum_ (Datum "true" _typ (Inl Done))) = PP.text "#t" -- prettyTerm (Datum_ (Datum "false" _typ (Inr (Inl Done)))) = PP.text "#f" prettyTerm (Datum_ d) = prettyDatum d prettyDatum :: (ABT Term abt) => Datum (abt '[]) t -> Doc prettyDatum (Datum hint _ d) = PP.parens $ PP.text "datum" <+> (PP.text (Text.unpack hint)) <+> (prettyDatumCode d) prettyDatumCode :: (ABT Term abt) => DatumCode xss (abt '[]) a -> Doc prettyDatumCode (Inr d) = PP.parens $ PP.text "inr" <+> (prettyDatumCode d) prettyDatumCode (Inl d) = PP.parens $ PP.text "inl" <+> (prettyDatumStruct d) prettyDatumStruct :: (ABT Term abt) => DatumStruct xs (abt '[]) a -> Doc prettyDatumStruct Done = PP.text "done" prettyDatumStruct (Et d1 d2) = PP.parens $ PP.text "et" <+> (prettyDatumFun d1) <+> (prettyDatumStruct d2) prettyDatumFun :: (ABT Term abt) => DatumFun x (abt '[]) a -> Doc prettyDatumFun (Konst a) = PP.parens $ PP.text "konst" <+> pretty a prettyDatumFun (Ident a) = PP.parens $ PP.text "ident" <+> pretty a prettyReducer :: (ABT Term abt) => Reducer abt xs a -> Doc prettyReducer (Red_Fanout red_a red_b) = PP.parens (PP.text "r_fanout" <+> prettyReducer red_a <+> prettyReducer red_b) prettyReducer (Red_Index i red_i red_a) = PP.parens (PP.text "r_index" <+> prettyViewABT i <+> prettyViewABT red_i <+> prettyReducer red_a) prettyReducer (Red_Split i red_a red_b) = PP.parens (PP.text "r_split" <+> prettyViewABT i <+> prettyReducer red_a <+> prettyReducer red_b) prettyReducer (Red_Nop) = PP.text "r_nop" prettyReducer (Red_Add _ a) = PP.parens (PP.text "r_add" <+> prettyViewABT a) prettyBranch :: (ABT Term abt) => Branch a abt b -> Doc prettyBranch (Branch pat e) = PP.parens $ prettyPattern pat <+> prettyViewABT e prettyPattern :: Pattern xs a -> Doc prettyPattern PWild = PP.text "*" prettyPattern PVar = PP.text "var" prettyPattern (PDatum hint c) = PP.parens $ PP.text "pdatum" <+> PP.text (Text.unpack hint) <+> goCode c goCode :: PDatumCode xss vars a -> Doc goCode c = PP.parens $ case c of (PInr d) -> PP.text "pc_inr" <+> goCode d (PInl s) -> PP.text "pc_inl" <+> goStruct s goStruct :: PDatumStruct xs vars a -> Doc goStruct s = PP.parens $ case s of (PDone) -> PP.text "ps_done" (PEt f s') -> PP.text "ps_et" <+> goFun f <+> goStruct s' goFun :: PDatumFun x vars a -> Doc goFun f = PP.parens $ case f of (PKonst p) -> PP.text "pf_konst" <+> prettyPattern p (PIdent p) -> PP.text "pf_ident" <+> prettyPattern p prettyViewABT :: (ABT Term abt) => abt xs a -> Doc prettyViewABT = prettyView . viewABT prettyView :: (ABT Term abt) => View (Term abt) xs a -> Doc prettyView (Bind x v) = PP.parens $ PP.text "bind" <+> prettyVariable x <+> prettyView v prettyView (Var x) = prettyVariable x prettyView (Syn t) = pretty (syn t) prettyShow :: (Show a) => a -> Doc prettyShow = PP.text . show prettyLiteral :: Literal a -> Doc prettyLiteral (LNat v) = PP.parens $ PP.text "nat_" <+> prettyShow v prettyLiteral (LInt i) = PP.parens $ PP.text "int_" <+> prettyShow i prettyLiteral (LProb p) = PP.parens $ PP.text "prob_" <+> PP.rational (fromNonNegativeRational p) prettyLiteral (LReal p) = PP.parens $ PP.text "real_" <+> PP.rational p prettyRatio :: (Show a, Integral a) => Ratio a -> Doc prettyRatio r | d == 1 = prettyShow n | otherwise = PP.parens $ PP.text "/" <+> prettyShow n <+> prettyShow d where d = denominator r n = numerator r prettyVariable :: Variable (a :: Hakaru) -> Doc prettyVariable x | Text.null (varHint x) = PP.text "_" <> (PP.int . fromNat .varID) x | otherwise = (PP.text . Text.unpack . varHint) x prettySCons :: (ABT Term abt) => SCon args a -> SArgs abt args -> Doc prettySCons Lam_ (e1 :* End) = caseBind e1 $ \x e1' -> PP.text "fn" <+> prettyVariable x <+> (prettyType $ typeOf e1') <+> pretty e1' prettySCons (PrimOp_ o) es = prettyPrimOp o es prettySCons (ArrayOp_ o) es = prettyArrayOp o es prettySCons (CoerceTo_ o) (e1 :* End) = PP.text (pCoerce o) <+> pretty e1 prettySCons (Summate _ _) (e1 :* e2 :* e3 :* End) = caseBind e3 $ \x e3' -> PP.text "summate" <+> PP.parens (prettyVariable x <+> pretty e1 <+> pretty e2) <+> pretty e3' prettySCons (Product _ _) (e1 :* e2 :* e3 :* End) = caseBind e3 $ \x e3' -> PP.text "product" <+> PP.parens (prettyVariable x <+> pretty e1 <+> pretty e2) <+> pretty e3' prettySCons App_ (e1 :* e2 :* End) = PP.text "app" <+> pretty e1 <+> pretty e2 prettySCons Let_ (e1 :* e2 :* End) = caseBind e2 $ \x e2' -> PP.text "let" <+> PP.parens (prettyVariable x <+> (prettyType $ typeOf e1) <+> pretty e1) <+> pretty e2' prettySCons (UnsafeFrom_ o) (e :* End) = PP.text (pUnsafeCoerce o) <+> pretty e prettySCons (MeasureOp_ o) es = prettyMeasureOp o es prettySCons Dirac (e1 :* End) = PP.text "dirac" <+> pretty e1 prettySCons MBind (e1 :* e2 :* End) = PP.text "mbind" <+> pretty e1 <+> prettyViewABT e2 prettySCons Plate (e1 :* e2 :* End) = PP.text "plate" <+> pretty e1 <+> prettyViewABT e2 prettySCons Chain (e1 :* e2 :* e3 :* End) = PP.text "chain" <+> pretty e1 <+> pretty e2 <+> prettyViewABT e3 prettySCons Integrate (e1 :* e2 :* e3 :* End) = PP.text "integrate" <+> pretty e1 <+> pretty e2 <+> prettyViewABT e3 prettySCons (Transform_ t) _ = PP.text $ Prelude.concat [ "SCons{", show t, "}: TODO" ] prettyMeasureOp :: (ABT Term abt, typs ~ UnLCs args, args ~ LCs typs) => MeasureOp typs a -> SArgs abt args -> Doc prettyMeasureOp Lebesgue = \(e1 :* e2 :* End) -> PP.text "lebesgue" <+> pretty e1 <+> pretty e2 prettyMeasureOp Counting = \End -> PP.text "counting" prettyMeasureOp Categorical = \(e1 :* End) -> PP.text "categorical" <+> pretty e1 prettyMeasureOp Uniform = \(e1 :* e2 :* End) -> PP.text "uniform" <+> pretty e1 <+> pretty e2 prettyMeasureOp Normal = \(e1 :* e2 :* End) -> PP.text "normal" <+> pretty e1 <+> pretty e2 prettyMeasureOp Poisson = \(e1 :* End) -> PP.text "poisson" <+> pretty e1 prettyMeasureOp Gamma = \(e1 :* e2 :* End) -> PP.text "gamma" <+> pretty e1 <+> pretty e2 prettyMeasureOp Beta = \(e1 :* e2 :* End) -> PP.text "beta" <+> pretty e1 <+> pretty e2 pUnsafeCoerce :: Coercion a b -> String pUnsafeCoerce (CCons (Signed HRing_Real) CNil) = "real2prob" pUnsafeCoerce (CCons (Signed HRing_Int) CNil) = "int2nat" pUnsafeCoerce c = "unsafeFrom_" ++ show c pCoerce :: Coercion a b -> String pCoerce (CCons (Signed HRing_Real) CNil) = "prob2real" pCoerce (CCons (Signed HRing_Int) CNil) = "nat2int" pCoerce (CCons (Continuous HContinuous_Real) CNil) = "int2real" pCoerce (CCons (Continuous HContinuous_Prob) CNil) = "nat2prob" pCoerce (CCons (Continuous HContinuous_Prob) (CCons (Signed HRing_Real) CNil)) = "nat2real" pCoerce (CCons (Signed HRing_Int) (CCons (Continuous HContinuous_Real) CNil)) = "nat2real" pCoerce c = "coerceTo_"++show c prettyNary :: (ABT Term abt) => NaryOp a -> Seq (abt '[] a) -> Doc prettyNary And es = PP.text "and" <+> foldMap pretty es prettyNary Or es = PP.text "or" <+> foldMap pretty es prettyNary Xor es = PP.text "xor" <+> foldMap pretty es prettyNary (Sum _) es = PP.text "+" <+> foldMap pretty es prettyNary (Prod _) es = PP.text "*" <+> foldMap pretty es prettyNary (Min _) es = PP.text "min" <+> foldMap pretty es prettyNary (Max _) es = PP.text "max" <+> foldMap pretty es prettyNary _ _ = error "Pretty.SExpression - prettyNary missing cases" prettyType :: Sing (a :: Hakaru) -> Doc prettyType SNat = PP.text "nat" prettyType SInt = PP.text "int" prettyType SProb = PP.text "prob" prettyType SReal = PP.text "real" prettyType (SArray a) = PP.parens $ PP.text "array" <+> prettyType a prettyType (SMeasure a) = PP.parens $ PP.text "measure" <+> prettyType a prettyType (SFun a b) = PP.parens $ prettyType a <+> PP.text "->" <+> prettyType b prettyType typ = case typ of SData (STyCon sym `STyApp` a `STyApp` b) _ | Just Refl <- jmEq1 sym sSymbol_Pair -> PP.parens $ PP.text "pair" <+> prettyType a <+> prettyType b | Just Refl <- jmEq1 sym sSymbol_Either -> PP.parens $ PP.text "either" <+> prettyType a <+> prettyType b SData (STyCon sym `STyApp` a) _ | Just Refl <- jmEq1 sym sSymbol_Maybe -> PP.parens $ PP.text "maybe" <+> prettyType a SData (STyCon sym) _ | Just Refl <- jmEq1 sym sSymbol_Bool -> PP.text "bool" | Just Refl <- jmEq1 sym sSymbol_Unit -> PP.text "unit" _ -> PP.text (showsPrec 11 typ "") prettyPrimOp :: (ABT Term abt, typs ~ UnLCs args, args ~ LCs typs) => PrimOp typs a -> SArgs abt args -> Doc prettyPrimOp Not (e1 :* End) = PP.text "not" <+> pretty e1 prettyPrimOp Pi End = PP.text "pi" prettyPrimOp Sin (e1 :* End) = PP.text "sin" <+> pretty e1 prettyPrimOp Cos (e1 :* End) = PP.text "cos" <+> pretty e1 prettyPrimOp Tan (e1 :* End) = PP.text "tan" <+> pretty e1 prettyPrimOp RealPow (e1 :* e2 :* End) = PP.text "realpow" <+> pretty e1 <+> pretty e2 prettyPrimOp Choose (e1 :* e2 :* End) = PP.text "choose" <+> pretty e1 <+> pretty e2 prettyPrimOp Exp (e1 :* End) = PP.text "exp" <+> pretty e1 prettyPrimOp Log (e1 :* End) = PP.text "log" <+> pretty e1 prettyPrimOp (Infinity _) End = PP.text "infinity" prettyPrimOp GammaFunc (e1 :* End) = PP.text "gammafunc" <+> pretty e1 prettyPrimOp BetaFunc (e1 :* e2 :* End) = PP.text "betafunc" <+> pretty e1 <+> pretty e2 prettyPrimOp (Equal _) (e1 :* e2 :* End) = PP.text "==" <+> pretty e1 <+> pretty e2 prettyPrimOp (Less _) (e1 :* e2 :* End) = PP.text "<" <+> pretty e1 <+> pretty e2 prettyPrimOp (NatPow _) (e1 :* e2 :* End) = PP.text "natpow" <+> pretty e1 <+> pretty e2 prettyPrimOp (Negate _) (e1 :* End) = PP.text "negate" <+> pretty e1 prettyPrimOp (Abs _) (e1 :* End) = PP.text "abs" <+> pretty e1 prettyPrimOp (Recip _) (e1 :* End) = PP.text "recip" <+> pretty e1 prettyPrimOp (NatRoot _) (e1 :* e2 :* End) = PP.text "root" <+> pretty e1 <+> pretty e2 prettyPrimOp Floor (e1 :* End) = PP.text "floor" <+> pretty e1 prettyPrimOp _ _ = error "prettyPrimop: a bunch of cases still need done!" prettyArrayOp :: (ABT Term abt, typs ~ UnLCs args, args ~ LCs typs) => ArrayOp typs a -> SArgs abt args -> Doc prettyArrayOp (Index _) (e1 :* e2 :* End) = PP.text "index" <+> pretty e1 <+> pretty e2 prettyArrayOp (Size _) (e1 :* End) = PP.text "size" <+> pretty e1 prettyArrayOp (Reduce _) _ = error "prettyArrayOp doesn't know how to print Reduce" prettyFile' :: [Char] -> [Char] -> IO () prettyFile' fname outFname = do fileText <- IO.readFile fname prettyText <- runPretty' fileText IO.writeFile outFname (Text.pack prettyText) print prettyText runPretty' :: Text -> IO String runPretty' prog = case parseAndInfer prog of Left _ -> return "err" Right (TypedAST _ ast) -> do summarised <- summary . expandTransformations $ ast return . render . pretty $ summarised fromAst :: Either Text (TypedAST (TrivialABT Term)) -> String fromAst prog = case prog of Left err -> Text.unpack err Right (TypedAST _ ast) -> render . pretty . expandTransformations $ ast