Theory BalancedBinary

section ‹Balanced Binary Trees›
theory BalancedBinary
imports
  "HOL-Library.Tree"
  Complex_Main
begin

text ‹
  @{text BalancedShape} captures the ∗‹shape› of a balancing scheme i.e. a rank function
  together with the constants @{text "cl"} and @{text "cu"}. The @{text balanced} predicate is
  defined relative to these, the locale then captures the emergent properties of the scheme
  as defined by cite‹blelloch2022joinable›.
  Here, ‹cl > 0› is demanded explicitly. With ‹cl = 0› the rank would
  never have to shrink towards the leaves, and the  cite‹blelloch2022joinable› itself divides 
  by ‹cl› in its Property 7.
  The actual guarantee that a scheme keeps its trees balanced is
  added by the @{text BalancedTree} locale further below, which connects an invariant to
  @{text balanced}.
›
locale BalancedShape =
fixes rank :: "('a * 'b) tree ⇒ real"
fixes cl cu :: real
assumes c_vals:"cl > 0 ∧ cl ≤ 1 ∧ cu ≥ 1"
assumes rule_empty: "rank Leaf = 0"
begin

lemma c_l_pos: "0 < cl"
  using c_vals by blast

lemma c_l_le_one: "cl ≤ 1"
  using c_vals by blast

lemma c_u_ge_one: "1 ≤ cu"
  using c_vals by blast

corollary c_l_nonneg: "0 ≤ cl"
  using c_l_pos by linarith

corollary c_u_pos: "0 < cu"
  using c_u_ge_one by linarith

corollary c_u_nonneg: "0 ≤ cu"
  using c_u_pos by linarith

subsection ‹Balance Predicate›
text "The central balance predicate is defined as a recursive function
over the structure of the tree. For each node it ensures that left and right subtrees
of any given parent node 
differ only by the balancing constants @{term cl} and @{term cu} in terms of @{term rank}
relative to it."
fun balanced :: "('a * 'b) tree ⇒ bool" where
"balanced Leaf = True" |
"balanced (Node l a r) = 
  (max (rank l) (rank r) + cl ≤ rank (Node l a r) ∧ 
   min (rank l) (rank r) + cu ≥ rank (Node l a r) ∧ 
   balanced l ∧ balanced r)"

text ‹It follows immediately that the ranks of the left/right subtrees of a parent node can 
differ at most by a constant, namely @{term  "cδ = cu - cl"} cite‹‹Property 4› in blelloch2022joinable›.›

definition "cδ = cu - cl"

(* Blelloch et al. property 4 *)
lemma dmax:"balanced (Node l a r) ⟹ ¦rank l - rank r¦ ≤ cδ"
  using cδ_def by force

(* useful dest rules *)
lemma balanced_children_explD:
  assumes "balanced (Node l a r)"
  shows "rank l ≤ rank (Node l a r) - cl"
    and "rank r ≤ rank (Node l a r) - cl"
    and "rank (Node l a r) - cu ≤ rank l"
    and "rank (Node l a r) - cu ≤ rank r"
using assms by fastforce+

subsection ‹Relationship between tree rank and height›

text ‹Rank is bounded below by $c_l \cdot \mathit{height~t}$ which is the 
  direction of cite‹‹Property 1› in blelloch2022joinable› needed for the 
  analysis.›

lemma rank_lower_height:"balanced t ⟹ cl * height t ≤ rank t"
proof(induction t)
  case Leaf
  then show ?case
    by (simp add: rule_empty)
next
  case (Node l a r)
  then have "cl * height l ≤ rank l"
    by simp
  moreover have "cl * height r ≤ rank r"
    using Node by simp
  ultimately have "cl * max (height l) (height r) ≤ max (rank l) (rank r)"
    by (simp add: mult_left_mono c_vals max_def)
  then have "cl * max (height l) (height r) + cl ≤ max (rank l) (rank r) + cl"
    by linarith
  also have "cl * max (height l) (height r) + cl = cl * (max (height l) (height r) + 1)"
    by (simp add: distrib_left)
  finally have "cl * (max (height l) (height r) + 1) ≤ max (rank l) (rank r) + cl"
    by linarith
  moreover have "cl * height (Node l a r) ≤ cl * (max (height l) (height r) + 1)"
    using c_vals by simp
  ultimately show ?case
    using Node by auto
