(*  Title:      RBT_Impl.thy
    Author:     Markus Reiter, TU Muenchen
    Author:     Alexander Krauss, TU Muenchen
*)

header {* Implementation of Red-Black Trees *}

theory RBT_Impl
imports Main "~~/src/HOL/Library/More_List" "~~/src/HOL/Library/Quickcheck_Narrowing"
begin

text {*
  For applications, you should use theory @{text RBT} which defines
  an abstract type of red-black tree obeying the invariant.
*}

subsection {* Datatype of RB trees *}

datatype color = R | B
datatype ('a, 'b) rbt = Empty | Branch color "('a, 'b) rbt" 'a 'b "('a, 'b) rbt"

lemma rbt_cases:
  obtains (Empty) "t = Empty" 
  | (Red) l k v r where "t = Branch R l k v r" 
  | (Black) l k v r where "t = Branch B l k v r"
proof (cases t)
  case Empty with that show thesis by blast
next
  case (Branch c) with that show thesis by (cases c) blast+
qed

subsection {* Tree properties *}

subsubsection {* Content of a tree *}

primrec entries :: "('a, 'b) rbt \<Rightarrow> ('a \<times> 'b) list"
where 
  "entries Empty = []"
| "entries (Branch _ l k v r) = entries l @ (k,v) # entries r"

abbreviation (input) entry_in_tree :: "'a \<Rightarrow> 'b \<Rightarrow> ('a, 'b) rbt \<Rightarrow> bool"
where
  "entry_in_tree k v t \<equiv> (k, v) \<in> set (entries t)"

definition keys :: "('a, 'b) rbt \<Rightarrow> 'a list" where
  "keys t = map fst (entries t)"

lemma keys_simps [simp, code]:
  "keys Empty = []"
  "keys (Branch c l k v r) = keys l @ k # keys r"
  by (simp_all add: keys_def)

lemma entry_in_tree_keys:
  assumes "(k, v) \<in> set (entries t)"
  shows "k \<in> set (keys t)"
proof -
  from assms have "fst (k, v) \<in> fst ` set (entries t)" by (rule imageI)
  then show ?thesis by (simp add: keys_def)
qed

lemma keys_entries:
  "k \<in> set (keys t) \<longleftrightarrow> (\<exists>v. (k, v) \<in> set (entries t))"
  by (auto intro: entry_in_tree_keys) (auto simp add: keys_def)


subsubsection {* Search tree properties *}

definition tree_less :: "'a\<Colon>order \<Rightarrow> ('a, 'b) rbt \<Rightarrow> bool"
where
  tree_less_prop: "tree_less k t \<longleftrightarrow> (\<forall>x\<in>set (keys t). x < k)"

abbreviation tree_less_symbol (infix "|\<guillemotleft>" 50)
where "t |\<guillemotleft> x \<equiv> tree_less x t"

definition tree_greater :: "'a\<Colon>order \<Rightarrow> ('a, 'b) rbt \<Rightarrow> bool" (infix "\<guillemotleft>|" 50) 
where
  tree_greater_prop: "tree_greater k t = (\<forall>x\<in>set (keys t). k < x)"

lemma tree_less_simps [simp]:
  "tree_less k Empty = True"
  "tree_less k (Branch c lt kt v rt) \<longleftrightarrow> kt < k \<and> tree_less k lt \<and> tree_less k rt"
  by (auto simp add: tree_less_prop)

lemma tree_greater_simps [simp]:
  "tree_greater k Empty = True"
  "tree_greater k (Branch c lt kt v rt) \<longleftrightarrow> k < kt \<and> tree_greater k lt \<and> tree_greater k rt"
  by (auto simp add: tree_greater_prop)

lemmas tree_ord_props = tree_less_prop tree_greater_prop

lemmas tree_greater_nit = tree_greater_prop entry_in_tree_keys
lemmas tree_less_nit = tree_less_prop entry_in_tree_keys

lemma tree_less_eq_trans: "l |\<guillemotleft> u \<Longrightarrow> u \<le> v \<Longrightarrow> l |\<guillemotleft> v"
  and tree_less_trans: "t |\<guillemotleft> x \<Longrightarrow> x < y \<Longrightarrow> t |\<guillemotleft> y"
  and tree_greater_eq_trans: "u \<le> v \<Longrightarrow> v \<guillemotleft>| r \<Longrightarrow> u \<guillemotleft>| r"
  and tree_greater_trans: "x < y \<Longrightarrow> y \<guillemotleft>| t \<Longrightarrow> x \<guillemotleft>| t"
  by (auto simp: tree_ord_props)

primrec sorted :: "('a::linorder, 'b) rbt \<Rightarrow> bool"
where
  "sorted Empty = True"
| "sorted (Branch c l k v r) = (l |\<guillemotleft> k \<and> k \<guillemotleft>| r \<and> sorted l \<and> sorted r)"

lemma sorted_entries:
  "sorted t \<Longrightarrow> List.sorted (List.map fst (entries t))"
by (induct t) 
  (force simp: sorted_append sorted_Cons tree_ord_props 
      dest!: entry_in_tree_keys)+

lemma distinct_entries:
  "sorted t \<Longrightarrow> distinct (List.map fst (entries t))"
by (induct t) 
  (force simp: sorted_append sorted_Cons tree_ord_props 
      dest!: entry_in_tree_keys)+


subsubsection {* Tree lookup *}

