Implementation of LF typechecking using canonical forms and hereditary substitution

lf.sml

(* uses some data structures from
 * https://github.com/standardml/cmlib/ *)

(* also the pretty printer uses some hacked up versions of some stuff
 * from tom7 that I don't actually include but will if anybody wants to
 * actually run this *)


signature VARIABLE =
sig
  type var
  type t = var
  val toStr : var -> string
  val eq : var * var -> bool
  val compare : var * var -> order
end


structure Variable : VARIABLE =
struct
  type var = string
  type t = var
  fun toStr s = s
  fun eq (v: var, v') = v = v'
  val compare = String.compare
end
structure VarDict = SplayDict(structure Key = Variable)
structure VarSet = SplaySet(structure Elem = Variable)

structure Const =
struct
  type const = string
  type t = const

  fun toStr (s: const) = s
  fun eq (x: const, y) = x = y
  val compare = String.compare
end
structure ConstDict = SplayDict(structure Key = Const)
structure ConstSet = SplaySet(structure Elem = Const)


structure LFSyntax =
struct

  type var = int * string
  type const = Const.const
  type binding = string


  datatype head = HVar of var
                | HConst of const

  datatype exp = EKind
               | EType
               | EProp (* ?? maybe not. *)
               | EPi of binding * exp * exp
               | ELam of binding * exp
               | EApp of head * spine
  (* Should spine just be a list? *)
       and spine = SNil
                 | SApp of exp * spine

  datatype entry_type = SgFamilyDecl | SgObjectDecl
  type sg_entry = entry_type * Const.const * exp
  type sg = sg_entry list

  val listToSpine = foldr SApp SNil
  fun spineToList SNil = nil
    | spineToList (SApp (e, s)) = e :: spineToList s
  (* welp. *)
  fun mapSpine f = listToSpine o map f o spineToList

end
structure LF = LFSyntax



signature LF_SUBST =
sig
  val hereditaryReduce : LF.exp -> LF.spine -> LF.exp
  val substExp : int -> LF.exp list -> int -> LF.exp -> LF.exp
  val liftExp : int -> LF.exp -> LF.exp
end


signature CONTEXT =
sig
  type ctx
  val empty : ctx
  val length : ctx -> int
  val sub : ctx -> int -> LF.exp
  val extend : ctx -> LF.exp -> ctx
end

structure LFSubst :> LF_SUBST =
struct

  open LFSyntax

  (* substXMain skip substs substs_len lift exp

     s = substs, n = substs_len, l = lift, m = skip

     if    |s| = n
     then  return exp[0 .. m-1 . s0[^m] .. sn-1[^m] . ^l+m]
   *)
  fun substExpMain skip substs substs_len lift exp =
      (case exp of
           EKind => EKind
         | EType => EType
         | EProp => EProp
         | ELam (b, e) =>
           let val e' = substExpMain (skip+1) substs substs_len lift e
           in ELam (b, e') end
         | EPi (b, e1, e2) =>
           let val e1' = substExpMain skip substs substs_len lift e1
               val e2' = substExpMain (skip+1) substs substs_len lift e2
           in EPi (b, e1', e2') end
         | EApp (head, spine) =>
           let val spine' = substSpineMain skip substs substs_len lift spine
           in (case head of
                   HConst c =>
                   EApp (HConst c, spine')
                 | HVar (i, s) =>
                   if i < skip then
                       EApp (HVar (i, s), spine')
                   else if i < skip+substs_len then
                       let val sub = List.nth (substs, i-skip)
                           val sub' = substExpMain 0 [] 0 skip sub
                       (* I'm *pretty* sure all of the index stuff is done.. *)
                       in hereditaryReduce sub' spine end
                   else
                       EApp (HVar (i-substs_len+lift, s), spine'))
           end
      )
  and substSpineMain skip substs substs_len lift spine =
      (case spine of
           SNil => SNil
         | SApp (exp, spine) =>
           let val exp' = substExpMain skip substs substs_len lift exp
               val spine' = substSpineMain skip substs substs_len lift spine
           in SApp (exp', spine') end)

  and hereditaryReduce head spine =
      let fun getBody (ELam (b, e)) n = getBody e (n+1)
            | getBody e n = (e, n)

           fun getSubsts SNil subs = subs
             | getSubsts (SApp (e, s)) subs = getSubsts s (e :: subs)

           val (body, count) = getBody head 0
           val subs = getSubsts spine []

           val () = if count = length subs then ()
                    else raise Fail "bogus application"

      in substExpMain 0 subs (length subs) 0 body end

  fun substExp skip substs lift exp =
      substExpMain skip substs (length substs) lift exp
  fun liftExp lift exp = substExp 0 [] lift exp
end

structure LFContext :> CONTEXT =
struct
  type ctx = LFSyntax.exp list
  val empty = nil
  val length = List.length
  fun sub G n =
      let val t = List.nth (G, n)
      in LFSubst.liftExp (n+1) t end
  fun extend G t = t :: G
end


signature SIGNATURE =
sig
  type sg
  val empty : sg
  val insert : sg -> LF.const -> LF.exp -> sg
  val lookup : sg -> LF.const -> LF.exp
end

structure Signature :> SIGNATURE =
struct
  val SignatureError = Fail

  type sg = LF.exp ConstDict.dict
  val empty = ConstDict.empty

  fun insert' sg c entry =
      if ConstDict.member sg c then
          raise SignatureError (Const.toStr c ^ " is already in signature")
      else ConstDict.insert sg c entry

  fun insert sg c typ = insert' sg c typ

  fun lookup sg c =
      (case ConstDict.find sg c of
           NONE => raise SignatureError (Const.toStr c ^ " not in signature")
         | SOME t => t)
end


structure PrettyLF =
struct
  local
      open LF

      structure L = Layout

      val $ = L.str
      val % = L.mayAlign
      val %% = L.freeStyleAlign
      val & = L.seq

      val WIDTH = 80
      val fmt = L.tostringex WIDTH
  in

  val look_good_but_be_wrong = true

  fun prettyConst s = s

  fun toLayoutHead (HVar (i, s)) =
      if look_good_but_be_wrong then $s
      else $(s ^ "/" ^ Int.toString i)
    | toLayoutHead (HConst s) = $(prettyConst s)
  fun toLayoutExp e =
      (case e of
           EKind => $"kind"
         | EType => $"type"
         | EProp => $"prop"
         (* Basing this on the variable being called "_" is a bit bogus. *)
         | EPi ("_", e1, e2) =>
           % [&[toLayoutTyParen e1, $" ->"],
              toLayoutExp e2]

         | EPi (b, e1, e2) =>
           % [piPartLayout b e1, toLayoutExp e2]
         | ELam (b, e) =>
           & [$"\\",
              % [ &[$b, $"."],
                  toLayoutExp e]
             ]
         | EApp (h, SNil) => toLayoutHead h
         | EApp (h, s) =>
           & [toLayoutHead h, $" ",
              toLayoutSpine s])
  and toLayoutSpine s =
      let val layouts = map toLayoutExpParen (spineToList s)
      in %% layouts end

  and toLayoutExpParen (e as ELam _) = L.paren (toLayoutExp e)
    | toLayoutExpParen (e as EApp (_, SApp _)) = L.paren (toLayoutExp e)
    | toLayoutExpParen e = toLayoutExp e
  and toLayoutTyParen (e as ELam _) = L.paren (toLayoutExp e)
    | toLayoutTyParen (e as EPi _) = L.paren (toLayoutExp e)
    | toLayoutTyParen e = toLayoutExp e

  and piPartLayout b e1 = &[$"pi ", $b, $" : ", toLayoutExp e1, $"."]

  (* A nicer layout for "toplevel" pis *)
  (* Another try that doesn't screw over ->s ?? *)
  fun toLayoutTop2 e =
      let fun loop l (e as EPi ("_", e1, e2)) = loop (&[toLayoutTyParen e1, $" ->"] :: l) e2
            | loop l (EPi (b, e1, e2)) = loop (piPartLayout b e1 :: l) e2
            | loop l e = %[%% (rev l), toLayoutExp e]
      in loop [] e end
  (* Actually, let's combine them. *)
  fun toLayoutTop1 e =
      let fun loop l (e as EPi ("_", _, _)) = %[% (rev l), toLayoutTop2 e]
            | loop l (EPi (b, e1, e2)) = loop (piPartLayout b e1 :: l) e2
            | loop l e = %[% (rev l), toLayoutTop2 e]
      in loop [] e end

  val toLayoutTop = toLayoutTop1

  fun prettyExp e = L.tostringex WIDTH (toLayoutTop e)

  fun prettyMsg msg e = fmt (&[$msg, toLayoutExp e])
  fun prettyMsg2 msg1 e1 sep msg2 e2 =
      fmt (%[ &[$msg1, toLayoutExp e1, $sep],
              &[$msg2, toLayoutExp e2]])

  fun prettyDecl (_, c, e) = fmt (&[$c, $": ", toLayoutTop e, $"."])

  fun prettySignature sg =
      String.concatWith "\n" (map prettyDecl sg)

  end
end


(* ***** OK Now the actual typechecker ****** *)
signature TYPE_CHECK_LF =
sig
  exception TypeError of string

  val expEquality : LF.exp -> LF.exp -> unit
  val checkExp : Signature.sg -> LFContext.ctx -> LF.exp -> LF.exp -> unit

  val checkSignatureEntry : Signature.sg -> LF.sg_entry -> Signature.sg

  val checkSignature : LF.sg -> Signature.sg

end


structure TypeCheckLF : TYPE_CHECK_LF =
struct

local
  open LFSyntax
  structure Ctx = LFContext
  structure Sig = Signature
in

  exception TypeError of string


  fun requireKind exp =
      if exp = EKind then () else raise TypeError "expected kind"
  fun requireAtomic exp =
      (case exp of
           EType => ()
         | EProp => ()
         | EApp _ => ()
         | _ => raise TypeError (PrettyLF.prettyMsg
                                     "required atomic type, got: "
                                     exp))

  fun expEquality e e' =
      (case (e, e') of
           (EKind, EKind) => ()
         | (EType, EType) => ()
         | (EProp, EProp) => ()
         | (EPi (_, e1, e2), EPi (_, e1', e2')) =>
           (expEquality e1 e1'; expEquality e2 e2')
         | (ELam (_, e), ELam (_, e')) =>
           expEquality e e'
         | (EApp (h, s), EApp (h', s')) =>
           (headEquality h h'; spineEquality s s')
         | _ => raise TypeError "exps not equal")
  and headEquality h h' =
      (case (h, h') of
           (HVar (i, _), HVar (i', _)) =>
           if i = i' then () else
           raise (TypeError ("bound variable mismatch: " ^
                             Int.toString i ^ " vs. " ^ Int.toString i'))
         | (HConst c, HConst c') =>
           if c = c' then () else
           raise (TypeError ("const mismatch: " ^ Const.toStr c ^ " vs. " ^ Const.toStr c'))
         | _ => raise TypeError "const vs. var mismatch")
  and spineEquality s s' =
      (case (s, s') of
           (SNil, SNil) => ()
         | (SApp (e, s), SApp (e', s')) =>
           (expEquality e e'; spineEquality s s')
         | _ => raise TypeError "spine length mismatch??")


  fun expEquality' e t t' =
      ((expEquality t t')
       handle TypeError s => raise TypeError (
          "equality failure: " ^ s ^ "\n" ^
          PrettyLF.prettyMsg2 "expected: " t ","  "got: " t' ^ "\n" ^
          PrettyLF.prettyMsg "while checking: " e))

  fun checkExp sg ctx exp typ =
      ((*print (PrettyLF.prettyMsg2 "checking: " exp "," "at: " typ ^ "\n");*)
       case exp of
           EKind => raise TypeError "kind is no classifier"
         | EType => requireKind typ
         | EProp => requireKind typ
         | EPi (b, e1, e2) =>
           (checkExp sg ctx e1 EType;
            checkExp sg (Ctx.extend ctx e1) e2 typ)

         | ELam (b, e) =>
           let val (t1, t2) =
               (case typ of EPi (_, t1, t2) => (t1, t2)
                          | _ => raise TypeError "lambda must have pi type")
           in checkExp sg (Ctx.extend ctx t1) e t2 end

         | EApp (h, spine) =>
           let val t = checkHead sg ctx h
               (*val () = print (PrettyLF.prettyMsg "head has type: " t ^ "\n")*)
               val t' = checkSpine sg ctx t spine
               val () = requireAtomic t'
           in expEquality' exp typ t' end)
  and checkHead _ ctx (HVar (n, _)) = Ctx.sub ctx n
    | checkHead sg _ (HConst c) = Sig.lookup sg c
  and checkSpine sg ctx typ SNil = typ
    | checkSpine sg ctx typ (SApp (e, s)) =
      let (*val () = print (PrettyLF.prettyMsg "checking at: " typ ^ "\n")*)
          val (t1, t2) =
               (case typ of EPi (_, t1, t2) => (t1, t2)
                          | _ => raise TypeError "lhs of app must be pi")
          val () = checkExp sg ctx e t1
          (* We could probably arrange to do one big substitution. *)
          val t2' = LFSubst.substExp 0 [e] 0 t2
      in checkSpine sg ctx t2' s end

  fun checkExp' sg ctx exp typ =
      ((checkExp sg ctx exp typ)
       handle TypeError s => (print ("Type error: " ^ s ^ "\n"); raise TypeError s))
  val checkExp = checkExp'
  fun expEquality'' e e' =
      ((expEquality e e')
       handle TypeError s => (print ("Type error: " ^ s ^ "\n"); raise TypeError s))
  val expEquality = expEquality''


  fun checkSignatureEntry sg ((entry_type, c, exp): LF.sg_entry) =
      let val classifier =
              (case entry_type of SgFamilyDecl => EKind
                                | SgObjectDecl => EType)
          val () = checkExp sg Ctx.empty exp classifier
      in Sig.insert sg c exp end

  fun checkSignature decls =
      foldl (fn (e, sg) => checkSignatureEntry sg e) Sig.empty decls

end
end