SProp sketch

SPropSketch.hs

{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE OverloadedStrings #-}

import Data.String (IsString(..))

type Ix = Int
type Lvl = Int
type ULvl = Int

data Sort = Type ULvl | Prop
  deriving (Show, Eq)

data SortType = SType | SProp
  deriving (Show, Eq)

sortType :: Sort -> SortType
sortType (Type _) = SType
sortType Prop = SProp

data Proj = Fst | Snd
  deriving (Eq)

instance Show Proj where
  show Fst = "fst"
  show Snd = "snd"

data Tm
  = Var Ix
  | Let Tm Tm Tm

  | Lam SortType Tm
  | App Tm Tm SortType
  | Pi Tm SortType Tm

  | Pair Tm SortType Tm SortType
  | Proj Proj Tm
  | Sigma Tm SortType SortType Tm

  | Sort Sort

  | Bot
  | Exfalso Tm

  | Id Tm SortType Tm Tm
  | Refl

instance Show Tm where
  show (Var i) = show i
  show (Let ty val body) = "(let " ++ show ty ++ " = " ++ show val ++ "; " ++ show body ++ ")"
  show (Lam _ body) = "(\\" ++ show body ++ ")"
  show (App fn arg _) = "(" ++ show fn ++ " " ++ show arg ++ ")"
  show (Pi ty _ body) = "(" ++ show ty ++ " -> " ++ show body ++ ")"
  show (Pair a _ b _) = "(" ++ show a ++ ", " ++ show b ++ ")"
  show (Proj p t) = "(" ++ show p ++ " " ++ show t ++ ")"
  show (Sigma ty _ _ body) = "(" ++ show ty ++ " ** " ++ show body ++ ")"
  show (Sort sort) = show sort
  show Bot = "Bot"
  show (Exfalso tm) = "(exfalso " ++ show tm ++ ")"
  show (Id _ _ a b) = "(" ++ show a ++ " = " ++ show b ++ ")"
  show Refl = "Refl"

umaxPi :: Sort -> Sort -> Sort
umaxPi _ Prop = Prop
umaxPi Prop (Type i) = Type i
umaxPi (Type i) (Type j) = Type (max i j)

umaxSigma :: Sort -> Sort -> Sort
umaxSigma Prop Prop = Prop
umaxSigma Prop (Type i) = Type i
umaxSigma (Type i) Prop = Type i
umaxSigma (Type i) (Type j) = Type (max i j)

-- values
type Env = [Val]
data Clos = Clos Tm Env

data Elim
  = EApp Val SortType
  | EProj Proj
  | EExfalso
type Sp = [Elim]

data Val
  = VNe Lvl Sp
  | VLam SortType Clos
  | VPi Val SortType Clos
  | VPair Val SortType Val SortType
  | VSigma Val SortType SortType Clos
  | VSort Sort
  | VBot
  | VId Val SortType Val Val
  | VRefl

pattern VVar l = VNe l []

-- evaluation
vinst :: Clos -> Val -> Val
vinst (Clos tm vs) v = eval (v : vs) tm

vapp :: Val -> Val -> SortType -> Val
vapp (VLam _ c) v s = vinst c v
vapp (VNe h sp) v s = VNe h (EApp v s : sp)
vapp _ _ _ = error "vapp"

vproj :: Proj -> Val -> Val
vproj Fst (VPair v _ _ _) = v
vproj Snd (VPair _ _ v _) = v
vproj p (VNe h sp) = VNe h (EProj p : sp)
vproj _ _ = error "vproj"

vexfalso :: Val -> Val
vexfalso (VNe h s) = VNe h (EExfalso : s)
vexfalso _ = error "vexfalso"

eval :: Env -> Tm -> Val
eval vs (Var i) = vs !! i
eval vs (Let _ v b) = eval (eval vs v : vs) b
eval vs (Lam s b) = VLam s (Clos b vs)
eval vs (App a b s) = vapp (eval vs a) (eval vs b) s
eval vs (Pi t s b) = VPi (eval vs t) s (Clos b vs)
eval vs (Pair a s b s') = VPair (eval vs a) s (eval vs b) s'
eval vs (Proj p t) = vproj p (eval vs t)
eval vs (Sigma t s s' b) = VSigma (eval vs t) s s' (Clos b vs)
eval vs (Sort s) = VSort s
eval vs Bot = VBot
eval vs (Exfalso tm) = vexfalso (eval vs tm)
eval vs (Id ty s a b) = VId (eval vs ty) s (eval vs a) (eval vs b)
eval vs Refl = VRefl

quoteElim :: Lvl -> Elim -> Tm -> Tm
quoteElim l (EApp v s) t = App t (quote l v) s
quoteElim l (EProj p) t = Proj p t
quoteElim l EExfalso t = Exfalso t

quote :: Lvl -> Val -> Tm
quote l (VNe hd sp) = foldr (quoteElim l) (Var (l - hd - 1)) sp
quote l (VLam s c) = Lam s (quote (l + 1) $ vinst c (VVar l))
quote l (VPi ty s c) = Pi (quote l ty) s (quote (l + 1) $ vinst c (VVar l))
quote l (VPair a s b s') = Pair (quote l a) s (quote l b) s'
quote l (VSigma ty s s' c) = Sigma (quote l ty) s s' (quote (l + 1) $ vinst c (VVar l))
quote _ (VSort s) = Sort s
quote _ VBot = Bot
quote l (VId ty s a b) = Id (quote l ty) s (quote l a) (quote l b)
quote l VRefl = Refl

nf :: Tm -> Tm
nf = quote 0 . eval []

-- conversion
convSp :: Lvl -> Sp -> Sp -> Bool
convSp l [] [] = True
convSp l (EApp v s : sp) (EApp v' _ : sp') = convSp l sp sp' && conv l s v v'
convSp l (EProj p : sp) (EProj p' : sp') = convSp l sp sp' && p == p'
convSp l (EExfalso : sp) (EExfalso : sp') = convSp l sp sp'
convSp _ _ _ = False

conv :: Lvl -> SortType -> Val -> Val -> Bool
conv _ SProp _ _ = True
conv l _ (VSort s) (VSort s') = s == s'
conv l _ VBot VBot = True
conv l _ VRefl _ = True
conv l _ _ VRefl = True
conv l _ (VId ty s a b) (VId ty' s' a' b') = s == s' && conv l SType ty ty' && conv l s a a' && conv l s b b'
conv l ss (VLam s c) (VLam _ c') = let v = VVar l in conv (l + 1) ss (vinst c v) (vinst c' v)
conv l ss (VLam s c) w = let v = VVar l in conv (l + 1) ss (vinst c v) (vapp w v s)
conv l ss w (VLam s c) = let v = VVar l in conv (l + 1) ss (vapp w v s) (vinst c v)
conv l _ (VPair a s b s') (VPair a' _ b' _) = conv l s a a' && conv l s' b b'
conv l _ (VPair a s b s') v = conv l s a (vproj Fst v) && conv l s' b (vproj Snd v)
conv l _ v (VPair a s b s') = conv l s (vproj Fst v) a && conv l s' (vproj Snd v) b
conv l ss (VPi ty s c) (VPi ty' s' c') = s == s' && conv l s ty ty' && (let v = VVar l in conv (l + 1) ss (vinst c v) (vinst c' v))
conv l ss (VSigma ty s _ c) (VSigma ty' s' _ c') = s == s' && conv l s ty ty' && (let v = VVar l in conv (l + 1) ss (vinst c v) (vinst c' v))
conv l _ (VNe hd sp) (VNe hd' sp') = hd == hd' && convSp l sp sp'
conv l _ _ _ = False

convTypes :: Lvl -> Val -> Val -> Bool
convTypes l a b = conv l SType a b

-- surface
type Name = String
data STm
  = SVar Name
  | SLet Name STm STm STm

  | SLam Name STm
  | SApp STm STm
  | SPi Name STm STm

  | SPair STm STm
  | SProj Proj STm
  | SSigma Name STm STm

  | SSort Sort

  | SBot
  | SExfalso STm

  | SId STm STm STm
  | SRefl

instance IsString STm where
  fromString = SVar

instance Show STm where
  show (SVar x) = x
  show (SLet x ty val body) = "(let " ++ x ++ " : " ++ show ty ++ " = " ++ show val ++ "; " ++ show body ++ ")"
  show (SLam x body) = "(\\" ++ x ++ ". " ++ show body ++ ")"
  show (SApp fn arg) = "(" ++ show fn ++ " " ++ show arg ++ ")"
  show (SPi "_" ty body) = "(" ++ show ty ++ " -> " ++ show body ++ ")"
  show (SPi x ty body) = "((" ++ x ++ " : " ++ show ty ++ ") -> " ++ show body ++ ")"
  show (SPair a b) = "(" ++ show a ++ ", " ++ show b ++ ")"
  show (SProj p t) = "(" ++ show p ++ " " ++ show t ++ ")"
  show (SSigma "_" ty body) = "(" ++ show ty ++ " ** " ++ show body ++ ")"
  show (SSigma x ty body) = "((" ++ x ++ " : " ++ show ty ++ ") ** " ++ show body ++ ")"
  show (SSort sort) = show sort
  show SBot = "Bot"
  show (SExfalso tm) = "(exfalso " ++ show tm ++ ")"
  show (SId _ a b) = "(" ++ show a ++ " = " ++ show b ++ ")"
  show SRefl = "Refl"

-- context
type Types = [(String, Val)]
data Ctx = Ctx { lvl :: Lvl, env :: [Val], types :: Types }

emptyCtx :: Ctx
emptyCtx = Ctx 0 [] []

bind :: Ctx -> Name -> Val -> Ctx
bind (Ctx lvl vs ts) x ty = Ctx (lvl + 1) (VVar lvl : vs) ((x, ty) : ts)

define :: Ctx -> Name -> Val -> Val -> Ctx
define (Ctx lvl vs ts) x ty val = Ctx (lvl + 1) (val : vs) ((x, ty) : ts)

lookupCtx :: Ctx -> Name -> Maybe (Int, Val)
lookupCtx ctx x = go 0 (types ctx)
  where
    go i [] = Nothing
    go i ((y, ty) : _) | x == y = Just (i, ty)
    go i (_ : rest) = go (i + 1) rest

-- typechecking
type TC t = Either String t

check :: Ctx -> STm -> Val -> TC Tm
check ctx (SLam x b) (VPi ty s c) = do
  eb <- check (bind ctx x ty) b (vinst c (VVar (lvl ctx)))
  return $ Lam s eb
check ctx (SPair a b) (VSigma ty s s' c) = do
  ea <- check ctx a ty
  eb <- check ctx b (vinst c (eval (env ctx) ea))
  return $ Pair ea s eb s'
check ctx (SLet x ty val b) ty' = do
  (ety, _) <- inferSort ctx ty
  let vt = eval (env ctx) ety
  ev <- check ctx val vt
  eb <- check (define ctx x vt (eval (env ctx) ev)) b ty'
  return $ Let ety ev eb
check ctx (SExfalso tm) ty = do
  etm <- check ctx tm VBot
  return $ Exfalso etm
check ctx SRefl idty@(VId ty s a b) =
  if conv (lvl ctx) s a b then
    return Refl
  else
    Left $ "refl failed: " ++ show (quote (lvl ctx) idty)
check ctx tm ty = do
  (etm, ty') <- infer ctx tm
  if convTypes (lvl ctx) ty' ty then
    return etm
  else
    Left $ "check failed (" ++ show tm ++ " : " ++ show (quote (lvl ctx) ty) ++ "), got " ++ show (quote (lvl ctx) ty') 

inferSort :: Ctx -> STm -> TC (Tm, Sort)
inferSort ctx ty = do
  (etm, ty') <- infer ctx ty
  case ty' of
    VSort s -> return (etm, s)
    _ -> Left $ "expected sort for " ++ show ty ++ " but got " ++ show (quote (lvl ctx) ty')

infer :: Ctx -> STm -> TC (Tm, Val)
infer ctx (SVar x) =
  case lookupCtx ctx x of
    Nothing -> Left $ "undefined var " ++ x
    Just (i, ty) -> return (Var i, ty)
infer ctx (SLet x ty val b) = do
  (ety, _) <- inferSort ctx ty
  let vt = eval (env ctx) ety
  ev <- check ctx val vt
  (eb, rty) <- infer (define ctx x vt (eval (env ctx) ev)) b
  return (Let ety ev eb, rty)
infer ctx tm@(SApp a b) = do
  (ea, ft) <- infer ctx a
  case ft of
    VPi ty s c -> do
      eb <- check ctx b ty
      return (App ea eb s, vinst c (eval (env ctx) eb))
    _ -> Left $ "expected pi in " ++ show tm ++ " but got " ++ show (quote (lvl ctx) ft)
infer ctx tm@(SProj p t) = do
  (et, pt) <- infer ctx t
  case pt of
    VSigma ty _ _ c ->
      case p of
        Fst -> return (Proj p et, ty)
        Snd -> return (Proj p et, vinst c (vproj Fst $ eval (env ctx) et))
    _ -> Left $ "expected sigma in " ++ show tm ++ " but got " ++ show (quote (lvl ctx) pt)
infer ctx (SPi x ty b) = do
  (ety, s) <- inferSort ctx ty
  (eb, s') <- inferSort (bind ctx x (eval (env ctx) ety)) b
  return (Pi ety (sortType s) eb, VSort (umaxPi s s'))
infer ctx (SSigma x ty b) = do
  (ety, s) <- inferSort ctx ty
  (eb, s') <- inferSort (bind ctx x (eval (env ctx) ety)) b
  return (Sigma ety (sortType s) (sortType s') eb, VSort (umaxSigma s s'))
infer ctx (SId ty a b) = do
  (ety, s) <- inferSort ctx ty
  let vty = eval (env ctx) ety
  ea <- check ctx a vty
  eb <- check ctx b vty
  return (Id ety (sortType s) ea eb, VSort Prop)
infer ctx (SSort Prop) = return (Sort Prop, VSort (Type 0))
infer ctx (SSort (Type l)) = return (Sort (Type l), VSort (Type (l + 1)))
infer ctx SBot = return (Bot, VSort Prop)
infer ctx tm = Left $ "cannot infer " ++ show tm

typecheck :: STm -> TC (Tm, Tm)
typecheck tm = do
  (etm, ty) <- infer emptyCtx tm
  return (etm, quote (lvl emptyCtx) ty)

-- testing
top :: STm
top = SPi "_" SBot SBot

tt :: STm
tt = SLet "tt" top (SLam "x" "x") "tt"

-- squash : 
squash :: STm -> STm
squash a = SPi "P" (SSort Prop) $ SPi "_" (SPi "_" a "P") "P"

sq :: STm -> STm -> STm
sq a x = SLet "squashed" (squash a) (SLam "P" $ SLam "f" $ SApp "f" x) "squashed"

unsq :: STm -> STm -> STm -> STm
unsq p f x = SApp (SApp x p) f

-- let x : (a b : Bot) -> (a = b) = \a b. Refl; x
example :: STm
example = SLet "x" (SPi "a" SBot $ SPi "b" SBot $ SId SBot "a" "b") (SLam "a" $ SLam "b" SRefl) "x"

-- let k : (ty : Type 0) (a b : ty) (q q' : a = b) -> q = q' = \ty a b q q'. Refl; k
example2 :: STm
example2 =
  SLet "k" (SPi "ty" (SSort $ Type 0) $ SPi "a" "ty" $ SPi "b" "ty" $ SPi "q" (SId "ty" "a" "b") $ SPi "q'" (SId "ty" "a" "b") $ SId ((SId "ty" "a" "b")) "q" "q'") (SLam "ty" $ SLam "a" $ SLam "b" $ SLam "q" $ SLam "q'" $ SRefl) "k"

-- let topeta : (x : Top) -> Id Top x tt = \_. Refl; topeta
example3 :: STm
example3 = SLet "topeta" (SPi "x" top $ SId top "x" tt) (SLam "_" SRefl) "topeta"

main :: IO ()
main = do
  let e = sq top tt
  putStrLn $ show e
  case typecheck e of
    Right (etm, ty) -> do
      putStrLn $ show ty
      putStrLn $ show etm
      putStrLn $ show $ nf etm
    Left err -> putStrLn err