-- |
-- Module      : LC
-- Description : Untyped lambda calculus
-- Stability   : experimental
--
-- An implementation of the untyped lambda calculus including let, letrec,
-- mutual letrec and let* expressions.
-- TODO: add example terms and fix Show instance
module LCLet where

import Rebound
import Rebound.Bind.Single
import qualified Rebound.Bind.Pat as Pat
import Rebound.Bind.PatN as PatN (BindN, bindN, instantiateN, getBodyN)
import Data.Fin
import Data.Vec qualified as Vec

-- | Datatype of well-scoped lambda-calculus expressions
data Exp (n :: Nat) where
  Var :: Fin n -> Exp n
  Lam :: Bind Exp Exp n -> Exp n
  App :: Exp n -> Exp n -> Exp n
  Let ::
    -- | single let expression
    -- "let x = e1 in e2" where x is bound in e2
    Exp n ->
    (Bind Exp Exp n) ->
    Exp n
  LetRec ::
    -- | "let rec x = e1 in e2" where x is bound in both e1 and e2
    Rec n ->
    Exp n
  LetTele ::
    -- | sequence of nested lets, where each one may depend on
    -- the previous binding
    -- "let x1 = e1 in x2 = e2 in ... in e" where x1 is bound
    -- in e2, e3 ... and e, x2 is bound in e3 and e, etc.
    Tele n ->
    Exp n
  LetMutRec ::
    -- | mutual recursive lets, where each one may depend on
    -- any other variable
    -- "let x1 = e1 in x2 = e2 in ... in e" where x1 ... xn
    -- are bound in e1, e2, e3 ... and e
    MutRec n ->
    Exp n

data Rec n =
  Rec { rec_rhs  :: Bind Exp Exp n,   -- single RHS
        rec_body :: Bind Exp Exp n }  -- body of let

data MutRec n = forall m. SNatI m =>
  MutRec { mutrec_rhss :: Vec m (BindN Exp Exp m n), -- Vector of RHSs
           mutrec_body :: BindN Exp Exp m n  -- body of let
           }

data Tele n where
  LetStar :: Exp n -> Bind Exp Tele n -> Tele n
  Body :: Exp n -> Tele n

----------------------------------------------
-- Example lambda-calculus expressions
----------------------------------------------

-- some variables
v0 :: Exp (S n)
v0 = Var f0
v1 :: Exp (S (S n))
v1 = Var f1
v2 :: Exp (S (S (S n)))
v2 = Var f2
-- | an application expression
(@@) :: Exp n -> Exp n -> Exp n
(@@) = App
-- | a lambda expression
lam :: Exp (S n) -> Exp n
lam = Lam . bind

letrec :: Exp (S n) -> Exp (S n) -> Exp n
letrec e1 e2 = LetRec (Rec (bind e1) (bind e2))

letstar :: Exp n -> Tele (S n) -> Tele n
letstar e t = LetStar e (bind t)

-- | The identity function "λ x. x".
-- With de Bruijn indices we write it as "λ. 0"
-- The `bind` function creates the binder
-- t0 :: Exp Z
t0 = lam v0

-- >>> t0
-- (λ. 0)

-- | A larger term "λ x. λy. x ((λ z. z) y)"
-- λ. λ. 1 ((λ. 0) 0))
t1 :: Exp Z
t1 = lam (lam (v1 @@ ((lam v0) @@ v0)))

-- >>> t1
-- (λ. (λ. (1 ((λ. 0) 0))))

-- let x = \y.y in x x
t2 :: Exp Z
t2 = Let t0 (bind (App v0 v0))

-- >>> t2
-- (let (λ. 0) in (0 0))

-- let rec fix = \f. f (fix f) in f
t3 :: Exp Z
t3 = letrec (lam (v0 @@(v1 @@ v0))) v0
-- >>> t3
-- (let rec (λ. (0 (1 0))) in 0)

-- let* x1 = \x.x ; x2 = x1 x1 ; x3 = x2 s1 in x3 x2 x1
t4 = LetTele
       (letstar t0
         (letstar (v0 @@ v0)
            (letstar (v0 @@ v1)
              (Body ((v0 @@ v1) @@ v2)))))

-- >>> t4
-- <let-tele>

----------------------------------------------
-- (Alpha-)Equivalence
----------------------------------------------

-- | The derivable equality instance
-- is alpha-equivalence
deriving instance (Eq (Exp n))

deriving instance (Eq (Tele n))

deriving instance (Eq (Rec n))

instance Eq (MutRec n) where
  (==) :: MutRec n -> MutRec n -> Bool
  MutRec { mutrec_rhss= (rhss1 :: Vec m1 t1), mutrec_body=body1} ==
    MutRec { mutrec_rhss= (rhss2 :: Vec m2 t2), mutrec_body=body2}
    = case testEquality (snat @m1) (snat @m2) of
       Just Refl -> Vec.all2 (==) rhss1 rhss2 && body1 == body2
       Nothing -> False
