Simple.hs

{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# OPTIONS_GHC -Wno-unticked-promoted-constructors #-}
module Simple where

import Data.Kind (Type)

data Nat
  = Zero
  | Succ Nat

data Fin :: Nat -> Type where
  FZero :: Fin (Succ n)
  FSucc :: Fin n -> Fin (Succ n)

deriving instance Eq (Fin n)

instance Show (Fin n) where
  show = show . toInt where
    toInt :: Fin n -> Int
    toInt FZero = 0
    toInt (FSucc f) = 1 + toInt f

data Var (n :: Nat)
  = Bound (Fin n)
  | Free String
  deriving Show

infixl 7 :@:
data Expr (n :: Nat)
  = Ref (Var n)
  | Lam (Expr (Succ n))
  | Expr n :@: Expr n
  deriving Show

var :: Fin n -> Expr n
var = Ref . Bound

iComb :: Expr (Succ n)
iComb = Lam (Ref (Bound FZero))

kComb :: Expr (Succ (Succ n))
kComb = Lam (Lam (Ref (Bound (FSucc FZero))))

sComb :: Expr (Succ (Succ (Succ n)))
sComb = Lam (Lam (Lam (var (FSucc (FSucc FZero)) :@: var FZero :@: (var (FSucc FZero) :@: var FZero))))

wComb :: Expr (Succ (Succ (Succ n)))
wComb = sComb :@: sComb :@: (sComb :@: kComb)

skk :: Expr (Succ (Succ (Succ n)))
skk = sComb :@: kComb :@: kComb

pretty :: Expr n -> String
pretty = \case
  Ref v -> case v of
    Bound f -> show f
    Free str -> str
  Lam e -> "\\" ++ pretty e
  e1 :@: e2 -> s1 ++ " " ++ s2
    where
      parens e = "(" ++ pretty e ++ ")"
      s1 = case e1 of
        Lam _ -> parens e1
        _ -> pretty e1
      s2 = case e2 of
        Lam _ -> parens e2
        _ :@: _ -> parens e2
        _ -> pretty e2

lower :: Fin (Succ n) -> Fin (Succ n) -> Maybe (Fin n)
lower i1 f1 = case i1 of
  FZero -> case f1 of
    FZero -> Nothing
    FSucc f2 -> Just f2
  FSucc i2 -> case f1 of
    FZero -> case i2 of
      FZero -> Just FZero
      FSucc _ -> Just FZero
    FSucc f2 -> case i2 of
      FZero -> fmap FSucc (lower FZero f2)
      FSucc _ -> fmap FSucc (lower i2 f2)

lift :: Expr n -> Expr (Succ n)
lift = go FZero
  where
    go :: Fin (Succ n) -> Expr n -> Expr (Succ n)
    go i = \case
      Ref r -> case r of
        Bound f -> Ref (Bound (raise i f))
        Free str -> Ref (Free str)
      Lam l -> Lam (go (FSucc i) l)
      e1 :@: e2 -> go i e1 :@: go i e2
    raise :: Fin (Succ n) -> Fin n -> Fin (Succ n)
    raise i1 f1 = case i1 of
      FZero -> FSucc f1
      FSucc i2 -> case f1 of
        FZero -> FZero
        FSucc f2 -> FSucc (raise i2 f2)

subst :: Fin (Succ n) -> Expr n -> Expr (Succ n) -> Expr n
subst i v = \case
  Ref r -> case r of
    Bound f -> case lower i f of
      Nothing -> v
      Just newF1 -> Ref (Bound newF1)
    Free str -> Ref (Free str)
  Lam e -> Lam (subst (FSucc i) (lift v) e)
  e1 :@: e2 -> subst i v e1 :@: subst i v e2

nf :: Expr n -> Expr n
nf = \case
  Ref v -> Ref v
  Lam e -> Lam (nf e)
  e1 :@: e2 -> case nf e1 of
    Lam l -> nf (subst FZero e2 l)
    _ -> nf e1 :@: nf e2

test :: Expr n -> IO ()
test e = do
  putStrLn "Before Normalization:"
  putStrLn (pretty e)
  putStrLn "After Normalization:"
  putStrLn (pretty (nf e))