qed

corollary rank_lower_height_inverse:"balanced t ⟹ height t ≤ rank t / cl"
  by (metis c_vals rank_lower_height pos_le_divide_eq mult_of_nat_commute)


text ‹Which also implies @{term rank} is never negative.›
corollary rank_pos:"balanced t ⟹ rank t ≥ 0"
  by (metis rank_lower_height c_vals of_nat_0_le_iff mult_nonneg_nonneg order_trans order_less_le)

text ‹The upper bound for @{term rank} holds as well and is added for sake of completeness.›

lemma rank_upper_height:"balanced t ⟹ rank t ≤ cu * height t"
proof(induction t)
  case Leaf
  then show ?case 
    by (simp add: rule_empty)
next
  case (Node l a r)
  then have "rank l ≤ cu * height l"
    by simp
  moreover have "rank r ≤ cu * height r"
    using Node by simp
  ultimately have "max (rank l) (rank r) ≤ cu * max (height l) (height r)"
    using c_vals by (auto simp: max_def intro: mult_left_mono order_trans)
  then have "max (rank l) (rank r) + cu ≤ cu * max (height l) (height r) + cu"
    using c_vals by argo
  also have "... ≤ cu * (max (height l) (height r) + 1)"
    using c_vals by (simp add: distrib_left)
  also have "... ≤ cu * height (Node l a r)"
    using c_vals by auto
  finally show ?case 
    using Node by auto
qed

subsection ‹Relationship between tree rank and size›

text ‹Size grows exponentially in rank cite‹‹Property 2› in blelloch2022joinable›.›
lemma pow_shift:
  assumes "cu ≠ 0"
  shows   "2 * 2 powr ((r - cu) / cu) = 2 powr (r / cu)"
proof -
  have "2 * 2 powr ((r - cu) / cu) = (2 powr 1) * 2 powr ((r - cu) / cu)"
    by simp
  also have "... = 2 powr (1 + (r - cu) / cu)"
    by (simp add: powr_add)
  also have "1 + (r - cu) / cu = r / cu"
    using assms by (simp add: field_simps)
  finally show ?thesis
    by simp
qed

lemma size1_ge_powr_rank:
  "balanced t ⟹ size1 t ≥ 2 powr (rank t / cu)"
proof (induction t)
  case Leaf
  show ?case
    by (simp add: rule_empty)
next
  case (Node l a r)
  then have 
    IHl: "size1 l ≥ 2 powr (rank l / cu)"
  and 
    IHr: "size1 r ≥ 2 powr (rank r / cu)"
    by auto

  have lower_l:
    "size1 l ≥ 2 powr ((rank (Node l a r) - cu) / cu)"
  proof - 
    have "rank (Node l a r) - cu ≤ rank l"
      using Node.prems by auto
    then have "2 powr ((rank (Node l a r) - cu) / cu) ≤ 2 powr (rank l / cu)"
      using c_vals divide_right_mono by fastforce
    then show ?thesis 
       using IHl by argo
  qed

  have lower_r:
    "size1 r ≥ 2 powr ((rank (Node l a r) - cu) / cu)"
  proof -
    have "rank (Node l a r) - cu ≤ rank r"
      using Node.prems by auto
    then have "2 powr ((rank (Node l a r) - cu) / cu) ≤ 2 powr (rank r / cu)"
      using c_vals divide_right_mono by fastforce
    then show ?thesis 
       using IHr by argo
  qed

  have "real (size1 (Node l a r)) = real (size1 l) + real (size1 r)"
    by simp
  also have "... ≥ 2 * (2 powr ((rank (Node l a r) - cu) / cu))"
    using lower_l lower_r by linarith
  finally show ?case 
    using pow_shift by fastforce
qed

corollary rank_le_cu_log_size1:
  assumes "balanced t"
  shows "rank t ≤ cu * log 2 (size1 t)"
proof -
  have "size1 t ≥ 2 powr (rank t / cu)"
    using size1_ge_powr_rank[OF assms] .
  then have "log 2 (size1 t) ≥ log 2 (2 powr (rank t / cu))"
    using le_log_iff by auto
  then show ?thesis
    using c_vals by (simp add: field_simps)
qed

subsection ‹Layers and rank-roots›