primrec lookup :: "('a\<Colon>linorder, 'b) rbt \<Rightarrow> 'a \<rightharpoonup> 'b"
where
  "lookup Empty k = None"
| "lookup (Branch _ l x y r) k = (if k < x then lookup l k else if x < k then lookup r k else Some y)"

lemma lookup_keys: "sorted t \<Longrightarrow> dom (lookup t) = set (keys t)"
  by (induct t) (auto simp: dom_def tree_greater_prop tree_less_prop)

lemma dom_lookup_Branch: 
  "sorted (Branch c t1 k v t2) \<Longrightarrow> 
    dom (lookup (Branch c t1 k v t2)) 
    = Set.insert k (dom (lookup t1) \<union> dom (lookup t2))"
proof -
  assume "sorted (Branch c t1 k v t2)"
  moreover from this have "sorted t1" "sorted t2" by simp_all
  ultimately show ?thesis by (simp add: lookup_keys)
qed

lemma finite_dom_lookup [simp, intro!]: "finite (dom (lookup t))"
proof (induct t)
  case Empty then show ?case by simp
next
  case (Branch color t1 a b t2)
  let ?A = "Set.insert a (dom (lookup t1) \<union> dom (lookup t2))"
  have "dom (lookup (Branch color t1 a b t2)) \<subseteq> ?A" by (auto split: split_if_asm)
  moreover from Branch have "finite (insert a (dom (lookup t1) \<union> dom (lookup t2)))" by simp
  ultimately show ?case by (rule finite_subset)
qed 

lemma lookup_tree_less[simp]: "t |\<guillemotleft> k \<Longrightarrow> lookup t k = None" 
by (induct t) auto

lemma lookup_tree_greater[simp]: "k \<guillemotleft>| t \<Longrightarrow> lookup t k = None"
by (induct t) auto

lemma lookup_Empty: "lookup Empty = empty"
by (rule ext) simp

lemma map_of_entries:
  "sorted t \<Longrightarrow> map_of (entries t) = lookup t"
proof (induct t)
  case Empty thus ?case by (simp add: lookup_Empty)
next
  case (Branch c t1 k v t2)
  have "lookup (Branch c t1 k v t2) = lookup t2 ++ [k\<mapsto>v] ++ lookup t1"
  proof (rule ext)
    fix x
    from Branch have SORTED: "sorted (Branch c t1 k v t2)" by simp
    let ?thesis = "lookup (Branch c t1 k v t2) x = (lookup t2 ++ [k \<mapsto> v] ++ lookup t1) x"

    have DOM_T1: "!!k'. k'\<in>dom (lookup t1) \<Longrightarrow> k>k'"
    proof -
      fix k'
      from SORTED have "t1 |\<guillemotleft> k" by simp
      with tree_less_prop have "\<forall>k'\<in>set (keys t1). k>k'" by auto
      moreover assume "k'\<in>dom (lookup t1)"
      ultimately show "k>k'" using lookup_keys SORTED by auto
    qed
    
    have DOM_T2: "!!k'. k'\<in>dom (lookup t2) \<Longrightarrow> k<k'"
    proof -
      fix k'
      from SORTED have "k \<guillemotleft>| t2" by simp
      with tree_greater_prop have "\<forall>k'\<in>set (keys t2). k<k'" by auto
      moreover assume "k'\<in>dom (lookup t2)"
      ultimately show "k<k'" using lookup_keys SORTED by auto
    qed
    
    {
      assume C: "x<k"
      hence "lookup (Branch c t1 k v t2) x = lookup t1 x" by simp
      moreover from C have "x\<notin>dom [k\<mapsto>v]" by simp
      moreover have "x\<notin>dom (lookup t2)" proof
        assume "x\<in>dom (lookup t2)"
        with DOM_T2 have "k<x" by blast
        with C show False by simp
      qed
      ultimately have ?thesis by (simp add: map_add_upd_left map_add_dom_app_simps)
    } moreover {
      assume [simp]: "x=k"
      hence "lookup (Branch c t1 k v t2) x = [k \<mapsto> v] x" by simp
      moreover have "x\<notin>dom (lookup t1)" proof
        assume "x\<in>dom (lookup t1)"
        with DOM_T1 have "k>x" by blast
        thus False by simp
      qed
      ultimately have ?thesis by (simp add: map_add_upd_left map_add_dom_app_simps)
    } moreover {
      assume C: "x>k"
      hence "lookup (Branch c t1 k v t2) x = lookup t2 x" by (simp add: less_not_sym[of k x])
      moreover from C have "x\<notin>dom [k\<mapsto>v]" by simp
      moreover have "x\<notin>dom (lookup t1)" proof
        assume "x\<in>dom (lookup t1)"
        with DOM_T1 have "k>x" by simp
        with C show False by simp
      qed
      ultimately have ?thesis by (simp add: map_add_upd_left map_add_dom_app_simps)
    } ultimately show ?thesis using less_linear by blast
  qed
  also from Branch have "lookup t2 ++ [k \<mapsto> v] ++ lookup t1 = map_of (entries (Branch c t1 k v t2))" by simp
  finally show ?case by simp
qed

lemma lookup_in_tree: "sorted t \<Longrightarrow> lookup t k = Some v \<longleftrightarrow> (k, v) \<in> set (entries t)"
  by (simp add: map_of_entries [symmetric] distinct_entries)

lemma set_entries_inject:
  assumes sorted: "sorted t1" "sorted t2" 
  shows "set (entries t1) = set (entries t2) \<longleftrightarrow> entries t1 = entries t2"
