TreesThatGrow.hs

{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE TypeFamilyDependencies #-}
module TreesThatGrow where

import Data.Void
import Data.Kind
import Data.Profunctor


------------
--- STLC ---
------------

data Exp
  = Lit Int
  | Var String
  | Ann Exp Typ
  | Abs String Exp
  | App Exp Exp
  
data Typ
  = TyInt
  | TyFun Typ Typ

-----------------------
--- Extensible STLC ---
-----------------------

data ExpX x
  = Litx (XLit x) Int
  | Varx (XVar x) String
  | Annx (XAnn x) (ExpX x) Typ
  | Absx (XAbs x) String (ExpX x)
  | Appx (XApp x) (ExpX x) (ExpX x)
  | Expx (XExp x)

type family XLit x 
type family XVar x 
type family XAnn x 
type family XAbs x 
type family XApp x 
type family XExp x
  
--------------------------------
--- Recovering Exp from ExpX ---
--------------------------------

data UD
type ExpUD = ExpX UD

void :: Void
void = undefined

type instance XLit UD = Void
type instance XVar UD = Void
type instance XAnn UD = Void
type instance XAbs UD = Void
type instance XApp UD = Void
type instance XExp UD = Void

type ISO a b = (a -> b, b -> a)

iso :: ISO Exp ExpUD
iso = (from, to)
  where
    from :: Exp -> ExpUD
    from = \case
      Lit i -> Litx void i
      Var str -> Varx void str
      Ann t1 ty -> Annx void (from t1) ty
      Abs bndr trm -> Absx void bndr (from trm)
      App t1 t2 -> Appx void (from t1) (from t2)

    to :: ExpUD -> Exp
    to = \case
      Litx _ i -> Lit i
      Varx _ str -> Var str
      Annx _ t1 ty -> Ann (to t1) ty
      Absx _ bndr trm -> Abs bndr (to trm)
      Appx _ t1 t2 -> App (to t1) (to t2)
    
incLit :: Exp -> Exp
incLit = \case
  Lit n -> Lit $ n + 1
  e -> e

incLitx :: ExpUD -> ExpUD
incLitx = \case
  Litx _ i -> Litx void $ i + 1
  e -> e

pattern LitUD :: Int -> ExpUD
pattern LitUD i <- Litx _ i
  where
    LitUD i = Litx void i

incLitx' :: ExpUD -> ExpUD
incLitx' = \case
  LitUD i -> LitUD $ i + 1
  e -> e

-----------------------------------------------
--- Decorating a field with new information ---
-----------------------------------------------

data TC
type ExpTC = ExpX TC

type instance XLit TC = Void
type instance XVar TC = Void
type instance XAnn TC = Void
type instance XAbs TC = Void
type instance XApp TC = Typ

pattern AppTC :: Typ -> ExpTC -> ExpTC -> ExpTC
pattern AppTC a l m = Appx a l m

---------------------------------
--- New Constructor Extension ---
---------------------------------

data Val = Val Bool

data PE
type ExpPE = ExpX PE

type instance XLit PE = Void
type instance XVar PE = Void
type instance XAnn PE = Void
type instance XAbs PE = Void
type instance XApp PE = Void
type instance XExp PE = Val

pattern ValPE :: Val -> ExpPE
pattern ValPE v = Expx v

------------------------------
--- Replacing Constructors ---
------------------------------
-- This assumes we do not export 'XApp'

-- data Exp = ... | App Exp [Exp]
   
data SA
type ExpSA = ExpX SA

type instance XLit SA = Void
type instance XVar SA = Void
type instance XAnn SA = Void
type instance XAbs SA = Void
type instance XApp SA = Void
type instance XExp SA = (ExpSA, [ExpSA])

pattern AppSA :: ExpSA -> [ExpSA] -> ExpSA
pattern AppSA l ms = Expx (l, ms)

----------------------------------------
--- Extensions Using Type Parameters ---
----------------------------------------
-- Consider a Variant of 'Exp' that is parametric at the type of variables:
-- data Exp α = . . . | Var α | Abs α (Exp α)

-- Also consider a variant of the above with additional let
-- expressions, as often introduced by passes such as let-insertion:
-- data Exp α = . . . | Let α (Exp α) (Exp α

data ExpX' x a
  = Litx' (XLit' x a) Int
  | Varx' (XVar' x a) String
  | Annx' (XAnn' x a) (ExpX' x a) Typ
  | Absx' (XAbs' x a) String (ExpX' x a)
  | Appx' (XApp' x a) (ExpX' x a) (ExpX' x a)
  | Expx' (XExp' x a)

type family XLit' x a 
type family XVar' x a 
type family XAnn' x a 
type family XAbs' x a 
type family XApp' x a 
type family XExp' x a

-- We can introduce a let binding:
data LE
type ExpLE a = ExpX' LE a

type instance XLit' LE a = Void
type instance XVar' LE a = Void
type instance XAnn' LE a = Void
type instance XAbs' LE a = Void
type instance XApp' LE a = Void
type instance XExp' LE a = (a, ExpLE a, ExpLE a)

pattern LetLE :: a -> ExpLE a -> ExpLE a -> ExpLE a
pattern LetLE bndr t1 t2 = Expx' (bndr, t1, t2)


------------------------------
--- Tagless Final Encoding ---
------------------------------
-- TODO as associated type
class Semantics repr where
  var  :: String -> repr String
  int  :: Int -> repr Int
  add  :: repr Int -> repr Int -> repr Int
  lam  :: (repr a -> repr b) -> repr (a -> b)
  app  :: repr (a -> b) -> repr a -> repr b
--  exp  :: (XExp x)

t1 :: Semantics repr => repr Int
t1 = add (int 1) (int 2)

t2 :: Semantics repr => repr (Int -> Int)
t2 = lam $ \x -> (add x x)

newtype R a = R { unR :: a }
  deriving Show

instance Semantics R where
  var a = R $ a
  int i = R $ i
  add (R x) (R y) = R $ x + y
  lam f = R $ dimap R unR f
  app (R f) (R a) = R $ f a

eval :: R a -> a
eval e = unR e