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

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

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

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

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

private

  -- 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 →
     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                           □