proof -
  from sorted have "distinct (map fst (entries t1))"
    "distinct (map fst (entries t2))"
    by (auto intro: distinct_entries)
  with sorted show ?thesis
    by (auto intro: map_sorted_distinct_set_unique sorted_entries simp add: distinct_map)
qed

lemma entries_eqI:
  assumes sorted: "sorted t1" "sorted t2" 
  assumes lookup: "lookup t1 = lookup t2"
  shows "entries t1 = entries t2"
proof -
  from sorted lookup have "map_of (entries t1) = map_of (entries t2)"
    by (simp add: map_of_entries)
  with sorted have "set (entries t1) = set (entries t2)"
    by (simp add: map_of_inject_set distinct_entries)
  with sorted show ?thesis by (simp add: set_entries_inject)
qed

lemma entries_lookup:
  assumes "sorted t1" "sorted t2" 
  shows "entries t1 = entries t2 \<longleftrightarrow> lookup t1 = lookup t2"
  using assms by (auto intro: entries_eqI simp add: map_of_entries [symmetric])

lemma lookup_from_in_tree: 
  assumes "sorted t1" "sorted t2" 
  and "\<And>v. (k\<Colon>'a\<Colon>linorder, v) \<in> set (entries t1) \<longleftrightarrow> (k, v) \<in> set (entries t2)" 
  shows "lookup t1 k = lookup t2 k"
proof -
  from assms have "k \<in> dom (lookup t1) \<longleftrightarrow> k \<in> dom (lookup t2)"
    by (simp add: keys_entries lookup_keys)
  with assms show ?thesis by (auto simp add: lookup_in_tree [symmetric])
qed


subsubsection {* Red-black properties *}

primrec color_of :: "('a, 'b) rbt \<Rightarrow> color"
where
  "color_of Empty = B"
| "color_of (Branch c _ _ _ _) = c"

primrec bheight :: "('a,'b) rbt \<Rightarrow> nat"
where
  "bheight Empty = 0"
| "bheight (Branch c lt k v rt) = (if c = B then Suc (bheight lt) else bheight lt)"

primrec inv1 :: "('a, 'b) rbt \<Rightarrow> bool"
where
  "inv1 Empty = True"
| "inv1 (Branch c lt k v rt) \<longleftrightarrow> inv1 lt \<and> inv1 rt \<and> (c = B \<or> color_of lt = B \<and> color_of rt = B)"

primrec inv1l :: "('a, 'b) rbt \<Rightarrow> bool" -- {* Weaker version *}
where
  "inv1l Empty = True"
| "inv1l (Branch c l k v r) = (inv1 l \<and> inv1 r)"
lemma [simp]: "inv1 t \<Longrightarrow> inv1l t" by (cases t) simp+

primrec inv2 :: "('a, 'b) rbt \<Rightarrow> bool"
where
  "inv2 Empty = True"
| "inv2 (Branch c lt k v rt) = (inv2 lt \<and> inv2 rt \<and> bheight lt = bheight rt)"

definition is_rbt :: "('a\<Colon>linorder, 'b) rbt \<Rightarrow> bool" where
  "is_rbt t \<longleftrightarrow> inv1 t \<and> inv2 t \<and> color_of t = B \<and> sorted t"

lemma is_rbt_sorted [simp]:
  "is_rbt t \<Longrightarrow> sorted t" by (simp add: is_rbt_def)

theorem Empty_is_rbt [simp]:
  "is_rbt Empty" by (simp add: is_rbt_def)


subsection {* Insertion *}

fun (* slow, due to massive case splitting *)
  balance :: "('a,'b) rbt \<Rightarrow> 'a \<Rightarrow> 'b \<Rightarrow> ('a,'b) rbt \<Rightarrow> ('a,'b) rbt"
where
  "balance (Branch R a w x b) s t (Branch R c y z d) = Branch R (Branch B a w x b) s t (Branch B c y z d)" |
  "balance (Branch R (Branch R a w x b) s t c) y z d = Branch R (Branch B a w x b) s t (Branch B c y z d)" |
  "balance (Branch R a w x (Branch R b s t c)) y z d = Branch R (Branch B a w x b) s t (Branch B c y z d)" |
  "balance a w x (Branch R b s t (Branch R c y z d)) = Branch R (Branch B a w x b) s t (Branch B c y z d)" |
  "balance a w x (Branch R (Branch R b s t c) y z d) = Branch R (Branch B a w x b) s t (Branch B c y z d)" |
  "balance a s t b = Branch B a s t b"

lemma balance_inv1: "\<lbrakk>inv1l l; inv1l r\<rbrakk> \<Longrightarrow> inv1 (balance l k v r)" 
  by (induct l k v r rule: balance.induct) auto

lemma balance_bheight: "bheight l = bheight r \<Longrightarrow> bheight (balance l k v r) = Suc (bheight l)"
  by (induct l k v r rule: balance.induct) auto

lemma balance_inv2: 
  assumes "inv2 l" "inv2 r" "bheight l = bheight r"
  shows "inv2 (balance l k v r)"
  using assms
  by (induct l k v r rule: balance.induct) auto

lemma balance_tree_greater[simp]: "(v \<guillemotleft>| balance a k x b) = (v \<guillemotleft>| a \<and> v \<guillemotleft>| b \<and> v < k)" 
  by (induct a k x b rule: balance.induct) auto

