Library UniMath.CategoryTheory.Inductives.Trees

Definition of binary trees (Tree) implemented similarily to lists as the initial algebra of the tree functor.
Written by: Anders Mörtberg (2016)

Binary trees

Section bintrees.

Variable A : HSET.

Local Open Scope cocont_functor_hset_scope.

The tree functor: F(X) = 1 + A * X * X
The type of binary trees
Definition Tree : HSET :=
  alg_carrier _ (InitialObject treeFunctor_Initial).

Let Tree_mor : HSETtreeFunctor Tree,Tree :=
  alg_map _ (InitialObject treeFunctor_Initial).

Let Tree_alg : algebra_ob treeFunctor :=
  InitialObject treeFunctor_Initial.

Definition leaf_map : HSETunitHSET,Tree.
Proof.
simpl; intro x.
use Tree_mor.
apply inl, x.
Defined.

Definition leaf : pr1 Tree := leaf_map tt.

Definition node_map : HSET(A × (Tree × Tree))%set,Tree.
Proof.
intros xs.
use Tree_mor.
exact (inr xs).
Defined.

Definition node : pr1 A × (pr1 Tree × pr1 Tree) pr1 Tree := node_map.

Get recursion/iteration scheme:
     x : X           f : A × X × X -> X
  ------------------------------------
        foldr x f : Tree A -> X
Definition mk_treeAlgebra (X : HSET) (x : pr1 X)
  (f : HSET(A × X × X)%set,X) : algebra_ob treeFunctor.
