TreeSort.aya

open inductive Nat | O | S Nat
open inductive Bool | True | False
open inductive List Type
| []
| A => infixr :> A (List A)

def isZero (a : Nat) : Bool
| 0 => True
| _ => False

open inductive Color | red | black
def Decider (A : Type) => Fn (x y : A) -> Bool

variable A : Type

open inductive RBTree (A : Type) : Type
| rbLeaf
| rbNode Color (RBTree A) A (RBTree A)

def rbTreeToList (rb : RBTree A) (r : List A) : List A elim rb
| rbLeaf => r
| rbNode x t1 a t2 => rbTreeToList t1 (a :> rbTreeToList t2 r)

def repaint (RBTree A) : RBTree A
| rbNode c l a r => rbNode black l a r
| rbLeaf => rbLeaf

def sub (x y : Nat) : Nat
| n, 0 => n
| 0, S _ => 0
| S n, S m => sub n m

def le (x y : Nat) => isZero (sub x y)

def balanceLeft Color (RBTree A) A (RBTree A) : RBTree A
| black, rbNode red (rbNode red a x b) y c, v, r =>
    rbNode red (rbNode black a x b) y (rbNode black c v r)
| black, rbNode red a x (rbNode red b y c), v, r =>
    rbNode red (rbNode black a x b) y (rbNode black c v r)
| c, a, v, r => rbNode c a v r

def balanceRight Color (RBTree A) A (RBTree A) : RBTree A
| black, l, v, rbNode red (rbNode red b y c) z d =>
    rbNode red (rbNode black l v b) y (rbNode black c z d)
| black, l, v, rbNode red b y (rbNode red c z d) =>
    rbNode red (rbNode black l v b) y (rbNode black c z d)
| c, l, v, b => rbNode c l v b

def insert (a : A) (node : RBTree A) (dec_le : Decider A) : RBTree A elim node
| rbLeaf => rbNode red rbLeaf a rbLeaf
| rbNode c l1 a1 l2 => match dec_le a1 a
  { True => balanceRight c l1 a1 (insert a l2 dec_le)
  | False => balanceLeft c (insert a l1 dec_le) a1 l2
  }

private def aux (ls : List A) (r : RBTree A) (dec_le : Decider A) : RBTree A elim ls
| [] => r
| a :> l => aux l (repaint (insert a r dec_le)) dec_le
def tree_sort (dec_le : Decider A) (l : List A) => rbTreeToList (aux l rbLeaf dec_le) []
def tree_sortNat (l : List Nat) => tree_sort le l