lemma balance_tree_less[simp]: "(balance a k x b |\<guillemotleft> v) = (a |\<guillemotleft> v \<and> b |\<guillemotleft> v \<and> k < v)"
  by (induct a k x b rule: balance.induct) auto

lemma balance_sorted: 
  fixes k :: "'a::linorder"
  assumes "sorted l" "sorted r" "l |\<guillemotleft> k" "k \<guillemotleft>| r"
  shows "sorted (balance l k v r)"
using assms proof (induct l k v r rule: balance.induct)
  case ("2_2" a x w b y t c z s va vb vd vc)
  hence "y < z \<and> z \<guillemotleft>| Branch B va vb vd vc" 
    by (auto simp add: tree_ord_props)
  hence "tree_greater y (Branch B va vb vd vc)" by (blast dest: tree_greater_trans)
  with "2_2" show ?case by simp
next
  case ("3_2" va vb vd vc x w b y s c z)
  from "3_2" have "x < y \<and> tree_less x (Branch B va vb vd vc)" 
    by simp
  hence "tree_less y (Branch B va vb vd vc)" by (blast dest: tree_less_trans)
  with "3_2" show ?case by simp
next
  case ("3_3" x w b y s c z t va vb vd vc)
  from "3_3" have "y < z \<and> tree_greater z (Branch B va vb vd vc)" by simp
  hence "tree_greater y (Branch B va vb vd vc)" by (blast dest: tree_greater_trans)
  with "3_3" show ?case by simp
next
  case ("3_4" vd ve vg vf x w b y s c z t va vb vii vc)
  hence "x < y \<and> tree_less x (Branch B vd ve vg vf)" by simp
  hence 1: "tree_less y (Branch B vd ve vg vf)" by (blast dest: tree_less_trans)
  from "3_4" have "y < z \<and> tree_greater z (Branch B va vb vii vc)" by simp
  hence "tree_greater y (Branch B va vb vii vc)" by (blast dest: tree_greater_trans)
  with 1 "3_4" show ?case by simp
next
  case ("4_2" va vb vd vc x w b y s c z t dd)
  hence "x < y \<and> tree_less x (Branch B va vb vd vc)" by simp
  hence "tree_less y (Branch B va vb vd vc)" by (blast dest: tree_less_trans)
  with "4_2" show ?case by simp
next
  case ("5_2" x w b y s c z t va vb vd vc)
  hence "y < z \<and> tree_greater z (Branch B va vb vd vc)" by simp
  hence "tree_greater y (Branch B va vb vd vc)" by (blast dest: tree_greater_trans)
  with "5_2" show ?case by simp
next
  case ("5_3" va vb vd vc x w b y s c z t)
  hence "x < y \<and> tree_less x (Branch B va vb vd vc)" by simp
  hence "tree_less y (Branch B va vb vd vc)" by (blast dest: tree_less_trans)
  with "5_3" show ?case by simp
next
  case ("5_4" va vb vg vc x w b y s c z t vd ve vii vf)
  hence "x < y \<and> tree_less x (Branch B va vb vg vc)" by simp
  hence 1: "tree_less y (Branch B va vb vg vc)" by (blast dest: tree_less_trans)
  from "5_4" have "y < z \<and> tree_greater z (Branch B vd ve vii vf)" by simp
  hence "tree_greater y (Branch B vd ve vii vf)" by (blast dest: tree_greater_trans)
  with 1 "5_4" show ?case by simp
qed simp+

lemma entries_balance [simp]:
  "entries (balance l k v r) = entries l @ (k, v) # entries r"
  by (induct l k v r rule: balance.induct) auto

lemma keys_balance [simp]: 
  "keys (balance l k v r) = keys l @ k # keys r"
  by (simp add: keys_def)

lemma balance_in_tree:  
  "entry_in_tree k x (balance l v y r) \<longleftrightarrow> entry_in_tree k x l \<or> k = v \<and> x = y \<or> entry_in_tree k x r"
  by (auto simp add: keys_def)

lemma lookup_balance[simp]: 
fixes k :: "'a::linorder"
assumes "sorted l" "sorted r" "l |\<guillemotleft> k" "k \<guillemotleft>| r"
shows "lookup (balance l k v r) x = lookup (Branch B l k v r) x"
by (rule lookup_from_in_tree) (auto simp:assms balance_in_tree balance_sorted)

primrec paint :: "color \<Rightarrow> ('a,'b) rbt \<Rightarrow> ('a,'b) rbt"
where
  "paint c Empty = Empty"
| "paint c (Branch _ l k v r) = Branch c l k v r"

lemma paint_inv1l[simp]: "inv1l t \<Longrightarrow> inv1l (paint c t)" by (cases t) auto
lemma paint_inv1[simp]: "inv1l t \<Longrightarrow> inv1 (paint B t)" by (cases t) auto
lemma paint_inv2[simp]: "inv2 t \<Longrightarrow> inv2 (paint c t)" by (cases t) auto
lemma paint_color_of[simp]: "color_of (paint B t) = B" by (cases t) auto
lemma paint_sorted[simp]: "sorted t \<Longrightarrow> sorted (paint c t)" by (cases t) auto
lemma paint_in_tree[simp]: "entry_in_tree k x (paint c t) = entry_in_tree k x t" by (cases t) auto
lemma paint_lookup[simp]: "lookup (paint c t) = lookup t" by (rule ext) (cases t, auto)
lemma paint_tree_greater[simp]: "(v \<guillemotleft>| paint c t) = (v \<guillemotleft>| t)" by (cases t) auto
lemma paint_tree_less[simp]: "(paint c t |\<guillemotleft> v) = (t |\<guillemotleft> v)" by (cases t) auto

