------------------------------------------------------------------------
-- A container for finite binary trees with information in internal
-- nodes
------------------------------------------------------------------------

{-# OPTIONS --without-K --safe #-}

module Container.Tree where

open import Container hiding (Shape; Position)
open import Container.List hiding (fold; fold-lemma)
open import Equality.Propositional
open import Prelude hiding (id; _∘_; List; []; _∷_)
import Tree

open import Bijection equality-with-J using (_↔_; module _↔_)
open import Function-universe equality-with-J

------------------------------------------------------------------------
-- The type

-- Shapes.

data Shape : Set where
  lf : Shape
  nd : Shape  Shape  Shape

-- Positions.

data Position : Shape  Set where
  root  :  {l r}  Position (nd l r)
  left  :  {l r}  Position l  Position (nd l r)
  right :  {l r}  Position r  Position (nd l r)

-- Trees.

Tree : Container lzero
Tree = Shape  Position

------------------------------------------------------------------------
-- An isomorphism

-- The type of shapes is isomorphic to Tree.Tree ⊤.
--
-- This lemma is included because it was mentioned in the paper "Bag
-- Equivalence via a Proof-Relevant Membership Relation".

Shape↔Tree-⊤ : Shape  Tree.Tree 
Shape↔Tree-⊤ = record
  { surjection = record
    { logical-equivalence = record
      { to   = to
      ; from = from
      }
    ; right-inverse-of = to∘from
    }
  ; left-inverse-of = from∘to
  }
  where
  to : Shape  Tree.Tree 
  to lf       = Tree.leaf
  to (nd l r) = Tree.node (to l) tt (to r)

  from : Tree.Tree   Shape
  from Tree.leaf          = lf
  from (Tree.node l tt r) = nd (from l) (from r)

  to∘from :  t  to (from t)  t
  to∘from Tree.leaf          = refl
  to∘from (Tree.node l tt r) =
    cong₂  l r  Tree.node l tt r) (to∘from l) (to∘from r)

  from∘to :  s  from (to s)  s
  from∘to lf       = refl
  from∘to (nd l r) = cong₂ nd (from∘to l) (from∘to r)

------------------------------------------------------------------------
-- Constructors

-- Leaves.

leaf : {A : Set}   Tree  A
leaf = (lf , λ ())

-- Internal nodes.

node : {A : Set}   Tree  A  A   Tree  A   Tree  A
node (l , lkup-l) x (r , lkup-r) =
  ( nd l r
  , λ { root       x
      ; (left  p)  lkup-l p
      ; (right p)  lkup-r p
      }
  )

-- Even if we don't assume extensionality we can prove that
-- intensionally distinct implementations of the constructors are bag
-- equivalent.

leaf≈ : {A : Set} {lkup : _  A} 
        _≈-bag_ {C₂ = Tree} leaf (lf , lkup)
leaf≈ _ = record
  { surjection = record
    { logical-equivalence = record
      { to   = λ { (() , _) }
      ; from = λ { (() , _) }
      }
    ; right-inverse-of = λ { (() , _) }
    }
  ; left-inverse-of = λ { (() , _) }
  }

node≈ :  {A : Set} {l r} {lkup : _  A} 
        _≈-bag_ {C₂ = Tree}
                (node (l , lkup  left) (lkup root) (r , lkup  right))
                (nd l r , lkup)
node≈ _ = record
  { surjection = record
    { logical-equivalence = record
      { to   = λ { (root    , eq)  (root    , eq)
                 ; (left p  , eq)  (left p  , eq)
                 ; (right p , eq)  (right p , eq)
                 }
      ; from = λ { (root    , eq)  (root    , eq)
                 ; (left p  , eq)  (left p  , eq)
                 ; (right p , eq)  (right p , eq)
                 }
      }
    ; right-inverse-of = λ { (root    , eq)  refl
                           ; (left p  , eq)  refl
                           ; (right p , eq)  refl
                           }
    }
  ; left-inverse-of = λ { (root    , eq)  refl
                        ; (left p  , eq)  refl
                        ; (right p , eq)  refl
                        }
  }

-- Any lemmas for the constructors.

Any-leaf :  {A : Set} (P : A  Set) 
           Any P leaf  ⊥₀
Any-leaf _ = record
  { surjection = record
    { logical-equivalence = record
      { to   = λ { (() , _) }
      ; from = λ ()
      }
    ; right-inverse-of = λ ()
    }
  ; left-inverse-of = λ { (() , _) }
  }

Any-node :  {A : Set} (P : A  Set) {l x r} 
           Any P (node l x r)  Any P l  P x  Any P r
