Theory Union_Find

(*
  File: Union_Find.thy
  Author: Bohua Zhan
*)

section ‹Union find›

theory Union_Find
  imports Partial_Equiv_Rel
begin

text ‹
  Development follows theory Union\_Find in \<^cite>‹"Separation_Logic_Imperative_HOL-AFP"›.
›

subsection ‹Representing a partial equivalence relation using rep\_of array›

function (domintros) rep_of where
  "rep_of l i = (if l ! i = i then i else rep_of l (l ! i))" by auto

setup ‹register_wellform_data ("rep_of l i", ["i < length l"])›
setup ‹add_backward_prfstep @{thm rep_of.domintros}›
setup ‹add_rewrite_rule @{thm rep_of.psimps}›
setup ‹add_prop_induct_rule @{thm rep_of.pinduct}›

definition ufa_invar :: "nat list ⇒ bool" where [rewrite]:
  "ufa_invar l = (∀i<length l. rep_of_dom (l, i) ∧ l ! i < length l)"

lemma ufa_invarD:
  "ufa_invar l ⟹ i < length l ⟹ rep_of_dom (l, i) ∧ l ! i < length l" by auto2
setup ‹add_forward_prfstep_cond @{thm ufa_invarD} [with_term "?l ! ?i"]›
setup ‹del_prfstep_thm_eqforward @{thm ufa_invar_def}›

lemma rep_of_id [rewrite]: "ufa_invar l ⟹ i < length l ⟹ l ! i = i ⟹ rep_of l i = i" by auto2

lemma rep_of_iff [rewrite]:
  "ufa_invar l ⟹ i < length l ⟹ rep_of l i = (if l ! i = i then i else rep_of l (l ! i))" by auto2
setup ‹del_prfstep_thm @{thm rep_of.psimps}›

lemma rep_of_min [rewrite]:
  "ufa_invar l ⟹ i < length l ⟹ l ! (rep_of l i) = rep_of l i"
@proof @prop_induct "rep_of_dom (l, i)" @qed

lemma rep_of_induct:
  "ufa_invar l ∧ i < length l ⟹
   ∀i<length l. l ! i = i ⟶ P l i ⟹
   ∀i<length l. l ! i ≠ i ⟶ P l (l ! i) ⟶ P l i ⟹ P l i"
@proof @prop_induct "rep_of_dom (l, i)" @qed
setup ‹add_prop_induct_rule @{thm rep_of_induct}›

lemma rep_of_bound [forward_arg1]:
  "ufa_invar l ⟹ i < length l ⟹ rep_of l i < length l"
@proof @prop_induct "ufa_invar l ∧ i < length l" @qed

lemma rep_of_idem [rewrite]:
  "ufa_invar l ⟹ i < length l ⟹ rep_of l (rep_of l i) = rep_of l i" by auto2

lemma rep_of_idx [rewrite]: 
  "ufa_invar l ⟹ i < length l ⟹ rep_of l (l ! i) = rep_of l i" by auto2

definition ufa_α :: "nat list ⇒ (nat × nat) set" where [rewrite]:
  "ufa_α l = {(x, y). x < length l ∧ y < length l ∧ rep_of l x = rep_of l y}"

lemma ufa_α_memI [backward, forward_arg]:
  "x < length l ⟹ y < length l ⟹ rep_of l x = rep_of l y ⟹ (x, y) ∈ ufa_α l"
  by (simp add: ufa_α_def)
  
lemma ufa_α_memD [forward]:
  "(x, y) ∈ ufa_α l ⟹ x < length l ∧ y < length l ∧ rep_of l x = rep_of l y"
  by (simp add: ufa_α_def)
setup ‹del_prfstep_thm @{thm ufa_α_def}›

lemma ufa_α_equiv [forward]: "part_equiv (ufa_α l)" by auto2

lemma ufa_α_refl [rewrite]: "(i, i) ∈ ufa_α l ⟷ i < length l" by auto2

subsection ‹Operations on rep\_of array›

definition uf_init_rel :: "nat ⇒ (nat × nat) set" where [rewrite]:
  "uf_init_rel n = ufa_α [0..<n]"

lemma ufa_init_invar [resolve]: "ufa_invar [0..<n]" by auto2

lemma ufa_init_correct [rewrite]:
  "(x, y) ∈ uf_init_rel n ⟷ (x = y ∧ x < n)"
@proof @have "ufa_invar [0..<n]" @qed

abbreviation ufa_union :: "nat list ⇒ nat ⇒ nat ⇒ nat list" where
  "ufa_union l x y ≡ l[rep_of l x := rep_of l y]"

lemma ufa_union_invar [forward_arg]:
  "ufa_invar l ⟹ x < length l ⟹ y < length l ⟹ l' = ufa_union l x y ⟹ ufa_invar l'"
@proof
  @have "∀i<length l'. rep_of_dom (l', i) ∧ l' ! i < length l'" @with
    @prop_induct "ufa_invar l ∧ i < length l"
  @end
@qed

lemma ufa_union_aux [rewrite]:
  "ufa_invar l ⟹ x < length l ⟹ y < length l ⟹ l' = ufa_union l x y ⟹
   i < length l' ⟹ rep_of l' i = (if rep_of l i = rep_of l x then rep_of l y else rep_of l i)"
@proof @prop_induct "ufa_invar l ∧ i < length l" @qed

text ‹Correctness of union operation.›
theorem ufa_union_correct [rewrite]:
  "ufa_invar l ⟹ x < length l ⟹ y < length l ⟹ l' = ufa_union l x y ⟹
   ufa_α l' = per_union (ufa_α l) x y"
@proof
  @have "∀a b. (a,b) ∈ ufa_α l' ⟷ (a,b) ∈ per_union (ufa_α l) x y" @with
    @case "(a,b) ∈ ufa_α l'" @with
      @case "rep_of l a = rep_of l x"
      @case "rep_of l a = rep_of l y"
    @end
  @end
@qed

abbreviation ufa_compress :: "nat list ⇒ nat ⇒ nat list" where
  "ufa_compress l x ≡ l[x := rep_of l x]"

lemma ufa_compress_invar [forward_arg]:
  "ufa_invar l ⟹ x < length l ⟹ l' = ufa_compress l x ⟹ ufa_invar l'"
@proof
  @have "∀i<length l'. rep_of_dom (l', i) ∧ l' ! i < length l'" @with
    @prop_induct "ufa_invar l ∧ i < length l"
  @end
@qed

lemma ufa_compress_aux [rewrite]:
  "ufa_invar l ⟹ x < length l ⟹ l' = ufa_compress l x ⟹ i < length l' ⟹
   rep_of l' i = rep_of l i"
@proof @prop_induct "ufa_invar l ∧ i < length l" @qed

text ‹Correctness of compress operation.›
theorem ufa_compress_correct [rewrite]:
  "ufa_invar l ⟹ x < length l ⟹ ufa_α (ufa_compress l x) = ufa_α l" by auto2

setup ‹del_prfstep_thm @{thm rep_of_iff}›

end