fun
  ins :: "('a\<Colon>linorder \<Rightarrow> 'b \<Rightarrow> 'b \<Rightarrow> 'b) \<Rightarrow> 'a \<Rightarrow> 'b \<Rightarrow> ('a,'b) rbt \<Rightarrow> ('a,'b) rbt"
where
  "ins f k v Empty = Branch R Empty k v Empty" |
  "ins f k v (Branch B l x y r) = (if k < x then balance (ins f k v l) x y r
                               else if k > x then balance l x y (ins f k v r)
                               else Branch B l x (f k y v) r)" |
  "ins f k v (Branch R l x y r) = (if k < x then Branch R (ins f k v l) x y r
                               else if k > x then Branch R l x y (ins f k v r)
                               else Branch R l x (f k y v) r)"

lemma ins_inv1_inv2: 
  assumes "inv1 t" "inv2 t"
  shows "inv2 (ins f k x t)" "bheight (ins f k x t) = bheight t" 
  "color_of t = B \<Longrightarrow> inv1 (ins f k x t)" "inv1l (ins f k x t)"
  using assms
  by (induct f k x t rule: ins.induct) (auto simp: balance_inv1 balance_inv2 balance_bheight)

lemma ins_tree_greater[simp]: "(v \<guillemotleft>| ins f k x t) = (v \<guillemotleft>| t \<and> k > v)"
  by (induct f k x t rule: ins.induct) auto
lemma ins_tree_less[simp]: "(ins f k x t |\<guillemotleft> v) = (t |\<guillemotleft> v \<and> k < v)"
  by (induct f k x t rule: ins.induct) auto
lemma ins_sorted[simp]: "sorted t \<Longrightarrow> sorted (ins f k x t)"
  by (induct f k x t rule: ins.induct) (auto simp: balance_sorted)

lemma keys_ins: "set (keys (ins f k v t)) = { k } \<union> set (keys t)"
  by (induct f k v t rule: ins.induct) auto

lemma lookup_ins: 
  fixes k :: "'a::linorder"
  assumes "sorted t"
  shows "lookup (ins f k v t) x = ((lookup t)(k |-> case lookup t k of None \<Rightarrow> v 
                                                       | Some w \<Rightarrow> f k w v)) x"
using assms by (induct f k v t rule: ins.induct) auto

definition
  insert_with_key :: "('a\<Colon>linorder \<Rightarrow> 'b \<Rightarrow> 'b \<Rightarrow> 'b) \<Rightarrow> 'a \<Rightarrow> 'b \<Rightarrow> ('a,'b) rbt \<Rightarrow> ('a,'b) rbt"
where
  "insert_with_key f k v t = paint B (ins f k v t)"

lemma insertwk_sorted: "sorted t \<Longrightarrow> sorted (insert_with_key f k x t)"
  by (auto simp: insert_with_key_def)

theorem insertwk_is_rbt: 
  assumes inv: "is_rbt t" 
  shows "is_rbt (insert_with_key f k x t)"
using assms
unfolding insert_with_key_def is_rbt_def
by (auto simp: ins_inv1_inv2)

lemma lookup_insertwk: 
  assumes "sorted t"
  shows "lookup (insert_with_key f k v t) x = ((lookup t)(k |-> case lookup t k of None \<Rightarrow> v 
                                                       | Some w \<Rightarrow> f k w v)) x"
unfolding insert_with_key_def using assms
by (simp add:lookup_ins)

definition
  insertw_def: "insert_with f = insert_with_key (\<lambda>_. f)"

lemma insertw_sorted: "sorted t \<Longrightarrow> sorted (insert_with f k v t)" by (simp add: insertwk_sorted insertw_def)
theorem insertw_is_rbt: "is_rbt t \<Longrightarrow> is_rbt (insert_with f k v t)" by (simp add: insertwk_is_rbt insertw_def)

lemma lookup_insertw:
  assumes "is_rbt t"
  shows "lookup (insert_with f k v t) = (lookup t)(k \<mapsto> (if k:dom (lookup t) then f (the (lookup t k)) v else v))"
using assms
unfolding insertw_def
by (rule_tac ext) (cases "lookup t k", auto simp:lookup_insertwk dom_def)

definition insert :: "'a\<Colon>linorder \<Rightarrow> 'b \<Rightarrow> ('a, 'b) rbt \<Rightarrow> ('a, 'b) rbt" where
  "insert = insert_with_key (\<lambda>_ _ nv. nv)"

lemma insert_sorted: "sorted t \<Longrightarrow> sorted (insert k v t)" by (simp add: insertwk_sorted insert_def)
theorem insert_is_rbt [simp]: "is_rbt t \<Longrightarrow> is_rbt (insert k v t)" by (simp add: insertwk_is_rbt insert_def)

lemma lookup_insert: 
  assumes "is_rbt t"
  shows "lookup (insert k v t) = (lookup t)(k\<mapsto>v)"
unfolding insert_def
using assms
by (rule_tac ext) (simp add: lookup_insertwk split:option.split)


subsection {* Deletion *}

lemma bheight_paintR'[simp]: "color_of t = B \<Longrightarrow> bheight (paint R t) = bheight t - 1"
by (cases t rule: rbt_cases) auto

