pi_sigma_self.ml

module Syntax = struct
  type var = string

  type term =
    | T_var of var
    | T_let of var * term * term
    | T_annot of term * term
    | T_with_self of term * term
    (* pi *)
    | T_forall of var * term * term
    | T_apply of term * term
    | T_lambda of var * term
    (* sigma *)
    | T_exists of var * term * term
    | T_fst of term
    | T_snd of term
    | T_pair of term * term
    (* self *)
    | T_self of var * term
    | T_unfold of term
    | T_fix of var * term
end

module Core = struct
  module Var = struct
    type t = { id : int; name : string option }

    (* TODO: global state *)
    let fresh =
      let counter = Atomic.make 0 in
      fun name ->
        let id = Atomic.fetch_and_add counter 1 in
        { id; name }

    let fresh_with_name name = fresh (Some name)
    let fresh () = fresh None

    let pp fmt var =
      let { id; name } = var in
      let name = Option.value ~default:"gen" name in
      Format.fprintf fmt "%s_%d" name id

    let equal a b = Int.equal a.id b.id

    let is_name var name =
      match var.name with
      | Some var_name -> String.equal var_name name
      | None -> false
  end

  type term =
    | T_univ
    | T_var of Var.t
    | T_let of Var.t * term * term
    | T_forall of Var.t * term * term
    | T_apply of term * term
    | T_lambda of Var.t * term
    | T_exists of Var.t * term * term
    | T_fst of term
    | T_snd of term
    | T_pair of term * term
    | T_self of Var.t * term
    | T_unfold of term
    | T_fix of Var.t * term

  type value =
    | V_univ
    | V_var of Var.t
    | V_forall of value * closure
    | V_apply of value * value
    | V_lambda of closure
    | V_exists of value * closure
    | V_fst of value
    | V_snd of value
    | V_pair of value * value
    | V_self of closure
    | V_unfold of value
    | V_fix of closure

  and env =
    (* _ *)
    | E_hole
    (* L[x = V; _] *)
    | E_let of Var.t * value * env

  and closure =
    (* L[x = _; K] *)
    | Closure of env * Var.t * term

  let rec lookup env var =
    match env with
    | E_hole -> failwith (Format.asprintf "lookup: unknown var %a" Var.pp var)
    | E_let (env_var, arg, env) -> (
        match Var.equal env_var var with true -> arg | false -> lookup env var)

  let rec eval env term =
    match term with
    | T_univ -> V_univ
    | T_var var -> lookup env var
    | T_let (var, arg, body) ->
        let arg = eval env arg in
        let env = E_let (var, arg, env) in
        eval env body
    | T_forall (var, param, body) ->
        let param = eval env param in
        let body = Closure (env, var, body) in
        V_forall (param, body)
    | T_apply (funct, arg) ->
        let funct = eval env funct in
        let arg = eval env arg in
        eval_apply funct ~arg
    | T_lambda (var, body) ->
        let body = Closure (env, var, body) in
        V_lambda body
    | T_exists (var, fst, snd) ->
        let fst = eval env fst in
        let snd = Closure (env, var, snd) in
        V_exists (fst, snd)
    | T_fst pair ->
        let pair = eval env pair in
        eval_fst pair
    | T_snd pair ->
        let pair = eval env pair in
        eval_snd pair
    | T_pair (fst, snd) ->
        let fst = eval env fst in
        let snd = eval env snd in
        V_pair (fst, snd)
    | T_self (var, body) ->
        let body = Closure (env, var, body) in
        V_self body
    | T_unfold fix ->
        let fix = eval env fix in
        eval_unfold fix
    | T_fix (var, body) ->
        let body = Closure (env, var, body) in
        V_fix body

  and eval_apply funct ~arg =
    match funct with
    | V_lambda body -> eval_with_arg body arg
    | funct -> V_apply (funct, arg)

  and eval_fst pair =
    match pair with V_pair (fst, _snd) -> fst | pair -> V_fst pair

  and eval_snd pair =
    match pair with V_pair (_fst, snd) -> snd | pair -> V_snd pair

  and eval_unfold fix =
    match fix with V_fix body -> eval_with_arg body fix | fix -> V_unfold fix

  and eval_with_arg body arg =
    let (Closure (env, var, term)) = body in
    let env = E_let (var, arg, env) in
    eval env term

  let rec equal received expected =
    match (received, expected) with
    | V_univ, V_univ -> ()
    | V_var received_var, V_var expected_var -> (
        match Var.equal received_var expected_var with
        | true -> ()
        | false ->
            failwith
              (Format.asprintf "equal: var clash, expected %a, received %a"
                 Var.pp expected_var Var.pp received_var))
    | ( V_forall (received_param, received_body),
        V_forall (expected_param, expected_body) ) ->
        equal received_param expected_param;
        equal_closure received_body expected_body
    | ( V_apply (received_funct, received_arg),
        V_apply (expected_funct, expected_arg) ) ->
        equal received_funct expected_funct;
        equal received_arg expected_arg
    | V_lambda received_body, V_lambda expected_body ->
        equal_closure received_body expected_body
    | ( V_exists (received_fst, received_snd),
        V_exists (expected_fst, expected_snd) ) ->
        equal received_fst expected_fst;
        equal_closure received_snd expected_snd
    | V_fst received_pair, V_fst expected_pair ->
        equal received_pair expected_pair
    | V_snd received_pair, V_snd expected_pair ->
        equal received_pair expected_pair
    | V_pair (received_fst, received_snd), V_pair (expected_fst, expected_snd)
      ->
        equal received_fst expected_fst;
        equal received_snd expected_snd
    | V_self received_body, V_self expected_body ->
        equal_closure received_body expected_body
    | V_unfold received_fix, V_unfold expected_fix ->
        equal received_fix expected_fix
    | V_fix received_body, V_fix expected_body ->
        equal_closure received_body expected_body
    | _received, _expected -> failwith "equal: type clash"

  and equal_closure received expected =
    let arg = V_var (Var.fresh ()) in
    let received = eval_with_arg received arg in
    let expected = eval_with_arg expected arg in
    equal received expected

  let split_univ univ =
    match univ with
    | V_univ -> ()
    | _value -> failwith "split_univ: expected type"

  let split_forall forall =
    match forall with
    | V_forall (param, body) -> (param, body)
    | _value -> failwith "split_forall: expected forall"

  let split_exists exists =
    match exists with
    | V_exists (fst, snd) -> (fst, snd)
    | _value -> failwith "split_exists: expected exists"

  let split_self self =
    match self with
    | V_self body -> body
    | _value -> failwith "split_self: expected self"
