A small typechecker using normalization by evaluation (NbE) for a Quantitative Dependently Typed Lambda Calculus (QDTLC).

#Typechecker for QDTLC

A small typechecker using normalization by evaluation (NbE) for a Quantitative Dependently Typed Lambda Calculus (QDTLC).

README.md

To run: ghc --make Main && ./Main.

Depending on the λ-term used, this will output either an error or <term> :: <type> if type-checking passed. Example execution:

>> Typechecking term @e@ against type @τ@ where
	e ≡ ((), ())
	τ ≡ (x :^1 𝟏) ⊗ 𝟏

>> Success!

((), ())
	:: (x :^1 𝟏) ⊗ 𝟏

This supports the following types:

• Π(x :ⁱ A). B is the quantitative dependent function type whose domain is A (with usage i) and codomain is B.

• ∀(x : A). B is the same as Π(x :⁰ A). B.

• (x : A) ⊸ B is the same as Π(x :¹ A). B.

• (x :ⁱ A) ⊗ B is the quantitative multiplicative dependent pair type.

• U is the type of universes of types. We have U :⁰ U in order to simplify the type system studied here.

• 𝟏 is the multiplicative unit type.

• ℕ is the type of natural numbers. and expressions:

• λ(x :ⁱ A). B is a function taking an x of type A and returning B which consumes x i times.

• (a, b) is the constructor of the quantitative multiplicative dependent pair type.

• () is the constructor for the multiplicative unit type.

• f x is the application of the argument x to the function f.

• x is a simple variable.

• let () as z = a in b is the general destructor of the multiplication unit type, where a should result in () and b is the return value of the expression.

• let x :ⁱ A = e₁ in e₂ locally binds the variable x with usage i and value e₁ for use within e₂.

• n denotes a simple naturel number (unsigned integer).

AST.hs