Proof.
set (x' := λ (_ : unit), x).
apply (tpair _ X (sumofmaps x' f) : algebra_ob treeFunctor).
Defined.

Definition foldr_map (X : HSET) (x : pr1 X) (f : HSET(A × X × X)%set,X) :
  algebra_mor _ Tree_alg (mk_treeAlgebra X x f).
Proof.
apply (InitialArrow treeFunctor_Initial (mk_treeAlgebra X x f)).
Defined.

Definition foldr (X : HSET) (x : pr1 X)
  (f : pr1 A × pr1 X × pr1 X pr1 X) : pr1 Tree pr1 X.
Proof.
apply (foldr_map _ x f).
Defined.

Lemma foldr_leaf (X : hSet) (x : X) (f : pr1 A × X × X X) : foldr X x f leaf = x.
Proof.
assert (F := maponpaths (λ x, BinCoproductIn1 _ (BinCoproductsHSET _ _) · x)
                        (algebra_mor_commutes _ _ _ (foldr_map X x f))).
apply (toforallpaths _ _ _ F tt).
Qed.

Lemma foldr_node (X : hSet) (x : X) (f : pr1 A × X × X X)
                 (a : pr1 A) (l1 l2 : pr1 Tree) :
  foldr X x f (node (a,,l1,,l2)) = f (a,,foldr X x f l1,,foldr X x f l2).
Proof.
assert (F := maponpaths (λ x, BinCoproductIn2 _ (BinCoproductsHSET _ _)· x)
                        (algebra_mor_commutes _ _ _ (foldr_map X x f))).
assert (Fal := toforallpaths _ _ _ F (a,,l1,,l2)).
clear F.
unfold compose in Fal.
simpl in Fal.
apply Fal.
Opaque foldr_map.
Qed. Transparent foldr_map.

This defines the induction principle for trees using foldr
Section tree_induction.

Variables (P : pr1 Tree UU) (PhSet : l, isaset (P l)).
Variables (P0 : P leaf)
          (Pc : (a : pr1 A) (l1 l2 : pr1 Tree), P l1 P l2 P (node (a,,l1,,l2))).

Let P' : UU := l, P l.
Let P0' : P' := (leaf,, P0).
Let Pc' : pr1 A × P' × P' P'.
Proof.
intros ap.
apply (tpair _ (node (pr1 ap,,pr1 (pr1 (pr2 ap)),,pr1 (pr2 (pr2 ap))))).
apply (Pc _ _ _ (pr2 (pr1 (pr2 ap))) (pr2 (pr2 (pr2 ap)))).
Defined.

Definition P'HSET : HSET.
Proof.
apply (tpair _ P').
abstract (apply (isofhleveltotal2 2); [ apply setproperty | intro x; apply PhSet ]).
Defined.

Opaque is_omega_cocont_treeFunctor.

Lemma isalghom_pr1foldr :
  is_algebra_mor _ Tree_alg Tree_alg (λ l, pr1 (foldr P'HSET P0' Pc' l)).
Proof.
apply BinCoproductArrow_eq_cor.
- apply funextfun; intro x; destruct x.
  apply (maponpaths pr1 (foldr_leaf P'HSET P0' Pc')).
- apply funextfun; intro x; destruct x as [a [l1 l2]].
  apply (maponpaths pr1 (foldr_node P'HSET P0' Pc' a l1 l2)).
Qed.

Transparent is_omega_cocont_treeFunctor.

Definition pr1foldr_algmor : algebra_mor _ Tree_alg Tree_alg :=
  tpair _ _ isalghom_pr1foldr.

Lemma pr1foldr_algmor_identity : identity _ = pr1foldr_algmor.
Proof.
now rewrite (@InitialEndo_is_identity _ treeFunctor_Initial pr1foldr_algmor).
Qed.

Lemma treeInd l : P l.
Proof.
assert (H : pr1 (foldr P'HSET P0' Pc' l) = l).
  apply (toforallpaths _ _ _ (!pr1foldr_algmor_identity) l).
rewrite <- H.
apply (pr2 (foldr P'HSET P0' Pc' l)).
Defined.

End tree_induction.

Lemma treeIndProp (P : pr1 Tree UU) (HP : l, isaprop (P l)) :
  P leaf ( a l1 l2, P l1 P l2 P (node (a,,l1,,l2))) l, P l.
Proof.
intros Pnil Pcons.
apply treeInd; try assumption.
intro l; apply isasetaprop, HP.
Defined.

End bintrees.

Some tests
Section nat_examples.

Local Open Scope nat_scope.

Definition size : pr1 (Tree natHSET) nat :=
  foldr natHSET natHSET 0 (λ x, S (pr1 (pr2 x) + pr2 (pr2 x))).

Lemma size_node a l1 l2 : size (node natHSET (a,,l1,,l2)) = 1 + size l1 + size l2.
Proof.
unfold size.
now rewrite foldr_node.
Qed.

Definition map (f : nat nat) (l : pr1 (Tree natHSET)) : pr1 (Tree natHSET) :=
  foldr natHSET (Tree natHSET) (leaf natHSET)
      (λ a, node natHSET (f (pr1 a),, pr1 (pr2 a),, pr2 (pr2 a))) l.

Lemma size_map (f : nat nat) : l, size (map f l) = size l.
Proof.
apply treeIndProp.
- intros l. apply isasetnat.
- now unfold map; rewrite foldr_leaf.
- intros a l1 l2 ih1 ih2; unfold map.
  now rewrite foldr_node, !size_node, <- ih1, <- ih2.
Qed.

Definition sum : pr1 (Tree natHSET) nat :=
  foldr natHSET natHSET 0 (λ x, pr1 x + pr1 (pr2 x) + pr2 (pr2 x)).

Definition testtree : pr1 (Tree natHSET).
Proof.
  use node_map; repeat split.
  - apply 5.
  - use node_map; repeat split.
    + apply 6.
    + exact (leaf_map _ tt).
    + exact (leaf_map _ tt).
  - exact (leaf_map _ tt).
Defined.

End nat_examples.

Flattening of a tree into a list

Local Notation "a :: b" := (cons _ a b).
Definition flatten (A : HSET) : pr1 (Tree A) List A.
Proof.
  intro t.
  use (foldr A).
  - apply nil.
  - intro all'.
    set (a := pr1 all').
    set (l := pr1 (pr2 all')).
    set (l' := pr2 (pr2 all')). cbn beta in l'.
    exact (concatenate _ l (concatenate _ (a :: nil _ ) l')).
  - exact t.
Defined.

Goal Lists.sum (flatten _ testtree) = sum testtree. reflexivity. Qed.