minimal MLTT implementation step-by-step

minimal-MLTT.ml

(* 这是我（预期，亦或是已经）在 ∞-type càfe summer school 上做的 talk，

     "从零教你手把手实现一个 MLTT 类型检查器"

   的内容稿。本 talk 计划以现场边讲解边手写代码的方式实现，
   所以虽然这份内容稿会尽量尝试还原 talk 的思路和逻辑，
   它的内容可能会与实际的 talk 有出入，建议有条件的人直接去听 talk 本身 *)

(* 本次 talk 将会使用 OCaml 来实现一个 MLTT 类型检查器。
   你可能不会写 OCaml，但这没有关系。本次 talk 只会使用以下的功能：
   - ADT（也就是 Rust 式 enum 或者说 tagged union）
   - 模式匹配（最基本的）
   我会尽可能在每个语言构造第一次出现时进行解释，
   这几项功能在 OCaml 中的语法也很直观。
   如果你有任何 OCaml 语法上的疑问，欢迎直接打断我提出。*)


(* [StringMap] 一个以 string 为 key 的映射。为什么这么写是无关的细节 *)
module StringMap = Map.Make(String)

(* 当我们在实现中遇到困难时，将会使用这句咒语来摆烂！*)
let wo_hai_bu_hui_xie () =
  failwith "你们懂吗"

(* 另一句咒语 *)
let wo_gai_tian_shen_me () =
  failwith "救命"

(* 本次 talk 实现的 MLTT 的语法树 *)
module Syntax = struct
  type expr =
    (* 变量 x, y, z *)
    | Var   of string
    (* “类型的类型”。
       由于直接令 [Type : Type] 会使语言不一致，
       我们需要一个 universe hierarchy：
       每个类型有一个 “宇宙层级”（这里的 [int]），
       小的宇宙活在大的宇宙中。*)
    | Type  of int
    (* 函数的类型，[FunType(x, A, B)] = "∏(x : A). B" *)
    | FunType of string * expr * expr
    (* 函数，[Fun(x, A, t)] = "λ(x : A). t" *)
    | Fun   of string * expr * expr
    (* 调用函数，[App f a] = "f a" *)
    | App   of expr * expr

  (* 在 dependent type 中，
     类型就是普通的表达式，表达式也能出现在类型中。
     所以，没有必要也很难在语法上区分类型和表达式，
     两者合而为一 *)
  type typ = expr

  (* 虽然这门语言极其简单，
     但我们已经可以利用它构造出大量有趣的例子了。
     诀窍在于，我们可以用 free variable 来表示常量。
     例如，在如下的 context 中：

       nat : Type 0, zero : nat, succ : nat -> nat

     我们可以构造出自然数。
     这种编码方式的缺点是无法编码计算规则，
     但本次 talk 要在短时间内从零写出一个完整的 dependent type 类型检查器，
     因此不会向语言添加更复杂的构造。
     如自然数之类的内建类型对类型检查实现的复杂度影响很小，
     因此后面将在举例子时假设某些常见的内建类型的存在 *)
end


