BidirComplete.hs

{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE PatternSynonyms #-}

import Data.Maybe
import Control.Monad.State
import Data.List

data Ty
  = Unit -- The unit type.
  | Univ String -- A polymorphic type variable. Introduced by a `ForAll`.
  | Exst String -- A temporary existential type variable.
  | ForAll String Ty -- Introduces a polymorphic variable.
  | Ty :->: Ty -- The function type.
  deriving (Show, Eq)

-- Useage: `replace ty (new, old)`
replace :: Ty -> (Ty, Ty) -> Ty
replace (ForAll x ty) delta = ForAll x (replace ty delta)
replace (a :->: b) delta = replace a delta :->: replace b delta
replace ty (new, old)
  | ty == old = new
  | otherwise = ty

occursIn :: Ty -> Ty -> Bool
(Exst x) `occursIn` Unit = False
(Exst x) `occursIn` (Univ _) = False
(Exst x) `occursIn` (Exst y) = x == y
(Exst x) `occursIn` (ForAll _ a) = (Exst x) `occursIn` a
(Exst x) `occursIn` (a :->: b) = (Exst x) `occursIn` a || (Exst x) `occursIn` b

type Ctx = [CtxElem]

-- Allows use of left-to-right list syntax.
pattern xs :& x = x : xs

data CtxElem
  = IntroUniv String
  | IntroAnn String Ty
  | IntroExst String
  | IntroSoln String Ty
  | IntroMarker String
  deriving (Show, Eq)

-- ctxHole Gamma Theta = (Gamma1, Gamma2) if Gamma == (Gamma1 ; Theta ; Gamma2)
ctxHole :: Ctx -> Ctx -> (Ctx, Ctx)
ctxHole [] hole@(_ :& _) = error ("Can't find" +++ code hole +++ "in context!")
ctxHole ctx hole =
  case stripPrefix hole ctx of
    Just ctx' -> (ctx', [])
    Nothing -> let
      ctx' :& top = ctx
      (before, after) = ctxHole ctx' hole
      in (before, after :& top)

fillHole :: (Ctx, Ctx) -> Ctx -> Ctx
fillHole (before, after) middle =
  -- This looks reversed because `Ctx` is a reversed-looking list.
  after <> middle <> before

ctxTwoHoles :: Ctx -> Ctx -> Ctx -> (Ctx, Ctx, Ctx)
ctxTwoHoles original hole1 hole2 = let
  (firstThird, lastTwoThirds) = ctxHole original hole1
  (middleThird, lastThird) = ctxHole lastTwoThirds hole2
  in (firstThird, middleThird, lastThird)

fillTwoHoles :: (Ctx, Ctx, Ctx) -> Ctx -> Ctx -> Ctx
fillTwoHoles (firstThird, middleThird, lastThird) repl1 repl2 = let
  firstTwoThirds = fillHole (firstThird, middleThird) repl1
  threeThirds = fillHole (firstTwoThirds, lastThird) repl2
  in threeThirds

mentions :: Ctx -> Ty -> Bool
mentions ctx ty = case (ctx, ty) of

  -- Rule `UVarWF`
  (ctx, Univ alpha) -> IntroUniv alpha `elem` ctx
    
  -- Rule `UnitWF`
  (_, Unit) -> True

  -- Rule `ArrowWF`
  (ctx, a :->: b) -> ctx `mentions` a && ctx `mentions` b

  -- Rule `ForallWF`
  (ctx, ForAll alpha a) -> (ctx :& IntroUniv alpha) `mentions` a

  -- Rule `EvarWF`
  (ctx, Exst alphaHat) -> IntroExst alphaHat `elem` ctx || hasSoln
    where hasSoln = isJust searchRes
          searchRes = find pred ctx
          pred = \case
            IntroSoln xHat _ -> xHat == alphaHat
            _ -> False

