Module UniSolver (T: TYPE_ALGEBRA).
(* The current set of constraints is represented by a record with two components: - [te_typ]: a partial map from variables to types - [te_equ]: a list of pairs [(x,y)] of variables, indicating that the type of [x] must be equal to the type of [y]. *)
Definition constraint : Type := (positive * positive)%type.
Record typenv : Type := Typenv {
te_typ: PTree.t T.t;       (**r mapping var -> typ *)
te_equ: list constraint    (**r additional equality constraints *)
}.
Definition initial : typenv := {| te_typ := PTree.empty _; te_equ := nil |}.
(** Add the constraint [T(x) = ty]. *)
Definition set (e: typenv) (x: positive) (ty: T.t) : res typenv :=
match e.(te_typ)!x with
| None =>
OK {| te_typ := PTree.set x ty e.(te_typ);
te_equ := e.(te_equ) |}
| Some ty' =>
if T.eq ty ty'
then OK e
else Error (MSG "bad definition/use of variable " :: POS x :: nil)
end.
Fixpoint set_list (e: typenv) (rl: list positive) (tyl: list T.t) {struct rl}: res typenv :=
match rl, tyl with
| nil, nil => OK e
| r1::rs, ty1::tys => do e1 <- set e r1 ty1; set_list e1 rs tys
| _, _ => Error (msg "arity mismatch")
end.
(** Add the constraint [T(x) = T(y)]. The boolean result is [true] if the types of [x] or [y] could be made more precise.  Otherwise, [te_typ] does not change and [false] is returned. *)
Definition move (e: typenv) (r1 r2: positive) : res (bool * typenv) :=
if peq r1 r2 then OK (false, e) else
match e.(te_typ)!r1, e.(te_typ)!r2 with
| None, None =>
OK (false, {| te_typ := e.(te_typ); te_equ := (r1, r2) :: e.(te_equ) |})
| Some ty1, None =>
OK (true, {| te_typ := PTree.set r2 ty1 e.(te_typ); te_equ := e.(te_equ) |})
| None, Some ty2 =>
OK (true, {| te_typ := PTree.set r1 ty2 e.(te_typ); te_equ := e.(te_equ) |})
| Some ty1, Some ty2 =>
if T.eq ty1 ty2
then OK (false, e)
else Error(MSG "ill-typed move from " :: POS r1 :: MSG " to " :: POS r2 :: nil)
end.
(** Solve the remaining subtyping constraints by iteration. *)
Fixpoint solve_rec (e: typenv) (changed: bool) (q: list constraint) : res (typenv * bool) :=
match q with
| nil =>
OK (e, changed)
| (r1, r2) :: q' =>
do (changed1, e1) <- move e r1 r2; solve_rec e1 (changed || changed1) q'
end.
(** Measuring the state *)
Definition weight_typenv (e: typenv) : nat := length e.(te_equ).
(** Iterative solving of the remaining constraints *)
Function solve_constraints (e: typenv) {measure weight_typenv e}: res typenv :=
match solve_rec {| te_typ := e.(te_typ); te_equ := nil |} false e.(te_equ) with
| OK(e', false) => OK e                   (**r no more changes, fixpoint reached *)
| OK(e', true)  => solve_constraints e'   (**r one more iteration *)
| Error msg => Error msg
end.
Definition typassign := positive -> T.t.
Definition makeassign (e: typenv) : typassign :=
fun x => match e.(te_typ)!x with Some ty => ty | None => T.default end.
Definition solve (e: typenv) : res typassign :=
do e' <- solve_constraints e; OK(makeassign e').
(** What it means to be a solution *)
Definition satisf (te: typassign) (e: typenv) : Prop :=
(forall x ty, e.(te_typ)!x = Some ty -> te x = ty)
/\ (forall x y, In (x, y) e.(te_equ) -> te x = te y).
(** Soundness proof *)
Hint Resolve set_incr: ty.
Hint Resolve set_list_incr: ty.
Hint Resolve move_incr: ty.

Lemma solve_rec_incr:
forall te q e changed e' changed',
solve_rec e changed q = OK(e', changed') -> satisf te e' -> satisf te e.