-- https://bentnib.org/docs/algebraic-typechecking-20150218.pdf

-- No values tracked. Just the types.
-- de-Bruijn index
data Typ
  = Nat
  | Boolean
  | Arrow Typ Typ
  deriving (Eq, Show)

data Term
  = Var Int
  | App Term Term
  | Lam Typ Term

-- Context here keeps track of the type of variables.
-- Not their values.
type Ctxt = [Typ] 

-- λ> [] !! 0
-- *** Exception: Prelude.!!: index too large
-- CallStack (from HasCallStack):
--   error, called at libraries/base/GHC/List.hs:1368:14 in base:GHC.List
--   tooLarge, called at libraries/base/GHC/List.hs:1378:50 in base:GHC.List
--   !!, called at <interactive>:140:4 in interactive:Ghci1

typecheck :: Term -> Ctxt -> Maybe Typ  
typecheck tm ctxt = case tm of
  Var i -> Just $ ctxt !! i
  App tm1 tm2 ->
    case typecheck tm1 ctxt of
      Just (Arrow typ1 typ2) ->
        case typecheck tm2 ctxt of
          Nothing -> Nothing
          Just typ1' ->
            if typ1 == typ1' then Just typ2
            else Nothing
      _ -> Nothing
  Lam ty1 tm ->
    case typecheck tm (ty1:ctxt) of
      Just ty2 -> Just $ Arrow ty1 ty2
      _ -> Nothing


eg1 = Lam Boolean (Var 0)  

-- Open terms possible too
eg2 = Lam Boolean (Var 2)  

eg3 = App (Lam Boolean (Var 2)) (Var 0)

-- λ> typecheck eg1 [Boolean, Nat, Arrow Nat Boolean]
-- Just (Arrow Boolean Boolean)
-- λ> typecheck eg2 [Boolean, Nat, Arrow Nat Boolean]
-- Just (Arrow Boolean Nat)
-- λ> typecheck eg3 [Boolean, Nat, Arrow Nat Boolean]
-- Just Nat


type TypeChecker = Ctxt -> Maybe Typ

-- Get typechecker that evaluates to type of given term
tm2Typechecker :: Term -> TypeChecker
tm2Typechecker tm = \ctxt -> case tm of
  Var i -> var i ctxt
  App tm1 tm2 -> app (tm2Typechecker tm1) (tm2Typechecker tm2) ctxt
  Lam ty tm -> lam ty (tm2Typechecker tm) ctxt

-- Inject a variable
var :: Int -> TypeChecker
var i ctxt = Just $ ctxt !! i

app :: TypeChecker -> TypeChecker -> TypeChecker
app tc1 tc2 = \ctxt -> case tc1 ctxt of
    Just (Arrow ty1 ty2) -> (case tc2 ctxt of
      Just ty' -> if ty1 == ty' then Just ty2 else Nothing
      _ -> Nothing)
    Nothing -> Nothing                 

lam :: Typ -> TypeChecker -> TypeChecker
lam ty tc = \ctxt -> case tc (ty:ctxt) of  -- TODO: Attention. Probably wrong
  Just ty' -> Just $ Arrow ty ty'
  _ -> Nothing
                             
-- Typechecker whose i-th element in context has type ty
have :: Int -> Typ -> TypeChecker -> TypeChecker
have i ty tc = \ctxt ->
  if ctxt!!i == ty then tc ctxt
  else Nothing

-- Typechecker that always fail
failure :: TypeChecker
failure = \ctxt -> Nothing

-- Typecheck that evaluates to a particular type
hasType :: Typ -> TypeChecker -> TypeChecker
hasType ty tc = \ctxt -> case tc ctxt of
  Just ty' -> if ty == ty' then Just ty else Nothing
  _ -> Nothing


-- λf:B -> N. λa:B. f a
tm1 :: Term
tm1 = Lam (Arrow Boolean Nat) $ Lam Boolean $ App (Var 1) (Var 0)

tc1 :: TypeChecker
tc1 = lam (Arrow Boolean Nat) (lam Boolean (app (var 1) (var 0)))

-- Term with a hole: λf:B -> N. λa:B. _
-- tm2 :: Term
-- tm2 = Lam (Arrow Boolean Nat) (Lam Boolean _)
tc2 :: forall x. TypeChecker
tc2 = \ctxt -> Nothing