subst :: Ctx -> Ty -> Ty
subst ctx Unit = Unit
subst [] ty = ty
subst ctx (ty1 :->: ty2) = subst ctx ty1 :->: subst ctx ty2
subst ctx (ForAll x ty) = ForAll x (subst ctx ty)
subst (ctx :& IntroSoln x ty') ty | Exst x == ty = ty'
subst (ctx :& _) ty = subst ctx ty

lookupVar :: Ctx -> String -> Ty
lookupVar [] x = error ("Unbound variable `" ++ x ++ "`")
lookupVar (ctx :& IntroAnn y ty) x | x == y = ty
lookupVar (ctx :& _) x = lookupVar ctx x

data Expr
  = Lit Ty -- `Lit ty` is a stand-in for a literal of type `ty`.
  | Var String -- A reference to a variable.
  | Lam String Expr -- Lambda function.
  | App Expr Expr -- Function application.
  | Ann Expr Ty -- Type-annotated expression.
  deriving (Show, Eq)

--------------------------------------------------------------------------------

data WithInCtx a = Ctx :|-: a
  deriving (Show, Eq)

data WithOutCtx a = a :-|: Ctx
  deriving (Show, Eq)

code :: Show a => a -> String 
code a = "`" ++ show a ++ "`"

(+++) :: String -> String -> String
a +++ b = a ++ " " ++ b

ensure :: Bool -> String -> TypeChecker ()
ensure True _ = return ()
ensure False msg = return $ error msg

--------------------------------------------------------------------------------

data Name = Name Char Integer
  deriving Eq

instance Show Name where
  show (Name ch n)
    | n == 0 = "?" ++ [ch]
    | otherwise = "?" ++ (ch : show n)

allNames :: [Name]
allNames = do
  n <- [0..]
  ch <- ['a'..'z']
  return $ Name ch n

instance Enum Name where
  toEnum i = allNames !! i

  fromEnum name = idx
    where Just idx = findIndex (==name) allNames

type TypeCheckerState = Name

type TypeChecker a = State TypeCheckerState a

initState :: TypeCheckerState
initState = toEnum 0

fresh :: TypeChecker String
fresh = do
  name <- get
  put $ succ name
  return $ show name

--------------------------------------------------------------------------------

infer :: WithInCtx Expr -> TypeChecker (WithOutCtx Ty)
infer = \case

  -- Rule `1I=>`
  ctx :|-: Lit ty -> do
    () <- ensure (ctx `mentions` ty) $ "1I=>:" +++ code ty +++ "not mentioned in" +++ code ctx
    return $ ty :-|: ctx

  -- Rule `Var`
  ctx :|-: Var x -> do
    let a = lookupVar ctx x
    return $ a :-|: ctx

  -- Rule `Anno`
  ctx :|-: Ann e a -> do
    () <- ensure (ctx `mentions` a) $ "Anno: Unknown type" +++ code a
    result@( _a :-|: ctx' ) <- check $ ctx :|-: (e, a)
    return result
  
  -- Rule `->E`
  ctx :|-: App e1 e2 -> do
    a :-|: ctx' <- infer (ctx :|-: e1)
    let a' = subst ctx' a
    result@( c :-|: ctx'' ) <- inferApp $ ctx' :|-: (a', e2)
    return result

  -- Rule `->I=>`
  ctx :|-: Lam x e -> do
    alphaHat <- fresh
    betaHat <- fresh
    let hypothesis = ctx :& IntroExst alphaHat
                         :& IntroExst betaHat
                         :& IntroAnn x (Exst alphaHat)
    _betaHat :-|: ctx' <- check $ hypothesis :|-: (e, Exst betaHat)
    let ctx'' = dropUntil ctx' (IntroAnn x (Exst alphaHat))
    return $ (Exst alphaHat :->: Exst betaHat) :-|: ctx''
    where
      dropUntil :: Ctx -> CtxElem -> Ctx
      dropUntil (ctx :& x) y
        | x == y = ctx
        | otherwise = dropUntil ctx y

--------------------------------------------------------------------------------

check :: WithInCtx (Expr, Ty) -> TypeChecker (WithOutCtx Ty)
check = \case

  -- Rule `1I`
  ctx :|-: (Lit ty1, ty2) | ty1 == ty2 -> return $ ty1 :-|: ctx

  -- Rule `Sub`
  ctx :|-: (e, b) -> do
    a :-|: ctx' <- infer $ ctx :|-: e
    let a' = subst ctx' a
    let b' = subst ctx' b
    () :-|: ctx'' <- isSubtype $ ctx' :|-: (a', b')
    return $ b :-|: ctx''

--------------------------------------------------------------------------------

inferApp :: WithInCtx (Ty, Expr) -> TypeChecker (WithOutCtx Ty)
inferApp = \case

  -- Rule `->App`
  ctx :|-: (a :->: c, e) -> do
    result@( _a :-|: ctx' ) <- check $ ctx :|-: (e, a)
    return result

  -- Rule `ForAll App`
  ctx :|-: (ForAll alpha a, e) -> do
    alphaHat <- fresh
    let hypothesis = ctx :& IntroExst alphaHat
    let a' = replace a (Univ alpha, Exst alphaHat)
    result@( c :-|: ctx' ) <- inferApp $ hypothesis :|-: (a', e)
    return result

  -- Rule `alphaHat App`
  ctx :|-: (Exst alphaHat, e) -> do
    let (before, after) = ctxHole ctx [IntroExst alphaHat]
    alphaHat1 <- fresh
    alphaHat2 <- fresh
    let middle = [] :& IntroExst alphaHat2
                    :& IntroExst alphaHat1
                    :& IntroSoln alphaHat (Exst alphaHat1 :->: Exst alphaHat2)
    let ctx' = fillHole (before, after) middle
    result@( _ :-|: ctx'' ) <- check $ ctx' :|-: (e, Exst alphaHat1)
    return result

--------------------------------------------------------------------------------

isSubtype :: WithInCtx (Ty, Ty) -> TypeChecker (WithOutCtx ())
isSubtype = \case

  -- Rule `<:Var`
  ctx :|-: (Univ alpha1, Univ alpha2) | alpha1 == alpha2 -> do
    () <- ensure (ctx `mentions` Univ alpha1) $ "<:Var: Couldn't find" ++ code (Univ alpha1)
    return $ () :-|: ctx

  -- Rule `<:Unit`
  ctx :|-: (Unit, Unit) -> do
    return $ () :-|: ctx

  -- Rule `<:Exvar`
  ctx :|-: (Exst alphaHat1, Exst alphaHat2) | alphaHat1 == alphaHat2 -> do
    () <- ensure (ctx `mentions` Exst alphaHat1) $ "<:Exvar: Couldn't find" ++ code (Exst alphaHat1)
    return $ () :-|: ctx

  -- Rule `<: ->`
  ctx :|-: (a1 :->: a2, b1 :->: b2) -> do
    () :-|: ctx' <- isSubtype $ ctx :|-: (b1, a1)
    let a2' = subst ctx' a2
    let b2' = subst ctx' b2
    result@( () :-|: ctx'' ) <- isSubtype $ ctx :|-: (a2', b2')
    return result

  -- Rule `<: ForAll L`
  ctx :|-: (ForAll alpha a, b) -> do
    alphaHat <- fresh
    let hypothesis = ctx :& IntroMarker alphaHat
                         :& IntroExst alphaHat
    let a' = replace a (Exst alphaHat, Univ alpha)
    () :-|: ctx' <- isSubtype $ hypothesis :|-: (a', b)
    let (before, _after) = ctxHole ctx' ([] :& IntroMarker alphaHat)
    let ctx'' = before
    return $ () :-|: ctx''

  -- Rule `<: ForAll R`
  ctx :|-: (a, ForAll alpha b) -> do
    let hypothesis = ctx :& IntroUniv alpha
    () :-|: ctx' <- isSubtype $ hypothesis :|-: (a, b)
    let (before, _after) = ctxHole ctx' ([] :& IntroUniv alpha)
    let ctx'' = before
    return $ () :-|: ctx''

  -- Rule `<: Instantiate L`
  ctx :|-: (Exst alphaHat, a) -> do
    () <- ensure (ctx `mentions` Exst alphaHat) $ "<:InstantiateL: Couldn't find" +++ code (Exst alphaHat)
    () <- ensure (not (Exst alphaHat `occursIn` a)) $ "<:InstantiateL: Occurs check failed for" +++ code (Exst alphaHat, a)
    result@( () :-|: ctx' ) <- instSubtype $ ctx :|-: (alphaHat, a)
    return result

  -- Rule `<: Instantiate R`
  ctx :|-: (a, Exst alphaHat) -> do
    () <- ensure (ctx `mentions` Exst alphaHat) $ "<:InstantiateR: Couldn't find" +++ code (Exst alphaHat)
    () <- ensure (not (Exst alphaHat `occursIn` a)) $ "<:InstantiateR: Occurs check failed for" +++ code (Exst alphaHat, a)
    result@( () :-|: ctx' ) <- instSupertype $ ctx :|-: (a, alphaHat)
    return result

--------------------------------------------------------------------------------

instSubtype :: WithInCtx (String, Ty) -> TypeChecker (WithOutCtx ())
instSubtype = \case

  -- Rule `InstLReach`
  ctx :|-: (alphaHat, Exst betaHat) -> do
    -- This corresponds to `Gamma[alphaHat][betaHat]` in the paper.
    let (before, middle, after) = ctxTwoHoles ctx [IntroExst alphaHat] [IntroExst betaHat]
    let ctx' = fillTwoHoles (before, middle, after) [IntroExst alphaHat] [IntroSoln betaHat (Exst alphaHat)]
    return $ () :-|: ctx'

  -- Rule `InstLArr`
  ctx :|-: (alphaHat, a1 :->: a2) -> do
    let (before, after) = ctxHole ctx [IntroExst alphaHat]
    alphaHat1 <- fresh
    alphaHat2 <- fresh
    let repl = [] :& IntroExst alphaHat2
                  :& IntroExst alphaHat1
                  :& IntroSoln alphaHat (Exst alphaHat1 :->: Exst alphaHat2)
    let ctx' = fillHole (before, after) repl
    () :-|: ctx'' <- instSupertype $ ctx' :|-: (a1, alphaHat1)
    let a2' = subst ctx'' a2
    result@( () :-|: ctx''' ) <- instSubtype $ ctx'' :|-: (alphaHat2, a2')
    return result

  -- Rule `InstLAllR`
  ctx :|-: (alphaHat, ForAll beta b) -> do
    () <- ensure (ctx `mentions` Exst alphaHat) $ "InstLAllR:" +++ code (Exst alphaHat, ctx)
    let hypothesis = ctx :& IntroUniv beta
    () :-|: ctx' <- instSubtype $ hypothesis :|-: (alphaHat, b)
    let (beforeBeta, afterBeta) = ctxHole ctx' [IntroUniv beta]
    return $ () :-|: beforeBeta

  -- Rule `InstLSolve`
  ctx :|-: (alphaHat, tau) -> do
    let (before, after) = ctxHole ctx [IntroExst alphaHat]
    () <- ensure (before `mentions` tau) $ "InstLSolve:" +++ code (alphaHat, tau)
    return $ () :-|: fillHole (before, after) [IntroSoln alphaHat tau]


--------------------------------------------------------------------------------

instSupertype :: WithInCtx (Ty, String) -> TypeChecker (WithOutCtx ())
instSupertype = \case

  -- Rule `InstRReach
  ctx :|-: (Exst betaHat, alphaHat) -> do
    let (before, middle, after) = ctxTwoHoles ctx [IntroExst alphaHat] [IntroExst betaHat]
    let ctx' = fillTwoHoles (before, middle, after) [IntroExst alphaHat] [IntroSoln betaHat (Exst alphaHat)]
    return $ () :-|: ctx'

  -- Rule `InstLArr`
  ctx :|-: (a1 :->: a2, alphaHat) -> do
    let (before, after) = ctxHole ctx [IntroExst alphaHat]
    alphaHat1 <- fresh
    alphaHat2 <- fresh
    let repl = [] :& IntroExst alphaHat2
                  :& IntroExst alphaHat1
                  :& IntroSoln alphaHat (Exst alphaHat1 :->: Exst alphaHat2)
    let ctx' = fillHole (before, after) repl
    () :-|: ctx'' <- instSubtype $ ctx' :|-: (alphaHat1, a1)
    let a2' = subst ctx'' a2
    result@( () :-|: ctx''' ) <- instSupertype $ ctx'' :|-: (a2', alphaHat2)
    return result

  -- Rule `InstRAllL`
  ctx :|-: (ForAll beta b, alphaHat) -> do
    () <- ensure (ctx `mentions` Exst alphaHat) $ "InstRAllL:" +++ code (ctx, Exst alphaHat)
    betaHat <- fresh
    let hypothesis = ctx :& IntroMarker betaHat :& IntroExst betaHat
    let b' = replace b (Exst betaHat, Univ beta)
    () :-|: ctx' <- instSupertype $ hypothesis :|-: (b', alphaHat)
    let (beforeMarker, afterMarker) = ctxHole ctx' [IntroMarker betaHat]
    return $ () :-|: beforeMarker

  -- Rule `InstRSolve`
  ctx :|-: (tau, alphaHat) -> do
    let (before, after) = ctxHole ctx [IntroExst alphaHat]
    () <- ensure (before `mentions` tau) $ "InstRSolve:" +++ code (tau, alphaHat)
    return $ () :-|: fillHole (before, after) [IntroSoln alphaHat tau]