{-# LANGUAGE DataKinds         #-}
{-# LANGUAGE DeriveFunctor     #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs             #-}
{-# LANGUAGE LambdaCase        #-}
{-# LANGUAGE PatternSynonyms   #-}
{-# LANGUAGE RankNTypes        #-}
{-# LANGUAGE TemplateHaskell   #-}
-- | Example implementation of untyped \(\lambda\)-calculus with the foil.
module Control.Monad.Free.Foil.Example where

import           Control.Monad.Foil
import           Control.Monad.Free.Foil
import           Data.Bifunctor.TH

-- $setup
-- >>> import Control.Monad.Foil

-- | Untyped \(\lambda\)-terms in scope @n@.
data ExprF scope term
  -- | Application of one term to another: \((t_1, t_2)\)
  = AppF term term
  -- | \(\lambda\)-abstraction introduces a binder and a term in an extended scope: \(\lambda x. t\)
  | LamF scope
  deriving ((forall a b. (a -> b) -> ExprF scope a -> ExprF scope b)
-> (forall a b. a -> ExprF scope b -> ExprF scope a)
-> Functor (ExprF scope)
forall a b. a -> ExprF scope b -> ExprF scope a
forall a b. (a -> b) -> ExprF scope a -> ExprF scope b
forall scope a b. a -> ExprF scope b -> ExprF scope a
forall scope a b. (a -> b) -> ExprF scope a -> ExprF scope b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall scope a b. (a -> b) -> ExprF scope a -> ExprF scope b
fmap :: forall a b. (a -> b) -> ExprF scope a -> ExprF scope b
$c<$ :: forall scope a b. a -> ExprF scope b -> ExprF scope a
<$ :: forall a b. a -> ExprF scope b -> ExprF scope a
Functor)
deriveBifunctor ''ExprF

pattern AppE :: AST binder ExprF n -> AST binder ExprF n -> AST binder ExprF n
pattern $mAppE :: forall {r} {binder :: S -> S -> *} {n :: S}.
AST binder ExprF n
-> (AST binder ExprF n -> AST binder ExprF n -> r)
-> ((# #) -> r)
-> r
$bAppE :: forall (binder :: S -> S -> *) (n :: S).
AST binder ExprF n -> AST binder ExprF n -> AST binder ExprF n
AppE x y = Node (AppF x y)

pattern LamE :: binder n l -> AST binder ExprF l -> AST binder ExprF n
pattern $mLamE :: forall {r} {binder :: S -> S -> *} {n :: S}.
AST binder ExprF n
-> (forall {l :: S}. binder n l -> AST binder ExprF l -> r)
-> ((# #) -> r)
-> r
$bLamE :: forall (binder :: S -> S -> *) (n :: S) (l :: S).
binder n l -> AST binder ExprF l -> AST binder ExprF n
LamE binder body = Node (LamF (ScopedAST binder body))

{-# COMPLETE Var, AppE, LamE #-}

type Expr = AST NameBinder ExprF

-- | Use 'ppExpr' to show \(\lambda\)-terms.
instance Show (Expr n) where
  show :: Expr n -> String
show = Expr n -> String
forall (n :: S). Expr n -> String
ppExpr

-- | Compute weak head normal form (WHNF) of a \(\lambda\)-term.
--
-- >>> whnf emptyScope (AppE (churchN 2) (churchN 2))
-- λx1. (λx0. λx1. (x0 (x0 x1)) (λx0. λx1. (x0 (x0 x1)) x1))
whnf :: Distinct n => Scope n -> Expr n -> Expr n
whnf :: forall (n :: S). Distinct n => Scope n -> Expr n -> Expr n
whnf Scope n
scope = \case
  AppE Expr n
fun Expr n
arg ->
    case Scope n -> Expr n -> Expr n
forall (n :: S). Distinct n => Scope n -> Expr n -> Expr n
whnf Scope n
scope Expr n
fun of
      LamE NameBinder n l
binder AST NameBinder ExprF l
body ->
        let subst :: Substitution (AST NameBinder ExprF) l n
subst = Substitution (AST NameBinder ExprF) n n
-> NameBinder n l
-> Expr n
-> Substitution (AST NameBinder ExprF) l n
forall (e :: S -> *) (i :: S) (o :: S) (i' :: S).
Substitution e i o -> NameBinder i i' -> e o -> Substitution e i' o
addSubst Substitution (AST NameBinder ExprF) n n
forall (e :: S -> *) (i :: S). InjectName e => Substitution e i i
identitySubst NameBinder n l
binder Expr n
arg
        in Scope n -> Expr n -> Expr n
forall (n :: S). Distinct n => Scope n -> Expr n -> Expr n
whnf Scope n
scope (Scope n
-> Substitution (AST NameBinder ExprF) l n
-> AST NameBinder ExprF l
-> Expr n
forall (sig :: * -> * -> *) (o :: S) (binder :: S -> S -> *)
       (i :: S).
(Bifunctor sig, Distinct o, CoSinkable binder, SinkableK binder) =>
Scope o
-> Substitution (AST binder sig) i o
-> AST binder sig i
-> AST binder sig o
substitute Scope n
scope Substitution (AST NameBinder ExprF) l n
subst AST NameBinder ExprF l
body)
      Expr n
fun' -> Expr n -> Expr n -> Expr n
forall (binder :: S -> S -> *) (n :: S).
AST binder ExprF n -> AST binder ExprF n -> AST binder ExprF n
AppE Expr n
fun' Expr n
arg
  Expr n
t -> Expr n
t

-- | Compute weak head normal form (WHNF) of a __closed__ \(\lambda\)-term.
--
-- >>> whnf' (AppE (churchN 2) (churchN 2))
-- λx1. (λx0. λx1. (x0 (x0 x1)) (λx0. λx1. (x0 (x0 x1)) x1))
whnf' :: Expr VoidS -> Expr VoidS
whnf' :: Expr 'VoidS -> Expr 'VoidS
whnf' = Scope 'VoidS -> Expr 'VoidS -> Expr 'VoidS
forall (n :: S). Distinct n => Scope n -> Expr n -> Expr n
whnf Scope 'VoidS
emptyScope

-- | Compute normal form (NF) of a \(\lambda\)-term.
--
-- >>> nf emptyScope (AppE (churchN 2) (churchN 2))
-- λx1. λx2. (x1 (x1 (x1 (x1 x2))))
nf :: Distinct n => Scope n -> Expr n -> Expr n
nf :: forall (n :: S). Distinct n => Scope n -> Expr n -> Expr n
nf Scope n
scope Expr n
expr = case Expr n
expr of
  LamE NameBinder n l
binder AST NameBinder ExprF l
body ->
    -- Instead of using 'assertDistinct',
    -- another option is to add 'Distinct l' constraint
    -- to the definition of 'LamE'.
    case NameBinder n l -> DistinctEvidence l
forall (n :: S) (pattern :: S -> S -> *) (l :: S).
(Distinct n, CoSinkable pattern) =>
pattern n l -> DistinctEvidence l
assertDistinct NameBinder n l
binder of
      DistinctEvidence l
Distinct ->
        let scope' :: Scope l
scope' = NameBinder n l -> Scope n -> Scope l
forall (n :: S) (l :: S). NameBinder n l -> Scope n -> Scope l
extendScope NameBinder n l
binder Scope n
scope
        in NameBinder n l -> AST NameBinder ExprF l -> Expr n
forall (binder :: S -> S -> *) (n :: S) (l :: S).
binder n l -> AST binder ExprF l -> AST binder ExprF n
LamE NameBinder n l
binder (Scope l -> AST NameBinder ExprF l -> AST NameBinder ExprF l
forall (n :: S). Distinct n => Scope n -> Expr n -> Expr n
nf Scope l
scope' AST NameBinder ExprF l
body)
  AppE Expr n
fun Expr n
arg ->
    case Scope n -> Expr n -> Expr n
forall (n :: S). Distinct n => Scope n -> Expr n -> Expr n
whnf Scope n
scope Expr n
fun of
      LamE NameBinder n l
binder AST NameBinder ExprF l
body ->
        let subst :: Substitution (AST NameBinder ExprF) l n
subst = Substitution (AST NameBinder ExprF) n n
-> NameBinder n l
-> Expr n
-> Substitution (AST NameBinder ExprF) l n
forall (e :: S -> *) (i :: S) (o :: S) (i' :: S).
Substitution e i o -> NameBinder i i' -> e o -> Substitution e i' o
addSubst Substitution (AST NameBinder ExprF) n n
forall (e :: S -> *) (i :: S). InjectName e => Substitution e i i
identitySubst NameBinder n l
binder Expr n
arg
        in Scope n -> Expr n -> Expr n
forall (n :: S). Distinct n => Scope n -> Expr n -> Expr n
nf Scope n
scope (Scope n
-> Substitution (AST NameBinder ExprF) l n
-> AST NameBinder ExprF l
-> Expr n
forall (sig :: * -> * -> *) (o :: S) (binder :: S -> S -> *)
       (i :: S).
(Bifunctor sig, Distinct o, CoSinkable binder, SinkableK binder) =>
Scope o
-> Substitution (AST binder sig) i o
-> AST binder sig i
-> AST binder sig o
substitute Scope n
scope Substitution (AST NameBinder ExprF) l n
subst AST NameBinder ExprF l
body)
      Expr n
fun' -> Expr n -> Expr n -> Expr n
forall (binder :: S -> S -> *) (n :: S).
AST binder ExprF n -> AST binder ExprF n -> AST binder ExprF n
AppE (Scope n -> Expr n -> Expr n
forall (n :: S). Distinct n => Scope n -> Expr n -> Expr n
nf Scope n
scope Expr n
fun') (Scope n -> Expr n -> Expr n
forall (n :: S). Distinct n => Scope n -> Expr n -> Expr n
nf Scope n
scope Expr n
arg)
  Expr n
t -> Expr n
t

-- | Compute normal form (NF) of a __closed__ \(\lambda\)-term.
--
-- >>> nf' (AppE (churchN 2) (churchN 2))
-- λx1. λx2. (x1 (x1 (x1 (x1 x2))))
nf' :: Expr VoidS -> Expr VoidS
nf' :: Expr 'VoidS -> Expr 'VoidS
nf' = Scope 'VoidS -> Expr 'VoidS -> Expr 'VoidS
forall (n :: S). Distinct n => Scope n -> Expr n -> Expr n
nf Scope 'VoidS
emptyScope

-- | Pretty print a name.
ppName :: Name n -> String
ppName :: forall (n :: S). Name n -> String
ppName Name n
name = String
"x" String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show (Name n -> Int
forall (l :: S). Name l -> Int
nameId Name n
name)

-- | Pretty-print a \(\lambda\)-term.
--
-- >>> ppExpr (churchN 3)
-- "\955x0. \955x1. (x0 (x0 (x0 x1)))"
ppExpr :: Expr n -> String
ppExpr :: forall (n :: S). Expr n -> String
ppExpr = \case
  Var Name n
name -> Name n -> String
forall (n :: S). Name n -> String
ppName Name n
name
  AppE Expr n
x Expr n
y -> String
"(" String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Expr n -> String
forall (n :: S). Expr n -> String
ppExpr Expr n
x String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
" " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Expr n -> String
forall (n :: S). Expr n -> String
ppExpr Expr n
y String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
")"
  LamE NameBinder n l
binder AST NameBinder ExprF l
body -> String
"λ" String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Name l -> String
forall (n :: S). Name n -> String
ppName (NameBinder n l -> Name l
forall (n :: S) (l :: S). NameBinder n l -> Name l
nameOf NameBinder n l
binder) String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
". " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> AST NameBinder ExprF l -> String
forall (n :: S). Expr n -> String
ppExpr AST NameBinder ExprF l
body

-- | A helper for constructing \(\lambda\)-abstractions.
lam :: Distinct n => Scope n -> (forall l. DExt n l => Scope l -> NameBinder n l -> Expr l) -> Expr n
lam :: forall (n :: S).
Distinct n =>
Scope n
-> (forall (l :: S).
    DExt n l =>
    Scope l -> NameBinder n l -> Expr l)
-> Expr n
lam Scope n
scope forall (l :: S). DExt n l => Scope l -> NameBinder n l -> Expr l
mkBody = Scope n
-> (forall {l :: S}. DExt n l => NameBinder n l -> Expr n)
-> Expr n
forall (n :: S) r.
Distinct n =>
Scope n -> (forall (l :: S). DExt n l => NameBinder n l -> r) -> r
withFresh Scope n
scope ((forall {l :: S}. DExt n l => NameBinder n l -> Expr n) -> Expr n)
-> (forall {l :: S}. DExt n l => NameBinder n l -> Expr n)
-> Expr n
forall a b. (a -> b) -> a -> b
$ \NameBinder n l
x ->
  let scope' :: Scope l
scope' = NameBinder n l -> Scope n -> Scope l
forall (n :: S) (l :: S). NameBinder n l -> Scope n -> Scope l
extendScope NameBinder n l
x Scope n
scope
   in NameBinder n l -> AST NameBinder ExprF l -> Expr n
forall (binder :: S -> S -> *) (n :: S) (l :: S).
binder n l -> AST binder ExprF l -> AST binder ExprF n
LamE NameBinder n l
x (Scope l -> NameBinder n l -> AST NameBinder ExprF l
forall (l :: S). DExt n l => Scope l -> NameBinder n l -> Expr l
mkBody Scope l
scope' NameBinder n l
x)

-- | Church-encoding of a natural number \(n\).
--
-- >>> churchN 0
-- λx0. λx1. x1
--
-- >>> churchN 3
-- λx0. λx1. (x0 (x0 (x0 x1)))
churchN :: Int -> Expr VoidS
churchN :: Int -> Expr 'VoidS
churchN Int
n =
  Scope 'VoidS
-> (forall {l :: S}.
    DExt 'VoidS l =>
    Scope l -> NameBinder 'VoidS l -> Expr l)
-> Expr 'VoidS
forall (n :: S).
Distinct n =>
Scope n
-> (forall (l :: S).
    DExt n l =>
    Scope l -> NameBinder n l -> Expr l)
-> Expr n
lam Scope 'VoidS
emptyScope ((forall {l :: S}.
  DExt 'VoidS l =>
  Scope l -> NameBinder 'VoidS l -> Expr l)
 -> Expr 'VoidS)
-> (forall {l :: S}.
    DExt 'VoidS l =>
    Scope l -> NameBinder 'VoidS l -> Expr l)
-> Expr 'VoidS
forall a b. (a -> b) -> a -> b
$ \Scope l
sx NameBinder 'VoidS l
nx ->
    Scope l
-> (forall {l :: S}.
    DExt l l =>
    Scope l -> NameBinder l l -> Expr l)
-> Expr l
forall (n :: S).
Distinct n =>
Scope n
-> (forall (l :: S).
    DExt n l =>
    Scope l -> NameBinder n l -> Expr l)
-> Expr n
lam Scope l
sx ((forall {l :: S}. DExt l l => Scope l -> NameBinder l l -> Expr l)
 -> Expr l)
-> (forall {l :: S}.
    DExt l l =>
    Scope l -> NameBinder l l -> Expr l)
-> Expr l
forall a b. (a -> b) -> a -> b
$ \Scope l
_sxy NameBinder l l
ny ->
      let x :: Expr l
x = Expr l -> Expr l
forall (e :: S -> *) (n :: S) (l :: S).
(Sinkable e, DExt n l) =>
e n -> e l
sink (Name l -> Expr l
forall (n :: S) (binder :: S -> S -> *) (sig :: * -> * -> *).
Name n -> AST binder sig n
Var (NameBinder 'VoidS l -> Name l
forall (n :: S) (l :: S). NameBinder n l -> Name l
nameOf NameBinder 'VoidS l
nx))
          y :: Expr l
y = Name l -> Expr l
forall (n :: S) (binder :: S -> S -> *) (sig :: * -> * -> *).
Name n -> AST binder sig n
Var (NameBinder l l -> Name l
forall (n :: S) (l :: S). NameBinder n l -> Name l
nameOf NameBinder l l
ny)
       in (Expr l -> Expr l) -> Expr l -> [Expr l]
forall a. (a -> a) -> a -> [a]
iterate (Expr l -> Expr l -> Expr l
forall (binder :: S -> S -> *) (n :: S).
AST binder ExprF n -> AST binder ExprF n -> AST binder ExprF n
AppE Expr l
x) Expr l
y [Expr l] -> Int -> Expr l
forall a. HasCallStack => [a] -> Int -> a
!! Int
n