Module ListUtils

Require Import Coqlib Tactics.
Require Import Psatz.
Require Import Permutation.

Section LISTNTH.

  Lemma list_nth_z_succeeds:
    forall {A} (l: list A) n,
      0 <= n < Z.of_nat (length l) ->
      exists x, list_nth_z l n = Some x.
Proof.
    induction l; simpl; intros; eauto.
    lia.
    destr. eauto.
    edestruct IHl; eauto. split. unfold Z.pred. lia.
    unfold Z.pred. lia.
  Qed.

  Lemma lnr_list_nth_z_diff:
    forall {A} (l: list A) (LNR: list_norepet l)
      n1 n2 (d: n1 <> n2)
      x1 x2
      (NTH1: list_nth_z l n1 = Some x1)
      (NTH2: list_nth_z l n2 = Some x2),
      x1 <> x2.
Proof.
    induction 1; simpl; intros; eauto. easy.
    - intro; subst. destr_in NTH1; destr_in NTH2.
      inv NTH1. eapply list_nth_z_in in NTH2. destr.
      inv NTH2. eapply list_nth_z_in in NTH1. destr.
      specialize (IHLNR (Z.pred n1) (Z.pred n2)). trim IHLNR. unfold Z.pred. lia.
      exploit IHLNR; eauto.
  Qed.

End LISTNTH.

Section FIND.

  Fixpoint find {A} (eq_dec: forall (a b: A), {a = b} + {a <> b}) x l :=
    match l with
      nil => None
    | a :: r => if eq_dec a x then Some O else option_map S (find eq_dec x r)
    end.

  Lemma find_in:
    forall {A} eq_dec l (x:A),
      In x l ->
      exists n, find eq_dec x l = Some n /\
           (n < length l)%nat.
Proof.
    induction l; simpl; intros; eauto. easy.
    destr. eexists; split; eauto. lia.
    edestruct IHl; eauto. des H. eauto.
    destruct H0. rewrite H0. simpl. eexists; split; eauto. lia.
  Qed.

  Lemma find_diff_diff:
    forall {A} E (a1 a2:A) (d: a1 <> a2) l n1 n2,
      find E a1 l = Some n1 ->
      find E a2 l = Some n2 ->
      n1 <> n2.
Proof.
    induction l; simpl; intros; eauto. congruence.
    destr_in H. inv H.
    destr_in H0. des (find E a2 l).
    des (find E a1 l). inv H. destr_in H0.
    des (find E a2 l). inv H0. specialize (IHl _ _ eq_refl eq_refl). congruence.
  Qed.

  Lemma find_some_range:
    forall {A} E (x:A) l n,
      find E x l = Some n ->
      0 <= Z.of_nat n < Z.of_nat (length l).
Proof.
    induction l; simpl; intros; eauto. congruence.
    destr_in H. inv H. lia.
    des (find E x l). specialize (IHl _ eq_refl). inv H. lia.
  Qed.

End FIND.

Section CONCAT.

  Definition concat {A} := fold_right (@app A) nil.

  Lemma in_concat:
    forall {A} l (x:A),
      In x (concat l) <-> exists l', In l' l /\ In x l'.
Proof.
    induction l; simpl; intros; eauto.
    - split; intros. easy. dex; destr.
    - rewrite in_app.
      split; intros.
      des H. eexists; split. left; eauto. auto.
      rewrite IHl in i. dex; repeat destr_and.
      eexists; split. right; eauto. auto.
      dex; destr. des H. des o. rewrite IHl. right; eexists; split; eauto.
  Qed.

End CONCAT.

Section REMOVEDBL.

  Fixpoint remove_dbl {A} (D: forall (a b: A), {a = b} + {a <> b}) l :=
    match l with
      nil => nil
    | a::r => if in_dec D a r then remove_dbl D r
             else a :: remove_dbl D r
    end.

  Lemma remove_dbl_In:
    forall {A} D (l: list A) a,
      (In a l <-> In a (remove_dbl D l)).
Proof.
    induction l; simpl; split; intros; auto.
    rewrite IHl in H. destr. des H. eauto.
    apply IHl; auto. destr_in H; eauto; destr;
                       rewrite IHl; eauto.
  Qed.

  Lemma remove_dbl_lnr:
    forall {A} D (l: list A),
      list_norepet (remove_dbl D l).
