equations_example.v

Require Import Equations.Equations.
Require Import Lia.
Require Import List.
Import ListNotations.

Fixpoint count_up n :=
  match n with
  | 0 => []
  | S m => 0 :: map S (count_up m)
  end.

Lemma count_up_lt : forall n k, In k (count_up n) -> k < n.
Proof.
  induction n; intros.
  - destruct H.
  - destruct H.
    + lia.
    + destruct k; rewrite in_map_iff in H;
      destruct H as [x [Hx1 Hx2]].
      * lia.
      * pose (IHn x Hx2); lia.
Qed.

Fixpoint add_pfs{X}(P : X -> Prop)(xs : list X) : (forall x, In x xs -> P x) -> list { x : X & P x } :=
  match xs with
  | [] => fun _ => []
  | x :: ys => fun pfs =>
    (existT _ x (pfs x (or_introl eq_refl))) :: add_pfs P ys (fun y pf => pfs y (or_intror pf))
  end.

Definition count_up_with_pfs : forall n, list { k : nat & k < n } :=
  fun n => add_pfs (fun k => k < n) (count_up n) (count_up_lt n).

Equations foo(n : nat) : nat by wf n :=
  foo n := S (list_sum (map (fun x => foo (projT1 x)) (count_up_with_pfs n))).

Compute foo 5.

Require Import Wf_nat.

Fixpoint exp n :=
  match n with
  | 0 => 1
  | S m => 2 * exp m
  end.

Lemma map_exp_S : forall xs, map exp (map S xs) = map (mult 2) (map exp xs).
Proof.
  induction xs.
  - reflexivity.
  - simpl.
    rewrite IHxs.
    reflexivity.
Qed.

Lemma list_sum_scale : forall n xs, list_sum (map (mult n) xs) = n * (list_sum xs).
Proof.
  induction xs.
  - simpl.
    apply mult_n_O.
  - simpl.
    rewrite IHxs.
    lia.
Qed.

Lemma geometric_series : forall n, exp n = S (list_sum (map exp (count_up n))).
Proof.
  induction n.
  - reflexivity.
  - unfold exp at 1; fold exp.
    simpl list_sum.
    rewrite map_exp_S.
    rewrite list_sum_scale.
    rewrite IHn.
    lia.
Qed.

Lemma add_remove : forall {X Y}(f : X -> Y)(P : X -> Prop)(xs : list X)(pfs : forall x, In x xs ->  P x),
  map (fun x => f (projT1 x)) (add_pfs P xs pfs) = map f xs.
Proof.
  induction xs.
  - intros; reflexivity.
  - intros; simpl; f_equal.
    apply IHxs.
Qed.

Lemma map_rewrite{X Y}(f g : X -> Y)(xs : list X) : (forall x, In x xs -> f x = g x)
  -> map f xs = map g xs.
Proof.
  induction xs; intros.
  - reflexivity.
  - simpl.
    rewrite H.
    + f_equal.
      apply IHxs.
      intros.
      apply H.
      right; auto.
    + left; auto.
Qed.

Lemma foo_correct : forall n, foo n = exp n.
Proof.
  intro n.
  induction (lt_wf n) as [n _ IHn].
  rewrite foo_equation_1.
  unfold count_up_with_pfs.
  rewrite add_remove.
  rewrite (map_rewrite foo exp).
  - symmetry; apply geometric_series.
  - intros.
    apply IHn.
    apply count_up_lt; auto.
Qed.