fun
  balance_left :: "('a,'b) rbt \<Rightarrow> 'a \<Rightarrow> 'b \<Rightarrow> ('a,'b) rbt \<Rightarrow> ('a,'b) rbt"
where
  "balance_left (Branch R a k x b) s y c = Branch R (Branch B a k x b) s y c" |
  "balance_left bl k x (Branch B a s y b) = balance bl k x (Branch R a s y b)" |
  "balance_left bl k x (Branch R (Branch B a s y b) t z c) = Branch R (Branch B bl k x a) s y (balance b t z (paint R c))" |
  "balance_left t k x s = Empty"

lemma balance_left_inv2_with_inv1:
  assumes "inv2 lt" "inv2 rt" "bheight lt + 1 = bheight rt" "inv1 rt"
  shows "bheight (balance_left lt k v rt) = bheight lt + 1"
  and   "inv2 (balance_left lt k v rt)"
using assms 
by (induct lt k v rt rule: balance_left.induct) (auto simp: balance_inv2 balance_bheight)

lemma balance_left_inv2_app: 
  assumes "inv2 lt" "inv2 rt" "bheight lt + 1 = bheight rt" "color_of rt = B"
  shows "inv2 (balance_left lt k v rt)" 
        "bheight (balance_left lt k v rt) = bheight rt"
using assms 
by (induct lt k v rt rule: balance_left.induct) (auto simp add: balance_inv2 balance_bheight)+ 

lemma balance_left_inv1: "\<lbrakk>inv1l a; inv1 b; color_of b = B\<rbrakk> \<Longrightarrow> inv1 (balance_left a k x b)"
  by (induct a k x b rule: balance_left.induct) (simp add: balance_inv1)+

lemma balance_left_inv1l: "\<lbrakk> inv1l lt; inv1 rt \<rbrakk> \<Longrightarrow> inv1l (balance_left lt k x rt)"
by (induct lt k x rt rule: balance_left.induct) (auto simp: balance_inv1)

lemma balance_left_sorted: "\<lbrakk> sorted l; sorted r; tree_less k l; tree_greater k r \<rbrakk> \<Longrightarrow> sorted (balance_left l k v r)"
apply (induct l k v r rule: balance_left.induct)
apply (auto simp: balance_sorted)
apply (unfold tree_greater_prop tree_less_prop)
by force+

lemma balance_left_tree_greater: 
  fixes k :: "'a::order"
  assumes "k \<guillemotleft>| a" "k \<guillemotleft>| b" "k < x" 
  shows "k \<guillemotleft>| balance_left a x t b"
using assms 
by (induct a x t b rule: balance_left.induct) auto

lemma balance_left_tree_less: 
  fixes k :: "'a::order"
  assumes "a |\<guillemotleft> k" "b |\<guillemotleft> k" "x < k" 
  shows "balance_left a x t b |\<guillemotleft> k"
using assms
by (induct a x t b rule: balance_left.induct) auto

lemma balance_left_in_tree: 
  assumes "inv1l l" "inv1 r" "bheight l + 1 = bheight r"
  shows "entry_in_tree k v (balance_left l a b r) = (entry_in_tree k v l \<or> k = a \<and> v = b \<or> entry_in_tree k v r)"
using assms 
by (induct l k v r rule: balance_left.induct) (auto simp: balance_in_tree)

fun
  balance_right :: "('a,'b) rbt \<Rightarrow> 'a \<Rightarrow> 'b \<Rightarrow> ('a,'b) rbt \<Rightarrow> ('a,'b) rbt"
where
  "balance_right a k x (Branch R b s y c) = Branch R a k x (Branch B b s y c)" |
  "balance_right (Branch B a k x b) s y bl = balance (Branch R a k x b) s y bl" |
  "balance_right (Branch R a k x (Branch B b s y c)) t z bl = Branch R (balance (paint R a) k x b) s y (Branch B c t z bl)" |
  "balance_right t k x s = Empty"

lemma balance_right_inv2_with_inv1:
  assumes "inv2 lt" "inv2 rt" "bheight lt = bheight rt + 1" "inv1 lt"
  shows "inv2 (balance_right lt k v rt) \<and> bheight (balance_right lt k v rt) = bheight lt"
using assms
by (induct lt k v rt rule: balance_right.induct) (auto simp: balance_inv2 balance_bheight)

lemma balance_right_inv1: "\<lbrakk>inv1 a; inv1l b; color_of a = B\<rbrakk> \<Longrightarrow> inv1 (balance_right a k x b)"
by (induct a k x b rule: balance_right.induct) (simp add: balance_inv1)+

lemma balance_right_inv1l: "\<lbrakk> inv1 lt; inv1l rt \<rbrakk> \<Longrightarrow>inv1l (balance_right lt k x rt)"
by (induct lt k x rt rule: balance_right.induct) (auto simp: balance_inv1)

lemma balance_right_sorted: "\<lbrakk> sorted l; sorted r; tree_less k l; tree_greater k r \<rbrakk> \<Longrightarrow> sorted (balance_right l k v r)"
apply (induct l k v r rule: balance_right.induct)
apply (auto simp:balance_sorted)
apply (unfold tree_less_prop tree_greater_prop)
by force+

lemma balance_right_tree_greater: 
  fixes k :: "'a::order"
  assumes "k \<guillemotleft>| a" "k \<guillemotleft>| b" "k < x" 
  shows "k \<guillemotleft>| balance_right a x t b"
