(**********************************************************************)
(* Equations                                                          *)
(* Copyright (c) 2009-2021 Matthieu Sozeau <matthieu.sozeau@inria.fr> *)
(**********************************************************************)
(* This file is distributed under the terms of the                    *)
(* GNU Lesser General Public License Version 2.1                      *)
(**********************************************************************)

Require Import Program.
From Equations Require Import Equations.
Require Import Utf8 Lia Arith.

Require Import List.
Import ListNotations.

Set Keyed Unification.

Equations map_In {A B : Type}
(l : list A) (f : (x : A), In x l B) : list B :=
| nil, _ := nil
| cons x xs, f := cons (f x _) (map_In xs (fun x Hf x _)).

Lemma map_In_spec {A B : Type} (f : A B) (l : list A) :
map_In l (fun (x : A) (_ : In x l) ⇒ f x) = List.map f l.
Proof.
funelim (map_In l _); rewrite ?H; trivial.
Qed.

Section list_size.
Context {A : Type} (f : A nat).
Equations list_size (l : list A) : nat :=
| nil := 0
| x :: xs := S (f x + list_size xs).

Lemma In_list_size:
x xs, In x xs f x < S (list_size xs).
Proof.
intros. funelim (list_size xs); simpl in ×. destruct H.
destruct H0.
× subst; lia.
× specialize (H _ H0). intuition. lia.
Qed.
End list_size.
Transparent list_size.

Module RoseTree.

Section roserec.
Context {A : Set}.

Inductive t : Set :=
| leaf (a : A) : t
| node (l : list t) : t.
Derive NoConfusion for t.

Fixpoint size (r : t) :=
match r with
| leaf a ⇒ 0
| node lS (list_size size l)
end.

Equations? elements (r : t) : list A by wf (size r) lt :=
elements (leaf a) := [a];
elements (node l) := concat (map_In l (fun x Helements x)).
Proof. red. now apply In_list_size. Qed.

Equations elements_def (r : t) : list A :=
elements_def (leaf a) := [a];
elements_def (node l) := concat (List.map elements_def l).
Lemma elements_equation (r : t) : elements r = elements_def r.
Proof.
funelim (elements r); simp elements_def; trivial. f_equal.
induction l; simpl; auto. simp map_In. rewrite H. rewrite IHl; auto.
intros. apply H. now constructor 2. now constructor.
Qed.

To solve measure subgoals
Hint Extern 4 (_ < _) ⇒ simpl; lia : Below.
Hint Extern 4 (MR _ _ _ _) ⇒ repeat red; simpl in *; lia : Below.

Obligation Tactic := simpl in *; program_simpl; try typeclasses eauto with Below subterm_relation.
(* Nested rec *)

Equations elements_acc (r : t) (acc : list A) : list A by wf (size r) lt :=
elements_acc (leaf a) acc := a :: acc;
elements_acc (node l) acc := aux l _
where aux (x : list t) (H : list_size size x < size (node l)) : list A by wf (list_size size x) lt :=
aux nil _ := acc;
aux (cons x xs) H := elements_acc x (aux xs _).

Definition elements2 (r : t) : list A := elements_acc r [].

Lemma elements2_equation r acc : elements_acc r acc = elements_def r ++ acc.
Proof.
revert r acc.
let t := constr:(fun_elim (f:=elements_acc)) in
apply (t (fun r acc resres = elements_def r ++ acc)
(fun r acc x H resres = concat (List.map elements_def x) ++ acc));
intros; simp elements; trivial.
rewrite H1. clear H1.
rewrite H0. simpl. now rewrite app_assoc.
Qed.

Equations elements' (r : t) : list A by wf r (MR lt size) :=
| leaf a := [a]
| node l := fn l hidebody

where fn (x : list t) (H : list_size size x < size (node l)) : list A
by wf x (MR lt (list_size size)) :=
| nil, _ := nil;
| cons x xs, _ := elements' x ++ fn xs hidebody.

Equations elements'_def (r : t) : list A :=
elements'_def (leaf a) := [a];
elements'_def (node l) := concat (List.map elements' l).

Lemma elements'_equation (r : t) : elements' r = elements'_def r.
Proof.
pose (fun_elim (f:=elements')).
apply (p (fun r ff = elements'_def r) (fun l x H rr = concat (List.map elements' x)));
clear p; intros; simp elements'_def; trivial.
simpl. f_equal. apply H1.
Qed.

End roserec.
Arguments t : clear implicits.

Section AltSize.
Context {A : Set}.

Let's use an alternative size criterion allowing to make recursive calls on non-strict subterms of the initial list: we just count the maximal depth of node constructors among all forests.
Equations alt_size (r : t A) : nat :=
{ alt_size (leaf _) ⇒ 0;
alt_size (node l) ⇒ S (max_size l) }
where max_size (l : list (t A)) : nat :=
{ max_size nil := 0;
max_size (cons a t) := Nat.max (alt_size a) (max_size t) }.

This has the property that the maximal size of two appended lists is the maximal size of the separate lists.
Lemma max_size_app l l' : max_size (l ++ l') = Nat.max (max_size l) (max_size l').
Proof.
induction l; simp max_size. reflexivity.
simpl. rewrite <- Nat.max_assoc. f_equal.
apply IHl.
Qed.

Context {B : Set} (f : A B).

It hence becomes possible to recurse on an arbitrary list as long as the depth decreases, for example by appending the subforest to itself in the node case. The same is possible with sized types where node has type j < i list^i (t^j) t^(S i).
Equations? map_t (r : t A) : t B by wf (alt_size r) lt :=
map_t (leaf a) := leaf (f a);
map_t (node l) := node (map_list (l ++ l) _)

where map_list (l' : list (t A)) (H : max_size l' max_size l) : list (t B) by struct l' :=
map_list nil _ := nil;
map_list (cons a t) Hl' := cons (map_t a) (map_list t _).
Proof. simp alt_size. apply le_lt_n_Sm. now apply Nat.max_lub_l in Hl'.
now apply Nat.max_lub_r in Hl'.
clear map_list. rewrite max_size_app. now rewrite Nat.max_id.
Defined.
End AltSize.

Section fns.
Context {A B : Set} (f : A B) (g : B A B) (h : A B B).

Equations map (r : t A) : t B :=
| leaf a := leaf (f a)
| node l := node (List.map map l).

Equations fold (acc : B) (r : t A) : B :=
fold acc (leaf a) := g acc a;
fold acc (node l) := List.fold_left fold l acc.

Equations fold_right (r : t A) (acc : B) : B :=
fold_right (leaf a) acc := h a acc;
fold_right (node l) acc := List.fold_right fold_right acc l.
End fns.

End RoseTree.