Proof.
    induction l; simpl; destr. constructor.
    constructor; auto.
    intro IN. apply n.
    eapply remove_dbl_In; eauto.
  Qed.

  Lemma lnr_remove:
    forall l,
      list_norepet l ->
      remove_dbl eq_nat_dec l = l.
Proof.
    induction 1; simpl; intros; eauto.
    destr.
  Qed.

  Lemma length_remove:
    forall l,
      (length (remove_dbl eq_nat_dec l) <= length l)%nat.
Proof.
    induction l; simpl; intros; eauto.
    destr. lia. lia.
  Qed.

  
End REMOVEDBL.


Section SUM.

  Definition sum l := fold_right Peano.plus O l.
  
  Lemma sum_permut:
    forall l l' (P: Permutation l l'),
      sum l = sum l'.
Proof.
    induction 1; simpl; intros; eauto.
    lia.
    lia.
  Qed.

  Lemma sum_app:
    forall l l' ,
      sum (l ++ l') = (sum l + sum l')%nat.
Proof.
    induction l; simpl; intros; eauto.
    rewrite IHl. lia.
  Qed.

  Lemma length_concat:
    forall {A} (l: list (list A)),
      length (concat l) = sum (map (@length A) l).
Proof.
    unfold concat. induction l; simpl; intros; eauto.
    rewrite app_length. auto.
  Qed.

End SUM.

