Library ATBR.DisjointSets

Require Import Common.
Require Import MyFSets.
Require Import BoolView.
Require Import Numbers.

Module Type DISJOINTSETS (N : NUM).
  Import N.

  Variable T: Type.
  Variable empty : T.

  Variable test_and_unify : T -> num -> num -> bool * T.
  Variable equiv : T -> num -> num -> bool * T.
  Variable union : T -> num -> num -> T.

  Variable class_of : T -> num -> NumSet.t.

  Parameter IsWF : T -> Prop.
  Class WF t: Prop := { wf: IsWF t }.

  Declare Instance empty_WF : WF empty.
  Declare Instance test_and_unify_WF `{WF} x y: WF (snd (test_and_unify t x y)).
  Declare Instance equiv_WF `{WF} x y : WF (snd (equiv t x y)).
  Declare Instance union_WF `{WF} x y : WF (union t x y).

  Parameter sameclass: T -> relation num.
  Hypothesis sameclass_equiv : forall `{WF} x y, (fst (equiv t x y) = true <-> sameclass t x y).
  Declare Instance sameclass_Equivalence `{WF}: Equivalence (sameclass t).

  Hypothesis sameclass_empty : forall x y, sameclass empty x y -> x = y.
  Hypothesis sameclass_union : forall `{WF} a b x y,
    sameclass (union t a b) x y <->
    (sameclass t x y \/ (sameclass t x a /\ sameclass t y b) \/ (sameclass t y a /\ sameclass t x b)).

  Hypothesis sameclass_class_of : forall `{WF} x y , sameclass t x y <-> NumSet.In x (class_of t y).

  Hypothesis test_and_unify_eq: forall `{WF} x y, test_and_unify t x y = (fst (equiv t x y), union t x y).

End DISJOINTSETS.

