module STLC where

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

{-
The language-style definition would look like:
 
 t ::= x | λx:T.t | (t t)
     | true | false | if t t t | t = t
     | # | t + t | t - t | t * t
     | nil T | cons t t | head t | tail t | isnil t
     | fix t
 
 T ::= 𝔹 | ℤ | T -> T | [T]
 
 v ::= λx:T.t | true | false | # | nil T | cons t t
 
-}

data Tm = Var String | Lam String Ty Tm | App Tm Tm 
        | Tru | Fls | If Tm Tm Tm | Equal Tm Tm
        | Num Int | Add Tm Tm | Sub Tm Tm | Mul Tm Tm
        | Nil Ty | Cons Tm Tm | Head Tm | Tail Tm | IsNil Tm
        | Fix Tm
          deriving (Show, Eq)

data Ty = TyBool | TyInt | Ty :->: Ty | TyList Ty deriving (Show, Eq)


-- Note: that :->: is an exampl of an "infix constructor". We could have used
-- "TyMap Ty Ty" instead, but it's nice to see the arrows like in our formal
-- presentations.

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

-- rather than use a separate datatype for values, we just define a predicate.
is_val :: Tm -> Bool
is_val (Lam _ _ _) = True
is_val Tru         = True
is_val Fls         = True
is_val (Num _)     = True
is_val (Nil _)     = True
is_val (Cons _ _)  = True
is_val (Fix t)     = False  -- remember, Fix is not a value! (this equation
                            -- could be deleted and be caught by the next)
is_val _           = False

--------------------------------------------------------------------------------
--------------------------------------------------------------------------------
{-
For the most part, evaluation is unchanged! Other than
pattern-matching including something for our lambdas' types, which
doesn't even need to be inspected, we have the same definition as
before.
-}

--------------------------------
eval :: Tm -> Tm
--------------------------------
eval t | is_val t = t
--------------------------------
eval (Var s) = error $ "found a free variable, whoops! ("++s++")"
eval (App t1 t2) | not (is_val t1) = eval (App (eval t1) t2)
                 | not (is_val t2) = eval (App t1 (eval t2))
eval (App (Lam s _ t) arg) = eval (subst s arg t)
--------------------------------
eval (If Tru b c) = eval b
eval (If Fls b c) = eval c
eval (If a b c) = eval (If (eval a) b c)
eval (Equal a b) = if (eval a) `alpha_equiv` (eval b) then Tru else Fls
--------------------------------
eval (Add a b) = case (eval a, eval b) of (Num a', Num b') -> Num (a' + b')
eval (Sub a b) = case (eval a, eval b) of (Num a', Num b') -> Num (a' - b')
eval (Mul a b) = case (eval a, eval b) of (Num a', Num b') -> Num (a' * b')
--------------------------------
eval (Head (Cons a b)) = a
eval (Head t) = eval (Head (eval t))
eval (Tail (Cons a b)) = b
eval (Tail t) = eval (Tail (eval t))
eval (IsNil (Nil _)) = Tru
eval (IsNil (Cons _ _)) = Fls
-- one example of trying to catch typecheck-failing things...
eval (IsNil t) = if (is_val t) then error "non-list isnil'd!" else eval (IsNil (eval t))
--------------------------------
eval (Fix (Lam s ty t)) = eval (subst s (Fix (Lam s ty t)) t)
eval (Fix t) = eval (Fix (eval t))
--------------------------------
-- this is mostly useful when we add extensions, as it shows what we forgot.
-- it shouldn't actually be used by (well-typed) terms, ever.
eval others = error $ "couldn't evaluate: "++show others

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

{- no change, other than leaving types alone when preserving structure. -}

--------------------------------
-- substitution helps us evaluate an application of a lambda with a value.
subst :: String -> Tm -> Tm -> Tm
subst v rep (Var s) | v==s      = rep
                    | otherwise = (Var s)