----------------------------------------------
-- Substitution
----------------------------------------------

-- To work with this library, we need two type class instances.

-- | Tell the library how to construct variables in the expression
-- type. This class is necessary to construct an indentity
-- substitution---one that maps each variable to itself.
instance SubstVar Exp where
  var :: Fin n -> Exp n
  var = Var

-- The library represents a substitution using an "Environment".
-- The type `Env Exp n m` is a substitution that can be applied to
-- indices bounded by n. It produces a result `Exp` with indices
-- bounded by m. The function `applyEnv` looks up a mapping in
-- an environment.

-- | The operation `applyE` applies an environment
-- (explicit substitution) to an expression.
--
-- The implementation of this operation applies the environment to
-- variable index in the variable case. All other caseas follow
-- via recursion. The library includes a type class instance for
-- the Bind type which handles the variable lifting needed under
-- the binder.
instance Shiftable Exp where
  shift = shiftFromApplyE @Exp
instance Subst Exp Exp where
  applyE :: Env Exp n m -> Exp n -> Exp m
  applyE r (Var x) = applyEnv r x
  applyE r (Lam b) = Lam (applyE r b)
  applyE r (App e1 e2) = App (applyE r e1) (applyE r e2)
  applyE r (Let e1 e2) = Let (applyE r e1) (applyE r e2)
  applyE r (LetRec e) = LetRec (applyE r e)
  applyE r (LetTele e) = LetTele (applyE r e)
  applyE r (LetMutRec e) = LetMutRec (applyE r e)

instance Subst Exp Rec where
  applyE r (Rec rhs body) = Rec (applyE r rhs) (applyE r body)

instance Shiftable Tele where
  shift = shiftFromApplyE @Exp
instance Subst Exp Tele where
  applyE r (Body e) = Body (applyE r e)
  applyE r (LetStar e1 e2) = LetStar (applyE r e1) (applyE r e2)

instance Subst Exp MutRec where
  applyE r (MutRec rhss body) =
    MutRec (fmap (applyE r) rhss) (applyE r body)

----------------------------------------------
-- Display (Show)
----------------------------------------------

-- | To show lambda terms, we use a simple recursive instance of
-- Haskell's `Show` type class. In the case of a binder, we use the `getBody`
-- operation to access the body of the lambda expression.
instance Show (Exp n) where
  showsPrec :: Int -> Exp n -> String -> String
  showsPrec _ (Var x) = shows x
  showsPrec d (App e1 e2) =
    showParen True $
      showsPrec 10 e1
        . showString " "
        . showsPrec 11 e2
  showsPrec d (Lam b) =
    showParen True $
      showString "λ. "
        . shows (getBody b)
  showsPrec d (Let e1 e2) =
    showParen True $
      showString "let "
        . showsPrec 10 e1
        . showString " in "
        . shows (getBody e2)
  showsPrec d (LetRec (Rec{rec_rhs=rhs,rec_body=body})) =
    showParen True $
      showString "let rec "
        . shows (getBody rhs)
        . showString " in "
        . shows (getBody body)
  showsPrec d (LetTele e) = showString "<let-tele>"
  showsPrec d (LetMutRec (MutRec {mutrec_rhss=rhss, mutrec_body=body})) =
    showParen True $
      showString "let rec "
        . showString " rhss " -- (getBodyN e1)
        . showString " in "
        . shows (getBodyN body)
-----------------------------------------------
-- (big-step) evaluation
-----------------------------------------------

-- >>> eval t1
-- (λ. (λ. (1 ((λ. 0) 0))))

-- >>> eval (t1 @@ t0)
-- (λ. ((λ. 0) ((λ. 0) 0)))

-- >>> eval t2
-- (λ. 0)

-- This one is an infinite loop
-- >>> eval t3
-- ProgressCancelledException

-- >>> eval t4
-- (λ. 0)

eval :: Exp n -> Exp n
eval (Var x) = Var x
eval (Lam b) = Lam b
eval (App e1 e2) =
  let v = eval e2
   in case eval e1 of
        Lam b -> eval (instantiate b v)
        t -> App t v
eval (Let e1 e2) =
  eval (instantiate e2 (eval e1))
eval (LetRec e) =
  -- use a Haskell recursive definition
  -- to tie the knot for a recursive definition
  -- in the object language
  let v = instantiate (rec_rhs e) v
   in eval (instantiate (rec_body e) v)
eval (LetTele e) = evalTele e
eval (LetMutRec (MutRec { mutrec_rhss = rhss, mutrec_body = body})) =
  -- use a Haskell recursive definition
  let vs = fmap (\b -> instantiateN b vs) rhss
  in eval (instantiateN body vs)

evalTele :: Tele n -> Exp n
evalTele (Body e) = eval e
evalTele (LetStar e t) =
  evalTele (instantiate t (eval e))