Use extension field in pretty-printing

lamagraph-compiler/src/Lamagraph/Compiler/PrettyAst.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE UndecidableInstances #-}
 {-# OPTIONS_GHC -Wno-orphans #-}
 
 {- | Orphan instances for pretty-printing AST
@@ -8,21 +9,46 @@ module Lamagraph.Compiler.PrettyAst () where
 
 import Relude
 
-import Control.Lens
+import Control.Lens hiding (Empty)
 import Prettyprinter
+import Prettyprinter.Internal
 
 import Lamagraph.Compiler.Extension
 import Lamagraph.Compiler.Parser.SrcLoc
 import Lamagraph.Compiler.Syntax
 
+----------------------
+-- Helper functions --
+----------------------
+
 spaceNumber :: Int
 spaceNumber = 1
 
+-- | Works like a normal 'vsep', but has stupid check for 'Empty' 'Doc' to help with empty lines
+smartVsep :: [Doc ann] -> Doc ann
+smartVsep = vsep . filter (\case Empty -> False; _ -> True)
+
+-- | Works like a normal 'hsep', but has stupid check for 'Empty' 'Doc' to help with redundant spaces
+smartHsep :: [Doc ann] -> Doc ann
+smartHsep = hsep . filter (\case Empty -> False; _ -> True)
+
 -- Default instance for 'Pretty (Maybe a)' prints empty string for 'Nothing'
 prettyMaybe :: (Pretty a) => Maybe a -> Doc ann
 prettyMaybe Nothing = parens "Nothing"
 prettyMaybe (Just a) = parens $ vsep ["Just", pretty a]
 
+-----------------------
+-- Utility instances --
+-----------------------
+
+instance Pretty DataConCantHappen where
+  pretty :: DataConCantHappen -> Doc ann
+  pretty _ = emptyDoc
+
+instance Pretty NoExtField where
+  pretty :: NoExtField -> Doc ann
+  pretty _ = emptyDoc
+
 instance Pretty RealSrcSpan where
   pretty :: RealSrcSpan -> Doc ann
   pretty rss = pretty _srcSpanFile <> ":" <> lineCol
@@ -65,84 +91,88 @@ instance {-# OVERLAPPING #-} Pretty (Located Text) where
    where
     inner = vsep ["L", pretty loc, dquotes $ pretty text]
 
+------------------------
+-- Main AST instances --
+------------------------
+
 instance Pretty RecFlag where
   pretty :: RecFlag -> Doc ann
   pretty Recursive = "Rec"
   pretty NonRecursive = "NonRec"
 
-instance Pretty (LmlDecl (LmlcPass pass)) where
+instance (ForallLmlDecl Pretty (LmlcPass pass)) => Pretty (LmlDecl (LmlcPass pass)) where
   pretty :: LmlDecl (LmlcPass pass) -> Doc ann
-  pretty (OpenD _ decl) = vsep ["OpenD", pretty decl]
-  pretty (ValD _ bindGroup) = vsep ["ValD", pretty bindGroup]
-  pretty (TyD _ decls) = vsep ["TyD", list $ map pretty (toList decls)]
+  pretty (OpenD ext decl) = smartVsep ["OpenD", pretty ext, pretty decl]
+  pretty (ValD ext bindGroup) = smartVsep ["ValD", pretty ext, pretty bindGroup]
+  pretty (TyD ext decls) = smartVsep ["TyD", pretty ext, list $ map pretty (toList decls)]
 
-instance Pretty (OpenDecl (LmlcPass pass)) where
+instance (ForallOpenDecl Pretty (LmlcPass pass)) => Pretty (OpenDecl (LmlcPass pass)) where
   pretty :: OpenDecl (LmlcPass pass) -> Doc ann
-  pretty (OpenDecl _ ident) = vsep ["OpenDecl", pretty ident]
+  pretty (OpenDecl ext ident) = smartVsep ["OpenDecl", pretty ext, pretty ident]
 
-instance Pretty (TyDecl (LmlcPass pass)) where
+instance (ForallTyDecl Pretty (LmlcPass pass)) => Pretty (TyDecl (LmlcPass pass)) where
   pretty :: TyDecl (LmlcPass pass) -> Doc ann
-  pretty (AliasDecl _ name vars ty) = vsep ["AliasDecl", pretty name, list $ map pretty vars, pretty ty]
-  pretty (DataDecl _ name vars constrs) = vsep ["DataDecl", pretty name, list $ map pretty vars, list $ map pretty constrs]
+  pretty (AliasDecl ext name vars ty) = smartVsep ["AliasDecl", pretty ext, pretty name, list $ map pretty vars, pretty ty]
+  pretty (DataDecl ext name vars constrs) = smartVsep ["DataDecl", pretty ext, pretty name, list $ map pretty vars, list $ map pretty constrs]
 
-instance Pretty (ConDecl (LmlcPass pass)) where
+instance (ForallConDecl Pretty (LmlcPass pass)) => Pretty (ConDecl (LmlcPass pass)) where
   pretty :: ConDecl (LmlcPass pass) -> Doc ann
-  pretty (ConDecl _ name args) = vsep ["ConDecl", pretty name, list (map pretty args)]
+  pretty (ConDecl ext name args) = smartVsep ["ConDecl", pretty ext, pretty name, list (map pretty args)]
 
-instance Pretty (LmlExpr (LmlcPass pass)) where
+instance (ForallLmlExpr Pretty (LmlcPass pass)) => Pretty (LmlExpr (LmlcPass pass)) where
   pretty :: LmlExpr (LmlcPass pass) -> Doc ann
-  pretty (LmlExprIdent _ ident) = vsep ["ExprIdent", pretty ident]
-  pretty (LmlExprConstant _ constant) = "ExprConstant" <+> pretty constant
-  pretty (LmlExprLet _ bindGroup expr) = vsep ["ExprLet", pretty bindGroup, pretty expr]
-  pretty (LmlExprFunction _ pat expr) = vsep ["ExprFunction", pretty pat, pretty expr]
-  pretty (LmlExprApply _ expr exprs) = vsep ["ExprApply", pretty expr, list $ map pretty (toList exprs)]
-  pretty (LmlExprMatch _ expr cases) = vsep ["ExprMatch", pretty expr, list $ map pretty (toList cases)]
-  pretty (LmlExprTuple _ expr exprs) = vsep ["ExprTuple", list $ map pretty (expr : toList exprs)]
-  pretty (LmlExprConstruct _ constr expr) = vsep ["ExprConstruct", pretty constr, prettyMaybe expr]
-  pretty (LmlExprIfThenElse _ cond t f) = vsep ["ExprITE", pretty cond, pretty t, pretty f]
-  pretty (LmlExprConstraint _ expr ty) = vsep ["ExprConstraint", pretty expr, pretty ty]
-
-instance Pretty (LmlBindGroup (LmlcPass pass)) where
+  pretty (LmlExprIdent ext ident) = smartVsep ["ExprIdent", pretty ext, pretty ident]
+  pretty (LmlExprConstant ext constant) = smartHsep ["ExprConstant", pretty ext, pretty constant]
+  pretty (LmlExprLet ext bindGroup expr) = smartVsep ["ExprLet", pretty ext, pretty bindGroup, pretty expr]
+  pretty (LmlExprFunction ext pat expr) = smartVsep ["ExprFunction", pretty ext, pretty pat, pretty expr]
+  pretty (LmlExprApply ext expr exprs) = smartVsep ["ExprApply", pretty ext, pretty expr, list $ map pretty (toList exprs)]
+  pretty (LmlExprMatch ext expr cases) = smartVsep ["ExprMatch", pretty ext, pretty expr, list $ map pretty (toList cases)]
+  pretty (LmlExprTuple ext expr exprs) = smartVsep ["ExprTuple", pretty ext, list $ map pretty (expr : toList exprs)]
+  pretty (LmlExprConstruct ext constr expr) = smartVsep ["ExprConstruct", pretty ext, pretty constr, prettyMaybe expr]
+  pretty (LmlExprIfThenElse ext cond t f) = smartVsep ["ExprITE", pretty ext, pretty cond, pretty t, pretty f]
+  pretty (LmlExprConstraint ext expr ty) = smartVsep ["ExprConstraint", pretty ext, pretty expr, pretty ty]
+
+instance (ForallLmlBindGroup Pretty (LmlcPass pass)) => Pretty (LmlBindGroup (LmlcPass pass)) where
   pretty :: LmlBindGroup (LmlcPass pass) -> Doc ann
-  pretty (LmlBindGroup _ recFlag binds) = vsep [pretty recFlag, list $ map pretty (toList binds)]
+  pretty (LmlBindGroup ext recFlag binds) = smartVsep [pretty ext, pretty recFlag, list $ map pretty (toList binds)]
 
-instance Pretty (LmlBind (LmlcPass pass)) where
+instance (ForallLmlBind Pretty (LmlcPass pass)) => Pretty (LmlBind (LmlcPass pass)) where
   pretty :: LmlBind (LmlcPass pass) -> Doc ann
-  pretty (LmlBind _ pat expr) = vsep ["Bind", pretty pat, pretty expr]
+  pretty (LmlBind ext pat expr) = smartVsep ["Bind", pretty ext, pretty pat, pretty expr]
 
-instance Pretty (LmlCase (LmlcPass pass)) where
+instance (ForallLmlCase Pretty (LmlcPass pass)) => Pretty (LmlCase (LmlcPass pass)) where
   pretty :: LmlCase (LmlcPass pass) -> Doc ann
-  pretty (LmlCase _ pat constraint expr) = vsep ["Case", pretty pat, prettyMaybe constraint, pretty expr]
+  pretty (LmlCase ext pat constraint expr) = smartVsep ["Case", pretty ext, pretty pat, prettyMaybe constraint, pretty expr]
 
-instance Pretty (LmlLit (LmlcPass pass)) where
+instance (ForallLmlLit Pretty (LmlcPass pass)) => Pretty (LmlLit (LmlcPass pass)) where
   pretty :: LmlLit (LmlcPass pass) -> Doc ann
-  pretty (LmlInt _ int) = pretty int
-  pretty (LmlChar _ char) = squotes $ pretty char
-  pretty (LmlString _ str) = dquotes $ pretty str
+  pretty (LmlInt ext int) = smartVsep [pretty ext, pretty int]
+  pretty (LmlChar ext char) = smartVsep [pretty ext, squotes $ pretty char]
+  pretty (LmlString ext str) = smartVsep [pretty ext, dquotes $ pretty str]
 
 instance Pretty Longident where
   pretty :: Longident -> Doc ann
   pretty (Longident idents) = dquotes $ hsep $ punctuate comma (map pretty (toList idents))
 
-instance Pretty (LmlPat (LmlcPass pass)) where
+instance (ForallLmlPat Pretty (LmlcPass pass)) => Pretty (LmlPat (LmlcPass pass)) where
   pretty :: LmlPat (LmlcPass pass) -> Doc ann
-  pretty (LmlPatAny _) = "PatAny"
-  pretty (LmlPatVar _ var) = vsep ["PatVar", pretty var]
-  pretty (LmlPatConstant _ constant) = vsep ["PatConstant", pretty constant]
-  pretty (LmlPatTuple _ pat pats) = vsep ["PatTuple", list $ map pretty (pat : toList pats)]
-  pretty (LmlPatConstruct _ constr pat) = vsep ["PatConstruct", pretty constr, prettyMaybe pat]
-  pretty (LmlPatOr _ pat1 pat2) = vsep ["ParOr", pretty pat1, pretty pat2]
-  pretty (LmlPatConstraint _ pat ty) = vsep ["PatConstraint", pretty pat, pretty ty]
-
-instance Pretty (LmlType (LmlcPass pass)) where
+  pretty (LmlPatAny ext) = smartVsep ["PatAny", pretty ext]
+  pretty (LmlPatVar ext var) = smartVsep ["PatVar", pretty ext, pretty var]
+  pretty (LmlPatConstant ext constant) = smartVsep ["PatConstant", pretty ext, pretty constant]
+  pretty (LmlPatTuple ext pat pats) = smartVsep ["PatTuple", pretty ext, list $ map pretty (pat : toList pats)]
+  pretty (LmlPatConstruct ext constr pat) = smartVsep ["PatConstruct", pretty ext, pretty constr, prettyMaybe pat]
+  pretty (LmlPatOr ext pat1 pat2) = smartVsep ["ParOr", pretty ext, pretty pat1, pretty pat2]
+  pretty (LmlPatConstraint ext pat ty) = smartVsep ["PatConstraint", pretty ext, pretty pat, pretty ty]
+
+instance (ForallLmlType Pretty (LmlcPass pass)) => Pretty (LmlType (LmlcPass pass)) where
   pretty :: LmlType (LmlcPass pass) -> Doc ann
-  pretty (LmlTyVar _ var) = vsep ["TyVar", pretty var]
-  pretty (LmlTyArrow _ ty1 ty2) = vsep ["TyArrow", pretty ty1, pretty ty2]
-  pretty (LmlTyTuple _ ty tys) = vsep ["TyTuple", list $ map pretty (ty : toList tys)]
-  pretty (LmlTyConstr _ constr tys) = vsep ["TyConstr", pretty constr, list (map pretty tys)]
+  pretty (LmlTyVar ext var) = smartVsep ["TyVar", pretty ext, pretty var]
+  pretty (LmlTyArrow ext ty1 ty2) = smartVsep ["TyArrow", pretty ext, pretty ty1, pretty ty2]
+  pretty (LmlTyTuple ext ty tys) = smartVsep ["TyTuple", pretty ext, list $ map pretty (ty : toList tys)]
+  pretty (LmlTyConstr ext constr tys) = smartVsep ["TyConstr", pretty ext, pretty constr, list (map pretty tys)]
 
-instance Pretty (LmlModule (LmlcPass pass)) where
+instance (ForallLmlModule Pretty (LmlcPass pass)) => Pretty (LmlModule (LmlcPass pass)) where
   pretty :: LmlModule (LmlcPass pass) -> Doc ann
-  pretty (LmlModule _ name decls) = parens $ nest spaceNumber inner
+  pretty (LmlModule ext name decls) = parens $ nest spaceNumber inner
    where
-    inner = vsep ["Module", prettyMaybe name, list $ map pretty decls]
+    inner = smartVsep ["Module", pretty ext, prettyMaybe name, list $ map pretty decls]

lamagraph-compiler/src/Lamagraph/Compiler/Syntax/Expr.hs

@@ -8,7 +8,7 @@ module Lamagraph.Compiler.Syntax.Expr (
   ForallLmlExpr,
   LLmlBindGroup,
   LmlBindGroup (..),
-  ForallLmlBindingGroup,
+  ForallLmlBindGroup,
   LLmlBind,
   LmlBind (..),
   ForallLmlBind,
@@ -97,11 +97,11 @@ data LmlBindGroup pass
   = LmlBindGroup (XLmlBindGroup pass) RecFlag (NonEmpty (LLmlBind pass))
   | XLmlBindGroup !(XXBindGroup pass)
 
-type ForallLmlBindingGroup (tc :: Type -> Constraint) pass =
+type ForallLmlBindGroup (tc :: Type -> Constraint) pass =
   (tc (XLmlBindGroup pass), tc (LLmlBind pass), tc (XXBindGroup pass))
 
-deriving instance (ForallLmlBindingGroup Show pass) => Show (LmlBindGroup pass)
-deriving instance (ForallLmlBindingGroup Eq pass) => Eq (LmlBindGroup pass)
+deriving instance (ForallLmlBindGroup Show pass) => Show (LmlBindGroup pass)
+deriving instance (ForallLmlBindGroup Eq pass) => Eq (LmlBindGroup pass)
 
 -- | Located let binder
 type LLmlBind pass = XLocated pass (LmlBind pass)