subst v rep (Lam x ty t) | v==x      = Lam x ty t
                         | otherwise = Lam x ty (subst v rep t)

subst v rep (App t1 t2) = App (subst v rep t1) (subst v rep t2)
--------------------------------
subst v rep Tru = Tru
subst v rep Fls = Fls
subst v rep (If a b c) = If (subst v rep a) (subst v rep b) (subst v rep c)
--------------------------------
subst _ _ (Num n) = Num n
subst v rep (Add a b) = Add (subst v rep a) (subst v rep b)
subst v rep (Sub a b) = Sub (subst v rep a) (subst v rep b)
subst v rep (Mul a b) = Mul (subst v rep a) (subst v rep b)
--------------------------------
subst v rep (Nil ty) = Nil ty
subst v rep (Cons a b) = Cons (subst v rep a) (subst v rep b)
subst v rep (Head a) = Head (subst v rep a)
subst v rep (Tail a) = Tail (subst v rep a)
subst v rep (IsNil a) = IsNil (subst v rep a)
--------------------------------
subst v rep (Fix a) = Fix (subst v rep a)
--------------------------------

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

{- no change, other than copying over types when preserving structure. -}

{-

how do we implement the equal evaluation over functions? We need "alpha equivalence",
meaning we need to know if the only difference in two functions is the chosen parameter
names (which doesn't really matter). Thus we need to be able to rename a variable (with
the rename function below), and then we can consider whether two terms are "alpha-equivalent"
(the same up to naming).

These two functions are only called in evaluation, thanks to "eval (Equal a b)".
We could delete them entirely if we didn't want the equals-operator to be a
language extension.

-}


-- for convenience, let's be able to change the variable used in a lambda.
rename sold snew (Lam s ty t) = Lam (if s==sold then snew else s) ty (rename sold snew t)
rename sold snew (Var s) = if s==sold then Var snew else Var s
rename sold snew (App t1 t2) = App (rename sold snew t1) (rename sold snew t2)
--
rename sold snew (If a b c) = If (rename sold snew a) (rename sold snew b) (rename sold snew c)
rename _ _ Tru = Tru
rename _ _ Fls = Fls
rename sold snew (Equal a b) = Equal (rename sold snew a) (rename sold snew b)
-- 
rename sold snew (Add a b) = Add (rename sold snew a) (rename sold snew b)
rename sold snew (Sub a b) = Sub (rename sold snew a) (rename sold snew b)
rename sold snew (Mul a b) = Mul (rename sold snew a) (rename sold snew b)
rename _ _ (Num n) = Num n
--
rename sold snew (Nil ty) = Nil ty
rename sold snew (Cons a b) = Cons (rename sold snew a) (rename sold snew b)
rename sold snew (Head  a) = Head (rename sold snew a)
rename sold snew (Tail  a) = Tail (rename sold snew a)
rename sold snew (IsNil a) = IsNil (rename sold snew a)
--
rename sold snew (Fix a) = Fix (rename sold snew a)
--

{- only change is requiring types to match in lambdas. The rest is just structural recursion. -}

-- if we re-name lambda variables throughout to match, are these two terms identical?
alpha_equiv (Var a) (Var b) = a==b
alpha_equiv lam1@(Lam s1 ty1 t1) lam2@(Lam s2 ty2 t2) = 
    if ty1 /= ty2  then False   else
        if s1==s2  then alpha_equiv t1 t2  else alpha_equiv lam1 (rename s2 s1 lam2)