using assms by (induct a x t b rule: balance_right.induct) auto

lemma balance_right_tree_less: 
  fixes k :: "'a::order"
  assumes "a |\<guillemotleft> k" "b |\<guillemotleft> k" "x < k" 
  shows "balance_right a x t b |\<guillemotleft> k"
using assms by (induct a x t b rule: balance_right.induct) auto

lemma balance_right_in_tree:
  assumes "inv1 l" "inv1l r" "bheight l = bheight r + 1" "inv2 l" "inv2 r"
  shows "entry_in_tree x y (balance_right l k v r) = (entry_in_tree x y l \<or> x = k \<and> y = v \<or> entry_in_tree x y r)"
using assms by (induct l k v r rule: balance_right.induct) (auto simp: balance_in_tree)

fun
  combine :: "('a,'b) rbt \<Rightarrow> ('a,'b) rbt \<Rightarrow> ('a,'b) rbt"
where
  "combine Empty x = x" 
| "combine x Empty = x" 
| "combine (Branch R a k x b) (Branch R c s y d) = (case (combine b c) of
                                      Branch R b2 t z c2 \<Rightarrow> (Branch R (Branch R a k x b2) t z (Branch R c2 s y d)) |
                                      bc \<Rightarrow> Branch R a k x (Branch R bc s y d))" 
| "combine (Branch B a k x b) (Branch B c s y d) = (case (combine b c) of
                                      Branch R b2 t z c2 \<Rightarrow> Branch R (Branch B a k x b2) t z (Branch B c2 s y d) |
                                      bc \<Rightarrow> balance_left a k x (Branch B bc s y d))" 
| "combine a (Branch R b k x c) = Branch R (combine a b) k x c" 
| "combine (Branch R a k x b) c = Branch R a k x (combine b c)" 

lemma combine_inv2:
  assumes "inv2 lt" "inv2 rt" "bheight lt = bheight rt"
  shows "bheight (combine lt rt) = bheight lt" "inv2 (combine lt rt)"
using assms 
by (induct lt rt rule: combine.induct) 
   (auto simp: balance_left_inv2_app split: rbt.splits color.splits)

lemma combine_inv1: 
  assumes "inv1 lt" "inv1 rt"
  shows "color_of lt = B \<Longrightarrow> color_of rt = B \<Longrightarrow> inv1 (combine lt rt)"
         "inv1l (combine lt rt)"
using assms 
by (induct lt rt rule: combine.induct)
   (auto simp: balance_left_inv1 split: rbt.splits color.splits)

lemma combine_tree_greater[simp]: 
  fixes k :: "'a::linorder"
  assumes "k \<guillemotleft>| l" "k \<guillemotleft>| r" 
  shows "k \<guillemotleft>| combine l r"
using assms 
by (induct l r rule: combine.induct)
   (auto simp: balance_left_tree_greater split:rbt.splits color.splits)

lemma combine_tree_less[simp]: 
  fixes k :: "'a::linorder"
  assumes "l |\<guillemotleft> k" "r |\<guillemotleft> k" 
  shows "combine l r |\<guillemotleft> k"
using assms 
by (induct l r rule: combine.induct)
   (auto simp: balance_left_tree_less split:rbt.splits color.splits)

lemma combine_sorted: 
  fixes k :: "'a::linorder"
  assumes "sorted l" "sorted r" "l |\<guillemotleft> k" "k \<guillemotleft>| r"
  shows "sorted (combine l r)"
using assms proof (induct l r rule: combine.induct)
  case (3 a x v b c y w d)
  hence ineqs: "a |\<guillemotleft> x" "x \<guillemotleft>| b" "b |\<guillemotleft> k" "k \<guillemotleft>| c" "c |\<guillemotleft> y" "y \<guillemotleft>| d"
    by auto
  with 3
  show ?case
    by (cases "combine b c" rule: rbt_cases)
      (auto, (metis combine_tree_greater combine_tree_less ineqs ineqs tree_less_simps(2) tree_greater_simps(2) tree_greater_trans tree_less_trans)+)
next
  case (4 a x v b c y w d)
  hence "x < k \<and> tree_greater k c" by simp
  hence "tree_greater x c" by (blast dest: tree_greater_trans)
  with 4 have 2: "tree_greater x (combine b c)" by (simp add: combine_tree_greater)
  from 4 have "k < y \<and> tree_less k b" by simp
  hence "tree_less y b" by (blast dest: tree_less_trans)
  with 4 have 3: "tree_less y (combine b c)" by (simp add: combine_tree_less)
  show ?case
  proof (cases "combine b c" rule: rbt_cases)
    case Empty
    from 4 have "x < y \<and> tree_greater y d" by auto
    hence "tree_greater x d" by (blast dest: tree_greater_trans)
    with 4 Empty have "sorted a" and "sorted (Branch B Empty y w d)" and "tree_less x a" and "tree_greater x (Branch B Empty y w d)" by auto
    with Empty show ?thesis by (simp add: balance_left_sorted)
  next
    case (Red lta va ka rta)
    with 2 4 have "x < va \<and> tree_less x a" by simp
    hence 5: "tree_less va a" by (blast dest: tree_less_trans)
    from Red 3 4 have "va < y \<and> tree_greater y d" by simp
    hence "tree_greater va d" by (blast dest: tree_greater_trans)
    with Red 2 3 4 5 show ?thesis by simp
  next
    case (Black lta va ka rta)
    from 4 have "x < y \<and> tree_greater y d" by auto
    hence "tree_greater x d" by (blast dest: tree_greater_trans)
    with Black 2 3 4 have "sorted a" and "sorted (Branch B (combine b c) y w d)" and "tree_less x a" and "tree_greater x (Branch B (combine b c) y w d)" by auto
    with Black show ?thesis by (simp add: balance_left_sorted)
  qed