{-# LANGUAGE PatternSynonyms #-}

module AST where

import Data.Bifunctor (first)
import Data.Map (Map)
import qualified Data.Map as Map
import Usage

type Name = String

data Term
  = -- | @λ(x :ᵖ T). e[x]@
    TLamda
      Name
      -- ^ @x@
      Usage
      -- ^ @p@
      Term
      -- ^ @T@
      Term
      -- ^ @e@
  | -- | @Π(x :ᵖ A). B[x]@
    TPi
      Name
      -- ^ @x@
      Usage
      -- ^ @p@
      Term
      -- ^ @A@
      Term
      -- ^ @B@
  | -- | @a b@
    TApplication
      Term
      -- ^ @a@
      Term
      -- ^ @b@
  | -- | The multiplicative dependent pair @(x :ᵖ A) ⊗ B[x]@
    TTensor
      Name
      -- ^ @x@
      Usage
      -- ^ @p@
      Term
      -- ^ @A@
      Term
      -- ^ @B@
  | -- | @(a, b)@
    TPair
      Term
      -- ^ @a@
      Term
      -- ^ @b@
  | -- @x@
    TVar
      Name
  | -- | @U@ (type universe, and we have @U : U@)
    TU
  | -- | The multiplicative unit type @𝟏@
    TOne
  | -- | @()@
    TUnit
  | -- | @let () as z = a in b@
    TUnitElim
      Name
      -- ^ @z@
      Term
      -- ^ @a@
      Term
      -- ^ @b@
  | -- | @let x :ᵖ T = a in b@
    TLet
      Name
      -- ^ @x@
      Usage
      -- ^ @p@
      Term
      -- ^ @T@
      Term
      -- ^ @a@
      Term
      -- ^ @b@
  | -- | @ℕ@
    TNatural
  | -- | @0@ and all other natural numbers
    TNumber
      Integer
  deriving (Show)

-- | @∀(x : τ). P@
pattern TForall ::
  -- | @x@
  Name ->
  -- | @τ@
  Term ->
  -- | @P@
  Term ->
  Term
pattern TForall x τ p = TPi x 0 τ p

-- | @(x : τ) ⊸ A@
pattern TLinear ::
  -- | @x@
  Name ->
  -- | @τ@
  Term ->
  -- | @A@
  Term ->
  Term
pattern TLinear x τ a = TPi x 1 τ a

-- | @(x : τ) → A@
pattern TFunction ::
  -- | @x@
  Name ->
  -- | @τ@
  Term ->
  -- | @A@
  Term ->
  Term
pattern TFunction x τ a = TPi x W τ a

prettyTerm :: Term -> String
prettyTerm (TVar x) = x
prettyTerm TU = "U"
prettyTerm (TLinear x a e) = "(" <> x <> " : " <> prettyTerm a <> ") ⊸ " <> prettyTerm e
prettyTerm (TForall x a e) = "∀(" <> x <> " : " <> prettyTerm a <> "). " <> prettyTerm e
prettyTerm (TFunction x a e) = "(" <> x <> " : " <> prettyTerm a <> ") → " <> prettyTerm e
prettyTerm (TLamda x p a e) = "λ(" <> x <> " :^" <> prettyUsage p <> " " <> prettyTerm a <> "). " <> prettyTerm e
prettyTerm (TPi x p a e) = "Π(" <> x <> " :^" <> prettyUsage p <> " " <> prettyTerm a <> "). " <> prettyTerm e
prettyTerm (TTensor x p a e) = "(" <> x <> " :^" <> prettyUsage p <> " " <> prettyTerm a <> ") ⊗ " <> prettyTerm e
prettyTerm (TPair a b) = "(" <> prettyTerm a <> ", " <> prettyTerm b <> ")"
prettyTerm TOne = "𝟏"
prettyTerm TUnit = "()"
prettyTerm (TUnitElim z a b) = "let " <> prettyTerm TUnit <> " as " <> z <> " = " <> prettyTerm a <> " in " <> prettyTerm b
prettyTerm (TLet x p t a b) = "let " <> x <> " :^" <> prettyUsage p <> " " <> prettyTerm t <> " = " <> prettyTerm a <> " in " <> prettyTerm b
prettyTerm TNatural = "ℕ"
prettyTerm (TNumber n) = show n
prettyTerm (TApplication f x) = maybeParens f <> " " <> maybeParens' x
  where
    maybeParens t@(TLamda _ _ _ _) = "(" <> prettyTerm t <> ")"
    maybeParens t@(TPi _ _ _ _) = "(" <> prettyTerm t <> ")"
    maybeParens t@(TTensor _ _ _ _) = "(" <> prettyTerm t <> ")"
    maybeParens t = prettyTerm t

    maybeParens' t@(TApplication _ _) = "(" <> prettyTerm t <> ")"
    maybeParens' t = maybeParens t

type Context = Map Name (Usage, Value)

type Environment = Map Name Value

data Closure = Clos Environment Term
  deriving (Show)

data Value
  = VU
  | VLamda Name Usage Value Closure
  | VPi Name Usage Value Closure
  | VApplication Value Value
  | VTensor Name Usage Value Closure
  | VPair Value Value
  | VVar Name
  | VOne
  | VUnit
  | VNatural
  | VNumber Integer
  deriving (Show)

(.**.) :: Usage -> Context -> Context
u .**. env = first (u .*.) <$> env

(.++.) :: Context -> Context -> Context
(.++.) = Map.unionWith add
  where
    add (u1, x) (u2, _) = (u1 + u2, x)

Eval.hs

{-# LANGUAGE LambdaCase #-}

module Eval where

import AST
import qualified Data.Map as Map
import Data.Maybe (fromJust)
import Debug.Trace (traceShow)
import Fresh
import GHC.Stack (HasCallStack)
import Usage

eval :: HasCallStack => Environment -> Term -> IO Value
eval δ (TVar x) = pure . fromJust $ Map.lookup x δ
eval δ (TLamda x p a b) = VLamda x p <$> eval δ a <*> pure (Clos δ b)
eval δ (TPi x p a b) = VPi x p <$> eval δ a <*> pure (Clos δ b)
eval δ (TApplication a b) =
  eval δ a >>= \case
    VLamda x _ _ c -> do
      b <- eval δ b
      applyClosure x b c
    e -> VApplication e <$> eval δ b
eval δ TU = pure VU
eval δ (TTensor x p a b) = VTensor x p <$> eval δ a <*> pure (Clos δ b)
eval δ (TPair a b) = VPair <$> eval δ a <*> eval δ b
eval δ TOne = pure VOne
eval δ TUnit = pure VUnit
eval δ (TUnitElim _ a b) = do
  eval δ a
  -- NOTE: is it really necessary? this is most likely not to be evaluated anyways because of Haskell's laziness
  eval δ b
eval δ (TLet x _ _ e r) = do
  e <- eval δ e
  eval (Map.insert x e δ) r
eval δ TNatural = pure VNatural
eval δ (TNumber n) = pure $ VNumber n

applyClosure :: HasCallStack => Name -> Value -> Closure -> IO Value
applyClosure x v (Clos δ b) = eval (Map.insert x v δ) b

quote :: HasCallStack => Environment -> Value -> IO Term
quote δ VU = pure TU
quote δ (VApplication a b) = TApplication <$> quote δ a <*> quote δ b
quote δ (VPi x p a clos) = do
  x' <- VVar <$> fresh δ

  let δ' = Map.insert x x' δ
  b <- quote δ' =<< applyClosure x x' clos
  TPi x p <$> quote δ' a <*> pure b
quote δ (VLamda x p a clos) = do
  x' <- VVar <$> fresh δ

  let δ' = Map.insert x x' δ
  b <- quote δ' =<< applyClosure x x' clos
  TLamda x p <$> quote δ' a <*> pure b
quote δ (VVar x) = pure $ TVar x
quote δ (VTensor x p a clos) = do
  x' <- VVar <$> fresh δ

  let δ' = Map.insert x x' δ
  b <- quote δ' =<< applyClosure x x' clos
  TTensor x p <$> quote δ' a <*> pure b
quote δ (VPair a b) = TPair <$> quote δ a <*> quote δ b
quote δ VOne = pure TOne
quote δ VUnit = pure TUnit
quote δ VNatural = pure TNatural
quote δ (VNumber n) = pure $ TNumber n

Fresh.hs

module Fresh where

import AST
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import qualified Data.Map as Map
import System.IO.Unsafe (unsafeDupablePerformIO)

freshCnt :: IORef Int
freshCnt = unsafeDupablePerformIO $ newIORef 0

-- | Generates a fresh name which is not bound in the current environment.
fresh :: Environment -> IO Name
fresh δ = do
  x <- readIORef freshCnt
  writeIORef freshCnt (x + 1)

  let name = "_" <> show x

  case Map.lookup name δ of
    Just _ -> fresh δ
    _ -> pure name

Main.hs

{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Main where

import AST
import Control.Exception (ErrorCall, catch, throwIO)
import System.Exit (exitSuccess)
import Typechecker
import Usage

main :: IO ()
main = do
  let expr = ex11Term
      typ = ex11Type

  putStrLn $ ">> Typechecking term @e@ against type @τ@ where"
  putStrLn $ "\te ≡ " <> prettyTerm expr
  putStrLn $ "\tτ ≡ " <> prettyTerm typ
  putStrLn ""

  tm <- catch (typecheck expr typ) \(e :: ErrorCall) -> do
    putStrLn ">> Error!"
    putStrLn ""

    throwIO e

  putStrLn ">> Success!"
  putStrLn ""
  putStrLn $ prettyTerm tm
  putStrLn $ "\t:: " <> prettyTerm typ

  exitSuccess

-- | The normal polymorphic identity function.
idTerm, idType :: Term
idTerm =
  -- λ(A :⁰ U). λ(x :¹ A). x
  TLamda "A" 0 TU $ TLamda "x" 1 (TVar "A") $ TVar "x"
idType =
  -- :: ∀(A : U). (x : A) ⊸ A
  TForall "A" TU $ TLinear "x" (TVar "A") $ TVar "A"

dupTerm, dupType :: Term
-- WARN: this is ill-typed because @x@ is linear and is used twice
dupTerm =
  -- λ(A :⁰ U). λ(x :¹ A). (x, x)
  TLamda "A" 0 TU $ TLamda "x" 1 (TVar "A") $ TPair (TVar "x") (TVar "x")
dupType =
  -- :: ∀(A : U). (x :¹ A) ⊸ (y :¹ A) ⊗ A
  TForall "A" TU $ TLinear "x" (TVar "A") $ TTensor "y" 1 (TVar "A") (TVar "A")

-- | A correct implementation of the duplication function @dup(x) = (x, x)@.
dup2Term, dup2Type :: Term
dup2Term =
  -- λ(A :⁰ U). λ(x :⁻ A). (x, x)
  TLamda "A" 0 TU $ TLamda "x" W (TVar "A") $ TPair (TVar "x") (TVar "x")
dup2Type =
  -- :: ∀(A : U). Π(x :⁻ A). (y :¹ A) ⊗ A
  TPi "A" 0 TU $ TPi "x" W (TVar "A") $ TTensor "y" 1 (TVar "A") (TVar "A")

-- | Implements the @;@ operator such that @(f ; g) ()@ is equivalent to @let () as _ = f () in g ()@.
seqTerm, seqType :: Term
seqTerm =
  -- λ(A :⁰ U). λ(f :¹ Π(_1 :¹ 𝟏) → 𝟏). λ(g :¹ Π(_2 :¹ 𝟏) → A). let () as z = f () in g ()
  TLamda "A" 0 TU $
    TLamda "f" 1 (TPi "_1" 1 TOne TOne) $
      TLamda "g" 1 (TPi "_2" 1 TOne $ TVar "A") $
        TUnitElim "z" (TApplication (TVar "f") TUnit) (TApplication (TVar "g") TUnit)
seqType =
  -- :: ∀(A : U). (f : (_1 : 𝟏) ⊸ 𝟏) ⊸ Π(g : (_2 : 𝟏) ⊸ A) ⊸ A
  TForall "A" TU $
    TLinear "f" (TLinear "_1" TOne TOne) $
      TLinear "g" (TLinear "_2" TOne $ TVar "A") $
        TVar "A"

-- | Usual dependent function composition operator, sometimes denoted @∘@.
composeTerm, composeType :: Term
composeTerm =
  -- λ(A :⁰ U). λ(B :⁰ ∀(a : A). U). λ(C :⁰ ∀(a : A). ∀(b : B a). U).
  --   λ(f :¹ (a : A) ⊸ (b : B a) ⊸ C a b). λ(g :¹ (a : A) ⊸ B a). λ(x :⁻ A).
  --     f x (g x)
  TLamda "A" 0 TU $
    TLamda "B" 0 (TForall "a" (TVar "A") TU) $
      TLamda "C" 0 (TForall "a" (TVar "A") $ TForall "b" (TApplication (TVar "B") (TVar "a")) TU) $
        TLamda "f" 1 (TLinear "a" (TVar "A") $ TLinear "b" (TApplication (TVar "B") (TVar "a")) $ TApplication (TApplication (TVar "C") (TVar "a")) (TVar "b")) $
          TLamda "g" 1 (TLinear "a" (TVar "A") $ TApplication (TVar "B") (TVar "a")) $
            TLamda "x" W (TVar "A") $
              TApplication (TApplication (TVar "f") (TVar "x")) (TApplication (TVar "g") (TVar "x"))
composeType =
  -- :: ∀(A : U). ∀(B : ∀(a : A). U). ∀(C : ∀(a : A). ∀(b : B a). U).
  --      (f : (a : A) ⊸ (b : B a) ⊸ C a b) ⊸ (g : (a : A) ⊸ B a) ⊸ Π(x :⁻ A).
  --        C x (g x)
  TForall "A" TU $
    TForall "B" (TForall "a" (TVar "A") TU) $
      TForall "C" (TForall "a" (TVar "A") $ TForall "b" (TApplication (TVar "B") (TVar "a")) TU) $
        TLinear "f" (TLinear "a" (TVar "A") $ TLinear "b" (TApplication (TVar "B") (TVar "a")) $ TApplication (TApplication (TVar "C") (TVar "a")) (TVar "b")) $
          TLinear "g" (TLinear "a" (TVar "A") $ TApplication (TVar "B") (TVar "a")) $
            TPi "x" W (TVar "A") $
              TApplication (TApplication (TVar "C") (TVar "x")) (TApplication (TVar "g") (TVar "x"))

----------- DUMB EXAMPLES -------------

ex1Term, ex1Type :: Term
ex1Term =
  -- (λ(A :⁰ U). λ(x :¹ A). x) (Π(A :⁰ U). Π(x :¹ A). A) (λ(A :⁰ U). λ(x :¹ A). x)
  TApplication (TApplication idTerm idType) idTerm
ex1Type =
  -- :: Π(A :⁰ U). Π(x :¹ A). A
  idType

ex2Term, ex2Type :: Term
ex2Term =
  -- ()
  TUnit
ex2Type =
  -- :: 𝟏
  TOne

ex3Term, ex3Type :: Term
ex3Term =
  -- let () as z = () in ()
  TUnitElim "z" TUnit TUnit
ex3Type =
  -- :: 𝟏
  TOne

ex4Term, ex4Type :: Term
-- WARN: this is ill-typed as @λ(x :¹ 𝟏).x@ is not of type @𝟏@
ex4Term =
  -- let () as z = λ(x :¹ 𝟏). 𝟏 in ()
  TUnitElim "z" (TLamda "x" 1 TOne (TVar "x")) TUnit
ex4Type =
  -- :: 𝟏
  TOne

ex5Term, ex5Type :: Term
ex5Term =
  -- let () as z = () in λ(x :¹ 𝟏). x
  TUnitElim "z" TUnit $ TLamda "x" 1 TOne (TVar "x")
ex5Type =
  -- :: (x : 𝟏) ⊸ 𝟏
  TLinear "x" TOne TOne

ex6Term, ex6Type :: Term
ex6Term =
  -- ((), ())
  TPair TUnit TUnit
ex6Type =
  -- :: (x :¹ 𝟏) ⊗ 𝟏
  TTensor "x" 1 TOne TOne

ex7Term, ex7Type :: Term
-- WARN: ill-typed because @A@ is not used linearly
ex7Term =
  -- λ(A :¹ U). λ(x :¹ A). x
  TLamda "A" 1 TU $ TLamda "x" 1 (TVar "A") $ TVar "x"
ex7Type =
  -- :: ∀(A : U). (x : A) ⊸ A
  TLinear "A" TU $ TLinear "x" (TVar "A") $ TVar "A"

ex8Term, ex8Type :: Term
-- WARN: ill-typed as @x@ should be linear
ex8Term =
  -- λ(f :¹ (_ : 𝟏) ⊸ 𝟏). λ(x :¹ 𝟏). let () as z = f x in x
  TLamda "f" 1 (TLinear "_" TOne TOne) $ TLamda "x" 1 TOne $ TUnitElim "z" (TApplication (TVar "f") (TVar "x")) (TVar "x")
ex8Type =
  -- :: (f : (_ : 𝟏) ⊸ 𝟏) ⊸ (x : 𝟏) ⊸ 𝟏
  TLinear "f" (TLinear "_" TOne TOne) $ TLinear "x" TOne TOne

ex9Term, ex9Type :: Term
ex9Term =
  -- let idType :⁰ U = ∀(A : U). (x : A) ⊸ A in
  -- let id :⁻ idType = λ(A :⁰ U). λ(x :¹ A). x in
  -- id idType id
  TLet "idType" 0 TU idType $
    TLet "id" W (TVar "idType") idTerm $
      TApplication (TApplication (TVar "id") (TVar "idType")) (TVar "id")
ex9Type =
  -- :: ∀(A : U). (x : A) ⊸ A
  idType

ex10Term, ex10Type :: Term
ex10Term =
  -- let id :¹ ∀(A : U). (x : A) ⊸ A = λ(A :⁰ U). λ(x :¹ A). x in
  -- id nat 0
  TLet "id" 1 idType idTerm $
    TApplication (TApplication (TVar "id") TNatural) (TNumber 0)
ex10Type =
  -- :: nat
  TNatural

ex11Term, ex11Type :: Term
ex11Term =
  -- (0, 1)
  TPair (TNumber 0) (TNumber 1)
ex11Type =
  -- :: (_ :¹ ℕ) ⊗ ℕ
  TTensor "_" 1 TNatural TNatural

Typechecker.hs

{-# LANGUAGE BlockArguments #-}

module Typechecker where

import AST
import Control.Monad (when)
import qualified Data.Map as Map
import Data.Maybe (fromMaybe)
import Debug.Trace (trace, traceShow)
import Eval
import Fresh
import GHC.Stack (HasCallStack)
import Usage

typecheck :: Term -> Term -> IO Term
typecheck e t = fst <$> (check ctx env e I =<< eval env t)
  where
    ctx = mempty
    env = mempty

-- | @check Γ Δ e p τ@ is the judgment @Γ ⊢ e ⇐ᵖ τ@.
check :: HasCallStack => Context -> Environment -> Term -> Usage -> Value -> IO (Term, Context)
check γ δ (TLamda x p a e) i (VPi _ p' a' b) = do
  when (p /= p') do
    error $ "Incorrect usage specification (needed " <> show p' <> " but got " <> show p <> ")"
  {-
     0Γ ⊢ A :⁰ U        Γ, x :ⁱᵖ A ⊢ e ⇐ⁱ B
    ---------------------------------------- [λ-I]
        Γ ⊢ λ(x :ᵖ A). e ⇐ⁱ Π(x :ᵖ A). B
  -}
  (a, _) <- check (0 .**. γ) δ a 0 VU
  a'' <- eval δ a
  unify δ a'' a'

  b <- applyClosure x (VVar x) b

  (γ, δ) <- pure (Map.insert x (i * p, a'') γ, Map.insert x (VVar x) δ)
  (e, γ) <- check γ δ e i b
  case (i * p, γ Map.! x) of
    -- If the usage was 1 and remains 1, then the linear variable has not been used.
    (1, (1, _)) -> error $ "Variable " <> x <> " not used in lambda body"
    _ -> pure ()

  pure (TLamda x p a e, Map.delete x γ)
check γ δ (TPair a b) i (VTensor x p a' b') = do
  {-
     0Γ ⊢ a ⇐⁰ A       Γ ⊢ b ⇐ⁱ B[x\a]       ip = 0
    ------------------------------------------------ [⊗-I₀]
              Γ ⊢ (a, b) ⇐ⁱ (x :ᵖ A) ⊗ B

      Γ₁ ⊢ a ⇐¹ A          Γ₂ ⊢ b ⇐ⁱ B[x\a]
    ---------------------------------------- [⊗-I₁]
       ipΓ₁ + Γ₂ ⊢ (a, b) ⇐ⁱ (x :ᵖ A) ⊗ B
  -}
  case i * p of
    0 -> do
      -- apply [⊗-I₀]
      (a, _) <- check (0 .**. γ) δ a 0 a'

      a'' <- eval δ a
      (b, γ) <- check γ δ b i =<< applyClosure x a'' b'

      pure (TPair a b, γ)
    _ -> do
      -- apply [⊗-I₁]
      (a, γ₁) <- check γ δ a 1 a'

      a'' <- eval δ a
      (b, γ₂) <- check γ δ b i =<< applyClosure x a'' b'

      γ <- checkLinearVariables γ ((i * p) .**. γ₁) γ₂

      pure (TPair a b, γ)
check γ δ (TUnitElim z a b) i τ = do
  {-
     Γ₁ ⊢ a ⇐ⁱ 𝟏        0Γ₁, z :⁰ 𝟏 ⊢ τ ⇐⁰ U         Γ₂ ⊢ b ⇐ⁱ τ[()\z]
    ------------------------------------------------------------------- [𝟏-E]
               Γ₁ + Γ₂ ⊢ let () as z = a in b ⇐ⁱ τ[a\z]
  -}

  (a, γ₁) <- check γ δ a i VOne
  a' <- eval δ a

  let δ' = Map.insert z VUnit δ

  τ <- quote δ' τ
  (τ, _) <- check (0 .**. Map.insert z (0, VOne) γ₁) δ' τ 0 VU
  τ <- eval δ' τ

  (b, γ₂) <- check (Map.insert z (i, a') γ) (Map.insert z a' δ) b i τ

  γ <- checkLinearVariables γ γ₁ γ₂

  pure (TUnitElim z a b, Map.delete z γ)
check γ δ (TLet x i a e r) p b = do
  {-
     0Γ ⊢ A ⇐⁰ U        Γ₁ ⊢ e ⇐ⁱ A           Γ₂, x :ⁱᵖ A ⊢ r ⇐ᵖ B
    --------------------------------------------------------------- [let-E]
                 Γ₁ + Γ₂ ⊢ let x :ⁱ A = e in r ⇐ᵖ B
  -}

  (a, _) <- check (0 .**. γ) δ a 0 VU
  a' <- eval δ a

  (e, γ₁) <- check γ δ e i a'
  e' <- eval δ e

  (r, γ₂) <- check (Map.insert x (i * p, a') γ) (Map.insert x e' δ) r p b

  γ <- checkLinearVariables γ γ₁ γ₂

  pure (TLet x i a e r, γ)
check γ δ e p b = do
  {-
     Γ ⊢ e ⇒ⁱ A        A ≅ B       i ⩽ p
    ------------------------------------- [≅]
                 Γ ⊢ e ⇐ᵖ B
  -}
  (e', i, a, γ) <- infer γ δ e
  if i <= p
    then (e', γ) <$ unify δ a b
    else error $ "Usage " <> show i <> " cannot be used in place of usage " <> show p

-- | @infer Γ Δ e@ is the judgment @Γ ⊢ e ⇒ᵖ τ@.
infer :: HasCallStack => Context -> Environment -> Term -> IO (Term, Usage, Value, Context)
infer γ δ (TVar x) =
  {-
    -------------------- [var-I]
     Γ, x :ᵖ τ ⊢ x ⇒ᵖ τ
  -}
  case Map.lookup x γ of
    -- whenever @x :¹ τ@, we have to have @x :⁰ τ@ after
    -- in order not to be able to use it at runtime after that
    Just (1, τ) -> pure (TVar x, 1, τ, Map.insert x (0, τ) γ)
    Just (p, τ) -> pure (TVar x, p, τ, γ)
    Nothing -> error $ "Name " <> x <> " not bound"
infer γ _ TU =
  {-
    ------------ [U-I]
     Γ ⊢ U :⁰ U
  -}
  pure (TU, 0, VU, γ)
infer γ _ TNatural =
  {-
    ------------ [ℕ-I]
     Γ ⊢ ℕ :⁰ U
  -}
  pure (TNatural, 0, VU, γ)
infer γ _ (TNumber n) =
  {-
     n is a natural number literal
    ------------------------------- [num-I]
             Γ ⊢ n :⁻ ℕ
  -}
  pure (TNumber n, W, VNatural, γ)
infer γ δ (TApplication f x) = do
  {-
     Γ ⊢ f ⇒ⁱ (y :ᵖ A) → B          0Γ ⊢ x ⇐⁰ A        ip = 0
    ---------------------------------------------------------- [λ-E₀]
                       Γ ⊢ f x ⇒ⁱ B[y\x]

     Γ₁ ⊢ f ⇒ⁱ (y :ᵖ A) → B          Γ₂ ⊢ x ⇐¹ A
    --------------------------------------------- [λ-E₁]
              Γ₁ + ipΓ₂ ⊢ f x ⇒ⁱ B[y\x]
  -}
  (f, i, τ, γ₁) <- infer γ δ f
  let (y, p, a, b) = case τ of
        VPi y p a b -> (y, p, a, b)
        _ -> error "Expected Π-type for function in application"

  case i * p of
    0 -> do
      -- apply [λ-E₀]
      (x, _) <- check (0 .**. γ₁) δ x 0 a

      b <- eval δ x >>= \x' -> applyClosure y x' b

      pure (TApplication f x, i, b, γ₁)
    _ -> do
      -- apply [λ-E₁]
      (x, γ₂) <- check γ δ x 1 a

      b <- eval δ x >>= \x' -> applyClosure y x' b

      γ <- checkLinearVariables γ γ₁ ((i * p) .**. γ₂)

      pure (TApplication f x, i, b, γ)
infer γ δ (TPi x p a b) = do
  {-
     0Γ ⊢ A ⇐⁰ U         0Γ, x :⁰ A ⊢ B ⇐⁰ U
    ----------------------------------------- [Π-F]
             Γ ⊢ Π(x :ᵖ A). B ⇒⁰ U
  -}
  (a, _) <- check (0 .**. γ) δ a 0 VU
  a' <- eval δ a

  (γ, δ) <- pure (Map.insert x (0, a') γ, Map.insert x (VVar x) δ)
  (b, _) <- check (0 .**. γ) δ b 0 VU

  pure (TPi x p a b, 0, VU, Map.delete x γ)
infer γ δ (TTensor x p a b) = do
  {-
     0Γ ⊢ A ⇐⁰ U         0Γ, x :⁰ A ⊢ B ⇐⁰ U
    ----------------------------------------- [⊗-F]
             Γ ⊢ (x :ᵖ A) ⊗ B ⇒⁰ U
  -}
  (a, _) <- check (0 .**. γ) δ a 0 VU
  a' <- eval δ a

  (γ, δ) <- pure (Map.insert x (0, a') γ, Map.insert x (VVar x) δ)
  (b, _) <- check (0 .**. γ) δ b 0 VU

  pure (TTensor x p a b, 0, VU, Map.delete x γ)
infer γ δ TOne = do
  {-
    ------------ [𝟏-F]
     Γ ⊢ 𝟏 ⇐⁰ U
  -}
  pure (TOne, 0, VU, γ)
infer γ δ TUnit = do
  {-
    ------------- [𝟏-I]
     Γ ⊢ () ⇐⁻ 𝟏
  -}
  pure (TUnit, W, VOne, γ)
infer γ δ (TLamda x p a e) = do
  {-
     0Γ ⊢ A ⇐⁰ U          Γ, x :ⁱᵖ A ⊢ e ⇒ⁱ B
    ------------------------------------------ [λ-I]
        Γ ⊢ λ(x :ᵖ A). e ⇒ⁱ Π(y :ᵖ A) → B
  -}
  (a, _) <- check (0 .**. γ) δ a 0 VU
  a' <- eval δ a

  (γ, δ) <- pure (Map.insert x (p * 1, a') γ, Map.insert x (VVar x) δ)
  (e, i, b, γ) <- infer γ δ e
  -- rec (e', i, b, γ') <- infer (Map.insert x (p * i, a') γ) (Map.insert v (VVar x) δ) e
  -- FIXME: this infinitely loops because `i` cannot be determined correctly
  -- Is there any way for us to determine it knowing the result of e.g. @i × p@? (we know @p@ for sure there)
  --
  -- There are multiple cases:
  -- - If @ip = 0@ then either @i = 0@ or @p = 0@:
  --   * When @p = 0@, then we don't know about @i@ (because it may be any usage).
  --   * When @i = 0@ then we get the desired result.
  --     However, this is a case where we actually cannot determine that @ip = 0@ because @p ∈ {1, ω}@.
  -- - If @ip = 1@ then @i = 1@ and @p = 1@. So we have the value of @i@.
  -- - If @ip = ω@ then either @i = ω@ or @p = ω@ or both:
  --   * When @p = ω@ then @i ∈ {1, ω}@ so we don't know @i@ for sure.
  --   * When @i = ω@ then we cannot really compute the value of @ip@ because @p ∈ {1, ω}@.
  --     So this case is also ill.
  --
  -- There are no case where @i@ can be known from @p@ only, so @i@ cannot be infered.
  --
  -- ============================================
  --
  -- This has a great implication, as infering the type of an application @f x@ relies on infering the type of @f@.
  -- However, if @f@ is a lambda (a simple example would be @(λx.x) (λx.x)@, assuming no types) then its type must be infered.
  --
  -- 1. In a way, we could simply disallow such case easily: throw an error when attempting to infer the type of a lambda.
  --    But this seems unsatisfactory as it is the only thing preventing us from infering the type.
  --
  -- 2. Another approach would be simply to infer a usage of @i = 1@ for such case only
  --    (as the function will be applied once because it is immediately applied).
  --    But this should work when targeting this specific case only (and not the general case, for example @let id = λx.x in id id@).
  --    In fact, is there any other case where we infer the type of a lambda?
  --
  -- We will try alternative 2. and see if it really works here, or if we need to move it to the specific direct application case.

  case (p, γ Map.! x) of
    -- If the usage was 1 and remains 1, then the linear variable has not been used.
    (1, (1, _)) -> error $ "Variable " <> x <> " is unused in lambda body"
    _ -> pure ()

  b <- Clos δ <$> quote δ b
  pure (TLamda x p a e, i, VPi x p a' b, Map.delete x γ)

-- | @unify A B@ is the relation @A ≅ B@.
unify :: HasCallStack => Environment -> Value -> Value -> IO ()
unify _ VU VU = pure ()
unify _ (VVar x) (VVar y) | x == y = pure ()
unify δ (VPi x _ a b) (VPi _ _ a' b') = do
  y <- VVar <$> fresh δ
  b <- applyClosure x y b
  b' <- applyClosure x y b'

  unify δ a a'
  unify δ b b'
unify δ (VLamda x _ a b) (VLamda _ _ a' b') = do
  y <- VVar <$> fresh δ
  b <- applyClosure x y b
  b' <- applyClosure x y b'

  unify δ a a'
  unify δ b b'
unify δ (VApplication a b) (VApplication a' b') = do
  unify δ a a'
  unify δ b b'
unify δ (VPair a b) (VPair a' b') = do
  unify δ a a'
  unify δ b b'
unify δ (VTensor x _ a b) (VTensor _ _ a' b') = do
  y <- VVar <$> fresh δ
  b <- applyClosure x y b
  b' <- applyClosure x y b'

  unify δ a a'
  unify δ b b'
unify _ VOne VOne = pure ()
unify _ VUnit VUnit = pure ()
unify _ VNatural VNatural = pure ()
unify _ (VNumber n1) (VNumber n2) | n1 == n2 = pure ()
unify δ t1 t2 = do
  t1 <- quote δ t1
  t2 <- quote δ t2
  error $ "Cannot unify terms @a@ and @b@ where\n\ta ≡ " <> prettyTerm t1 <> "\n\tb ≡ " <> prettyTerm t2 <> ""

checkLinearVariables :: HasCallStack => Context -> Context -> Context -> IO Context
checkLinearVariables γ γ₁ γ₂ =
  flip Map.traverseMaybeWithKey γ \x (u, τ) -> case u of
    -- If our variable was linear, and it has been used in both contexts γ₁ and γ₂
    -- then the variable has actually been used multiple times.
    1 -> case (γ₁ Map.!? x, γ₂ Map.!? x) of
      (Just (0, _), Just (0, _)) -> error $ "Variable " <> x <> " has been used non-linearly"
      (Just (1, _), Just (0, _)) -> pure $ Just (0, τ)
      (Just (0, _), Just (1, _)) -> pure $ Just (0, τ)
      (Just (1, _), Just (1, _)) -> pure $ Just (1, τ)
      (Just (i, _), Nothing) -> pure $ Just (i, τ)
      (Nothing, Just (i, _)) -> pure $ Just (i, τ)
      (Nothing, Nothing) -> pure $ Just (u, τ)
      (_, _) -> pure $ Just (W, τ)
    -- In any other case, we don't really care.
    _ -> pure $ Just (u, τ)

Usage.hs

{-# LANGUAGE PatternSynonyms #-}

module Usage where

data Usage
  = -- | 0
    Erased
  | -- | 1
    Linear
  | -- | ω
    Unrestricted
  deriving (Show, Eq)

pattern O, I, W :: Usage
pattern O = Erased
pattern I = Linear
pattern W = Unrestricted

instance Num Usage where
  fromInteger 0 = O
  fromInteger 1 = I
  fromInteger _ = W

  (+) = (.+.)
  (*) = (.*.)

-- | Addition of usage, defined by the following table:
--
-- +-+-+-+-+
-- |+|0|1|ω|
-- +=+=+=+=+
-- |0|0|1|ω|
-- +-+-+-+-+
-- |1|1|ω|ω|
-- +-+-+-+-+
-- |ω|ω|ω|ω|
-- +-+-+-+-+
(.+.) :: Usage -> Usage -> Usage
O .+. u = u
u .+. O = u
_ .+. _ = W

-- | Multiplication of usage, defined by the following table:
--
-- +-+-+-+-+
-- |×|0|1|ω|
-- +=+=+=+=+
-- |0|0|0|0|
-- +-+-+-+-+
-- |1|0|1|ω|
-- +-+-+-+-+
-- |ω|0|ω|ω|
-- +-+-+-+-+
(.*.) :: Usage -> Usage -> Usage
O .*. _ = O
_ .*. O = O
I .*. I = I
_ .*. _ = W

instance Ord Usage where
  -- ω ⩽ _
  -- 1 ⩽ 1
  -- 0 ⩽ 0
  W <= _ = True
  I <= I = True
  I <= _ = False
  O <= O = True
  O <= _ = False

prettyUsage :: Usage -> String
prettyUsage O = "0"
prettyUsage I = "1"
prettyUsage W = "ω"