module Language.Lambda.SystemF.Expression
  ( SystemFExpr(..),
    Ty(..),
    prettyPrint,
    upperLambda
  ) where

import Data.Monoid
import Prettyprinter
import Prettyprinter.Render.Text (renderStrict)
import RIO

data SystemFExpr name ty
  -- | Variable: `x`
  = Var name
  -- | Function application: `x y`
  | App (SystemFExpr name ty) (SystemFExpr name ty)
  -- | Lambda abstraction: `\x: X. x`
  | Abs name (Ty ty) (SystemFExpr name ty)
  -- | Type Abstraction: `\X. body`
  | TyAbs ty (SystemFExpr name ty)                  
  -- | Type Application: `x [X]`
  | TyApp (SystemFExpr name ty) (Ty ty)
  deriving (SystemFExpr name ty -> SystemFExpr name ty -> Bool
(SystemFExpr name ty -> SystemFExpr name ty -> Bool)
-> (SystemFExpr name ty -> SystemFExpr name ty -> Bool)
-> Eq (SystemFExpr name ty)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall name ty.
(Eq name, Eq ty) =>
SystemFExpr name ty -> SystemFExpr name ty -> Bool
/= :: SystemFExpr name ty -> SystemFExpr name ty -> Bool
$c/= :: forall name ty.
(Eq name, Eq ty) =>
SystemFExpr name ty -> SystemFExpr name ty -> Bool
== :: SystemFExpr name ty -> SystemFExpr name ty -> Bool
$c== :: forall name ty.
(Eq name, Eq ty) =>
SystemFExpr name ty -> SystemFExpr name ty -> Bool
Eq, Int -> SystemFExpr name ty -> ShowS
[SystemFExpr name ty] -> ShowS
SystemFExpr name ty -> String
(Int -> SystemFExpr name ty -> ShowS)
-> (SystemFExpr name ty -> String)
-> ([SystemFExpr name ty] -> ShowS)
-> Show (SystemFExpr name ty)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall name ty.
(Show name, Show ty) =>
Int -> SystemFExpr name ty -> ShowS
forall name ty.
(Show name, Show ty) =>
[SystemFExpr name ty] -> ShowS
forall name ty.
(Show name, Show ty) =>
SystemFExpr name ty -> String
showList :: [SystemFExpr name ty] -> ShowS
$cshowList :: forall name ty.
(Show name, Show ty) =>
[SystemFExpr name ty] -> ShowS
show :: SystemFExpr name ty -> String
$cshow :: forall name ty.
(Show name, Show ty) =>
SystemFExpr name ty -> String
showsPrec :: Int -> SystemFExpr name ty -> ShowS
$cshowsPrec :: forall name ty.
(Show name, Show ty) =>
Int -> SystemFExpr name ty -> ShowS
Show)

data Ty name
  = TyVar name                  -- ^ Type variable (T)
  | TyArrow (Ty name) (Ty name) -- ^ Type arrow    (T -> U)
  | TyForAll name (Ty name)     -- ^ Universal type (forall T. X)
  deriving (Ty name -> Ty name -> Bool
(Ty name -> Ty name -> Bool)
-> (Ty name -> Ty name -> Bool) -> Eq (Ty name)
forall name. Eq name => Ty name -> Ty name -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Ty name -> Ty name -> Bool
$c/= :: forall name. Eq name => Ty name -> Ty name -> Bool
== :: Ty name -> Ty name -> Bool
$c== :: forall name. Eq name => Ty name -> Ty name -> Bool
Eq, Int -> Ty name -> ShowS
[Ty name] -> ShowS
Ty name -> String
(Int -> Ty name -> ShowS)
-> (Ty name -> String) -> ([Ty name] -> ShowS) -> Show (Ty name)
forall name. Show name => Int -> Ty name -> ShowS
forall name. Show name => [Ty name] -> ShowS
forall name. Show name => Ty name -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Ty name] -> ShowS
$cshowList :: forall name. Show name => [Ty name] -> ShowS
show :: Ty name -> String
$cshow :: forall name. Show name => Ty name -> String
showsPrec :: Int -> Ty name -> ShowS
$cshowsPrec :: forall name. Show name => Int -> Ty name -> ShowS
Show)

instance (Pretty name, Pretty ty) => Pretty (SystemFExpr name ty) where
  pretty :: SystemFExpr name ty -> Doc ann
