-- |
-- Module      : LC
-- Description : Untyped lambda calculus
-- Stability   : experimental
--
-- An implementation of the untyped lambda calculus including evaluation
-- and small-step reduction.
--
-- This module demonstrates the use of well-scoped lambda calculus terms using `Rebound`.
-- The natural number index `n` is the scoping level -- a bound on the number
-- of free variables that can appear in the term. If `n` is 0, then the
-- term must be closed.
module LC where

import Data.Fin
import Data.Vec qualified
import Rebound
import Rebound.Bind.Single
import Data.Fin
import Data.Vec qualified
import qualified Data.Maybe as Maybe

-- | Datatype of well-scoped lambda-calculus expressions
--
-- The `Var` constructor of this datatype takes an index that must
-- be strictly less than the bound. Note that the type `Fin (S n)`
-- has `n` different elements.

-- The `Lam` constructor binds a variable, using the the type `Bind`
-- from the library. The type arguments state that the binder is
-- for a single expression variable, inside an expression term, that may
-- have at most `n` free variables.
data Exp (n :: Nat) where
  Var :: Fin n -> Exp n
  Lam :: Bind Exp Exp n -> Exp n
  App :: Exp n -> Exp n -> Exp n
    deriving (Generic1)
-- Enable generic derivation for substitution
-- deriving instance (Generic1 Exp)

 

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

-- To make it easier to construct lambda calculus
-- expressions, we'll first define some helper
-- definitions

-- | a lambda expression
lam :: Exp (S n) -> Exp n
lam = Lam . bind
-- | an application expression
(@@) :: Exp n -> Exp n -> Exp n
(@@) = App
-- | variable with index 0
v0 :: Exp (S n)
v0 = Var f0
-- | variable with index 1
v1 :: Exp (S (S n))
v1 = Var f1


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

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


-- For example, we can write
-- (λx. ((x ((λy. y) x)) (λz. z)))
-- using this term with de Bruijn indices
-- (λ. ((0 ((λ. 0) 0)) (λ. 0)))
-- and then construct it with the definitions above
t :: Exp Z
t = lam ((v0 @@ ((lam v0) @@ v0)) @@ (lam v0))

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

-- The nice thing about de Bruijn indices is that
-- we can use structural equality as alpha equivalence.
-- The built-in Eq instance for Bind, makes sure that 
-- the delayed substitutions are not observable here.
deriving instance Eq (Exp n)


----------------------------------------------
-- Substitution
----------------------------------------------

-- To work with this library, we need two type class instances.
-- First, we 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

-- Second, the operation `applyE` applies an environment
-- (explicit substitution) to an expression, and can be
-- automatically generated by the `Subst` type class, as
-- long as it can identify the variable constructor.
-- (Insead of generic programming, this operation can also
-- be written explicitly.)

instance Subst Exp Exp where
  isVar (Var x) = Just (Refl, x)
  isVar _ = Nothing


----------------------------------------------
-- 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)

-----------------------------------------------
-- (big-step) evaluation
-----------------------------------------------

-- | Calculate the value of a lambda-calculus expression
-- This function looks like it uses call-by-value evaluation:
-- in an application it evaluates the argument `e2` before
-- using the `instantiate` function from the library to substitute
-- the bound variable of `Bind` by v. However, this is Haskell,
-- a lazy language, so that result won't be evaluated unless the
-- function actually uses its argument.
eval :: Exp Z -> Exp Z
eval (Var x) = case x of {}
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


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

-- >>> eval (t `App` t0)
-- (λ. 0)

-- ((λ. (λ. 1)) ((λ. 0) (λ. 0)))

t2 = App (Lam (bind (Lam (bind (Var f1)))))
         (App (Lam (bind (Var f0))) (Lam (bind (Var f0))))

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

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

----------------------------------------------
-- small-step evaluation
----------------------------------------------

-- | Do one step of evaluation, if possible
-- If the function is already a value or is stuck
-- this function returns `Nothing`
step :: Exp n -> Maybe (Exp n)
step (Var x) = Nothing
step (Lam b) = Nothing
step (App (Lam b) e2) = Just (instantiate b e2)
step (App e1 e2)
  | Just e1' <- step e1 = Just (App e1' e2)
  | Just e2' <- step e2 = Just (App e1 e2')
  | otherwise = Nothing

-- | Evaluate the term as much as possible
eval' :: Int -> Exp n -> Maybe (Exp n)
eval' 0 e = Nothing
eval' k e = case step e of
              Just e' -> eval' (k - 1) e'
              Nothing -> Just e

-- >>> step (t0 `App` t0)
-- Just (λ. 0)

-- >>> eval' 5 (t `App` t0)
-- Just (λ. 0)


--------------------------------------------------------
-- full normalization
--------------------------------------------------------

-- | Calculate the normal form of a lambda expression. This
-- is like evaluation except that it also reduces underneath
-- the binders of `Lam` expressions. There, we must first `getBody`
-- the binder and then rebind when finished
nf :: Exp n -> Exp n
nf (Var x) = Var x
nf (Lam b) = Lam (bind (nf (getBody b)))
nf (App e1 e2) =
  case nf e1 of
    Lam b -> nf (instantiate b e2)
    t -> App t (nf e2)

--------------------------------------------------------
-- weak-head normalization / full reduction
--------------------------------------------------------

nf1 :: Exp n -> Exp n
nf1 (Var x) = Var x
nf1 (Lam b) = Lam (bind (nf1 (getBody b)))
nf1 (App e1 e2) =
  case whnf e1 of
    Lam b -> nf1 (instantiate b (whnf e2))
    t -> App t (nf e2)

whnf :: Exp n -> Exp n
whnf (Var x) = Var x
whnf (Lam b) = Lam b
whnf (App e1 e2) =
  case nf e1 of
    Lam b -> nf (instantiate b (whnf e2))
    t -> App t (nf e2)

--------------------------------------------------------
-- environment based evaluation / normalization
--------------------------------------------------------

-- invariant: expressions in the range of the environment are in whnf
whnfEnv :: Env Exp m n -> Exp m -> Exp n
whnfEnv r (Var x) = applyEnv r x
whnfEnv r (Lam b) = applyE r (Lam b)
whnfEnv r (App f a) =
  case whnfEnv r f of
     Lam b ->
       instantiateWith b (whnfEnv r a) whnfEnv
        -- unbindWith b (\r' e' -> whnfEnv (whnfrEnv r a .: r') e')
     f' -> App f' (applyE r a)

-- >>> whnfEnv zeroE t     -- start with "empty environment"
-- (λ. ((0 ((λ. 0) 0)) (λ. 0)))

-- For full reduction, we need to normalize under the binder too.
nfEnv :: Exp n -> Exp n
nfEnv (Var x) = Var x
nfEnv (Lam b) = Lam (bind (nfEnv (getBody b)))
nfEnv (App f a) =
   case whnfEnv idE f of
        Lam b -> nfEnv (instantiate b (whnfEnv idE a))
        f' -> App (nfEnv f') (nfEnv a)



----------------------------------------------------------------