Section PERMUTATIONS.

  Lemma perm_concat:
    forall {A} (l l': list (list A)) (P: list_forall2 (@Permutation A) l l'),
      Permutation (concat l) (concat l').
Proof.
    induction 1; simpl; intros; eauto. apply Permutation_app; auto.
  Qed.

  Lemma partition_l:
    forall {A} (l: list A) f ll lr,
      partition f l = (ll,lr) ->
      ll = filter_map (fun b => if f b then Some b else None) l.
Proof.
    induction l; simpl; intros; eauto. inv H; auto.
    repeat destr_in H; inv H; f_equal; eauto.
  Qed.
  
End PERMUTATIONS.

Lemma partition_app_perm:
  forall {A: Type} f (l: list A) a b,
    partition f l = (a,b) ->
    Permutation l (a ++ b).
Proof.
  induction l; simpl; intros; eauto.
  - inv H. simpl; constructor.
  - repeat destr_in H.
    + inv H. specialize (IHl _ _ eq_refl).
      simpl. constructor. auto.
    + inv H. specialize (IHl _ _ eq_refl).
      eapply perm_trans. apply perm_skip. eauto.
      apply Permutation_middle.
Qed.




Lemma lnr_nin:
  forall {A: Type} (f: A -> bool) l a b,
    partition f l = (a, b) ->
    forall x, ~ In x l ->
         ~ In x a /\ ~ In x b.
Proof.
  induction l; simpl; intros; eauto.
  - inv H; destr.
  - repeat destr_in H; inv H; specialize (IHl _ _ eq_refl); destr.
    trim (IHl x). destr. split; destr.
    trim (IHl x). destr. split; destr.
Qed.

Lemma lnr_partition:
  forall {A: Type} (f: A -> bool) l a b,
    partition f l = (a, b) ->
    list_norepet l ->
    list_norepet a /\ list_norepet b.
Proof.
  induction l; simpl; intros; eauto.
  - inv H. split; constructor.
  - repeat destr_in H; inv H; specialize (IHl _ _ eq_refl); inv H0; split; trim IHl; auto.
    constructor; auto.
    eapply lnr_nin in H2; eauto. destr.
    destr. destr. destr.
    constructor; auto.
    eapply lnr_nin in H2; eauto. destr. destr.
Qed.

Lemma in_partition_l:
  forall {A: Type} (f: A -> bool) l a b,
    partition f l = (a, b) ->
    forall x,
      In x a <->
      In x l /\ f x = true.
Proof.
  induction l; simpl; intros; eauto.
  - inv H. simpl; tauto.
  - repeat destr_in H; inv H; specialize (IHl _ _ eq_refl).
    simpl.
    + split; intros.
      destruct H; auto. destr. apply IHl in H. destr.
      destruct H. destruct H; auto. right; rewrite IHl. destr.
    + split; intros.
      apply IHl in H. destr.
      destruct H. rewrite IHl. destruct H; auto. congruence.
Qed.




Lemma partition_filter:
  forall {A} (f: A -> bool) l l1 l2,
    partition f l = (l1,l2) ->
    l1 = filter f l.
Proof.
  induction l; simpl; intros.
  inv H; auto.
  destr_in H. specialize (IHl _ _ eq_refl). subst.
  destr_in H; inv H; auto.
Qed.

Lemma partition_filter_r:
  forall {A} (f: A -> bool) l l1 l2,
    partition f l = (l1,l2) ->
    l2 = filter (fun b => negb (f b)) l.
Proof.
  induction l; simpl; intros.
  inv H; auto.
  destr_in H. specialize (IHl _ _ eq_refl). subst.
  destr_in H; inv H; auto.
Qed.

Lemma filter_app:
  forall {A} f (l1 l2: list A),
    filter f (l1 ++ l2) = filter f l1 ++ filter f l2.
Proof.
  induction l1; simpl; intros; eauto.
  rewrite IHl1. destr.
Qed.

Arguments pair {A B} _ _.
Arguments fst {A B} _.
Arguments snd {A B} _.

Lemma lnr_nin':
  forall {A B: Type} (f: A * B -> bool) (l: list (A*B)) a b,
    partition f l = (a, b) ->
    forall x, ~ In x (map fst l) ->
         ~ In x (map fst a) /\ ~ In x (map fst b).
Proof.
  induction l; simpl; intros; eauto.
  - inv H; destr.
  - repeat destr_in H; inv H; specialize (IHl _ _ eq_refl); destr.
    trim (IHl x). destr. split; destr.
    trim (IHl x). destr. split; destr.
Qed.


Lemma lnr_partition':
  forall {A B: Type} (f: A * B -> bool) (l: list (A * B)) a b,
    partition f l = (a, b) ->
    list_norepet (map fst l) ->
    list_norepet (map fst a) /\ list_norepet (map fst b).
Proof.
  induction l; simpl; intros; eauto.
  - inv H. split; constructor.
  - repeat destr_in H; inv H; specialize (IHl _ _ eq_refl); inv H0; split; trim IHl; auto; simpl.
    constructor; auto.
    eapply lnr_nin' in H2; eauto. destr.
    destr. destr. destr.
    constructor; auto.
    eapply lnr_nin' in H2; eauto. destr. destr.
Qed.


Lemma permut_filter_map:
  forall {A B} (f: A -> option B) (l l': list A),
    Permutation l l' ->
    Permutation (filter_map f l) (filter_map f l').
Proof.
  induction 1; simpl; intros; eauto.
  - destr; auto.
  - repeat destr; auto. apply perm_swap.
  - rewrite IHPermutation1; auto.
Qed.

Lemma in_partition_r:
  forall {A: Type} (f: A -> bool) l a b,
    partition f l = (a, b) ->
    forall x,
      In x b <->
      In x l /\ f x = false.
Proof.
  induction l; simpl; intros; eauto.
  - inv H. simpl; tauto.
  - repeat destr_in H; inv H; specialize (IHl _ _ eq_refl).
    simpl.
    + split; intros.
      apply IHl in H. destr.
      destruct H. rewrite IHl. destr.
    + split; intros.
      destruct H; auto. destr. apply IHl in H. destr.
      destruct H. destr. destruct H; auto. right; rewrite IHl. destr.
Qed.



Section FILTER_MAP.

  Lemma length_filters:
    forall {A} f (l: list A),
      length l = length (filter f l ++ filter (fun x => negb (f x)) l).
Proof.
    intros.
    des (partition f l).
    exploit (partition_app_perm (A:=A)). eauto.
    exploit (partition_filter (A:= A)). eauto. intro; subst.
    exploit (partition_filter_r (A:=A)). eauto. intro; subst.
    intros.
    apply Permutation_length. auto.
  Qed.

  Remark filter_charact:
    forall (A: Type) (f: A -> bool) x l,
      In x (List.filter f l) <-> In x l /\ f x = true.
Proof.
    induction l; simpl. tauto.
    destruct (f a) eqn:?.
    simpl. rewrite IHl. intuition congruence.
    intuition congruence.
  Qed.

  
  Lemma filter_map_ext:
    forall {A B} (f f': A -> option B) l
      (EXT: forall x, In x l -> f x = f' x),
      filter_map f l = filter_map f' l.
Proof.
    induction l; simpl; intros; eauto.
    rewrite EXT; auto.
    trim IHl. intros; apply EXT; auto. destr.
  Qed.

  Lemma map_filter_map:
    forall {A B C D}
      (f: B -> C) (g: A -> option B)
      (f': A -> D) (g': D -> option C)
      (EQ: forall a b, g a = Some b -> exists c, g' (f' a) = Some c /\ f b = c)
      (EQ': forall a c, g' (f' a) = Some c -> exists b, g a = Some b)
      (l: list A),
      map f (filter_map g l) = filter_map g' (map f' l).
Proof.
    induction l; simpl; intros; eauto.
    repeat destr.
    - exploit ((EQ a b)). eauto. rewrite Heqo0. intros (c0 & ? & ?). inv H. f_equal; auto.
    - apply EQ in Heqo. dex; destr.
    - apply EQ' in Heqo0. dex; destr.
  Qed.

  Lemma map_filter_map':
    forall {A B C} (g: A -> option B) (f: B -> C) l,
      map f (filter_map g l) = filter_map (fun x => match g x with
                                                   Some gx => Some (f gx)
                                                 | None => None
                                                 end) l.
Proof.
    induction l; simpl; intros; eauto.
    destr.
  Qed.

  Lemma filter_map_app:
    forall {A B:Type} (f: A -> option B) a b,
      filter_map f (a ++ b) = filter_map f a ++ filter_map f b.
Proof.
    induction a; simpl; intros; eauto.
    destr.
  Qed.

  Lemma in_map_filter:
    forall {A B} b0 (f: A -> option B) l,
      In b0 (filter_map f l) <->
      (exists a0, In a0 l /\ f a0 = Some b0).
Proof.
    induction l; simpl; intros. split; intros; dex; easy.
    des (f a); eauto.
    - split; intros.
      + destruct H; subst; eauto.
        rewrite IHl in H. dex.
        exists a0; destr.
      + dex. destruct H as [[C|C] D].
        subst; auto. left; congruence. right; rewrite IHl.
        exists a0; destr.
    - rewrite IHl.
      split; intros [a0 [C D]]; exists a0; destr.
  Qed.

  Lemma list_norepet_filter_map:
    forall {A B} (f: A -> option B)
      (MAP: forall a1 a2 b1 b2,
          f a1 = Some b1 ->
          f a2 = Some b2 ->
          a1 <> a2 -> b1 <> b2
      ),
    forall (l: list A), list_norepet l ->
                   list_norepet (filter_map f l).
Proof.
    induction 2; simpl; intros; eauto. constructor.
    destr.
    constructor; auto.
    intro IN.
    rewrite in_map_filter in IN.
    destruct IN as (a0 & IN & EQ).
    eapply MAP. apply Heqo. apply EQ. congruence. auto.
  Qed.

  Lemma in_fold_left_filter:
    forall {A B} (f: A -> B -> bool) l lacc x ,
      In x (fold_left (fun acc b => filter (f b) acc) l lacc) ->
      In x lacc /\ Forall (fun y => f y x = true) l.
Proof.
    induction l; simpl; intros; eauto.
    apply IHl in H. rewrite filter_In in H.
    des H. des a0. repeat constructor; auto.
  Qed.

  Lemma filter_filter:
    forall {A} f f' (l: list A),
      filter f (filter f' l) = filter (fun x => f x && f' x) l.
Proof.
    induction l; simpl; intros; eauto.
    destr. destr. rewrite andb_false_r. auto.
  Qed.

  Lemma filter_ext:
    forall {A} f f' (EXT: forall x, f x = f' x) (l: list A),
      (filter f l) = (filter f' l).
Proof.
    induction l; simpl; intros; eauto. rewrite EXT; destr.
  Qed.

  Lemma filter_true:
    forall {A} (l: list A),
      filter (fun x => true) l = l.
Proof.
    induction l; destr.
  Qed.

  
End FILTER_MAP.