pretty (Var name
name) = name -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty name
name
  pretty (App SystemFExpr name ty
e1 SystemFExpr name ty
e2) = SystemFExpr name ty -> SystemFExpr name ty -> Doc ann
forall name ty a.
(Pretty name, Pretty ty) =>
SystemFExpr name ty -> SystemFExpr name ty -> Doc a
prettyApp SystemFExpr name ty
e1 SystemFExpr name ty
e2
  pretty (Abs name
name Ty ty
ty SystemFExpr name ty
body) = name -> Ty ty -> SystemFExpr name ty -> Doc ann
forall name ty ann.
(Pretty name, Pretty ty) =>
name -> Ty ty -> SystemFExpr name ty -> Doc ann
prettyAbs name
name Ty ty
ty SystemFExpr name ty
body
  pretty (TyAbs ty
ty SystemFExpr name ty
body) = ty -> SystemFExpr name ty -> Doc ann
forall name ty ann.
(Pretty name, Pretty ty) =>
ty -> SystemFExpr name ty -> Doc ann
prettyTyAbs ty
ty SystemFExpr name ty
body
  pretty (TyApp SystemFExpr name ty
expr Ty ty
ty) = SystemFExpr name ty -> Ty ty -> Doc ann
forall name ty ann.
(Pretty name, Pretty ty) =>
SystemFExpr name ty -> Ty ty -> Doc ann
prettyTyApp SystemFExpr name ty
expr Ty ty
ty

instance Pretty name => Pretty (Ty name) where
  pretty :: Ty name -> Doc ann
pretty = Bool -> Ty name -> Doc ann
forall name ann. Pretty name => Bool -> Ty name -> Doc ann
prettyTy Bool
False

prettyPrint :: Pretty pretty => pretty -> Text
prettyPrint :: pretty -> Text
prettyPrint pretty
expr = SimpleDocStream Any -> Text
forall ann. SimpleDocStream ann -> Text
renderStrict SimpleDocStream Any
docStream
  where docStream :: SimpleDocStream Any
docStream = LayoutOptions -> Doc Any -> SimpleDocStream Any
forall ann. LayoutOptions -> Doc ann -> SimpleDocStream ann
layoutPretty LayoutOptions
defaultLayoutOptions (pretty -> Doc Any
forall a ann. Pretty a => a -> Doc ann
pretty pretty
expr)

upperLambda :: Char
upperLambda :: Char
upperLambda = Char
'Λ'

prettyApp
  :: (Pretty name, Pretty ty)
  => SystemFExpr name ty
  -> SystemFExpr name ty
  -> Doc a
prettyApp :: SystemFExpr name ty -> SystemFExpr name ty -> Doc a
prettyApp e1 :: SystemFExpr name ty
e1@Abs{} e2 :: SystemFExpr name ty
e2@Abs{} = Doc a -> Doc a
forall ann. Doc ann -> Doc ann
parens (SystemFExpr name ty -> Doc a
forall a ann. Pretty a => a -> Doc ann
pretty SystemFExpr name ty
e1) Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc a -> Doc a
forall ann. Doc ann -> Doc ann
parens (SystemFExpr name ty -> Doc a
forall a ann. Pretty a => a -> Doc ann
pretty SystemFExpr name ty
e2)
prettyApp e1 :: SystemFExpr name ty
e1@Abs{} SystemFExpr name ty
e2 = Doc a -> Doc a
forall ann. Doc ann -> Doc ann
parens (SystemFExpr name ty -> Doc a
forall a ann. Pretty a => a -> Doc ann
pretty SystemFExpr name ty
e1) Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> SystemFExpr name ty -> Doc a
forall a ann. Pretty a => a -> Doc ann
pretty SystemFExpr name ty
e2
prettyApp SystemFExpr name ty
e1 e2 :: SystemFExpr name ty
e2@Abs{} = SystemFExpr name ty -> Doc a
forall a ann. Pretty a => a -> Doc ann
pretty SystemFExpr name ty
e1 Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc a -> Doc a
forall ann. Doc ann -> Doc ann
parens (SystemFExpr name ty -> Doc a
forall a ann. Pretty a => a -> Doc ann
pretty SystemFExpr name ty
e2)
prettyApp SystemFExpr name ty
e1 e2 :: SystemFExpr name ty
e2@App{} = SystemFExpr name ty -> Doc a
forall a ann. Pretty a => a -> Doc ann
pretty SystemFExpr name ty
e1 Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc a -> Doc a
forall ann. Doc ann -> Doc ann
parens (SystemFExpr name ty -> Doc a
forall a ann. Pretty a => a -> Doc ann
pretty SystemFExpr name ty
e2)
prettyApp SystemFExpr name ty
e1 SystemFExpr name ty
e2 = SystemFExpr name ty -> Doc a
forall a ann. Pretty a => a -> Doc ann
pretty SystemFExpr name ty
e1 Doc a -> Doc a -> Doc a
forall ann. Doc ann -> Doc ann -> Doc ann
<+> SystemFExpr name ty -> Doc a
forall a ann. Pretty a => a -> Doc ann
pretty SystemFExpr name ty
e2