alpha_equiv (App t11 t12) (App t21 t22) = alpha_equiv t11 t12 && alpha_equiv t21 t22
--
alpha_equiv Tru Tru = True
alpha_equiv Fls Fls = True
alpha_equiv (If  a b c) (If  x y z) = (alpha_equiv a x) && (alpha_equiv b y) && (alpha_equiv c z)
alpha_equiv (Equal a b) (Equal x y) = (alpha_equiv a x) && (alpha_equiv b y)
--
alpha_equiv (Num a)   (Num b)   = a==b
alpha_equiv (Add a b) (Add x y) = (alpha_equiv a x) && (alpha_equiv b y)
alpha_equiv (Sub a b) (Sub x y) = (alpha_equiv a x) && (alpha_equiv b y)
alpha_equiv (Mul a b) (Mul x y) = (alpha_equiv a x) && (alpha_equiv b y)
--
alpha_equiv (Nil tya) (Nil tyb)       = tya==tyb
alpha_equiv (Cons a1 b1) (Cons a2 b2) = (alpha_equiv a1 a2) && (alpha_equiv b1 b2)
alpha_equiv (Head  a1) (Head  a2)     =  alpha_equiv a1 a2
alpha_equiv (Tail  a1) (Tail  a2)     =  alpha_equiv a1 a2
alpha_equiv (IsNil a1) (IsNil a2)     =  alpha_equiv a1 a2
--
alpha_equiv (Fix a1) (Fix a2) = alpha_equiv a1 a2
--
alpha_equiv a b = False

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

-- driver function that starts up the process with an empty environment, [].
-- the real work happens in the tc function below.
typecheck :: Tm -> Ty
typecheck tm = tc [] tm

-- NOTE: it's quite common to always have exactly one typechecking
-- rule per term.  we don't need to repeatedly substitute, reduce,
-- collapse structures with multiple rules as we did in evaluation, we
-- just recurse on sub-terms and use their types to describe/find our
-- own.

-- by the way, Haskell accepts unicode! Alas, Γ is an upper-case letter so
-- we can't use it for pattern-matching...but it could be the type synonym.

type Γ     = [(String, Ty)]   -- valid! But we'll use Gamma throughout the file.
type Gamma = [(String,Ty)]
--------------------------------
--------------------------------
-- tc :: Γ  -> Tm -> Ty   This is valid!
tc :: Gamma -> Tm -> Ty
--------------------------------
tc gamma (Var s) =
  case lookup s gamma of
    Just ty -> ty
    Nothing -> error $ "found free variable ("++s++")."

tc gamma (Lam s ty tm) =
  let ty_range = tc ((s,ty):gamma) tm
  in ty :->: ty_range

tc gamma (App t1 t2) =
  let ty1 = tc gamma t1
      ty2 = tc gamma t2
  in case ty1 of
    (tyD :->: tyR) -> if tyD==ty2 
                      then tyR
                      else error $ "incorrect type of argument.\n\n"
                                 ++"\texpected:   "++show tyD++"\n"
                                 ++"\tfound:      "++show ty2++"\n"
    _ -> error $ "needed function type in application. (got "++show ty1++")"
--------------------------------
tc _ Tru = TyBool
tc _ Fls = TyBool

tc gamma (If a b c) =
  let tyA = tc gamma a
      tyB = tc gamma b
      tyC = tc gamma c
  in if tyA==TyBool
     then if tyB==tyC
          then tyB
          else error "if-branches' types didn't match."
     else error ("if-expr's guard wasn't a boolean. (found "++(show tyA)++".)")

tc gamma (Equal t1 t2) = if tc gamma t1 == tc gamma t2
                         then TyBool
                         else error "can't equate differently-typed things."
--------------------------------
tc gamma (Num _) = TyInt

tc gamma (Add a b) =
  case (tc gamma a, tc gamma b) of
    (TyInt, TyInt) -> TyInt
    _ -> error $ "saw add of something other than two numbers."

tc gamma (Sub a b) =
  case (tc gamma a, tc gamma b) of
    (TyInt, TyInt) -> TyInt
    _ -> error $ "saw sub of something other than two numbers."

tc gamma (Mul a b) =
  case (tc gamma a, tc gamma b) of
    (TyInt, TyInt) -> TyInt
    _ -> error $ "saw mul of something other than two numbers."

                                       
--------------------------------
tc gamma (Nil ty) = TyList ty

