Files
rlp/src/Core/Syntax.hs
2024-02-26 14:59:37 -07:00

562 lines
18 KiB
Haskell

{-|
Module : Core.Syntax
Description : Core ASTs and the like
-}
-- {-# OPTIONS_GHC -freduction-depth=0 #-}
{-# LANGUAGE PatternSynonyms, OverloadedStrings #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE TemplateHaskell #-}
-- for recursion-schemes
{-# LANGUAGE DeriveTraversable, TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE QuantifiedConstraints #-}
module Core.Syntax
(
-- * Core AST
ExprF(..), ExprF'
, ScDef(..), ScDef'
, Program(..), Program'
, Type(..), Kind, pattern (:->), pattern TyInt
, Alter(..), Alter', AltCon(..)
, pattern Binding, pattern Alter
, Rec(..), Lit(..)
, Pragma(..)
-- ** Variables and identifiers
, Name, Var(..), Tag
, Binding(..), pattern (:=), pattern (:$)
, type Binding'
-- ** Working with the fixed point of ExprF
, Expr, Expr'
, pattern Con, pattern Var, pattern App, pattern Lam, pattern Let
, pattern Case, pattern Type, pattern Lit
-- * Misc
, Pretty(pretty)
-- * Optics
, programScDefs, programTypeSigs, programDataTags
, formalising
, HasRHS(_rhs), HasLHS(_lhs)
, HasBinders(binders)
)
where
----------------------------------------------------------------------------------
import Data.Coerce
import Data.Pretty
import Data.List (intersperse)
import Data.Function ((&))
import Data.String
import Data.HashMap.Strict (HashMap)
import Data.HashMap.Strict qualified as H
import Data.Hashable
import Data.Foldable (traverse_)
import Data.Functor
import Data.Functor.Classes
import Data.Text qualified as T
import Data.Char
import Data.These
import GHC.Generics (Generic, Generically(..))
import Text.Show.Deriving
import Data.Eq.Deriving
import Data.Fix hiding (cata, ana)
import Data.Bifunctor (Bifunctor(..))
import Data.Bifoldable (bifoldr, Bifoldable(..))
import Data.Bifunctor.TH
import Data.Bitraversable
import Data.Functor.Foldable
import Data.Functor.Foldable.TH (makeBaseFunctor)
-- Lift instances for the Core quasiquoters
import Misc
import Misc.Lift1
import Control.Lens
----------------------------------------------------------------------------------
data ExprF b a = VarF Name
| ConF Tag Int -- ^ Con Tag Arity
| CaseF a [AlterF b a]
| LamF [b] a
| LetF Rec [BindingF b a] a
| AppF a a
| LitF Lit
| TypeF Type
deriving (Functor, Foldable, Traversable)
type Expr b = Fix (ExprF b)
instance IsString (ExprF b a) where
fromString = VarF . fromString
instance (IsString (f (Fix f))) => IsString (Fix f) where
fromString = Fix . fromString
data Type = TyFun
| TyVar Name
| TyApp Type Type
| TyCon Name
| TyForall Var Type
| TyKindType
deriving (Show, Eq, Lift)
type Kind = Type
-- data TyCon = MkTyCon Name Kind
-- deriving (Eq, Show, Lift)
data Var = MkVar Name Type
deriving (Eq, Show, Lift)
instance Hashable Var where
hashWithSalt s (MkVar n _) = hashWithSalt s n
pattern Con :: Tag -> Int -> Expr b
pattern Con t a = Fix (ConF t a)
pattern Var :: Name -> Expr b
pattern Var b = Fix (VarF b)
pattern App :: Expr b -> Expr b -> Expr b
pattern App f x = Fix (AppF f x)
pattern Lam :: [b] -> Expr b -> Expr b
pattern Lam bs e = Fix (LamF bs e)
pattern Let :: Rec -> [Binding b] -> Expr b -> Expr b
pattern Let r bs e = Fix (LetF r bs e)
pattern Case :: Expr b -> [Alter b] -> Expr b
pattern Case e as = Fix (CaseF e as)
pattern Type :: Type -> Expr b
pattern Type t = Fix (TypeF t)
pattern Lit :: Lit -> Expr b
pattern Lit t = Fix (LitF t)
pattern TyInt :: Type
pattern TyInt = TyCon "Int#"
infixr 1 :->
pattern (:->) :: Type -> Type -> Type
pattern a :-> b = TyApp (TyApp TyFun a) b
data BindingF b a = BindingF b (ExprF b a)
deriving (Functor, Foldable, Traversable)
type Binding b = BindingF b (Fix (ExprF b))
type Binding' = Binding Name
-- collapse = foldFix embed
pattern Binding :: b -> Expr b -> Binding b
pattern Binding k v <- BindingF k (wrapFix -> v)
where Binding k v = BindingF k (unwrapFix v)
{-# COMPLETE (:=) #-}
{-# COMPLETE Binding #-}
infixl 1 :=
pattern (:=) :: b -> Expr b -> Binding b
pattern k := v = Binding k v
infixl 2 :$
pattern (:$) :: Expr b -> Expr b -> Expr b
pattern f :$ x = App f x
data AlterF b a = AlterF AltCon [b] (ExprF b a)
deriving (Functor, Foldable, Traversable)
pattern Alter :: AltCon -> [b] -> Expr b -> Alter b
pattern Alter con bs e <- AlterF con bs (wrapFix -> e)
where Alter con bs e = AlterF con bs (unwrapFix e)
type Alter b = AlterF b (Fix (ExprF b))
type Alter' = Alter Name
-- pattern Alter :: AltCon -> [b] -> Expr b -> Alter b
-- pattern Alter con bs e <- Fix (AlterF con bs (undefined -> e))
-- where Alter con bs e = Fix (AlterF con bs undefined)
newtype Pragma = Pragma [T.Text]
data Rec = Rec
| NonRec
deriving (Show, Eq, Lift)
data AltCon = AltData Name
| AltTag Tag
| AltLit Lit
| AltDefault
deriving (Show, Eq, Lift)
newtype Lit = IntL Int
deriving (Show, Eq, Lift)
type Name = T.Text
type Tag = Int
data ScDef b = ScDef b [b] (Expr b)
-- unliftScDef :: ScDef b -> Expr b
-- unliftScDef (ScDef _ as e) = Lam as e
data Module b = Module (Maybe (Name, [Name])) (Program b)
data Program b = Program
{ _programScDefs :: [ScDef b]
, _programTypeSigs :: HashMap b Type
, _programDataTags :: HashMap Name (Tag, Int)
-- ^ map constructors to their tag and arity
}
deriving (Generic)
deriving (Semigroup, Monoid)
via Generically (Program b)
makeLenses ''Program
-- makeBaseFunctor ''Expr
pure []
-- this is a weird optic, stronger than Lens and Prism, but weaker than Iso.
programTypeSigsP :: (Hashable b) => Prism' (Program b) (HashMap b Type)
programTypeSigsP = prism
(\b -> mempty & programTypeSigs .~ b)
(Right . view programTypeSigs)
type ExprF' = ExprF Name
type Program' = Program Name
type Expr' = Expr Name
type ScDef' = ScDef Name
-- type Alter' = Alter Name
-- type Binding' = Binding Name
-- instance IsString (Expr b) where
-- fromString = Var . fromString
----------------------------------------------------------------------------------
class HasRHS s t a b | s -> a, t -> b, s b -> t, t a -> s where
_rhs :: Lens s t a b
instance HasRHS (AlterF b a) (AlterF b a') (ExprF b a) (ExprF b a') where
_rhs = lens
(\ (AlterF _ _ e) -> e)
(\ (AlterF t as _) e' -> AlterF t as e')
instance HasRHS (ScDef b) (ScDef b) (Expr b) (Expr b) where
_rhs = lens
(\ (ScDef _ _ e) -> e)
(\ (ScDef n as _) e' -> ScDef n as e')
instance HasRHS (BindingF b a) (BindingF b' a') (ExprF b a) (ExprF b' a')
class HasLHS s t a b | s -> a, t -> b, s b -> t, t a -> s where
_lhs :: Lens s t a b
instance HasLHS (ScDef b) (ScDef b) (b, [b]) (b, [b]) where
_lhs = lens
(\ (ScDef n as _) -> (n,as))
(\ (ScDef _ _ e) (n',as') -> ScDef n' as' e)
-- instance HasLHS (Binding b) (Binding b) b b where
-- _lhs = lens
-- (\ (k := _) -> k)
-- (\ (_ := e) k' -> k' := e)
-- | This is not a valid isomorphism for expressions containing lambdas whose
-- bodies are themselves lambdas with multiple arguments:
--
-- >>> [coreExpr|\x -> \y z -> x|] ^. from (from formalising)
-- Lam ["x"] (Lam ["y"] (Lam ["z"] (Var "x")))
-- >>> [coreExpr|\x -> \y z -> x|]
-- Lam ["x"] (Lam ["y","z"] (Var "x"))
--
-- For this reason, it's best to consider 'formalising' as if it were two
-- unrelated unidirectional getters.
formalising :: Iso (Expr a) (Expr b) (Expr a) (Expr b)
formalising = iso sa bt where
sa :: Expr a -> Expr a
sa = ana \case
Lam [b] e -> LamF [b] e
Lam (b:bs) e -> LamF [b] (Lam bs e)
x -> project x
bt :: Expr b -> Expr b
bt = cata \case
LamF [b] (Lam bs e) -> Lam (b:bs) e
x -> embed x
--------------------------------------------------------------------------------
instance (Hashable b, Pretty b) => Pretty (Program b) where
pretty p = (datatags <> "\n")
$+$ defs
where
datatags = ifoldrOf (programDataTags . ifolded) cataDataTag mempty p
defs = vlinesOf (programJoinedDefs . to prettyGroup) p
programJoinedDefs :: Fold (Program b) (These (b, Type) (ScDef b))
programJoinedDefs = folding $ \p ->
foldMapOf programTypeSigs thisTs p
`u` foldMapOf programScDefs thatSc p
where u = H.unionWith unionThese
thisTs = ifoldMap @b @(HashMap b)
(\n t -> H.singleton n (This (n,t)))
thatSc = foldMap $ \sc ->
H.singleton (sc ^. _lhs . _1) (That sc)
prettyGroup :: These (b, Type) (ScDef b) -> Doc
prettyGroup = bifoldr vcatWithSemi vcatWithSemi mempty
. bimap prettyTySig pretty
vcatWithSemi a b = (a <+> ";") $$ b
prettyTySig (n,t) = hsep [ttext n, "::", pretty t]
unionThese (This a) (That b) = These a b
unionThese (That b) (This a) = These a b
unionThese (These a b) _ = These a b
cataDataTag n (t,a) acc = prettyDataTag n t a $+$ acc
prettyDataTag n t a =
hsep ["{-#", "PackData", ttext n, ttext t, ttext a, "#-}"]
instance Pretty Type where
prettyPrec _ (TyVar n) = ttext n
prettyPrec _ TyFun = "(->)"
prettyPrec _ (TyCon n) = ttext n
prettyPrec p (a :-> b) = maybeParens (p>0) $
hsep [prettyPrec 1 a, "->", prettyPrec 0 b]
prettyPrec p (TyApp f x) = maybeParens (p>1) $
prettyPrec 1 f <+> prettyPrec 2 x
instance (Pretty b) => Pretty (ScDef b) where
pretty sc = hsep [name, as, "=", hang empty 1 e]
where
name = ttext $ sc ^. _lhs . _1
as = sc & hsepOf (_lhs . _2 . each . to ttext)
e = pretty $ sc ^. _rhs
instance (Pretty b) => Pretty (Expr b) where
-- prettyPrec _ (Var n) = ttext n
-- prettyPrec _ (Con t a) = "Pack{" <> (ttext t <+> ttext a) <> "}"
-- prettyPrec _ (Lam bs e) = hsep ["λ", hsep (prettyPrec 1 <$> bs), "->", pretty e]
-- prettyPrec _ (Let r bs e) = hsep [word, explicitLayout bs]
-- $+$ hsep ["in", pretty e]
-- where word = if r == Rec then "letrec" else "let"
-- prettyPrec p (App f x) = maybeParens (p>0) $
-- prettyPrec 0 f <+> prettyPrec 1 x
-- prettyPrec _ (Lit l) = pretty l
-- prettyPrec p (Case e as) = maybeParens (p>0) $
-- "case" <+> pretty e <+> "of"
-- $+$ nest 2 (explicitLayout as)
instance (Pretty b, Pretty a) => Pretty (AlterF b a) where
-- pretty (Alter c as e) =
-- hsep [pretty c, hsep (pretty <$> as), "->", pretty e]
instance Pretty AltCon where
pretty (AltData n) = ttext n
pretty (AltLit l) = pretty l
pretty (AltTag t) = ttext t
pretty AltDefault = "_"
instance Pretty Lit where
pretty (IntL n) = ttext n
instance (Pretty b, Pretty a) => Pretty (BindingF b a) where
-- pretty (k := v) = hsep [pretty k, "=", pretty v]
explicitLayout :: (Pretty a) => [a] -> Doc
explicitLayout as = vcat inner <+> "}" where
inner = zipWith (<+>) delims (pretty <$> as)
delims = "{" : repeat ";"
instance Pretty Var where
prettyPrec p (MkVar n t) = maybeParens (p>0) $
hsep [pretty n, ":", pretty t]
--------------------------------------------------------------------------------
-- instance Functor Alter where
-- fmap f (Alter con bs e) = Alter con (f <$> bs) e'
-- where
-- e' = foldFix (embed . bimap' f id) e
-- bimap' = $(makeBimap ''ExprF)
-- instance Foldable Alter where
-- instance Traversable Alter where
-- instance Functor Binding where
-- instance Foldable Binding where
-- instance Traversable Binding where
liftShowsPrecExpr :: (Show b)
=> (Int -> a -> ShowS)
-> ([a] -> ShowS)
-> Int -> ExprF b a -> ShowS
liftShowsPrecExpr = $(makeLiftShowsPrec ''ExprF)
showsPrec1Expr :: (Show b, Show a)
=> Int -> ExprF b a -> ShowS
showsPrec1Expr = $(makeShowsPrec1 ''ExprF)
instance (Show b) => Show1 (AlterF b) where
liftShowsPrec sp spl d (AlterF con bs e) =
showsTernaryWith showsPrec showsPrec (liftShowsPrecExpr sp spl)
"AlterF" d con bs e
instance (Show b) => Show1 (BindingF b) where
liftShowsPrec sp spl d (BindingF k v) =
showsBinaryWith showsPrec (liftShowsPrecExpr sp spl)
"BindingF" d k v
instance (Show b, Show a) => Show (BindingF b a) where
showsPrec d (BindingF k v)
= showParen (d > 10)
$ showString "BindingF" . showChar ' '
. showsPrec 11 k . showChar ' '
. showsPrec1Expr 11 v
instance (Show b, Show a) => Show (AlterF b a) where
showsPrec d (AlterF con bs e)
= showParen (d > 10)
$ showString "AlterF" . showChar ' '
. showsPrec 11 con . showChar ' '
. showsPrec 11 bs . showChar ' '
. showsPrec1Expr 11 e
deriveShow1 ''ExprF
deriving instance (Show b, Show a) => Show (ExprF b a)
-- deriving instance (Show b, Show a) => Show (BindingF b a)
-- deriving instance (Show b, Show a) => Show (AlterF b a)
deriving instance Show b => Show (ScDef b)
deriving instance Show b => Show (Program b)
bimapExpr :: (b -> b') -> (a -> a')
-> ExprF b a -> ExprF b' a'
bimapExpr = $(makeBimap ''ExprF)
bifoldrExpr :: (b -> c -> c)
-> (a -> c -> c)
-> c -> ExprF b a -> c
bifoldrExpr = $(makeBifoldr ''ExprF)
bitraverseExpr :: Applicative f
=> (b -> f b')
-> (a -> f a')
-> ExprF b a -> f (ExprF b' a')
bitraverseExpr = $(makeBitraverse ''ExprF)
instance Bifunctor AlterF where
bimap f g (AlterF con bs e) = AlterF con (f <$> bs) (bimapExpr f g e)
instance Bifunctor BindingF where
bimap f g (BindingF k v) = BindingF (f k) (bimapExpr f g v)
instance Bifoldable AlterF where
bifoldr f g z (AlterF con bs e) = bifoldrExpr f g z' e where
z' = foldr f z bs
instance Bitraversable AlterF where
bitraverse f g (AlterF con bs e) =
AlterF con <$> traverse f bs <*> bitraverseExpr f g e
instance Bifoldable BindingF where
bifoldr f g z (BindingF k v) = bifoldrExpr f g (f k z) v
instance Bitraversable BindingF where
bitraverse f g (BindingF k v) =
BindingF <$> f k <*> bitraverseExpr f g v
deriveBifunctor ''ExprF
deriveBifoldable ''ExprF
deriveBitraversable ''ExprF
instance Lift b => Lift1 (ExprF b) where
lift1 (VarF k) = liftCon 'VarF (lift k)
lift1 (AppF f x) = liftCon2 'AppF (lift f) (lift x)
lift1 (LamF b e) = liftCon2 'LamF (lift b) (lift e)
lift1 (LetF r bs e) = liftCon3 'LetF (lift r) (lift bs) (lift e)
lift1 (CaseF e as) = liftCon2 'CaseF (lift e) (lift as)
lift1 (TypeF t) = liftCon 'TypeF (lift t)
lift1 (LitF l) = liftCon 'LitF (lift l)
lift1 (ConF t a) = liftCon2 'ConF (lift t) (lift a)
deriving instance (Lift b, Lift a) => Lift (ExprF b a)
deriving instance (Lift b, Lift a) => Lift (BindingF b a)
deriving instance (Lift b, Lift a) => Lift (AlterF b a)
deriving instance Lift b => Lift (ScDef b)
deriving instance Lift b => Lift (Program b)
--------------------------------------------------------------------------------
class HasBinders s t a b | s -> a, t -> b, s b -> t, t a -> s where
binders :: Traversal s t a b
-- instance (HasBinders a a' b b')
-- => HasBinders (ExprF b a) (ExprF b' a') b b' where
-- binders k = undefined
instance HasBinders a a' b b'
=> HasBinders (ExprF b a) (ExprF b' a') b b' where
binders :: forall f. (Applicative f)
=> LensLike f (ExprF b a) (ExprF b' a') b b'
binders k = go where
go :: ExprF b a -> f (ExprF b' a')
go (LamF bs e) = LamF <$> traverse k bs <*> binders k e
go (CaseF e as) = CaseF <$> binders k e <*> eachbind as
go (LetF r bs e) = LetF r <$> eachbind bs <*> binders k e
go f = bitraverse k (binders k) f
eachbind :: forall p. Bitraversable p => [p b a] -> f [p b' a']
eachbind bs = bitraverse k (binders k) `traverse` bs
instance HasBinders a a b b'
=> HasBinders (AlterF b a) (AlterF b' a) b b' where
binders k (AlterF con bs e) =
AlterF con <$> traverse k bs <*> traverseOf binders k e
instance HasBinders (BindingF b a) (BindingF b' a) b b' where
binders = undefined
instance (HasBinders (f b (Fix (f b))) (f b' (Fix (f b'))) b b')
=> HasBinders (Fix (f b)) (Fix (f b')) b b' where
binders k (Fix f) = Fix <$> binders k f
liftEqExpr :: (Eq b)
=> (a -> a' -> Bool)
-> ExprF b a -> ExprF b a' -> Bool
liftEqExpr = $(makeLiftEq ''ExprF)
instance (Eq b, Eq a) => Eq (BindingF b a) where
BindingF ka va == BindingF kb vb =
ka == kb && va `eq` vb
where eq = liftEqExpr (==)
instance (Eq b, Eq a) => Eq (AlterF b a) where
AlterF cona bsa ea == AlterF conb bsb eb =
cona == conb && bsa == bsb && ea `eq` eb
where eq = liftEqExpr (==)
instance (Eq b) => Eq1 (AlterF b) where
liftEq f (AlterF cona bsa ea) (AlterF conb bsb eb) =
cona == conb && bsa == bsb && ea `eq` eb
where eq = liftEqExpr f
instance (Eq b) => Eq1 (BindingF b) where
liftEq f (BindingF ka va) (BindingF kb vb) =
ka == kb && va `eq` vb
where eq = liftEqExpr f
deriveEq1 ''ExprF
deriving instance (Eq b, Eq a) => Eq (ExprF b a)