module Language.HERMIT.PrettyPrinter.Clean
(
corePrettyH
)
where
import Control.Arrow hiding ((<+>))
import Control.Applicative ((<$>))
import Data.List (intersperse)
import Data.Char (isSpace)
import Data.Set (notMember)
import GhcPlugins (TyCon(..), Coercion(..), Var(..), Expr(..))
import qualified GhcPlugins as GHC
import Language.HERMIT.Context
import Language.HERMIT.Core
import Language.HERMIT.GHC
import Language.HERMIT.Kure
import Language.HERMIT.Monad
import Language.HERMIT.Syntax
import Language.HERMIT.PrettyPrinter.Common
import Language.HERMIT.Primitive.Common
import Pair
import Text.PrettyPrint.MarkedHughesPJ as PP
data RetExpr
= RetLam PathH [DocH] PathH DocH
| RetLet PathH [DocH] PathH DocH
| RetApp DocH [RetExpr]
| RetForAll PathH [DocH] PathH DocH
| RetArrowType [DocH]
| RetExpr DocH
| RetAtom DocH
| RetEmpty
retApp :: RetExpr -> RetExpr -> RetExpr
retApp f RetEmpty = f
retApp RetEmpty e = e
retApp (RetApp f es) e = RetApp f (es ++ [e])
retApp f e = RetApp (parenExpr f) [e]
retApps :: RetExpr -> [RetExpr] -> RetExpr
retApps = foldl retApp
retLam :: PathH -> RetExpr -> RetExpr -> RetExpr
retLam _ RetEmpty e = e
retLam p v (RetLam _ vs pb e) = RetLam p (parenExpr v : vs) pb e
retLam p v e = RetLam p [parenExpr v] (p ++ [Lam_Body]) (normalExpr e)
retLet :: PathH -> RetExpr -> RetExpr -> RetExpr
retLet _ RetEmpty body = body
retLet p bnd (RetLet _ bnds pb body) = RetLet p (normalExpr bnd : bnds) pb body
retLet p bnd body = RetLet p [normalExpr bnd] (p ++ [Let_Body]) (normalExpr body)
retForAll :: PathH -> Crumb -> RetExpr -> RetExpr -> RetExpr
retForAll _ _ RetEmpty ty = ty
retForAll p _ v (RetForAll _ vs pb ty) = RetForAll p (parenExpr v : vs) pb ty
retForAll p cr v ty = RetForAll p [parenExpr v] (p ++ [cr]) (normalExpr ty)
retArrowType :: RetExpr -> RetExpr -> RetExpr
retArrowType ty1 (RetArrowType tys) = RetArrowType (parenExprExceptApp ty1 : tys)
retArrowType ty1 ty2 = RetArrowType [parenExprExceptApp ty1, parenExprExceptApp ty2]
isEmptyR :: RetExpr -> Bool
isEmptyR RetEmpty = True
isEmptyR _ = False
isAtom :: RetExpr -> Bool
isAtom (RetAtom _) = True
isAtom _ = False
normalExpr :: RetExpr -> DocH
normalExpr RetEmpty = empty
normalExpr (RetAtom e) = e
normalExpr (RetExpr e) = e
normalExpr (RetLam p vs pb e) = hang (attrP p (specialSymbol LambdaSymbol) <+> hsep vs <+> attrP pb (specialSymbol RightArrowSymbol)) 2 e
normalExpr (RetLet p vs pb e) = sep [ attrP p (keyword "let") <+> vcat vs, attrP pb (keyword "in") <+> e ]
normalExpr (RetApp f es) = let (atoms,exprs) = span isAtom es
in sep [ hsep (f : map normalExpr atoms)
, nest 2 (sep $ map parenExpr exprs) ]
normalExpr (RetForAll p vs pb ty) = attrP p (specialSymbol ForallSymbol) <+> hsep vs <+> attrP pb (symbol '.') <+> ty
normalExpr (RetArrowType tys) = hsep (intersperse typeArrow tys)
parenExpr :: RetExpr -> DocH
parenExpr RetEmpty = empty
parenExpr (RetAtom e) = e
parenExpr (RetApp d []) = d
parenExpr other = ppParens (normalExpr other)
parenExprExceptApp :: RetExpr -> DocH
parenExprExceptApp e@(RetApp _ _) = normalExpr e
parenExprExceptApp e = parenExpr e
attrPAtomExpr :: PathH -> RetExpr -> RetExpr
attrPAtomExpr p (RetAtom d) = RetAtom (attrP p d)
attrPAtomExpr p (RetExpr d) = RetExpr (attrP p d)
attrPAtomExpr _ e = e
ppParens :: DocH -> DocH
ppParens p = symbol '(' <> p <> symbol ')'
specialSymbol :: SpecialSymbol -> DocH
specialSymbol = markColor SyntaxColor . specialFont . char . renderSpecial
symbol :: Char -> DocH
symbol = markColor SyntaxColor . char
keyword :: String -> DocH
keyword = markColor KeywordColor . text
coChar :: Char -> DocH
coChar = coercionColor . char
coSymbol :: SpecialSymbol -> DocH
coSymbol = coercionColor . specialFont . char . renderSpecial
castSymbol :: DocH
castSymbol = coSymbol CastSymbol
coercionSymbol :: DocH
coercionSymbol = coSymbol CoercionSymbol
coercionBindSymbol :: DocH
coercionBindSymbol = coSymbol CoercionBindSymbol
coKeyword :: String -> DocH
coKeyword = coText
tySymbol :: SpecialSymbol -> DocH
tySymbol = typeColor . specialFont . char . renderSpecial
typeSymbol :: DocH
typeSymbol = tySymbol TypeSymbol
typeBindSymbol :: DocH
typeBindSymbol = tySymbol TypeBindSymbol
typeArrow :: DocH
typeArrow = tySymbol RightArrowSymbol
hasType :: DocH -> DocH -> DocH
e `hasType` ty = e <+> tySymbol TypeOfSymbol <+> ty
corePrettyH :: PrettyOptions -> PrettyH CoreTC
corePrettyH opts =
let
ppSDoc :: GHC.Outputable a => PrettyH a
ppSDoc = do dynFlags <- constT GHC.getDynFlags
arr (toDoc . GHC.showSDoc dynFlags . GHC.ppr)
where toDoc s | any isSpace s = parens (text s)
| otherwise = text s
ppVar :: Translate PrettyC HermitM Var RetExpr
ppVar = readerT $ \ v -> GHC.varName ^>> ppName (varColor v)
varColor :: Var -> SyntaxForColor
varColor var | GHC.isTyVar var = TypeColor
| GHC.isCoVar var = CoercionColor
| otherwise = IdColor
ppName :: SyntaxForColor -> Translate PrettyC HermitM GHC.Name RetExpr
ppName color = (GHC.nameOccName >>> GHC.occNameString) ^>> arr (\ name -> let doc = markColor color (text name)
in RetAtom $ if all isScriptInfixIdChar name
then ppParens doc
else doc
)
ppLitTy :: Translate PrettyC HermitM TyLit RetExpr
ppLitTy = arr $ \case
NumTyLit i -> RetAtom $ tyText (show i)
StrTyLit fs -> RetAtom $ tyText (show fs)
ppTyCon :: Translate PrettyC HermitM TyCon RetExpr
ppTyCon = GHC.getName ^>> ppName TypeColor
ppTyConCo :: Translate PrettyC HermitM TyCon RetExpr
ppTyConCo = GHC.getName ^>> ppName CoercionColor
ppBinderMode :: Translate PrettyC HermitM Var RetExpr
ppBinderMode = do v <- idR
if
| GHC.isTyVar v -> case po_exprTypes opts of
Omit -> return (RetEmpty)
Abstract -> return (RetAtom typeBindSymbol)
_ -> ppVar
| GHC.isCoVar v -> case po_coercions opts of
Omit -> return (RetEmpty)
Abstract -> return (RetAtom coercionBindSymbol)
Show -> ppVar
Kind -> do pCoKind <- ppCoKind <<^ CoVarCo
return (RetExpr (coercionBindSymbol `hasType` pCoKind))
| otherwise -> ppVar
ppModGuts :: PrettyH GHC.ModGuts
ppModGuts = do name <- ppSDoc <<^ GHC.mg_module
vtys <- mapT (ppVar &&& (ppType <<^ GHC.idType)) <<< modGutsT progIdsT (\ _ ids -> ids)
let defs = [ normalExpr v `hasType` ty | (v,ty) <- vtys ]
return $ hang (keyword "module" <+> name <+> keyword "where") 2 (vcat defs)
ppCoreProg :: PrettyH CoreProg
ppCoreProg = progConsT ppCoreBind ppCoreProg ($+$) <+ progNilT empty
ppCoreBind :: PrettyH GHC.CoreBind
ppCoreBind = (nonRecT idR (ppCoreExprR &&& ppTypeSig) (,) >>> ppDef)
<+ (do p <- rootPathT
recT (const ppCoreDef) (\ bnds -> attrP p (keyword "rec") <+> vcat bnds)
)
ppCoreAlt :: PrettyH GHC.CoreAlt
ppCoreAlt = do p <- rootPathT
altT (readerT $ \case
GHC.DataAlt dcon -> return (GHC.dataConName dcon) >>> ppName IdColor >>^ parenExpr
GHC.LitAlt lit -> return lit >>> ppSDoc
GHC.DEFAULT -> return (symbol '_')
)
(\ _ -> ppBinderMode)
ppCoreExpr
(\ con vs e -> hang (con <+> hsep (map parenExpr vs) <+> attrP p (specialSymbol RightArrowSymbol)) 2 e)
ppCoreDef :: PrettyH CoreDef
ppCoreDef = defT idR (ppCoreExprR &&& ppTypeSig) (,) >>> ppDef
ppDef :: PrettyH (Var,(RetExpr,DocH))
ppDef = do p <- rootPathT
(v,(e,ty)) <- idR
case po_coercions opts of
Omit | GHC.isCoVar v -> return empty
Kind | GHC.isCoVar v -> return $ case po_exprTypes opts of
Show -> (coercionBindSymbol `hasType` ty) $+$ (coercionBindSymbol <+> symbol '=' <+> coercionSymbol)
_ -> coercionBindSymbol <+> attrP p (symbol '=') <+> normalExpr e
_ -> do pv <- normalExpr ^<< ppBinderMode <<< return v
let pre = pv <+> symbol '='
body = case e of
RetLam p' vs pb e0 -> hang (pre <+> attrP p' (specialSymbol LambdaSymbol) <+> hsep vs <+> attrP pb (specialSymbol RightArrowSymbol)) 2 e0
_ -> hang pre 2 (normalExpr e)
return $ case po_exprTypes opts of
Omit | GHC.isTyVar v -> empty
Show -> (pv `hasType` ty) $+$ body
_ -> body
ppCoreExpr :: PrettyH GHC.CoreExpr
ppCoreExpr = ppCoreExprR >>^ normalExpr
ppCoreExprR :: Translate PrettyC HermitM GHC.CoreExpr RetExpr
ppCoreExprR = rootPathT >>= ppCoreExprPR
where
ppCoreExprPR :: PathH -> Translate PrettyC HermitM GHC.CoreExpr RetExpr
ppCoreExprPR p =
lamT ppBinderMode ppCoreExprR (retLam p)
<+ letT (RetExpr <$> ppCoreBind) ppCoreExprR (retLet p)
<+ (acceptR (\ e -> case e of
App (Type _) (Lam {}) | po_exprTypes opts == Omit -> True
App (App (Var _) (Type _)) (Lam {}) | po_exprTypes opts == Omit -> True
_ -> False) >>>
(appT ppCoreExprR ppCoreExprR (\ (RetAtom e1) (RetLam p' vs pb e0) ->
RetExpr $ hang (e1 <+>
symbol '(' <>
attrP p' (specialSymbol LambdaSymbol) <+>
hsep vs <+>
attrP pb (specialSymbol RightArrowSymbol)) 2 (e0 <> symbol ')')))
)
<+ appT ppCoreExprR ppCoreExprR retApp
<+ caseT ppCoreExpr ppVar ppTypeModeR (const ppCoreAlt) (\ s w ty alts -> RetExpr $ attrP p ((keyword "case" <+> s <+> keyword "of" <+> parenExpr w <+> parenExpr ty) $$ nest 2 (vcat alts)))
<+ varT (attrPAtomExpr p <$> (do (c,i) <- exposeT
if (GHC.isLocalId i) && (i `notMember` boundVars c)
then GHC.varName ^>> ppName WarningColor
else ppVar
)
)
<+ litT ((RetAtom . attrP p) <$> ppSDoc)
<+ typeT (attrPAtomExpr p <$> ppTypeModeR)
<+ coercionT (attrPAtomExpr p <$> ppCoercionModeR)
<+ castT ppCoreExprR ppCoercionModeR (\ e co -> if isEmptyR co
then e
else RetExpr (parenExprExceptApp e <+> attrP p castSymbol <+> parenExprExceptApp co)
)
<+ tickT ppSDoc ppCoreExprR (\ tk e -> RetExpr $ attrP p (text "Tick" $$ nest 2 (tk <+> parenExpr e)))
ppType :: PrettyH Type
ppType = ppTypeR >>^ normalExpr
ppTypeModeR :: Translate PrettyC HermitM Type RetExpr
ppTypeModeR = case po_exprTypes opts of
Omit -> return RetEmpty
Abstract -> return (RetAtom typeSymbol)
_ -> ppTypeR
ppTypeR :: Translate PrettyC HermitM Type RetExpr
ppTypeR = rootPathT >>= ppTypePR
where
ppTypePR :: PathH -> Translate PrettyC HermitM Type RetExpr
ppTypePR p =
tyVarT (attrPAtomExpr p <$> ppVar)
<+ litTyT (attrPAtomExpr p <$> ppLitTy)
<+ appTyT ppTypeR ppTypeR retApp
<+ funTyT ppTypeR ppTypeR retArrowType
<+ forAllTyT ppVar ppTypeR (retForAll p ForAllTy_Body)
<+ tyConAppT (ppTyCon &&& idR) (const ppTypeR)
(\ (pCon,tyCon) tys -> if | GHC.isFunTyCon tyCon && length tys == 2 -> let [ty1,ty2] = tys in retArrowType ty1 ty2
| tyCon == GHC.listTyCon -> RetAtom $ attrP p $ tyText "[" <> (case tys of
[] -> empty
t:_ -> normalExpr t)
<> tyText "]"
| GHC.isTupleTyCon tyCon -> RetAtom $ attrP p $ tyText "(" <> (if null tys
then empty
else foldr1 (\ ty r -> ty <> tyText "," <+> r) (map normalExpr tys)
)
<> tyText ")"
| otherwise -> retApps (attrPAtomExpr p pCon) tys
)
ppCoercion :: PrettyH Coercion
ppCoercion = ppCoercionR >>^ normalExpr
ppCoercionModeR :: Translate PrettyC HermitM Coercion RetExpr
ppCoercionModeR = case po_coercions opts of
Omit -> return RetEmpty
Abstract -> return (RetAtom coercionSymbol)
Show -> ppCoercionR
Kind -> ppCoKind >>^ (\ k -> RetExpr (coercionSymbol `hasType` k))
ppCoercionR :: Translate PrettyC HermitM Coercion RetExpr
ppCoercionR = rootPathT >>= ppCoercionPR
where
ppCoercionPR :: PathH -> Translate PrettyC HermitM Coercion RetExpr
ppCoercionPR p =
reflT (ppTypeModeR >>^ \ ty -> RetAtom $ attrP p $ if isEmptyR ty then coText "refl" else coText "<" <> normalExpr ty <> coText ">")
<+ coVarCoT (attrPAtomExpr p <$> ppVar)
<+ symCoT (ppCoercionR >>^ \ co -> RetExpr (attrP p (coKeyword "sym") <+> parenExpr co))
<+ forAllCoT ppBinderMode ppCoercionR (retForAll p ForAllCo_Body)
<+ transCoT ppCoercionR ppCoercionR (\ co1 co2 -> RetExpr (parenExprExceptApp co1 <+> attrP p (coChar ';') <+> parenExprExceptApp co2))
<+ unsafeCoT ppTypeModeR ppTypeModeR (\ ty1 ty2 -> (if isEmptyR ty1 && isEmptyR ty2 then RetAtom else RetExpr)
(attrP p (coKeyword "unsafe") <+> parenExpr ty1 <+> parenExpr ty2)
)
<+ nthCoT idR ppCoercionR (\ n co -> RetExpr (attrP p (coKeyword "nth") <+> attrP p (coText $ show n) <+> parenExpr co))
<+ instCoT ppCoercionR ppTypeModeR (\ co ty -> if isEmptyR ty
then RetExpr (attrP p (coText "inst") <+> parenExpr co)
else RetExpr (parenExprExceptApp co <+> attrP p (coChar '@') <+> parenExprExceptApp ty)
)
<+ tyConAppCoT (attrPAtomExpr p <$> ppTyConCo) (const ppCoercionR) retApps
<+ appCoT ppCoercionR ppCoercionR retApp
#if __GLASGOW_HASKELL__ > 706
<+ axiomInstCoT (coAxiomName ^>> ppName CoercionColor) (RetAtom <$> ppSDoc) (const ppCoercionR) (\ ax idx coes -> RetExpr (attrP p (coText "axiomInst") <+> attrP p (parenExpr ax) <+> attrP p (parenExpr idx) <+> sep (map parenExpr coes)))
<+ lrCoT (coercionColor <$> ppSDoc) ppCoercionR (\ lr co -> RetExpr (attrP p lr <+> parenExpr co))
#else
<+ axiomInstCoT (coAxiomName ^>> ppName CoercionColor) (const ppCoercionR) (\ ax coes -> RetExpr (attrP p (coText "axiomInst") <+> attrP p (parenExpr ax) <+> sep (map parenExpr coes)))
#endif
ppCoKind :: PrettyH Coercion
ppCoKind = (GHC.coercionKind >>> unPair) ^>> (ppTypeModeR *** ppTypeModeR) >>^ ( \(ty1,ty2) -> parenExprExceptApp ty1 <+> coText "~#" <+> parenExprExceptApp ty2)
ppTypeSig :: PrettyH GHC.CoreExpr
ppTypeSig = coercionT ppCoKind <+ (GHC.exprType ^>> ppType)
in promoteT (ppCoreExpr :: PrettyH GHC.CoreExpr)
<+ promoteT (ppCoreProg :: PrettyH CoreProg)
<+ promoteT (ppCoreBind :: PrettyH GHC.CoreBind)
<+ promoteT (ppCoreDef :: PrettyH CoreDef)
<+ promoteT (ppModGuts :: PrettyH GHC.ModGuts)
<+ promoteT (ppCoreAlt :: PrettyH GHC.CoreAlt)
<+ promoteT (ppType :: PrettyH GHC.Type)
<+ promoteT (ppCoercion :: PrettyH Coercion)