text ‹Layer $i$ of a tree collects all nodes whose rank falls into $[i, i+1)$ i.e. 
its rank lies in the unit band $[i, i+1)$ cite‹‹Definition 3› in blelloch2022joinable›.
In essence, ‹layer› discretizes the real valued ‹rank› into finite layers, enabling
the summation arguments used in the asymptotic analysis.›

abbreviation in_band :: "nat ⇒ real ⇒ bool" where
  "in_band i x ≡ real i ≤ x ∧ x < real i + 1"

fun layer :: "('a * 'b) tree ⇒ nat ⇒ (('a * 'b) tree) list" where
"layer Leaf _ = []" |
"layer (Node l x r) i =
     ((if in_band i (rank (Node l x r)) then [(Node l x r)] else []) @
   layer l i
   @ layer r i)"

lemma layer_empty_if_rank_lt:
  assumes Bt: "balanced t"
  assumes lt: "rank t < i"
  shows "layer t i = []"
using Bt lt c_vals by(induction t) auto

text ‹Where @{term layer} recurses the entire tree, @{term rank_roots} stops at the first
node whose rank falls into layer $i$, which is called the ‹rank root› of that layer
cite‹‹Definition 4› in blelloch2022joinable›.
A layer then decomposes as the concatenation of sub-layers rooted at these nodes
(cf. cite‹‹Definition 5› in blelloch2022joinable›).›

fun rank_roots :: "('a * 'b) tree ⇒ nat ⇒ (('a * 'b) tree) list" where
"rank_roots Leaf _ = []" |
"rank_roots (Node l x r) i =
     (if in_band i (rank (Node l x r)) then [(Node l x r)] else
   rank_roots l i
   @ rank_roots r i)"

lemma rank_rank_roots: "r ∈ set(rank_roots t i) ⟹ in_band i (rank r)"
  by(induction t) (auto split: if_splits)

text ‹Balance is inherited.›

lemma balanced_rank_roots: "⟦balanced t; r ∈ set(rank_roots t i)⟧ ⟹ balanced r"
  by(induction t) (auto split: if_splits)

lemma layer_eq_concat_rank_roots:
  "layer T i = concat (map (λr. layer r i) (rank_roots T i))"
proof (induction T)
  case Leaf
  then show ?case by simp
next
  case (Node l x r)
  then show ?case
    by (cases "in_band i (rank (Node l x r))") auto
qed

subsection ‹Size bound on layer and root nodes›

text ‹
  The work bounds for union, intersection and difference will ultimately be obtained by
  decomposing the total work into layers and bounding each layer's contribution
  separately. Two main ingredients make this work:

  ▪ ‹Across› layers, the number of rank roots decays geometrically i.e. layer
    @{term i} has at most $\mathit{size1~t}/2^{i/c_u}$ rank roots. This sharpens
    cite‹‹Lemma 1› in blelloch2022joinable›, where the exponent is $i/(1 + c_u)$.

  ▪ ‹Within› a layer, the number of nodes belonging to each rank root 
    is bounded by the constant $C = 2^{\lceil 1/c_l \rceil} - 1$. This absorbs intra-layer
    overhead into a constant factor  cite‹‹Property 7› in blelloch2022joinable››
text‹
  Both facts are established in this subsection and are consumed separately by the
  main theory, as the decay bound on rank roots and the constant bound per rank root.
  The combined geometric layer bound closing this subsection is not needed
  downstream, but is still shown to make it clear that it is an emergent property
  of balanced trees.
  First, each rank root in layer $i$ has rank at least $i$, so the exponential
  size lower bound applies.›

corollary rank_roots_size_lower:
  assumes "balanced t"
  shows "r ∈ set (rank_roots t i) ⟹ size1 r ≥ 2 powr (i / cu)"
proof -
  fix r
  assume r_in: "r ∈ set (rank_roots t i)"

  have "i ≤ rank r"
    using rank_rank_roots[OF r_in] by auto
  then have "(2::real) powr (real i / cu) ≤ (2::real) powr (rank r / cu)"
    using c_vals by(simp add: divide_simps)
  then show "size1 r ≥ 2 powr (real i / cu)"
    using size1_ge_powr_rank[OF balanced_rank_roots[OF assms r_in]] by linarith
