-- An implementation of tree sort, formally proved to return a
-- permutation of the input

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

open import Prelude hiding (id; _∘_; List; module List; []; _∷_)

module Container.Tree-sort
  {A : Set}
  (_≤_ : A  A  Bool) -- A comparison function.

open import Container
open import Container.List as List
open import Container.Tree as Tree
open import Equality.Propositional
open import Logical-equivalence using (module _⇔_)

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

-- Boring lemmas


  -- The following lemma is easy to prove automatically (for
  -- instance by using a ring solver).

  lemma₁ : (A B C D : Set)  A  B  C  D  C  A  B  D
  lemma₁ A B C D =
    A  B  C  D      ↔⟨ id ⊎-cong ⊎-assoc 
    A  (B  C)  D    ↔⟨ id ⊎-cong ⊎-comm ⊎-cong id 
    A  (C  B)  D    ↔⟨ ⊎-assoc 
    (A  C  B)  D    ↔⟨ ⊎-assoc ⊎-cong id 
    ((A  C)  B)  D  ↔⟨ (⊎-comm ⊎-cong id) ⊎-cong id 
    ((C  A)  B)  D  ↔⟨ inverse ⊎-assoc ⊎-cong id 
    (C  A  B)  D    ↔⟨ inverse ⊎-assoc 
    C  (A  B)  D    ↔⟨ id ⊎-cong inverse ⊎-assoc 
    C  A  B  D      

  lemma₂ : {A B C D : Set} (b : Bool) 
           T b × ((A  B)  C  D)  T (not b) × (B  C  A  D) 
           A  B  C  D
  lemma₂ = if-lemma  _  _) (inverse ⊎-assoc) (lemma₁ _ _ _ _)

-- Insertion into trees

-- Inserts an element into the tree.

insert : A   Tree  A   Tree  A
insert x = Tree.fold
  (singleton x)
   l y r x+l x+r 
     if x  y then node x+l y r
              else node l y x+r)

-- The insert function inserts.

Any-insert :  (P : A  Set) x t 
             Any P (insert x t)  P x  Any P t
Any-insert P x = Tree.fold-lemma
   t t′  Any P t′  P x  Any P t)

   t₁ t₂ t₁≈t₂ t hyp 
     Any P t         ↔⟨ hyp 
     P x  Any P t₁  ↔⟨ id ⊎-cong _⇔_.to (∼⇔∼″ t₁ t₂) t₁≈t₂ P 
     P x  Any P t₂  )

  (Any P (singleton x)  ↔⟨ Any-singleton P 
   P x                  ↔⟨ inverse ⊎-right-identity 
   P x                ↔⟨ id ⊎-cong inverse (Any-leaf P) 
   P x  Any P leaf     )

   l y r l′ r′ ih-l ih-r 
     Any P (if x  y then node l′ y r else node l y r′)   ↔⟨ Any-if P (node l′ y r) (node l y r′) (x  y) 

     T (x  y) × Any P (node l′ y r) 
       T (not (x  y)) × Any P (node l y r′)              ↔⟨ id ×-cong Any-node P ⊎-cong
                                                             id ×-cong Any-node P 
     T (x  y) × (Any P l′  P y  Any P r) 
       T (not (x  y)) × (Any P l  P y  Any P r′)       ↔⟨ id ×-cong (ih-l ⊎-cong id) ⊎-cong
                                                             id ×-cong (id ⊎-cong id ⊎-cong ih-r) 
     T (x  y) × ((P x  Any P l)  P y  Any P r) 
       T (not (x  y)) × (Any P l  P y  P x  Any P r)  ↔⟨ lemma₂ (x  y) 

     P x  Any P l  P y  Any P r                        ↔⟨ id ⊎-cong inverse (Any-node P) 

     P x  Any P (node l y r)                             )

-- Turning a list into a search tree

-- Converts the list to a search tree.

to-search-tree :  List  A   Tree  A
to-search-tree = List.fold leaf  x _ t  insert x t)

-- No elements are added or removed.

to-search-tree≈ :  xs  to-search-tree xs ≈-bag xs
to-search-tree≈ = List.fold-lemma
   xs t  t ≈-bag xs)

   xs ys xs≈ys t t≈xs z 
     z  t   ↔⟨ t≈xs z 
     z  xs  ↔⟨ xs≈ys z 
     z  ys  )

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

   x xs t t≈xs z 
     z  insert x t  ↔⟨ Any-insert  x  z  x) _ _ 
     z  x  z  t   ↔⟨ id ⊎-cong t≈xs z 
     z  x  z  xs  ↔⟨ inverse $ Any-∷  x  z  x) 
     z  x  xs      )

-- Sorting

-- Sorts a list.

sort :  List  A   List  A
sort = flatten  to-search-tree

-- The result is a permutation of the input.

sort≈ :  xs  sort xs ≈-bag xs
sort≈ xs = λ z 
  z  sort xs                      ↔⟨⟩
  z  flatten (to-search-tree xs)  ↔⟨ flatten≈ (to-search-tree xs) _ 
  z  to-search-tree xs            ↔⟨ to-search-tree≈ xs _ 
  z  xs