content: -- |
  -- 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."
  