qed

text ‹Further, the combined size of all rank roots in a layer never exceeds the size of the whole tree.›

lemma rank_roots_size_sum:
"sum_list (map size1 (rank_roots t i)) ≤ size1 t"
  by(induction t) auto

text ‹Combining both gives the desired bound on the number of rank roots.›
lemma rank_roots_count_powr:
  assumes "balanced t"
  shows "length (rank_roots t i) * 2 powr (i / cu) ≤ size1 t"
proof -
  have "length (rank_roots t i) * 2 powr (i / cu)
      = sum_list (map (λ_. 2 powr (i / cu)) (rank_roots t i))"
    by (simp add: sum_list_triv)
  also have "... ≤ sum_list (map (real o size1) (rank_roots t i))"
    using rank_roots_size_lower[OF assms] by (simp add: sum_list_mono)
  also have "... ≤ real (size1 t)"
    using rank_roots_size_sum[of t i]
    by (metis list.map_comp of_nat_le_iff sum_list_of_nat)
  finally show ?thesis .
qed

text ‹Restated via division.›

corollary rank_roots_count_powr_le:
  assumes "balanced t"
  shows "length (rank_roots t i)
        ≤ size1 t / (2 powr (i / cu))"
proof -
  have "length (rank_roots t i) * 2 powr (i / cu) ≤ size1 t"
    using rank_roots_count_powr[OF assms] .
  moreover have "0 < 2 powr (real i / cu)"
    by simp
  ultimately show ?thesis
    by (simp add: pos_le_divide_eq)
qed

text ‹
  A balanced tree whose rank overshoots layer $i$ by at most $d$ has at most
  $2^{\left\lceil \frac{d}{c_l} \right\rceil} - 1$ nodes in that layer.
  The budget $d$ shrinks by @{term cl} at every level, giving the exponential bound.
  The parameter $d$ is generalized beyond the natural choice $d = 1$
  because the induction step requires applying the lemma to children 
  with a reduced budget.›

lemma layer_len_budget:
  fixes d :: real
  assumes "balanced t"
  assumes "rank t < real i + d"
  assumes "0 ≤ d"
  shows "length (layer t i) ≤ 2^(nat (ceiling (d / cl))) - 1"
using assms
proof (induction t arbitrary: d)
  case Leaf
  then show ?case
    by simp
next
  case (Node l x r)

  let ?t = "Node l x r"
  let ?m = "nat (ceiling (d / cl))"

  have Bl: "balanced l" and Br: "balanced r"
    using Node by auto

  show ?case
  proof (cases "?m = 0")
    case True
    then show ?thesis
      using layer_empty_if_rank_lt[OF Node.prems(1)] Node.prems(2,3) c_l_pos 
      by (auto simp add: divide_le_0_iff)
  next
    case False
    have mpos: "?m ≥ 1"
      using False by linarith
    have child_cut_l: "rank l < real i + (d - cl)"
      using Node.prems(1,2) by force
    have child_cut_r: "rank r < real i + (d - cl)"
      using Node.prems(1,2) by force
    have ceil_step:
      "nat (ceiling ((d - cl) / cl)) ≤ ?m - 1"
      using c_vals by (auto simp: diff_divide_distrib)

    have "length (layer ?t i) ≤ 1 + length (layer l i) + length (layer r i)"
      by simp
    also have "... ≤ 1 + (2^(?m - 1) - 1) + (2^(?m - 1) - 1)"
    proof - 
      have "length (layer l i) ≤ 2^(?m - 1) - 1"
      proof -
        have "length (layer l i) ≤ 2^(nat (ceiling ((d - cl) / cl))) - 1"
          using Bl Node.IH(1) child_cut_l layer_empty_if_rank_lt by force
        also have "... ≤ 2^(?m - 1) - 1"
          using ceil_step by (metis diff_le_mono le_add_same_cancel1 one_add_one power_increasing zero_le_one)
        finally show ?thesis .
      qed

      moreover have "length (layer r i) ≤ 2^(?m - 1) - 1"
      proof -
        have "length (layer r i) ≤ 2^(nat (ceiling ((d - cl) / cl))) - 1"
          using Br Node.IH(2) child_cut_r layer_empty_if_rank_lt by force
        also have "... ≤ 2^(?m - 1) - 1"
          using ceil_step by (metis diff_le_mono le_add_same_cancel1 one_add_one power_increasing zero_le_one)
        finally show ?thesis .
      qed
      
      ultimately show ?thesis by linarith
    qed
    also have "... = 2 * 2^(?m - 1) - 1"
        using one_le_power by fastforce
    also have "... = 2^?m - 1"
      using mpos False by (metis power_eq_if)
    finally show ?thesis .
  qed