module Try_To_Implement_Typecheck = struct
  open Syntax

  (* 由于语法中有充足的类型标注，可以直接推导出表达式的类型 *)
  let rec infer (env : typ StringMap.t) (expr : expr) =
    (* 这个辅助函数检查一个表达式是否是一个合法的类型。
       如果是，返回它的宇宙层级 *)
    let check_expr_is_type env expr =
      match infer env expr with
      | Type universe ->
          universe
      | _ ->
          (* Dependent type 中类型可以依赖于值。
             有宇宙的情况下，类型自身也可能是计算的结果。
             如果 [expr] 的类型不直接是 [Type u]，
             而是要经过一些计算才能变成 [Type u] 怎么办？*)
          failwith "not a type"
    in
    match expr with
    | Var name ->
        (* 变量的类型直接在上下文中查询 *)
        StringMap.find name env
    | Type universe ->
        (* 类型生活在更高一层的类型中 *)
        Type (universe + 1)
    | FunType(param_name, param_type, return_type) ->
        let universe_of_param_type  = check_expr_is_type env param_type in
        let universe_of_return_type = check_expr_is_type (StringMap.add param_name param_type env) return_type in

        Type (max universe_of_param_type universe_of_return_type)
    | Fun(param_name, param_type, body) ->
        (*    Γ, x : A |- t : B
           --------------------------------
          Γ |- λ(x : A). t : ∏(x : A). B
        *)
        (* [param_type] 应当是一个合法的类型 *)
        let _ = check_expr_is_type env param_type in
        let return_type = infer (StringMap.add param_name param_type env) body in
        FunType(param_name, param_type, return_type)
    | App(func, arg) ->
        (* Γ |- f : ∏(x : A). B     Γ |- a : A
           ------------------------------------
              Γ |- f a : B[x ↦  a]
        *)
        let func_type = infer env func in
        let arg_type  = infer env arg  in
        match func_type with
        | FunType(param_name, param_type, return_type) ->
            (* 如果 [param_type] 和 [arg_type] 不是直接相等，
               而是经过计算才相等，例如 [1 + 1] v.s. [2] 怎么办？*)
            if param_type <> arg_type then 
              failwith "type mismatch";
            (* [result_type] 应当通过
               把 [return_type] 中的 [param_name] 替换为 [arg] 得到。
               为了实现类型检查，我们需要 substitution 这一求值常用的操作，
               这是 dependent type 的必然代价！*)
            let result_type = wo_hai_bu_hui_xie () in
            result_type
      | _ ->
          (* 如果 [func_type] 不*直接*是 ∏(x : A). B，
             而是需要计算才能变成一个 "∏" 怎么办？*)
          failwith "not a function"

  (* 总结：由于 dependent type 的特性，实现普通类型检查的思路不再适用。
     我们需要一些其他操作，具体来说：*)

  (* 判断两个类型（根据 definitinal equality/经过计算）是否相等 *)
  let types_are_equal (typ1 : typ) (typ2 : typ) : bool =
    wo_hai_bu_hui_xie ()

  (* 在表达式（类型）中替换变量 *)
  let subst (var : string) (map_to : expr) (target : expr) : expr =
    wo_hai_bu_hui_xie ()
end