tc gamma (Cons a b) =
   case (tc gamma b) of
     TyList bty -> if bty == tc gamma a
                   then TyList bty
                   else error $ "can't cons a "++show (tc gamma a)
                        ++" onto a list of "++show bty++"."
     _   -> error "second cons argument must be a list."

tc gamma (Head t) = case tc gamma t of
  TyList aty -> aty
  _ -> error "can't call head of non-list."

tc gamma (Tail t) = case tc gamma t of
  TyList aty -> TyList aty
  _ -> error "can't call tail of non-list."

tc gamma (IsNil t) = case tc gamma t of
  TyList anyType -> TyBool
  _ -> error "can't call isnil of non-list."
--------------------------------
tc gamma (Fix t) = case tc gamma t of
  tyA :->: tyB -> if tyA==tyB
                  then tyA
                  else error "fix needs argument with type of shape T->T."
  _ -> error "can't call fix on non-function."

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

-- force (typecheck tm) to be evaluated (be strict about it), but then
-- throw away the result and just evaluate it. Without seq, typecheck
-- tm's result would never be needed and we would just lazily skip it!
-- :-o
evaluate tm = typecheck tm `seq` eval tm

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

id_bool = Lam "b" TyBool (Var "b")
id_num  = Lam "n" TyInt (Var "n")

e1 = App id_bool Tru
e2 = App id_num (Num 4)
f0 = App id_num Tru  -- should fail with our own error msg.
{-
try it out:

    *STLC> typecheck e1
    TyBool
    *STLC> typecheck e2
    TyInt
    *STLC> typecheck f0
    *** Exception: incorrect type of argument.
    
    	expected:   TyInt
    	found:      TyBool
    
    CallStack (from HasCallStack):
      error, called at STLC.hs:240:28 in main:STLC
    *STLC> 



    *STLC> evaluate e1
    Tru
    *STLC> evaluate e2
    Num 4
    *STLC> evaluate f0
    *** Exception: incorrect type of argument.
    
    	expected:   TyInt
    	found:      TyBool
    
    CallStack (from HasCallStack):
      error, called at STLC.hs:240:28 in main:STLC
    *STLC> 

-}

-- remember, we can create Tm values that aren't actually in our
-- language, but our typecheck function will weed out the rest of
-- these impostors.
f1 = App id_bool (Num 1)
f2 = App id_num Tru

-- if you only call eval, and not tc/typecheck, it'll ignore types and
-- happen to do something.

{-
try it out:
    typecheck f1  -- fails
    typecheck f2  -- fails
    eval f1       -- returns (Num 2). Task failed successfully.
    eval f2       -- return Tru.      Task failed successfully.      
-}


inc = Lam "x" TyInt (Add (Num 1) (Var "x"))
mul5 = Lam "v" TyInt (Mul (Num 5) (Var "v"))

maybeAdd5 = Lam "b" TyBool $ Lam "n" TyInt $ If (Var "b")    
                                                (Add (Var"n") (Num 5))
                                                (Var"n")

-- example recursive program. Not much trickery after the first line
-- (finding the type can take practice, but then it's simple.)
sumints = Fix $ Lam "self" (TyList TyInt :->: TyInt)
              $ Lam "xs" (TyList TyInt)
              $ If (IsNil (Var "xs"))
                   (Num 0)
                   (Add (Head (Var "xs")) (App (Var "self") (Tail (Var "xs"))))


-- [True, False, True] :: [TyBool]
bools3 = Cons Tru (Cons Fls (Cons Tru (Nil TyBool)))

-- [2,4,6,8,10] :: [TyInt]
nums5  = Cons (Num 2)
         (Cons (Num 4)
          (Cons (Num 6)
           (Cons (Num 8)
            (Cons (Num 10)
             (Nil TyInt)))))


{-
try it out:
    *STLC> typecheck (App sumints nums5)
    TyInt
    *STLC> eval (App sumints nums5)
    Num 30
    *STLC> evaluate (App sumints nums5)
    Num 30
-}