Module DSUtils (Import N : NUM) (Import M: DISJOINTSETS N).
  Module Import NU := NumUtils N.

  Definition le (s t : T): Prop := inclusion _ (sameclass s) (sameclass t).

  Instance le_PreOrder : PreOrder le.

  Definition eq s t := same_relation _ (sameclass s) (sameclass t).
  Instance eq_Equivalence : Equivalence eq.
  Instance sameclass_compat : Proper (eq ==> same_relation num) sameclass.

  Instance sameclass_compat' : Proper (eq ==> @Logic.eq num ==> @Logic.eq num ==> iff) sameclass.

  Module Notations.
    Notation "{{ t }}" := (sameclass t).
    Notation "x [=T=] y" := (eq x y) (at level 75).
    Notation "x <T= y" := (le x y) (at level 75).
  End Notations.

  Import Notations.

  Lemma le_union `{WF}: forall x y, t <T= union t x y.

  Lemma eq_union `{WF}: forall x y, {{t}} x y -> t [=T=] union t x y.

  Lemma le_incr `{WF} {t'} `{WF t'}: forall x y,
    t <T= t' ->
    union t x y <T= union t' x y.

  Lemma compat_union_eq `{WF} {t'} `{WF t'} :
    t [=T=] t' ->
    forall x y, union t x y [=T=] union t' x y.

  Lemma class_of_empty: forall x, NumSet.Equal (class_of empty x) (NumSet.singleton x).

  Lemma class_of_union: forall `{WF} x y,
    NumSet.Equal (class_of (union t x y) x) (NumSet.union (class_of t x) (class_of t y)).

  Lemma class_of_union_2: forall `{WF} x y z,
    ~ sameclass t x z ->
    ~ sameclass t y z ->
    NumSet.Equal (class_of (union t x y) z) (class_of t z).

  Lemma in_class_of_self: forall `{WF} x, NumSet.In x (class_of t x).

  Instance class_of_compat `{WF} : Proper (sameclass t ==> NumSet.Equal) (class_of t).

  Lemma class_of_compat' : forall `{WF} {t'} `{WF t'} x,
    eq t t' ->
    NumSet.Equal (class_of t x) (class_of t' x).

  Lemma class_of_union_1: forall `{WF} x y z,
    {{t}} x z \/ {{t}} y z ->
    NumSet.Equal (class_of (union t x y) z) (NumSet.union (class_of t x) (class_of t y)).

  Lemma class_of_injective: forall `{WF} x y,
    NumSet.Equal (class_of t x) (class_of t y) -> {{t}} x y.

  Lemma sameclass_spec `{WF} x y: reflect ({{t}} x y) (fst (equiv t x y)).

End DSUtils.

Module PosDisjointSets <: DISJOINTSETS Positive.

  Import PositiveUtils.
  Module Z := FMapFacts.OrdProperties NumMap.   Section protect.

  Open Scope nat_scope. Notation S := Datatypes.S.
  Definition num := num.
  Definition eq_num := @eq num.
  Notation "M '[' key ']' " := (NumMap.find key M)(at level 0, no associativity).
  Notation "M '[' key '>-' v ']' " := (NumMap.add key v M)(at level 0, no associativity).
  Notation "'[]'" :=(NumMap.empty _).

  Structure UF :=
    {
      p : NumMap.t num;
      rank : NumMap.t num;
      size : nat
    }.

  Definition T := UF.

  Definition empty : T :=
    Build_UF
    []
    []
    (S O).

  Fixpoint find_aux n x t :=
    match n with
      | (S n) => match t [x] with
                   | None => (x, t)
                   | (Some y) =>
                     let (ry, t) := find_aux n y t in
                       (ry, t [x>- ry] )
                 end
      | O => (x,t)
    end.

  Definition find t x :=
    let (rx,t') := find_aux (size t) x (p t) in
      (rx, Build_UF t' (rank t) (size t)).


  Definition get_rank x rank :=
    match rank [x] with
      | None => 1%positive
      | Some rx => rx
    end.

  Definition link t x y :=
    let rank := rank t in
    let rx := get_rank x rank in
    let ry := get_rank y rank in
    let p := p t in
      match (compare rx ry) with
           | Lt =>
             Build_UF
             (p [x >- y])
             rank
             (S (size t))
           | Eq =>
             Build_UF
             (p [x >- y])
             (rank [y >- Positive.S ry] )
             (S (size t))
           | Gt =>
             Build_UF
             (p [y >- x])
             rank
             (S (size t))
         end
     .

  Definition union t x y :=
    let (rx,t) := find t x in
    let (ry,t) := find t y in
      if Positive.eqb rx ry then t else link t rx ry.

  Definition equiv t x y :=
    let (rx,t) := find t x in
    let (ry,t) := find t y in
     (Positive.eqb rx ry,t).

  Definition test_and_unify t x y :=
    let (rx,t) := find t x in
    let (ry,t) := find t y in
      if Positive.eqb rx ry then (true,t) else (false, link t rx ry).

  Import NumMapProps.

  Inductive repr (t: NumMap.t num) : num -> num -> Prop :=
  | rzero : forall i, ~ NumMap.In i t -> repr t i i
  | rsucc : forall i j r, NumMap.MapsTo i j t -> repr t j r -> repr t i r
    .

  Inductive D (t: NumMap.t num) : num -> nat -> Prop :=
  | DO : forall i, ~NumMap.In i t -> D t i O
  | DS : forall i j n, NumMap.MapsTo i j t -> D t j n -> D t i (S n).

  Definition Equiv t x y := exists2 z, repr t x z & repr t y z.

  Definition FEquiv t t' := forall x y, repr t x y <-> repr t' x y .

  Definition bounded t n:= forall x, exists2 m, D t x m & m < n.

  Section repr.
    Variable t : NumMap.t num.

    Lemma repr_zero_inv : forall i x, i <> x -> ~NumMap.In i t -> ~ repr t i x.

    Lemma repr_inj_right : forall a b, repr t a b -> forall c, repr t a c -> b = c.

    Lemma repr_inv_In : forall a b, repr t a b -> ~ NumMap.In b t.

    Lemma repr_inj_idem : forall i j , repr t i j -> forall r, repr t j r -> j = r.
    Lemma DO_inv_In : forall x , D t x O -> ~ NumMap.In x t.

    Hint Resolve repr_inv_In DO_inv_In.

    Lemma D_repr : forall x, D t x O <-> repr t x x.

    Lemma D_inv_n : forall x y n, D t y O -> D t x (S n) -> x <> y.

    Lemma repr_idem : forall x rx, repr t x rx -> repr t rx rx.

    Lemma D_succ : forall x y n, NumMap.MapsTo x y t -> D t x n -> exists2 m, D t y m & n = S m.
    Lemma D_inj : forall x n , D t x n -> forall m, D t x m -> n = m.

    Lemma repr_D : forall i j, repr t i j -> exists n, D t i n.

    Lemma update_In : forall i a b, repr t a b -> a <> b -> ~NumMap.In i t -> ~NumMap.In i (t [a >- b]).
  End repr.

  Lemma D_update : forall x y t, x <> y -> D t y 0 -> D (t [x>- y]) y 0.

  Hint Resolve repr_idem DO_inv_In D_inv_n repr_D D_update update_In repr_inv_In.
  Lemma FEquiv_D : forall t t', FEquiv t t' -> (forall x, D t x O <-> D t' x O).

  Instance FEquiv_eq : Equivalence FEquiv.

  Instance repr_compat : Proper (FEquiv ==> (@eq num) ==> (@eq num)==> iff) repr.

  Section update.
    Variable t : NumMap.t num.
    Variable a b : num.
    Hypothesis a_diff_b: a <> b.
    Hypothesis a_repr_b: repr t a b.

    Lemma update_equiv : FEquiv t (t [a >- b ]).

    Lemma D_update_ex : forall x n, D t x n -> exists2 m, D (t [a>-b]) x m & m <= n.


    Lemma boundage : forall n, bounded t n ->
      bounded (t [a>- b]) n.
  End update.

  Section link.
    Variable t : NumMap.t num.
    Variables a b : num.
    Hypothesis a_diff_b : a <> b.
    Hypothesis repr_a : repr t a a.
    Hypothesis repr_b : repr t b b.
    Lemma D_link_ex : forall x n, D t x n -> exists2 m, D (t [a>-b]) x m & m <= S n.

    Lemma boundage_link : forall n, bounded t n -> bounded (t [a>- b]) (S n).

    Lemma repr_update_fwd :
      forall y, repr t y b -> repr (t[ b>- a]) y a.

    Lemma repr_x_inv_upd :
      forall x,
        repr t x a -> repr (t[ b>- a]) x a.

    Lemma link_lemma_1 :
      forall z,
        repr (t[ b>- a]) z a -> (repr t z a \/ repr t z b).


    Lemma link_lemma_2 :
      forall z rz, repr (t[ b>- a]) z rz -> rz<>a ->
        repr t z rz.

    Lemma repr_inv_upd :
      forall z rz, repr t z rz -> rz<>b ->
        repr (t[ b>- a]) z rz.
    Hint Resolve repr_x_inv_upd repr_inv_upd link_lemma_1 link_lemma_2 repr_update_fwd.

    Lemma link_main_lemma :
      forall x y,
        repr t x a -> repr t y b ->
        forall u v,
          (Equiv (t[ b>- a]) u v <->
            (Equiv t u v \/ ((Equiv t u x /\ Equiv t v y) \/ (Equiv t v x /\ Equiv t u y)))).
  End link.

  Section find_aux.
    Lemma find_equiv : forall n s t x y t' px,
      find_aux n x t = (y,t') ->
      D t x px ->
      px < n ->
      bounded t s ->
      (
        D t' y O /\
        (x = y \/ NumMap.MapsTo x y t') /\
        FEquiv t t' /\
        bounded t' s
      ).

  End find_aux.

  Section WF.
    Definition IsWF t : Prop := bounded (p t) (size t).
    Class WF t: Prop := { wf : IsWF t }.
    Hint Constructors WF.

    Inductive find_spec_decl t x : (num * T) -> Type :=
    | find_spec_1 : forall rx t'
      (Heq:FEquiv (p t) (p t'))
      (Hrepr: repr (p t') x rx)
      (Hwf: WF t'),
      find_spec_decl t x (rx,t').

    Lemma find_spec : forall `{WF} x, find_spec_decl t x (find t x).

    Global Instance empty_WF : WF empty.

    Global Instance find_WF `{WF} x: WF (snd (find t x )).

    Global Instance equiv_WF `{WF} x y: WF (snd (equiv t x y)).

    Instance link_WF `{WF} x y : repr (p t) x x -> repr (p t) y y -> x<> y -> WF (link t x y).

    Global Instance union_WF `{WF} x y: WF (union t x y).

    Global Instance test_and_unify_WF `{WF} x y: WF (snd (test_and_unify t x y)).
  End WF.

  Lemma repr_helper `{WF} x: exists z, repr (p t) x z.

  Section Equiv.
    Definition sameclass t x y := Equiv (p t) x y.

    Lemma Equiv_empty : forall m s t,NumMap.Empty m -> Equiv m s t -> s = t.

    Lemma sameclass_empty : forall s t, sameclass empty s t -> s = t.

    Instance Equiv_compat : Proper (FEquiv ==> @eq num ==> @eq num ==> iff ) Equiv.

    Global Instance sameclass_Equivalence `{WF} : Equivalence (sameclass t).

    Lemma sameclass_union : forall {uf} `{WF uf} a b,
      forall s t, sameclass (union uf a b) s t <->
        (sameclass uf s t \/ (sameclass uf s a /\ sameclass uf t b) \/ (sameclass uf t a /\ sameclass uf s b)).

    Lemma sameclass_equiv : forall `{WF} x y, (fst (equiv t x y) = true <-> sameclass t x y).

    Lemma sameclass_find : forall `{WF} x, sameclass t (fst (find t x)) x.

  End Equiv.

  Lemma test_and_unify_eq : forall {uf} `{WF uf} x y, test_and_unify uf x y = (fst (equiv uf x y), union uf x y).

  Definition class_of' (t : T) x :=
    NumMap.fold (fun a b acc =>
      if fst (equiv t x a) && fst (equiv t x b)
        then (NumSet.add a (NumSet.add b acc))
        else acc
    )
    (p t)
    NumSet.empty.

  Lemma sameclass_equiv' : forall `{WF} y z,
    sameclass t y z -> forall x, @fst bool UF (equiv t y x) = fst (equiv t z x).


  Lemma NumMap_fold_compat A t f g (x: A):
    (forall acc (k e : num) , f k e acc = g k e acc) ->
    NumMap.fold f t x = NumMap.fold g t x.

  Lemma MapsTo_sameclass : forall `{WF} x z, NumMap.MapsTo x z (p t) -> sameclass t x z.

  Lemma equiv_reflexive_false : forall `{WF} x, ~ fst (equiv t x x) = false .

  Definition class_of t x := NumSet.add x (class_of' t x).

  Lemma class_of_Equal_MapsTo : forall `{WF} x z, NumMap.MapsTo x z (p t) ->
    NumSet.Equal (class_of' t z) (class_of' t x) .

  Lemma class_of_Equal_repr : forall `{WF} x z, repr (p t) x z -> NumSet.Equal (class_of' t z) (class_of' t x).

  Lemma class_of_In_MapsTo : forall `{WF} x z, NumMap.MapsTo x z (p t) -> NumSet.In x (class_of t z).

  Lemma class_of_In_MapsTo' : forall `{WF} x z, NumMap.MapsTo x z (p t) -> NumSet.In z (class_of t x).

  Lemma class_of_In_repr : forall `{WF} x z, repr (p t) x z -> NumSet.In x (class_of t z).

  Lemma class_of_In_repr' : forall `{WF} x z, repr (p t) x z -> NumSet.In z (class_of t x).

  Lemma NumMap_Equal_Empty_empty : forall elt (t : NumMap.t elt) , NumMap.Empty t -> NumMap.Equal t (NumMap.empty _).

  Lemma sameclass_class_of : forall `{WF} x y , sameclass t x y <-> NumSet.In x (class_of t y).

  End protect.
End PosDisjointSets.