qed

text ‹Instantiating with $d=1$ begets the actual constant.›
definition C :: nat where "C = 2^(nat (ceiling (1 / cl))) - 1"

lemma C_pos: "real C ≥ 0"
  by force

corollary layer_bound_const:
  assumes "balanced t" "rank t < real i + 1"
  shows "length (layer t i) ≤ C"
  using layer_len_budget[OF assms(1) assms(2)] C_def by simp


text ‹
  Thus the total number of nodes in layer $i$ is at most @{term C} times the
  number of rank roots.›

lemma length_concat_map_le_sum:
  "length (concat (map f xs)) ≤ sum_list (map (λx. length (f x)) xs)"
  by (induction xs) auto

lemma layer_bound_via_rank_roots:
  assumes BT: "balanced T"
  shows "length (layer T i)
       ≤ C * length (rank_roots T i)"
proof -
  have decomp: "layer T i = concat (map (λr. layer r i) (rank_roots T i))"
    using layer_eq_concat_rank_roots .
  then have "length (layer T i)
      = length (concat (map (λr. layer r i) (rank_roots T i)))"
    by simp
  also have "... ≤ sum_list (map (λr. length (layer r i)) (rank_roots T i))"
    using length_concat_map_le_sum by metis
  also have "... ≤ sum_list (map (λ_. C) (rank_roots T i))"
    using balanced_rank_roots[OF BT] rank_rank_roots layer_bound_const C_def by (meson sum_list_mono)
  also have "sum_list (map (λ_. C) (rank_roots T i)) = C * (length (rank_roots T i))"
    by (simp add: sum_list_triv)
  finally show ?thesis 
   by simp
qed


text ‹Together, this yields the fully combined geometric layer bound.› 

corollary layer_bound_geometric:
  assumes "balanced t"
  shows "(length (layer t i)) ≤ C * (size1 t / (2 powr (i / cu)))"
proof -
  have
    "length (layer t i) ≤ C * length (rank_roots t i)"
    using layer_bound_via_rank_roots[OF assms] .
  also have
    " ... ≤ C * (size1 t / (2 powr (i / cu)))"
    using rank_roots_count_powr_le[OF assms] by (metis mult_left_mono of_nat_0_le_iff of_nat_mult)
  finally show ?thesis 
    by simp
qed

end

subsection ‹Balanced schemes›

text ‹
  A @{text BalancedTree} is a @{text BalancedShape} together with an invariant
  that implies balance. This records the balancing rule of 
  cite‹blelloch2022joinable›.
›
locale BalancedTree = BalancedShape rank cl cu
  for rank :: "('a × 'b) tree ⇒ real"
  and cl cu :: real
  +
  fixes inv :: "('a × 'b) tree ⇒ bool"
  assumes rule_bal: "inv t ⟹ balanced t"
begin

declare rule_bal[intro]

corollary rank_pos'[simp]: "inv t ⟹ 0 ≤ rank t"
  using rank_pos rule_bal by blast

corollary dmax': "inv (Node l a r) ⟹ ¦rank l - rank r¦ ≤ cδ"
  using dmax rule_bal by blast

corollary rank_le_cu_log_size1': "inv t ⟹ rank t ≤ cu * log 2 (size1 t)"
  using rank_le_cu_log_size1 rule_bal by blast

corollary rank_lower_height_inverse': "inv t ⟹ height t ≤ rank t / cl"
  using rank_lower_height_inverse rule_bal by blast

corollary inv_children_explD:
  assumes "inv (Node l a r)"
  shows "rank l ≤ rank (Node l a r) - cl"
    and "rank r ≤ rank (Node l a r) - cl"
    and "rank (Node l a r) - cu ≤ rank l"
    and "rank (Node l a r) - cu ≤ rank r"
  using balanced_children_explD[OF rule_bal[OF assms]] by blast+
end

end