Any-node _ {l = _ , _} {r = _ , _} = record
  { surjection = record
    { logical-equivalence = record
      { to   = λ { (root    , eq)  inj₂ (inj₁ eq)
                 ; (left p  , eq)  inj₁ (p , eq)
                 ; (right p , eq)  inj₂ (inj₂ (p , eq))
                 }
      ; from = λ { (inj₁ (p , eq))         (left p  , eq)
                 ; (inj₂ (inj₁ eq))        (root    , eq)
                 ; (inj₂ (inj₂ (p , eq)))  (right p , eq)
                 }
      }
    ; right-inverse-of = λ { (inj₁ (p , eq))         refl
                           ; (inj₂ (inj₁ eq))        refl
                           ; (inj₂ (inj₂ (p , eq)))  refl
                           }
    }
  ; left-inverse-of = λ { (root    , eq)  refl
                        ; (left p  , eq)  refl
                        ; (right p , eq)  refl
                        }
  }

------------------------------------------------------------------------
-- More functions

-- Singleton trees.

singleton : {A : Set}  A   Tree  A
singleton x = node leaf x leaf

-- Any lemma for singleton.

Any-singleton :  {A : Set} (P : A  Set) {x} 
                Any P (singleton x)  P x
Any-singleton P {x} =
  Any P (singleton x)            ↔⟨⟩
  Any P (node leaf x leaf)       ↔⟨ Any-node P 
  Any P leaf  P x  Any P leaf  ↔⟨ Any-leaf P ⊎-cong id ⊎-cong Any-leaf P 
    P x                      ↔⟨ ⊎-left-identity 
  P x                          ↔⟨ ⊎-right-identity 
  P x                            

-- For the design considerations underlying the inclusion of fold and
-- fold-lemma, see Container.List.fold/fold-lemma.

-- A fold for trees. (Well, this is not a catamorphism, it is a
-- paramorphism.)

fold : {A B : Set} 
       B  ( Tree  A  A   Tree  A  B  B  B)   Tree  A  B
fold {A} {B} fl fn = uncurry fold′
  where
  fold′ : (s : Shape)  (Position s  A)  B
  fold′ lf       lkup = fl
  fold′ (nd l r) lkup =
    fn (l , lkup  left )
       (lkup root)
       (r , lkup  right)
       (fold′ l (lkup  left ))
       (fold′ r (lkup  right))

-- A lemma which can be used to prove properties about fold.
--
-- The "respects bag equivalence" argument could be omitted if
-- equality of functions were extensional.

fold-lemma :  {A B : Set}
               {fl : B} {fn :  Tree  A  A   Tree  A  B  B  B}
             (P :  Tree  A  B  Set) 
             (∀ t₁ t₂  t₁ ≈-bag t₂   b  P t₁ b  P t₂ b) 
             P leaf fl 
             (∀ l x r b₁ b₂ 
                P l b₁  P r b₂  P (node l x r) (fn l x r b₁ b₂)) 
              t  P t (fold fl fn t)
fold-lemma {A} {fl = fl} {fn} P resp P-le P-no = uncurry fold-lemma′
  where
  fold-lemma′ : (s : Shape) (lkup : Position s  A) 
                P (s , lkup) (fold fl fn (s , lkup))
  fold-lemma′ lf       lkup = resp _ _ leaf≈ _ P-le
  fold-lemma′ (nd l r) lkup = resp _ _ node≈ _ $
    P-no _ _ _ _ _
       (fold-lemma′ l (lkup  left ))
       (fold-lemma′ r (lkup  right))

-- Inorder flattening of a tree.

flatten : {A : Set}   Tree  A   List  A
flatten = fold []  _ x _ xs ys  xs ++ x  ys)

-- Flatten does not add or remove any elements.

flatten≈ : {A : Set} (t :  Tree  A)  flatten t ≈-bag t
flatten≈ = fold-lemma
   t xs  xs ≈-bag t)

   t₁ t₂ t₁≈t₂ xs xs≈t₁ z 
     z  xs  ↔⟨ xs≈t₁ z 
     z  t₁  ↔⟨ t₁≈t₂ z 
     z  t₂  )

   z 
   z  []    ↔⟨ Any-[]  x  z  x) 
            ↔⟨ inverse $ Any-leaf  x  z  x) 
   z  leaf  )

   l x r xs ys xs≈l ys≈r z 
     z  xs ++ x  ys         ↔⟨ Any-++  x  z  x) _ _ 
     z  xs  z  x  ys      ↔⟨ id ⊎-cong Any-∷  x  z  x) 
     z  xs  z  x  z  ys  ↔⟨ xs≈l z ⊎-cong id ⊎-cong ys≈r z 
     z  l  z  x  z  r    ↔⟨ inverse $ Any-node  x  z  x) 
     z  node l x r           )