end

module Typer = struct
  open Core

  type context =
    (* _ *)
    | C_hole
    (* C[x : A; _]*)
    | C_def of context * Var.t * value
    (* C[x = (M : A); _] *)
    | C_let of context * Var.t * value * value

  let rec untype_context ctx =
    match ctx with
    | C_hole -> E_hole
    | C_def (ctx, var, _type_) ->
        let env = untype_context ctx in
        E_let (var, V_var var, env)
    | C_let (ctx, var, arg, _type) ->
        let env = untype_context ctx in
        E_let (var, arg, env)

  let rec lookup ctx name =
    match ctx with
    | C_hole -> failwith (Format.asprintf "lookup: unknown var %S" name)
    | C_def (ctx, var, type_) -> (
        match Var.is_name var name with
        | true -> (var, type_)
        | false -> lookup ctx name)
    | C_let (ctx, var, _arg, type_) -> (
        match Var.is_name var name with
        | true -> (var, type_)
        | false -> lookup ctx name)

  let eval ctx term =
    let env = untype_context ctx in
    eval env term

  (* TODO: make arg, snd and fix lazy *)
  let rec infer ctx term =
    match (term : Syntax.term) with
    | T_var name ->
        let var, type_ = lookup ctx name in
        (T_var var, type_)
    (* TODO: let is duplicated on infer and check *)
    | T_let (var, arg, body) ->
        let var = Var.fresh_with_name var in
        let arg, arg_type = infer ctx arg in
        let ctx =
          let arg = eval ctx arg in
          C_let (ctx, var, arg, arg_type)
        in
        let body, type_ = infer ctx body in
        (T_let (var, arg, body), type_)
    | T_annot (term, annot) ->
        let annot = check ctx annot V_univ in
        let annot = eval ctx annot in
        let term = check ctx term annot in
        (term, annot)
    | T_with_self (term, self) ->
        let self, type_ = infer ctx self in
        let self = eval ctx self in
        let term = check_with_self ctx term type_ (Some self) in
        let () =
          let term = eval ctx term in
          equal term self
        in
        (term, type_)
    | T_forall (var, param, body) ->
        let var = Var.fresh_with_name var in
        let param = check ctx param V_univ in
        let body =
          let param = eval ctx param in
          let ctx = C_def (ctx, var, param) in
          check ctx body V_univ
        in
        (T_forall (var, param, body), V_univ)
    | T_apply (funct, arg) ->
        let funct, forall = infer ctx funct in
        let param_type, body_type = split_forall forall in
        let arg = check ctx arg param_type in
        let type_ =
          let arg = eval ctx arg in
          eval_with_arg body_type arg
        in
        (T_apply (funct, arg), type_)
    | T_exists (var, fst, snd) ->
        let var = Var.fresh_with_name var in
        let fst = check ctx fst V_univ in
        let snd =
          let fst = eval ctx fst in
          let ctx = C_def (ctx, var, fst) in
          check ctx snd V_univ
        in
        (T_exists (var, fst, snd), V_univ)
    | T_fst pair ->
        let pair, exists = infer ctx pair in
        let fst_type, _snd_type = split_exists exists in
        (T_fst pair, fst_type)
    | T_snd pair ->
        let pair, exists = infer ctx pair in
        let _fst_type, snd_type = split_exists exists in
        let snd_type =
          let pair = eval ctx pair in
          let fst = eval_fst pair in
          eval_with_arg snd_type fst
        in
        (T_snd pair, snd_type)
    | T_unfold fix ->
        let fix, self = infer ctx fix in
        let body_type = split_self self in
        let type_ =
          let fix = eval ctx fix in
          eval_with_arg body_type fix
        in
        (T_unfold fix, type_)
    | T_lambda _ | T_pair _ | T_self _ | T_fix _ ->
        failwith "infer: cannot infer, need annotation"

  and check ctx term expected =
    (* an alternative to None here would be to introduce a skolem *)
    check_with_self ctx term expected None

  and check_with_self ctx term expected self =
    match (term : Syntax.term) with
    (* TODO: this is duplicated *)
    | T_let (var, arg, body) ->
        let var = Var.fresh_with_name var in
        let arg, arg_type = infer ctx arg in
        let ctx =
          let arg = eval ctx arg in
          C_let (ctx, var, arg, arg_type)
        in
        let body = check_with_self ctx body expected self in
        T_let (var, arg, body)
    | T_lambda (var, body) ->
        let param_type, body_type = split_forall expected in
        let var = Var.fresh_with_name var in
        let body =
          let arg = V_var var in
          let ctx = C_def (ctx, var, param_type) in
          let expected = eval_with_arg body_type arg in
          let self = Option.map (fun self -> eval_apply self ~arg) self in
          check_with_self ctx body expected self
        in
        T_lambda (var, body)
    | T_pair (fst, snd) ->
        let fst_type, snd_type = split_exists expected in
        let fst =
          let self = Option.map eval_fst self in
          check_with_self ctx fst fst_type self
        in
        let snd =
          let fst = eval ctx fst in
          let snd_type = eval_with_arg snd_type fst in
          let self = Option.map eval_snd self in
          check_with_self ctx snd snd_type self
        in
        T_pair (fst, snd)
    | T_self (var, body) ->
        let () = split_univ expected in
        let self =
          match self with
          | Some self -> self
          | None -> failwith "check_with_self: missing self"
        in
        let var = Var.fresh_with_name var in
        let body =
          let ctx = C_def (ctx, var, self) in
          (* intentionally drop self *)
          check ctx body V_univ
        in
        T_self (var, body)
    | T_fix (var, body) ->
        let body_type = split_self expected in
        let var = Var.fresh_with_name var in
        let body =
          let ctx, self =
            match self with
            | Some self ->
                let ctx = C_let (ctx, var, self, expected) in
                (ctx, self)
            | None ->
                let ctx = C_def (ctx, var, expected) in
                (ctx, V_var var)
          in
          let expected = eval_with_arg body_type self in
          let self = eval_unfold self in
          check_with_self ctx body expected (Some self)
        in
        T_fix (var, body)
    | T_var _ | T_annot _ | T_with_self _ | T_forall _ | T_apply _ | T_exists _
    | T_fst _ | T_snd _ | T_unfold _ ->
        let term, received = infer ctx term in
        let () = equal received expected in
        term
end