(* 让我们首先尝试直接通过 substitution 来实现 MLTT 的类型检查 *)
module Subst = struct
  open Syntax

  let rec subst var map_to expr =
    match expr with
    | Var name ->
        if var = name
        then map_to
        else Var name
    | Type universe ->
        Type universe
    | FunType(param_name, param_type, return_type) ->
        (* 如果 [param_name] 在 [map_to] 里出现了怎么办？
           我们需要把 param_name 重命名成新的名字，非常麻烦！

           这里为了节省时间，摆烂不处理了。
           这样实现出来的东西是有 bug 的，实践中请勿模仿，
           好好处理变量捕获问题！*)
        let param_type' = subst var map_to param_type in
        if param_name = var
        then FunType(param_name, param_type', return_type)
        else FunType(param_name, param_type', subst var map_to return_type)
    | Fun(param_name, param_type, body) ->
        (* 函数同样有意外捕获的问题！ *)
        let param_type' = subst var map_to param_type in
        if param_name = var
        then Fun(param_name, param_type', body)
        else Fun(param_name, param_type', subst var map_to body)
    | App(func, arg) ->
        App(subst var map_to func, subst var map_to arg)


  (* 我们要如何判断两个类型（表达式）是否（根据 definitional equality）相等呢？
     我们需要一个 "normalization" 函数。
     它会把一个表达式化简成某种最简的形式。
     如果两个表达式根据 definitional equality 相等，它们的最简形式就应该直接相等 *)
  let types_are_equal (typ1 : expr) (typ2 : expr) =
    let normalize expr = wo_hai_bu_hui_xie () in
    (* 我们还需要考虑变量名字的问题。
       例如 "λ(x:A). x" 和 "λ(y:A). y" 就应当被看作同一类型。
       允许替换变量名字意义下的相等叫作 ɑ-equivalence *)
    let is_alpha_equivalent expr1 expr2 = wo_hai_bu_hui_xie () in
    is_alpha_equivalent (normalize typ1) (normalize typ2)

  (* 首先是如何实现 normalization 来化简表达式。
     我们直接找出所有的可以进行 "计算" 的子表达式，
     如 "(λ(x : A). t) a"，并不断将它们化简，直到没有能化简的子表达式即可

     这里的实现是比较低效的，这是为了节省时间。
     事实上基于 substitution 的实现再怎么优化也快不到哪去，所以我就摆烂了

     另外，normalization 的实现很简单，但证明这个实现的正确性和停机性则困难许多。
     是研究一个 dependent type 系统的性质时最困难的问题之一 *)
  let rec normalize expr =
    match expr with
    | Var name ->
        Var name
    | Type universe ->
        Type universe
    | FunType(param_name, param_type, return_type) ->
        FunType(param_name, normalize param_type, normalize return_type)
    | Fun(param_name, param_type, body) ->
        Fun(param_name, normalize param_type, normalize body)
    | App(func, arg) ->
        match normalize func with
        | Fun(param_name, _, body) ->
            (* 可以进行计算，把参数代入函数体 *)
            normalize (subst param_name arg body)
        | func' ->
            App(func', normalize arg)


  (* 接下来我们需要实现 alpha equivalence 的检查 *)
  let is_alpha_equivalent expr1 expr2 =
    (* 我们维护一个映射 [var_subst]，
       记录 [expr1] 中的哪些变量被映射到了 [expr2] 中的哪些变量 *)
    let rec loop (var_subst : string StringMap.t) e1 e2 =
      match e1, e2 with
      | Var name1, Var name2 ->
          (match StringMap.find_opt name1 var_subst with
          | Some name1' -> name1' = name2
          (* 我们在 [var_subst] 中只记录 [expr1] 和 [expr2] 中被 "λ" 或 "∏" 中定义的变量，
             但 [expr1] 和 [expr2] 中还可能有 free variable *)
          | None -> name1 = name2)
      | Type universe1, Type universe2 ->
          universe1 = universe2
      | FunType(param_name1, param_type1, return_type1),
          FunType(param_name2, param_type2, return_type2) ->
          loop var_subst param_type1 param_type2
          (* 在判断两个 [return_type] 是否 ɑ-equivalent 时，
             把 [param_name1] 映射到 [param_name2] *)
          && loop (StringMap.add param_name1 param_name2 var_subst) return_type1 return_type2
      | Fun(param_name1, param_type1, body1),
          Fun(param_name2, param_type2, body2) ->
          loop var_subst param_type1 param_type2
          (* 同理 *)
          && loop (StringMap.add param_name1 param_name2 var_subst) body1 body2
      | App(func1, arg1), App(func2, arg2) ->
          (* [is_alpha_equivalent] 不考虑计算，它假设输入已经被 [normalize] 处理过了。
             所以遇到的 [App] 一定是无法继续化简的，直接比较子表达式即可 *)
          loop var_subst func1 func2
          && loop var_subst arg1 arg2
      | _ ->
          false
    in
    loop StringMap.empty expr1 expr2

  (* 现在我们可以（不借助摆烂咒语地）实现 [types_are_equal] 了 *)
  let types_are_equal typ1 typ2 =
    is_alpha_equivalent (normalize typ1) (normalize typ2)


  (* 然后完整的类型检查也能实现出来了。
     我用 "!!!!" 标记出了和最开始的尝试不一样的地方 *)
  let rec infer (env : typ StringMap.t) (expr : expr) =
    let check_expr_is_type env expr =
      match (* !!!! *) normalize (infer env expr) with
      | Type universe ->
          universe
      | _ ->
          failwith "not a type"
    in
    match expr with
    | Var name ->
        StringMap.find name env
    | Type universe ->
        Type (universe + 1)
    | FunType(param_name, param_type, return_type) ->
        let universe_of_param_type  = check_expr_is_type env param_type in
        let universe_of_return_type = check_expr_is_type (StringMap.add param_name param_type env) return_type in

        Type (max universe_of_param_type universe_of_return_type)
    | Fun(param_name, param_type, body) ->
        let return_type = infer (StringMap.add param_name param_type env) body in
        FunType(param_name, param_type, return_type)
    | App(func, arg) ->
        let func_type = infer env func in
        let arg_type  = infer env arg  in
        match (* !!!! *) normalize func_type with
        | FunType(param_name, param_type, return_type) ->
            if (* !!!! *) not (types_are_equal param_type arg_type) then 
              failwith "type mismatch";
            let result_type = (* !!!! *) subst param_name arg return_type in
            result_type
        | _ ->
            failwith "not a function"
end


(* 好！我们已经实现出了一个完整的 MLTT type checker！
   本次 talk 圆满结束！

   ......才怪了！

   这个基于 substitution 的实现：
   - 有 BUG（没处理意外捕获变量的问题）
   - 非常慢！

   那么，有没有更好的实现方式呢？
   答案是：有！
   这种实现方式叫作 Normalization By Evaluation (NBE)。
   下面就让我们来一探究竟。*)

module NBE = struct
  open Syntax

  (* NBE 的大体直觉是，既然 normalization 需要进行计算，
     为什么我们不直接用一个求值器来实现 normalization 呢？
     为了做到这一点，我们需要解决两个问题：

     - normalization 中可能出现 free variable，求值器应当如何处理它们？
     - 如何把求值的结果转回一个表达式，方便比较 *)

  (* 第一个问题，要通过求值的“值”的设计来解决。
     free variable 会使所有计算 “卡住”，所以我们把它们也当成一种特殊的 “值”。

     free variable 会卡住更大的表达式的计算，如 "x a" 这个函数调用。
     些表达式同样会卡住计算。所以我们也把它们当作值。

     这种 "被 free variable 卡住" 的值，称为 neutral value（中性值）。
     在我们的语言中，neutral value（和其他普通的值）的定义如下：*)
  type neutral =
    | Ne_Var of string
    (* 无论参数是什么，只要函数被卡住了，整个函数调用就会被卡住 *)
    | Ne_App of neutral * value

  and value =
    | V_Ne of neutral
    | V_Type of int
    (* 我们使用 Higher Order Abstract Syntax (HOAS) 的方式，
       直接用宿主语言 OCaml 的函数来表示 MLTT 中的一个函数值。

       除了 HOAS，也可以直接用一个 closure

         value StringMap.t * expr

       （函数定义处的环境，加上函数体）来表示函数值。

       由于把值转换为函数时需要一个参数类型的标注，这里把参数类型也一并存储。
       在一种叫作 "type-directed NBE" 的 NBE 变种中，可以不存储参数类型。
       这里由于时间原因不做介绍 *)
    | V_Fun of value * (value -> value)
    (* 在函数类型 "∏(x : A). B" 中，"B" 依赖于变量 "x"。
       而在类型检查时，我们也需要把 "B" 中的 "x" 替换为别的东西，“调用” "B"。
       所以，这里把 "B" 用 HOAS 也表示为一个 OCaml 函数 *)
    | V_FunType of value * (value -> value)


  (* 接下来，我们可以定义一个 "求值器" 了。
     [env] 把表达式中的 free variable 映射到它们的值 *)
  let rec eval (env : value StringMap.t) (expr : expr) : value =
    match expr with
    | Var name ->
        StringMap.find name env
    | Type universe ->
        V_Type universe
    | FunType(param_name, param_type, return_type) ->
        (* 描述传入参数的值之后，能得到怎样的返回类型的值 *)
        let value_of_return_type arg_value =
          eval (StringMap.add param_name arg_value env) return_type
        in
        V_FunType(eval env param_type, value_of_return_type)
    | Fun(param_name, param_type, body) ->
        let value_function arg_value =
          eval (StringMap.add param_name arg_value env) body
        in
        V_Fun(eval env param_type, value_function)
    | App(func, arg) ->
        match eval env func with
        | V_Fun(_, f) ->
            (* 如果 [func] 求值后是一个函数，直接调用 *)
            f (eval env arg)
        | V_Ne neutral ->
            (* 如果 [func] 的求值被卡住了，整个函数调用也会被卡住 *)
            V_Ne (Ne_App(neutral, eval env arg))
        | _ ->
            failwith "runtime error"

  (* 下一步是把值转换回表达式，让我们首先进行一次尝试 *)
  let rec value_to_expr value : expr =
    match value with
    | V_Ne neutral ->
        neutral_to_expr neutral
    | V_Type universe ->
        Type universe
    | V_FunType(param_type_value, return_type_value) ->
        (* 我们需要提供一个函数参数的名字，我们应该选什么名字？*)
        let param_name = wo_gai_tian_shen_me () in
        FunType(
          param_name,
          value_to_expr param_type_value,
          (* 我们需要给 [return_type_value] 提供一个参数，
             这个参数应该是函数的参数，也就是 [param_name]。
             幸运的是，我们可以直接把自由变量当成值 *)
          value_to_expr (return_type_value (V_Ne (Ne_Var param_name))))
    | V_Fun(param_type_value, function_value) ->
        (* 同样的问题 *)
        let param_name = wo_gai_tian_shen_me () in
        Fun(
          param_name,
          value_to_expr param_type_value,
          (* 同理 *)
          value_to_expr (function_value (V_Ne (Ne_Var param_name))))

  and neutral_to_expr neutral =
    match neutral with
    | Ne_Var name ->
        (* 自由变量不需要特别处理 *)
        Var name
    | Ne_App(func, arg) ->
        App(neutral_to_expr func, value_to_expr arg)

  (* 可以看到，我们缺少的是一种生成新的变量名字的机制。
     但这并不是一个简单的问题：在基于 substitution 的实现中，
     就遇到了很麻烦的变量重名和捕获的问题！

     所以，我们要想办法保证新生成的名字和 “所有其他名字” 都不冲突。
     那么，“所有其他名字” 有哪些呢？
     我们只需要保证同一个表达式中的变量互不重名即可。
     因此，在将值转换回表达式时，我们维护 “当前有哪些变量” 的列表，
     并保证生成的新变量和列表中的变量不重名即可 *)

  (* [generate_name other_names] 生成一个和 [other_names] 中的名字不同的新名字 *)
  let generate_name other_names =
    (* failwith "我懒得写" *)
    (* 好吧还是要让写的代码跑得起来才健全 *)
    let rec try_xi i =
      let name = "x" ^ string_of_int i in
      if List.mem name other_names
      then try_xi (i + 1)
      else name
    in
    try_xi 0

  (* 可以完整实现了值到表达式的转换了。不同的地方用 "!!!!" 标出 *)
  let rec value_to_expr (* !!!! *) (names : string list) value =
    match value with
    | V_Ne neutral ->
        neutral_to_expr names neutral
    | V_Type universe ->
        Type universe
    | V_FunType(param_type_value, return_type_value) ->
        let param_name = (* !!!! *) generate_name names in
        FunType(
          param_name,
          value_to_expr names param_type_value,
          (* 把 [param_name] 加到已使用的变量列表里 *)
          value_to_expr (* !!!! *) (param_name :: names)
            (return_type_value (V_Ne (Ne_Var param_name))))
    | V_Fun(param_type_value, function_value) ->
        let param_name = generate_name names in
        Fun(
          param_name,
          value_to_expr names param_type_value,
          value_to_expr (* !!!! *) (param_name :: names)
            (function_value (V_Ne (Ne_Var param_name))))

  and neutral_to_expr names neutral =
    match neutral with
    | Ne_Var name ->
        Var name
    | Ne_App(func, arg) ->
        App(neutral_to_expr names func, value_to_expr names arg)


  (* 现在可以利用上面的函数实现出 normalization 了。
     这里为了让所有东西都更好理解，实现的方式可能有重复计算的地方。
     知道原理之后实际实现的时候可以做很多小优化来解决问题 *)
  let normalize (env : value StringMap.t) expr =
    let value = eval env expr in
    (* 获得 [env] 里的所有名字。具体的实现用的 API 无需在意 *)
    let names_in_env = List.map fst (StringMap.bindings env) in
    value_to_expr names_in_env value

  (* 虽然我们已经成功使用 NBE 实现了 normalization，
     但利用 NBE 实现类型检查时其实并不需要 normalization 函数。
     让我们考虑这么一个问题：
     在类型检查时，推导出的 “类型” 应该用什么形式表达？
     用表达式表达固然可行，但我们才使用类型时，
     往往想要的是它计算后的结果：

     - 检查函数调用时，我们关心函数的类型*经过计算*是否是一个 "∏"
     - 检查函数参数的类型时，我们关心预期类型和实际类型*计算/normalization 后的结果*是否一致

     所以，如果用表达式表示类型，我们需要常常调用 normalization 来化简它们。
     但如果我们直接用*值*来表示类型，就不需要 normalization 了：值本身已经是最简的了。
     按照这一思路，我们需要实现一个 “比较值是否相等” 的函数。*)
  let types_are_equal names (value1 : value) (value2 : value) =
    (* 这里不需要使用 alpha-equivalence，因为生成新名字的方式是确定性的。
       因此直接比较表达式是否相等即可 *)
    value_to_expr names value1 = value_to_expr names value2

  (* 剩下的问题是 [subst] 如何实现。
     但很快我们会看到，如果用值来表示类型，不需要实现 [subst] 也能实现类型检查 *)

  (* 基于 NBE 的类型检查使用的上下文 *)
  type env =
    { typs   : value StringMap.t (* 变量的类型 *)
    ; values : value StringMap.t (* 变量的值，提供给 eval 用于求值表达式 *)
    ; names  : string list (* 上下文中的名字列表，用于 [value_to_expr] 等 *) }

  let add_var name typ env =
    (* 在类型检查时，我们不知道一个变量的值。
       因此，我们直接把这个变量自身作为值存入上下文 *)
    let value = V_Ne (Ne_Var name) in
    { typs   = StringMap.add name typ env.typs
    ; values = StringMap.add name value env.values
    ; names  = name :: env.names }

  (* 变化的地方较多，建议从头重新阅读 *)
  let rec infer env expr : value =
    let check_expr_is_type env expr =
      match infer env expr with
      | V_Type universe ->
          universe
      | _ ->
          failwith "not a type"
    in
    match expr with
    | Var name ->
        StringMap.find name env.typs
    | Type universe ->
        V_Type (universe + 1)
    | FunType(param_name, param_type, return_type) ->
        let universe_of_param_type  = check_expr_is_type env param_type in
        (* 由于用类型表示值，需要对 [param_type] 进行求值，才能把它放入 [env] 中 *)
        let param_type_value = eval env.values param_type in
        let universe_of_return_type =
          check_expr_is_type (add_var param_name param_type_value env) return_type
        in

        V_Type (max universe_of_param_type universe_of_return_type)
    | Fun(param_name, param_type, body) ->
        let _ = check_expr_is_type env param_type in
        let param_type_value = eval env.values param_type in
        let body_env = add_var param_name param_type_value env in
        (* [return_type_value] 是一个值，
           在推断出它时，[param_name] 对应的值是 [V_Ne (Ne_Var param_name)] *)
        let return_type_value = infer body_env body in
        (* 但在构建 [V_FunType] 时，返回值的类型应该是一个函数。
           所以，我们需要先把 [return_type_value] 转回表达式，
           再通过 [eval] 得到一个函数 *)
        let return_type = value_to_expr body_env.names return_type_value in
        V_FunType( param_type_value
                 , fun param_value -> eval (StringMap.add param_name param_value env.values) return_type)
    | App(func, arg) ->
        let func_type = infer env func in
        let arg_type  = infer env arg  in
        match func_type with
        | V_FunType(param_type, return_type) ->
          (* 由于用值表示类型，不需要进行 normalization，得到的类型已经经过了化简 *)
            if not (types_are_equal env.names param_type arg_type) then 
              failwith "type mismatch";
            (* 为了得到具体的返回类型，把实际的参数求值后喂给 [return_type] *)
            return_type (eval env.values arg)
        | _ ->
            failwith "not a function"
end

(* 一些广告/延伸阅读内容：

   - 实现 dependent type type-checker 的各种例子和技巧：https://github.com/AndrasKovacs/elaboration-zoo

   - 如果除了简单的类型检查，还想实现更多实用功能，例如类型推导、隐式参数等，应该如何实现？
     这些功能底层往往通过一个叫作 meta variable 的机制实现。除了 elboration zoo 中相关的例子，
     也可以看我写的一份教程：https://github.com/Guest0x0/pruning-tutor

   - NBE 比 substitution 快多少？这是一个我做的 benchmark：https://github.com/Guest0x0/normalization-bench。
     里面有一个项目就是比较基于 substitution 和 NBE 的 normalization。NBE 的性能碾压 substitution。

   - NBE 算法有很多不同的写法，而且往往都有非常有意思的结构。
     对这方面的内容感兴趣的话可以看这篇教程：https://guest0x0.xyz/NBE-intro.pdf

   - 对 NBE 的形式化和证明感兴趣的话，可以看 Abel 的一篇经典综述：https://www.cse.chalmers.se/~abela/habil.pdf
*)



(* 用于测试的内建定义 *)
open Syntax
let (@=>) (name, param_type) return_type = FunType(name, param_type, return_type)
let fn (param, param_type) body = Fun(param, param_type, body)
let app f args = List.fold_left (fun f a -> App(f, a)) (Var f) args


let predef =
  [ "Nat", Type 0
  ; "zero", Var "Nat"
  ; "succ", ("n", Var "Nat") @=> Var "Nat"

  ; "Bool", Type 0
  ; "tt", Var "Bool"
  ; "ff", Var "Bool"
  ; "if",
    ("A", Type 0) @=>
      ("cond", Var "Bool") @=>
      ("conseq", Var "A") @=>
      ("alter", Var "A") @=> Var "A"

  ; "List", ("A", Type 0) @=> ("n", Var "Nat") @=> Type 0
  ; "nil", ("A", Type 0) @=> app "List" [Var "A"; Var "zero"]
  ; "cons",
    ("A", Type 0) @=> ("n", Var "Nat") @=>
      ("hd", Var "A") @=>
      ("tl", app "List" [Var "A"; Var "n"]) @=>
      app "List" [Var "A"; app "succ" [Var "n"]]

  ; "eq", ("A", Type 0) @=> ("lhs", Var "A") @=> ("rhs", Var "A") @=> Type 0
  ; "refl", ("A", Type 0) @=> ("a", Var "A") @=> app "eq" [Var "A"; Var "a"; Var "a"] ]


(* 进行测试 *)
let infer_subst =
  let env = List.fold_left
      (fun env (name, typ) -> StringMap.add name typ env)
      StringMap.empty predef
  in
  fun expr -> Subst.normalize (Subst.infer env expr)

let infer_nbe =
  let open NBE in
  let env = List.fold_left
      (fun env (name, typ) -> add_var name (eval env.values typ) env)
      { typs = StringMap.empty; values = StringMap.empty; names = [] }
      predef
  in
  fun expr -> value_to_expr env.names (infer env expr)


type test_result =
  { typ_subst : Syntax.expr
  ; typ_nbe   : Syntax.expr
  ; two_impl_agree : bool }

let test expr =
  let typ_subst = infer_subst expr in
  let typ_nbe   = infer_nbe   expr in
  assert (Subst.is_alpha_equivalent typ_subst typ_nbe);
  (typ_subst, typ_nbe)



(* boolean and *)
let ex1 =
  fn ("x", Var "Bool") @@
  fn ("y", Var "Bool") @@
  app "if" [
    Var "Bool";
    Var "x";
    Var "y";
    Var "ff"
  ]

(* function application *)
let ex2 =
  fn ("A", Type 0) @@
  fn ("B", ("x", Var "A") @=> Type 0) @@
  fn ("f", ("x", Var "A") @=> app "B" [Var "x"]) @@
  fn ("a", Var "A") @@
  app "f" [Var "a"]


let appE f args = List.fold_left (fun f a -> App(f, a)) f args
let check_expr_has_type typ expr =
  let f = fn ("x", typ) @@ Var "x" in
  App(f, expr)

(* 测试表达式相等的判定 *)
let ex3 =
  let id = fn ("x", Var "Nat") @@ Var "x" in
  fn ("n", Var "Nat") @@
    check_expr_has_type
      (app "eq" [Var "Nat"; appE id [Var "n"]; Var "n"])
      (app "refl" [Var "Nat"; Var "n"])

(* 测试函数之间的相等判定 *)
let ex4 =
  fn ("A", Type 0) (
    let typ = ("a", Var "A") @=> Var "A" in
    let id = fn ("x", Var "A") @@ Var "x" in
    check_expr_has_type
      (* eq (A -> A) id (apply A (fun _ -> A) id) *)
      (app "eq" [
            typ; id;
            appE ex2 [
              Var "A";
              fn ("a", Var "A") @@ Var "A";
              id
            ]
          ])
      (* refl (A -> A) id *)
      (app "refl" [typ; id])
  )