(*** Library of finite maps indexed by positive binary integers ***)

Require Export ZArith.
Require Import Bool.

Lemma peq: forall (n1 n2: positive), {n1=n2}+{n1<>n2}.
Proof.
  decide equality.
Defined.

Section MAPS.

Variable A: Set.
Variable Aeq: forall (x y: A), {x=y} + {x<>y}.

Inductive map: Set :=
  | Leaf: map
  | Node: map -> option A -> map -> map.

Lemma map_eq:
  forall (m1 m2: map), {m1=m2} + {m1<>m2}.
Proof.
  assert (forall (o1 o2: option A), {o1=o2}+{o1<>o2}). decide equality.
  decide equality. 
Defined.

Definition empty := Leaf.

Fixpoint get (m: map) (i: positive) {struct i} : option A :=
  match m with
  | Leaf => None
  | Node l o r =>
      match i with
      | xH => o
      | xO ii => get l ii
      | xI ii => get r ii
      end
  end.

Fixpoint set (m: map) (i: positive) (v: A) {struct i} : map :=
  match m with
  | Leaf =>
      match i with
      | xH => Node Leaf (Some v) Leaf
      | xO ii => Node (set Leaf ii v) None Leaf
      | xI ii => Node Leaf None (set Leaf ii v)
      end
  | Node l o r =>
      match i with
      | xH => Node l (Some v) r
      | xO ii => Node (set l ii v) o r
      | xI ii => Node l o (set r ii v)
      end
  end.

Lemma get_empty:
  forall (n: positive), get empty n = None.
Proof.
  induction n; simpl; auto.
Qed.

Lemma get_set_same: 
  forall (n: positive) (m: map) (x: A),
  get (set m n x) n = Some x.
Proof.
  induction n; destruct m; intros; simpl; auto.
Qed.  

Remark get_leaf:
  forall (n: positive), get Leaf n = None.
Proof get_empty.

Lemma get_set_other:
  forall (n1 n2: positive) (m: map) (x: A),
  n1 <> n2 -> get (set m n1 x) n2 = get m n2.
Proof.
  induction n1; intros; destruct n2; destruct m; simpl;
  try rewrite <- (get_leaf n2); auto; try apply IHn1; try congruence.
Qed.

Lemma get_set:
  forall (m: map) (n1 n2: positive) (x: A),
  get (set m n1 x) n2 =
  if peq n1 n2 then Some x else get m n2.
Proof.
  intros. case (peq n1 n2); intro.
  subst n2. apply get_set_same. 
  apply get_set_other; auto.
Qed.

End MAPS. 

Notation "a # b" := (get _ a b) (at level 1).
Notation "a # b <- c" := (set _ a b c) (at level 1, b at next level).
 
Section FORALL2.

Variables A B: Set.
Variable pred: A -> option B -> bool.

Definition pred_opt (x: option A) (y: option B) :=
  match x with
  | None => true
  | Some z => pred z y
  end.

Fixpoint map_forall2 (m1: map A) (m2: map B) {struct m1}: bool :=
  match m1, m2 with
  | Leaf, _ => true
  | Node l1 o1 r1, Leaf =>
      map_forall2 l1 (Leaf B) && pred_opt o1 None && map_forall2 r1 (Leaf B)
  | Node l1 o1 r1, Node l2 o2 r2 =>
      map_forall2 l1 l2 && pred_opt o1 o2 && map_forall2 r1 r2
  end.

Lemma map_forall2_correct:
  forall m1 m2 i x,
  map_forall2 m1 m2 = true -> 
  m1#i = Some x -> pred x m2#i = true. 
Proof.
  induction m1; simpl; intros until x.
  rewrite get_leaf. congruence.
  destruct m2;
  intro C; elim (andb_prop _ _ C); intros C1 C2;
  elim (andb_prop _ _ C1); intros C3 C4; clear C; clear C1;
  destruct i; simpl.
  rewrite <- (get_leaf B i). auto.
  rewrite <- (get_leaf B i). auto.
  intro. subst o. exact C4.
  auto.
  auto.
  intro. subst o. exact C4.
Qed.

Lemma map_forall2_complete:
  forall m1 m2,
  (forall i x, m1#i = Some x -> pred x m2#i = true) ->
  map_forall2 m1 m2 = true.
Proof.
  induction m1; intros; simpl.
  auto.
  destruct m2.
  rewrite IHm1_1. rewrite IHm1_2. rewrite <- (get_leaf B 1). 
  unfold pred_opt. destruct o. rewrite H. reflexivity. reflexivity. 
  reflexivity.
  intros. rewrite get_leaf. rewrite <- (get_leaf B (xI i)). apply H. assumption.
  intros. rewrite get_leaf. rewrite <- (get_leaf B (xO i)). apply H. assumption.
  rewrite IHm1_1. rewrite IHm1_2.  
  unfold pred_opt. destruct o. change o0 with (get B (Node B m2_1 o0 m2_2) 1).
  rewrite H. reflexivity. reflexivity. 
  reflexivity.
  intros. change (m2_2#i) with (get B (Node B m2_1 o0 m2_2) (xI i)). apply H. assumption.
  intros. change (m2_1#i) with (get B (Node B m2_1 o0 m2_2) (xO i)). apply H. assumption.
Qed.

End FORALL2.