prettyAbs
  :: (Pretty name, Pretty ty)
  => name
  -> Ty ty
  -> SystemFExpr name ty
  -> Doc ann
prettyAbs :: name -> Ty ty -> SystemFExpr name ty -> Doc ann
prettyAbs name
name Ty ty
ty SystemFExpr name ty
body
  = Doc ann
forall ann. Doc ann
lambda
    Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
hsep (((name, Ty ty) -> Doc ann) -> [(name, Ty ty)] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map ((name -> Ty ty -> Doc ann) -> (name, Ty ty) -> Doc ann
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry name -> Ty ty -> Doc ann
forall name ty ann.
(Pretty name, Pretty ty) =>
name -> Ty ty -> Doc ann
prettyArg) [(name, Ty ty)]
names)
    Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
dot
    Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> SystemFExpr name ty -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty SystemFExpr name ty
body'
  where ([(name, Ty ty)]
names, SystemFExpr name ty
body') = name
-> Ty ty
-> SystemFExpr name ty
-> ([(name, Ty ty)], SystemFExpr name ty)
forall n t.
n -> Ty t -> SystemFExpr n t -> ([(n, Ty t)], SystemFExpr n t)
uncurryAbs name
name Ty ty
ty SystemFExpr name ty
body

prettyTyAbs :: (Pretty name, Pretty ty) => ty -> SystemFExpr name ty -> Doc ann
prettyTyAbs :: ty -> SystemFExpr name ty -> Doc ann
prettyTyAbs ty
name SystemFExpr name ty
body = Doc ann
forall ann. Doc ann
upperLambda' Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
hsep ((ty -> Doc ann) -> [ty] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map ty -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty [ty]
names) Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
dot
    Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> SystemFExpr name ty -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty SystemFExpr name ty
body'
  where ([ty]
names, SystemFExpr name ty
body') = ty -> SystemFExpr name ty -> ([ty], SystemFExpr name ty)
forall t n. t -> SystemFExpr n t -> ([t], SystemFExpr n t)
uncurryTyAbs ty
name SystemFExpr name ty
body
prettyTyApp :: (Pretty name, Pretty ty) => SystemFExpr name ty -> Ty ty -> Doc ann
prettyTyApp :: SystemFExpr name ty -> Ty ty -> Doc ann
prettyTyApp SystemFExpr name ty
expr Ty ty
ty = SystemFExpr name ty -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty SystemFExpr name ty
expr Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
brackets (Ty ty -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty Ty ty
ty)

prettyTy :: Pretty name => Bool -> Ty name -> Doc ann
prettyTy :: Bool -> Ty name -> Doc ann
prettyTy Bool
_ (TyVar name
name) = name -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty name
name
prettyTy Bool
compact (TyArrow Ty name
t1 Ty name
t2) = Bool -> Ty name -> Ty name -> Doc ann
forall name ann.
Pretty name =>
Bool -> Ty name -> Ty name -> Doc ann
prettyTyArrow Bool
compact Ty name
t1 Ty name
t2
prettyTy Bool
compact (TyForAll name
name Ty name
ty) = Bool -> name -> Ty name -> Doc ann
forall name ann. Pretty name => Bool -> name -> Ty name -> Doc ann
prettyTyForAll Bool
compact name
name Ty name
ty

prettyTyArrow :: Pretty name => Bool -> Ty name -> Ty name -> Doc ann
prettyTyArrow :: Bool -> Ty name -> Ty name -> Doc ann
prettyTyArrow Bool
compact (TyArrow Ty name
t1 Ty name
t2) Ty name
t3
  = Bool -> Doc ann -> Doc ann -> Doc ann
forall ann. Bool -> Doc ann -> Doc ann -> Doc ann
prettyTyArrow' Bool
compact Doc ann
compositeTy (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$ Bool -> Ty name -> Doc ann
forall name ann. Pretty name => Bool -> Ty name -> Doc ann
prettyTy Bool
compact Ty name
t3
  where compositeTy :: Doc ann
compositeTy = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$ Bool -> Ty name -> Ty name -> Doc ann
forall name ann.
Pretty name =>
Bool -> Ty name -> Ty name -> Doc ann
prettyTyArrow Bool
compact Ty name
t1 Ty name
t2

prettyTyArrow Bool
compact Ty name
t1 Ty name
t2
  = Bool -> Doc ann -> Doc ann -> Doc ann
forall ann. Bool -> Doc ann -> Doc ann -> Doc ann
prettyTyArrow' Bool
compact (Bool -> Ty name -> Doc ann
forall name ann. Pretty name => Bool -> Ty name -> Doc ann
prettyTy Bool
compact Ty name
t1) (Bool -> Ty name -> Doc ann
forall name ann. Pretty name => Bool -> Ty name -> Doc ann
prettyTy Bool
compact Ty name
t2)

prettyTyForAll :: Pretty name => Bool -> name -> Ty name -> Doc ann
prettyTyForAll :: Bool -> name -> Ty name -> Doc ann
prettyTyForAll Bool
compact name
name Ty name
ty
  = Doc ann
"forall"
  Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> name -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty name
name Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
dot
  Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Bool -> Ty name -> Doc ann
forall name ann. Pretty name => Bool -> Ty name -> Doc ann
prettyTy Bool
compact Ty name
ty

lambda :: Doc ann
lambda :: Doc ann
lambda = Char -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty Char
'λ'

prettyArg :: (Pretty name, Pretty ty) => name -> Ty ty -> Doc ann
prettyArg :: name -> Ty ty -> Doc ann
prettyArg name
name (TyArrow Ty ty
t1 Ty ty
t2)
  = name -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty name
name Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
colon Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Bool -> Ty ty -> Ty ty -> Doc ann
forall name ann.
Pretty name =>
Bool -> Ty name -> Ty name -> Doc ann
prettyTyArrow Bool
True Ty ty
t1 Ty ty
t2)
prettyArg name
name Ty ty
ty = name -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty name
name Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
colon Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Ty ty -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty Ty ty
ty

upperLambda' :: Doc ann
upperLambda' :: Doc ann
upperLambda' = Char -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty Char
upperLambda

prettyTyArrow' :: Bool -> Doc ann -> Doc ann -> Doc ann
prettyTyArrow' :: Bool -> Doc ann -> Doc ann -> Doc ann
prettyTyArrow' Bool
compact Doc ann
doc1 Doc ann
doc2 = Doc ann
doc1 Doc ann -> Doc ann -> Doc ann
`add'` Doc ann
"->" Doc ann -> Doc ann -> Doc ann
`add'` Doc ann
doc2
  where add' :: Doc ann -> Doc ann -> Doc ann
add'
          | Bool
compact = Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
(<>) 
          | Bool
otherwise = Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
(<+>)

uncurryAbs :: n -> Ty t -> SystemFExpr n t -> ([(n, Ty t)], SystemFExpr n t)
uncurryAbs :: n -> Ty t -> SystemFExpr n t -> ([(n, Ty t)], SystemFExpr n t)
uncurryAbs n
name Ty t
ty = [(n, Ty t)] -> SystemFExpr n t -> ([(n, Ty t)], SystemFExpr n t)
forall a ty.
[(a, Ty ty)]
-> SystemFExpr a ty -> ([(a, Ty ty)], SystemFExpr a ty)
uncurry' [(n
name, Ty t
ty)] 
  where uncurry' :: [(a, Ty ty)]