next
  case (5 va vb vd vc b x w c)
  hence "k < x \<and> tree_less k (Branch B va vb vd vc)" by simp
  hence "tree_less x (Branch B va vb vd vc)" by (blast dest: tree_less_trans)
  with 5 show ?case by (simp add: combine_tree_less)
next
  case (6 a x v b va vb vd vc)
  hence "x < k \<and> tree_greater k (Branch B va vb vd vc)" by simp
  hence "tree_greater x (Branch B va vb vd vc)" by (blast dest: tree_greater_trans)
  with 6 show ?case by (simp add: combine_tree_greater)
qed simp+

lemma combine_in_tree: 
  assumes "inv2 l" "inv2 r" "bheight l = bheight r" "inv1 l" "inv1 r"
  shows "entry_in_tree k v (combine l r) = (entry_in_tree k v l \<or> entry_in_tree k v r)"
using assms 
proof (induct l r rule: combine.induct)
  case (4 _ _ _ b c)
  hence a: "bheight (combine b c) = bheight b" by (simp add: combine_inv2)
  from 4 have b: "inv1l (combine b c)" by (simp add: combine_inv1)

  show ?case
  proof (cases "combine b c" rule: rbt_cases)
    case Empty
    with 4 a show ?thesis by (auto simp: balance_left_in_tree)
  next
    case (Red lta ka va rta)
    with 4 show ?thesis by auto
  next
    case (Black lta ka va rta)
    with a b 4  show ?thesis by (auto simp: balance_left_in_tree)
  qed 
qed (auto split: rbt.splits color.splits)

subsection {* Custom Generator for Red--Black Trees *}

subsubsection {* Dedicated Type for Red--Black Trees *}

typedef (open) ('a, 'b) abstract_rbt = "{t :: ('a\<Colon>linorder, 'b) rbt. is_rbt t}"
  morphisms impl_of RBT
proof -
  have "Empty \<in> {t. is_rbt t}" by simp
  then show ?thesis ..
qed

lemma rbt_eq_iff:
  "t1 = t2 \<longleftrightarrow> impl_of t1 = impl_of t2"
  by (simp add: impl_of_inject)

lemma rbt_eqI:
  "impl_of t1 = impl_of t2 \<Longrightarrow> t1 = t2"
  by (simp add: rbt_eq_iff)

lemma is_rbt_impl_of [simp, intro]:
  "is_rbt (impl_of t)"
  using impl_of [of t] by simp

lemma RBT_impl_of [simp, code abstype]:
  "RBT (impl_of t) = t"
  by (simp add: impl_of_inverse)

subsubsection {* Lifting operations *}
  
definition empty_rbt :: "('a\<Colon>linorder, 'b) abstract_rbt" where
  "empty_rbt = RBT Empty"

lemma impl_of_empty [code abstract]:
  "impl_of empty_rbt = Empty"
  by (simp add: empty_rbt_def RBT_inverse)

definition insert_rbt :: "'a\<Colon>linorder \<Rightarrow> 'b \<Rightarrow> ('a, 'b) abstract_rbt \<Rightarrow> ('a, 'b) abstract_rbt" where
  "insert_rbt k v t = RBT (insert k v (impl_of t))"

lemma impl_of_insert [code abstract]:
  "impl_of (insert_rbt k v t) = insert k v (impl_of t)"
  by (simp add: insert_rbt_def RBT_inverse)

subsubsection {* Defining the Test Data Generator *}

instantiation abstract_rbt :: ("{full_exhaustive, linorder}", full_exhaustive) full_exhaustive
begin

fun full_exhaustive_abstract_rbt :: "(('a, 'b) abstract_rbt \<times> (unit \<Rightarrow> term) \<Rightarrow> term list option) \<Rightarrow> code_numeral \<Rightarrow> term list option"
where
  "full_exhaustive_abstract_rbt f d = Quickcheck_Exhaustive.orelse (f (empty_rbt,
    \<lambda>u \<Colon> unit. Code_Evaluation.Const (STR ''RBT_Impl.empty_rbt'') (Typerep.Typerep (STR ''RBT_Impl.abstract_rbt'') [Typerep.typerep (TYPE('a)), Typerep.typerep (TYPE('b))])))
    (if d = 0 then None else (full_exhaustive (%(k, kt). full_exhaustive (%(v, vt). full_exhaustive_abstract_rbt (%(t, tt). f (insert_rbt k v t,
      \<lambda>u \<Colon> unit. Code_Evaluation.App (Code_Evaluation.App (Code_Evaluation.App (Code_Evaluation.Const (STR ''RBT_Impl.insert_rbt'') (
         let A = Typerep.typerep (TYPE('a));
             B = Typerep.typerep (TYPE('b));
             R = Typerep.Typerep (STR ''RBT_Impl.abstract_rbt'') [A, B];
             fun = (\<lambda>T U. Typerep.Typerep (STR ''fun'') [T, U]) 
         in fun A (fun B (fun R R)))) (kt ())) (vt ())) (tt ()))) (d - 1)) d) d))"

instance ..

end

hide_const (open) R B



end