open Bwd
open Bwd.Infix

module D = Domain

exception Unequal

module Internal =
struct
  (** Context size *)
  type env = int

  module Eff = Algaeff.Reader.Make (struct type nonrec t = env end)

  let bind f =
    let arg = D.var (Eff.read()) in
    Eff.scope (fun size -> size + 1) (fun () -> f arg)

  let rec equate v1 v2 = match v1, v2 with
    | D.Neutral ne, v | v, D.Neutral ne -> equate_ne ne v
    | D.Pi (base1, fam1), D.Pi (base2, fam2) ->
      equate base1 base2;
      equate_clo fam1 fam2
    | D.Lam clo1, D.Lam clo2 -> equate_clo clo1 clo2
    | D.Sg (base1, fam1), D.Sg (base2, fam2) ->
      equate base1 base2;
      equate_clo fam1 fam2
    | D.Pair (fst1, snd1), D.Pair (fst2, snd2) ->
      equate fst1 fst2;
      equate snd1 snd2
    | Type, Type -> ()
    | Bool, Bool -> ()
    | True, True -> ()
    | False, False -> ()
    | _ -> raise Unequal

  and equate_clo clo1 clo2 = bind @@ fun arg ->
    equate (Eval.inst_clo clo1 arg) (Eval.inst_clo clo2 arg)

  and equate_ne_head (D.Var lvl1) (D.Var lvl2) =
    if lvl1 = lvl2 then () else raise Unequal

  and equate_frm frm1 frm2 = match frm1, frm2 with
    | D.App arg1, D.App arg2 -> equate arg1 arg2
    | D.Fst, D.Fst -> ()
    | D.Snd, D.Snd -> ()
    | D.BoolElim { cmot = mot1; vtrue = t1; vfalse = f1 },
      D.BoolElim { cmot = mot2; vtrue = t2; vfalse = f2 } ->
      equate_clo mot1 mot2;
      equate t1 t2;
      equate f1 f2;
    | _ -> raise Unequal

  and equate_spine sp1 sp2 = match sp1, sp2 with
    | Emp, Emp -> ()
    | Snoc (sp1, frm1), Snoc (sp2, frm2) ->
      equate_frm frm1 frm2;
      equate_spine sp1 sp2
    | _ -> raise Unequal

  and equate_ne (hd, sp) v = match v with
    | D.Neutral (hd2, sp2) ->
      equate_ne_head hd hd2;
      equate_spine sp sp2
    (* eta expansion *)
    | D.Lam clo -> bind @@ fun arg ->
      equate_ne (hd, sp <: D.App arg) (Eval.inst_clo clo arg)
    | D.Pair (fst, snd) ->
      equate_ne (hd, sp <: D.Fst) fst;
      equate_ne (hd, sp <: D.Snd) snd
    | _ -> raise Unequal
end

let equate ~size v1 v2 = Internal.Eff.run ~env:size @@ fun () ->
  Internal.equate v1 v2