-> SystemFExpr a ty -> ([(a, Ty ty)], SystemFExpr a ty)
uncurry' [(a, Ty ty)]
ns (Abs a
n' Ty ty
t' SystemFExpr a ty
body') = [(a, Ty ty)]
-> SystemFExpr a ty -> ([(a, Ty ty)], SystemFExpr a ty)
uncurry' ((a
n', Ty ty
t')(a, Ty ty) -> [(a, Ty ty)] -> [(a, Ty ty)]
forall a. a -> [a] -> [a]
:[(a, Ty ty)]
ns) SystemFExpr a ty
body'
        uncurry' [(a, Ty ty)]
ns SystemFExpr a ty
body'             = ([(a, Ty ty)] -> [(a, Ty ty)]
forall a. [a] -> [a]
reverse [(a, Ty ty)]
ns, SystemFExpr a ty
body')

uncurryTyAbs :: t -> SystemFExpr n t -> ([t], SystemFExpr n t)
uncurryTyAbs :: t -> SystemFExpr n t -> ([t], SystemFExpr n t)
uncurryTyAbs t
ty = [t] -> SystemFExpr n t -> ([t], SystemFExpr n t)
forall a name.
[a] -> SystemFExpr name a -> ([a], SystemFExpr name a)
uncurry' [t
ty]
  where uncurry' :: [a] -> SystemFExpr name a -> ([a], SystemFExpr name a)
uncurry' [a]
ts (TyAbs a
t' SystemFExpr name a
body') = [a] -> SystemFExpr name a -> ([a], SystemFExpr name a)
uncurry' (a
t'a -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
ts) SystemFExpr name a
body'
        uncurry' [a]
ts SystemFExpr name a
body'            = ([a] -> [a]
forall a. [a] -> [a]
reverse [a]
ts, SystemFExpr name a
body')