module Choc.SimplyTyped where

import Data.Functor (($>))
import Control.Monad (guard)

import Choc.Parser

----------------------
-- Type Definitions --
----------------------

data Type = Bass
          | Arrow Type Type
          deriving (Eq)

data Term = Lambda Type Term
          | App Term Term
          | Var Integer
          | Const
          deriving (Eq)

instance Show Type where
  show Bass = "B"
  show (Arrow b1 b2) = "(-> " <> show b1 <> " " <> show b2 <> ")"

instance Show Term where
  show Const = "C"
  show (Var i) = show i
  show (App t1 t2) = "(" <> show t1 <> " " <> show t2 <> ")"
  show (Lambda ty te) = "(L " <> show ty <> " " <> show te <> ")"

------------------
-- Type Checker --
------------------

check :: [Type] -> Term -> Maybe Type
check _ Const = return Bass
check env (Var v)
  | fromInteger v < length env = Just (env !! fromInteger v)
  | otherwise = Nothing
check env (Lambda tp te) = Arrow tp <$> check (tp : env) te
check env (App tm tm') = do
  Arrow i o <- check env tm
  i' <- check env tm'
  guard (i == i')
  return o

-------------------------------
-- Evaluation of Expressions --
-------------------------------

eval :: Term -> Term
eval (App t1 t2) =
  case eval t1 of
    Lambda _ body -> eval (subst 0 t2 body)
    _ -> error "Cannot apply non function."
eval v = v

subst :: Integer -> Term -> Term -> Term
subst _ _ Const = Const
subst n tm (Var m) =
  case compare m n of
    LT -> Var m
    EQ -> tm
    GT -> Var (m - 1)
subst n tm (Lambda t body) = Lambda t (subst (n + 1) tm body)
subst n tm (App tm' tm'') = App (subst n tm tm') (subst n tm tm'')

evalMay :: Term -> Maybe Term
evalMay t = eval t <$ check [] t

------------
-- Parser --
------------

parseLambda :: Parser Term
parseLambda = parens $ do
  _ <- reserved "L"
  Lambda <$> parseType <*> parseTerm

parseApp :: Parser Term
parseApp = parens $ App <$> parseTerm <*> parseTerm

parseVar :: Parser Term
parseVar = Var <$> integer

parseType :: Parser Type
parseType =
  reserved "B" $> Bass <|>
  parens (reservedOp "->" >> (Arrow <$> parseType <*> parseType))

parseTerm :: Parser Term
parseTerm = try parseLambda <|> parseApp <|> parseVar <|> (reserved "C" $> Const)

parseL :: String -> Term
parseL a =
  case runParser parseTerm () "" a of
    Right b -> b
    Left _ -> error "Could not parse expression."