-- |
-- Module       : LinLC
-- Description  : Linear simply typed lambda calculus
-- Stability    : experimental
--
-- A typechecker for a linear lambda calculus. This module demonstrates the
-- usage of the `ScopedState` monad, which can be used when a typing context
-- has to be threaded during typechecking.
module LinLC where

import Control.Monad
import Control.Monad.Except
import Data.Fin
import Data.Vec as Vec hiding (bind)
import Rebound hiding (rescope)
import Rebound.Bind.Single hiding (rescope)
import Rebound.MonadScoped
import Prelude

-- | Run a monadic computation, then another, and return the result of the
-- first. Or, put another way: the reverse of `<<`.
(<<) :: (Monad m) => m a -> m b -> m a
l << r = do
  vl <- l
  r
  return vl

--------------------------------------------------------------------------------
--- Basic definitions
--------------------------------------------------------------------------------

-- | Represent the (current) usage of a variable.
data Usage where
  Unused :: Usage
  Used :: Usage
  deriving (Show, Eq)

-- | Representation of types.
data Ty where
  TyUnit :: Ty
  TyArrow :: Ty -> Ty -> Ty
  deriving (Show, Eq)

-- | Representation of terms.
data Exp (n :: Nat) where
  Var :: Fin n -> Exp n
  CUnit :: Exp n
  Lam :: Bind Exp Exp n -> Exp n
  App :: Exp n -> Exp n -> Exp n
  deriving (Eq, Generic1)

-- Some instances required by rebound. See LC.hs for more explanations.

instance SubstVar Exp where var = Var

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

--------------------------------------------------------------------------------
--- Some helper-constructors
--------------------------------------------------------------------------------

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

-- | Notation for the arrow type
(~>) :: Ty -> Ty -> Ty
(~>) = TyArrow

infixr 8 ~>

--------------------------------------------------------------------------------
--- Typechecking infrastructure
--------------------------------------------------------------------------------

-- | A typing environment has to keep track of two things about variables:
-- 1. What their type is (`types`).
-- 2. Whether they've already been used or not (`usages`).
data TCEnv n = TCEnv
  { types :: Vec n Ty,
    usages :: Vec n Usage
  }

-- | The typechecking monad.
--
-- Unlike other calculi, the typing environment is not local. When two
-- sub-expressions have to be typechecked, the second sub-expression has to be
-- typechecked in the environment generated by the first sub-expression. Note
-- that this new environment will not include any new binding, but typing the
-- first sub-expression may have altered the usage of variables. This means that
-- `ScopedReader` cannot be used, as any changes done to the context are
-- forgotten when the scope is reverted. Hence, the `ScopedState` monad must be
-- used instead.
type TC n a = ScopedStateT TCEnv (Except String) n a

-- | Attempt to "consume" a variable, returning its type. Fails if the variable
-- was already used.
consumeVar :: Fin n -> TC n Ty
consumeVar i = setUsage i >> getsS ((! i) . types)
  where
    -- \| Set a variable to `Used`. Fails if the variable was already used.
    setUsage :: Fin n -> TC n ()
    setUsage i = do
      current <- getsS usages
      let (new, old) = set i Used current
      unless (old == Unused) $ throwError "Variable has already been used."
      modifyS $ \s -> s {usages = new}

    -- \| Set a value in the vector. Return the updated vector as well as the
    -- previous value.
    set :: Fin n -> t -> Vec n t -> (Vec n t, t)
    set FZ v (h ::: t) = (v ::: t, h)
    set (FS i) v (h ::: t) =
      let (t', v') = set i v t
       in (h ::: t', v')

-- | Add a variable to the scope.
--
-- Since the typing environment is threaded rather than passed down, as would be
-- the case with `ScopedReader`'s `localS`, we need to provide functions to both
-- enter a new scope, as with `localS`, and to leave it. Additionally, we need
-- to check that variables were used when they go out of scope.
addBinder :: Ty -> TC (S n) a -> TC n a
addBinder ty m = rescope enter leave (m << checkUsed FZ)
  where
    checkUsed :: Fin n -> TC n ()
    checkUsed i = do
      u <- getsS ((! i) . usages)
      unless (u == Used) $ throwError "Variable was not used."

    -- \| Add the binding in the scope.
    enter :: TCEnv n -> TCEnv (S n)
    enter e = e {types = ty ::: types e, usages = Unused ::: usages e}

    -- \| And remove it.
    leave :: TCEnv (S n) -> TCEnv n
    leave e = e {types = Vec.tail $ types e, usages = Vec.tail $ usages e}

-- | Run a computation in the typechecking monad, assuming that all free
-- variables need to be used.
runTC :: forall n a. (SNatI n) => Vec n Ty -> TC n a -> Either String a
runTC ts c = runExcept $ evalScopedStateT (c << checkAllUsed) initEnv
  where
    -- \| The initial environment, in which all variables are unused.
    initEnv :: TCEnv n
    initEnv = TCEnv {types = ts, usages = Vec.tabulate $ const Unused}

    -- \| Checks that all values are tagged as used.
    checkAllUsed :: TC n ()
    checkAllUsed = do
      us <- getsS usages
      let r = Vec.foldr (\u acc -> u == Used && acc) True us
      unless r $ throwError "Some variables in the initial scope were not used."

--------------------------------------------------------------------------------
--- Bi-directional typechecking
--------------------------------------------------------------------------------

-- Most of the unusual stuff about linear type systems is contained in the
-- handling of variables, which was implemented above. Hence the following code
-- should be rather unsurprising.

inferType :: Exp n -> TC n Ty
inferType (Var i) = consumeVar i
inferType CUnit = return TyUnit
inferType _ = throwError "Cannot infer type of this construct."

checkType :: Exp n -> Ty -> TC n ()
checkType (Lam bnd) ty = do
  let t = unbindl bnd
  (xTy, tTy) <- ensureArrow ty
  addBinder xTy $ checkType t tTy
  where
    ensureArrow :: Ty -> TC n (Ty, Ty)
    ensureArrow (TyArrow l r) = return (l, r)
    ensureArrow _ = throwError "Type is not arrow."
checkType (App f a) rTy = do
  aTy <- inferType a
  checkType f (TyArrow aTy rTy)
checkType t ty = do
  ty' <- inferType t
  unless (ty == ty') $ throwError "Inferred type does not match expected type."

-- >>> runTC empty $ checkType (lam $ v0) (TyUnit ~> TyUnit)
-- Right ()

-- >>> runTC empty $ checkType (lam $ lam $ v0 @@ v1) (TyUnit ~> (TyUnit ~> TyUnit) ~> TyUnit)
-- Right ()

-- >>> runTC empty $ checkType (lam $ lam $ v1) (TyUnit ~> (TyUnit ~> TyUnit) ~> TyUnit)
-- Left "Variable was not used."

-- >>> runTC empty $ checkType (lam $ lam $ v0) (TyUnit ~> (TyUnit ~> TyUnit) ~> TyUnit)
-- Left "Inferred type does not match expected type."

-- >>> runTC empty $ checkType (lam $ lam $ v0) (TyUnit ~> (TyUnit ~> TyUnit) ~> TyUnit ~> TyUnit)
-- Left "Variable was not used."

-- >>> runTC (TyUnit ::: (TyUnit ~> TyUnit) ::: TyUnit ::: Vec.empty) $ checkType (v1 @@ v0) (TyUnit)
-- Left "Some variables in the initial scope were not used."