Section BTree.
Parameter d : nat.
Hypothesis d_min : (1 <= d)%nat.
(* could use bin_search here instead *)
Inductive find_result := Found (i : nat) | NotFound (i : nat).
Fixpoint find_index tgt keys :=
match keys with
| [] => NotFound O
| k :: rest => if tgt <? k then NotFound O
else if tgt =? k then Found O
else match find_index tgt rest with
| Found i => Found (S i)
| NotFound i => NotFound (S i)
end
end.
Fixpoint find_path tgt t :=
let search_nth := fix f children i :=
match children, i with
| [], _ => None
| c :: _, O => find_path tgt c
| _ :: rest, S i => f rest i
end in
match t with BNode keys children =>
match find_index tgt keys with
| NotFound i => match search_nth children i with
| Some p => Some (i :: p)
| None => None
end
| Found i => Some [i]
end
end.
Definition search tgt t := match find_path tgt t with Some _ => true | None => false end.
Definition replace {A} i (x : A) l := firstn i l ++ x :: skipn (S i) l.
Fixpoint find_leaf tgt t :=
let find_nth := fix f children i :=
match children, i with
| [], _ => Some []
| c :: _, O => find_leaf tgt c
| _ :: rest, S i => f rest i
end in
match t with BNode keys children =>
match find_index tgt keys with
| NotFound i => match find_nth children i with
| Some p => Some (i :: p)
| None => None
end
| Found _ => None
end
end.
Definition split parent i :=
match parent with BNode keys children =>
match nth_error children i with
| Some (BNode keys' children') =>
let mid := (length keys' / 2)%nat in
BNode (insert i (nth mid keys' 0) keys)
(insert i (BNode (firstn mid keys') (firstn (S mid) children'))
(replace i (BNode (skipn (S mid) keys') (skipn (S mid) children')) children))
| None => parent
end
end.
Fixpoint insert_at tgt t path :=
match path with
| [] => None
| [i] => match t with BNode keys children => Some (BNode (insert i tgt keys) children)
end
| [i; j] => match t with BNode keys children =>
match nth_error children i with
| Some (BNode keys' []) =>
let t' := BNode keys (replace i (BNode (insert j tgt keys') []) children) in
if lt_dec (length keys') (2 * d) then Some t'
else Some (split t' i)
| _ => None
end end
| i :: rest => match t with BNode keys children =>
match nth_error children i with
| Some c => match insert_at tgt c rest with
| Some (BNode keys' children') =>
let t' := BNode keys (replace i (BNode keys' children') children) in
if le_dec (length keys') (2 * d) then Some t'
else Some (split t' i)
| None => None
end
| _ => None
end end
end.
Definition b_insert tgt t :=
match find_leaf tgt t with
| Some p => match insert_at tgt t p with
| Some t' => match t' with BNode keys children =>
if le_dec (length keys) (2 * d) then Some t'
else Some (split (BNode [] [t']) 0)
end
| None => None
end
| None => None
end.
Definition remove_at {A} i (l : list A) := firstn i l ++ skipn (S i) l.
Fixpoint del_max t :=
let del_last := fix f children :=
match children with
| [] => (BNode [] [], 0, None)
| [c] => del_max c
| _ :: rest => f rest
end in
match t with
| BNode keys [] => (BNode (removelast keys) [], last keys 0,
if le_dec (length keys) d then Some [] else None)
| BNode keys children => let '(t', v, p) := del_last children in
(BNode keys (replace (length children - 1) t' children),
v, match p with Some p => Some ((length children - 1)%nat :: p)
| None => None end)
end.
Definition merge t i j := match t with BNode keys children =>
let sep := nth i keys 0 in
match nth i children (BNode [] []) with BNode keysl childrenl =>
match nth j children (BNode [] []) with BNode keysr childrenr =>
let l' := BNode (keysl ++ sep :: keysr) (childrenl ++ childrenr) in
if eq_dec (length keys) (S O) then l'
else BNode (remove_at i keys) (remove_at j (replace i l' children))
end end end.
Definition tryleft t i :=
match i with
| O => merge t O (S O)
| S i' => match t with BNode keys children =>
match nth_error children i' with
| Some (BNode keysl childrenl) =>
if lt_dec d (length keysl) then
let sep := nth i' keys 0 in
match nth_error children i with
| Some (BNode keys' children') =>
let c' := BNode (sep :: keys')
(match childrenl with [] => children' |
_ => last childrenl (BNode [] []) :: children' end) in
BNode (replace i' (last keysl 0) keys)
(replace i' (BNode (removelast keysl) (removelast childrenl))
(replace i c' children))
| None => t
end
else merge t i' i
| None => t
end end
end.
Definition rebalance_local t i := match t with BNode keys children =>
match nth_error children (S i) with
| Some (BNode (k :: keysr) childrenr) =>
if lt_dec d (length keysr) then
let sep := nth i keys 0 in
match nth_error children i with
| Some (BNode keys' children') =>
let c' := BNode (keys' ++ [sep])
(match childrenr with [] => children' |
c :: _ => c :: children' end) in
BNode (replace i k keys)
(replace (S i) (BNode keysr (tl childrenr)) (replace i c' children))
| None => t
end
else tryleft t i
| _ => tryleft t i
end end.
Fixpoint rebalance t path {struct path} :=
match path with
| [] => t
| [i] => rebalance_local t i
| i :: rest => match t with BNode keys children =>
match nth_error children i with
| Some c => match rebalance c rest with BNode keys' children' =>
let t' := BNode keys (replace i (BNode keys' children') children) in
if lt_dec (length keys') d then rebalance_local t' i else t' end
| None => t
end
end end.
Fixpoint delete_at t path :=
let delete_nth := fix f children i path :=
match children, i with
| [], _ => []
| c :: rest, O => delete_at c path :: rest
| c :: rest, S i => c :: f rest i path
end in
match path with
| [] => t
| [i] => let '(t', p') := match t with
| BNode keys [] => (BNode (remove_at i keys) [],
if le_dec (length keys) d then Some [] else None)
| BNode keys children =>
match nth_error children i with
| Some c => let '(c', v, p) := del_max c in
(BNode (replace i v keys) (replace i c' children), p)
| None => (t, None)
end end in
match p' with Some p' => rebalance t' p' | None => t' end
| i :: rest => match t with BNode keys children =>
BNode keys (delete_nth children i rest) end
end.
Definition delete tgt t := option_map (delete_at t) (find_path tgt t).
Fixpoint contents t :=
let contents_list := fix f lt lv :=
match lt, lv with
| [], _ => lv
| t :: lt, v :: lv => contents t ++ v :: f lt lv
| _, _ => []
end in
match t with BNode keys children => contents_list children keys end.
Fixpoint wf_btree_aux t :=
let wf_btrees := fix f lt :=
match lt with
| [] => True
| t :: rest => wf_btree_aux t /\ f rest
end in
match t with BNode keys children =>
(d <= length keys <= 2 * d)%nat /\
(children = [] \/ length children = S (length keys) /\ wf_btrees children)
end.
Definition wf_btree t := match t with BNode keys children =>
(length keys <= 2 * d /\ (children = [] \/
(2 <= length children /\ length children = S (length keys) /\
Forall wf_btree_aux children)))%nat end.
Definition b_empty := BNode [] [].
Fixpoint leaf_path t p {struct p} :=
match p with
| [] => match t with BNode keys children => children = [] end
| i :: rest => match t with BNode keys children =>
match nth_error children i with
| Some c => leaf_path c rest
| None => False
end end
end.
Arguments Nat.div x y : simpl never.
Arguments skipn A n l : simpl nomatch.
Arguments firstn A n l : simpl nomatch.
Fixpoint wf_btree_over t :=
match t with BNode keys children =>
(length keys = 2 * d + 1)%nat /\
(children = [] \/ length children = S (length keys) /\ Forall wf_btree_aux children)
end.

Lemma In_Znth : forall A (l : list A) x d,
In x l ->
exists i, 0 <= i < Zlength l /\ Znth i l d = x.