RBTree.agda

module RBTree where

open import Agda.Primitive public
  using (Level; lzero; lsuc; _⊔_)

data Σ {l k} (A : Set l) (B : A → Set k) : Set (l ⊔ k) where
  _,_ : (a : A) → (b : B a) → Σ A B
infixr 10 _,_

fst : ∀ {l k} {A : Set l} {B : A → Set k} -> Σ A B -> A
fst (a , _) = a

snd : ∀ {l k} {A : Set l} {B : A → Set k} -> (σ : Σ A B) -> (B (fst σ))
snd (_ , b) = b

_o_ : ∀ {a b c} {A : Set a} {B : Set b} {C : Set c}
      → (B → C) → (A → B) → (A → C)
f o g = λ x → f (g x)

infixr 9 _o_

data Two : Set where
  tt : Two
  ff : Two

data Comp : Set where
  LT : Comp
  EQ : Comp
  GT : Comp

record Ord {l} (X : Set l) : Set l where
  field
    compare : X -> X -> Comp

  _≤_ _≥_ _≡_ : X -> X -> Two

  a ≤ b with compare a b
  a ≤ b | GT = ff
  a ≤ b | _  = tt

  a ≥ b with compare a b
  a ≥ b | LT = ff
  a ≥ b | _  = tt

  a ≡ b with compare a b
  a ≡ b | EQ = tt
  a ≡ b | _  = ff

open Ord {{...}} public

data Nat : Set where
  Z :       Nat
  S : Nat → Nat

instance ordNat : Ord Nat
ordNat = record { compare = cmpNat }
  where
    cmpNat : Nat -> Nat -> Comp
    cmpNat  Z     Z    = EQ
    cmpNat  Z    (S b) = LT
    cmpNat (S a)  Z    = GT
    cmpNat (S a) (S b) = cmpNat a b

{-# BUILTIN NATURAL Nat #-}

data List (X : Set) : Set where
  []   : List X
  _::_ : (x : X) → (xs : List X) → List X
infixr 9 _::_

data Colour : Nat → Nat → Set where
  R : forall {n : Nat} → Colour n    n
  B : forall {n : Nat} → Colour n (S n)

data RBT (X : Set) : Nat → Set where
  <>       : RBT X 1

  _<_[_]>_ :
    ∀ {m n : Nat}
    → (l : RBT X m)
    → (c : Colour m n)
    → (v : X)
    → (r : RBT X m)
    →  RBT X n


balance-left : ∀ {X n} → RBT X n → RBT X n
balance-left (l < R [ x ]> r) = l < R [ x ]> r
balance-left ((t₁ < R [ x ]> (t₂ < R [ y ]> t₃)) < B [ z ]> t₄) =
  (t₁ < B [ x ]> t₂) < R [ y ]> (t₃ < B [ z ]> t₄)
balance-left (((t₁ < R [ x ]> t₂) < R [ y ]> t₃) < B [ z ]> t₄) =
  (t₁ < B [ x ]> t₂) < R [ y ]> (t₃ < B [ z ]> t₄)
balance-left t = t

balance-right : ∀ {X n} → RBT X n → RBT X n
balance-right (l < R [ x ]> r) = l < R [ x ]> r
balance-right (t₁ < B [ x ]> (t₂ < R [ y ]> (t₃ < R [ z ]> t₄))) =
        (t₁ < B [ x ]> t₂) < R [ y ]> (t₃ < B [ z ]> t₄)
balance-right (t₁ < B [ x ]> ((t₂ < R [ z ]> t₃) < R [ y ]> t₄)) =
        (t₁ < B [ x ]> t₂) < R [ y ]> (t₃ < B [ z ]> t₄)
balance-right t = t

insert' :
  ∀ {X n}
    {{_ : Ord X}}
  → X
  → RBT X n
  → RBT X n
insert' x <> = <> < R [ x ]> <>
insert' x (l < c [ v ]> r) with compare x v
insert' x (l < c [ v ]> r) | LT = balance-left ((insert' x l) < c [ v ]> r)
insert' x (l < c [ v ]> r) | EQ = l < c [ v ]> r
insert' x (l < c [ v ]> r) | GT = balance-right (l < c [ v ]> (insert' x r))

blacken-root :
  ∀ {X n}
  → RBT X n
  → Σ Nat (RBT X)
blacken-root <> = 1 , <>
blacken-root (l < _ [ v ]> r) = _ , l < B [ v ]> r

insert :
  ∀ {X n}
    {{_ : Ord X}}
  → X
  → RBT X n
  → Σ Nat (RBT X)

insert x t = blacken-root (insert' x t)

fromList :
  ∀ {X}
    {{_ : Ord X}}
  → List X
  → Σ Nat (RBT X)
fromList [] = _ , <>
fromList (x :: xs) with fromList xs
fromList (x :: xs) | _ , rbt = insert x rbt

module Laws {X : Set} {{_ : Ord X}} where
  mutual
    data _is-balanced : ∀ {n : Nat} → RBT X n → Set₁ where
      R-Bal :
        ∀ {n} {t : RBT X n}
        → (p : t is-R-balanced)
        → t is-balanced

      B-Bal :
        ∀ {n} {t : RBT X n}
        → (p : t is-B-balanced)
        → t is-balanced





    data _is-R-balanced : ∀ {n : Nat} → RBT X n → Set₁ where
      N-Bal :
        ∀ {n} {l r : RBT X n} {v : X}
        → (lp : l is-B-balanced)
        → (rp : r is-B-balanced)
        → (l < R [ v ]> r) is-R-balanced





    data _is-B-balanced : ∀ {n : Nat} → RBT X n → Set₁ where
      L-Bal : <> is-B-balanced

      N-Bal :
        ∀ {n} {l r : RBT X n} {v : X}
        → (lp : l is-balanced)
        → (rp : r is-balanced)
        → (l < B [ v ]> r) is-B-balanced

  bal-=>-l-bal :
    ∀ {n m} {x : X} {c : Colour n m} {l r : RBT X n}
    → ((l < c [ x ]> r) is-balanced)
    → l is-balanced
  bal-=>-l-bal (R-Bal (N-Bal lp rp)) = B-Bal lp
  bal-=>-l-bal (B-Bal (N-Bal lp rp)) = lp

  bal-=>-r-bal :
    ∀ {n m} {x : X} {c : Colour n m} {l r : RBT X n}
    → ((l < c [ x ]> r) is-balanced)
    → r is-balanced
  bal-=>-r-bal (R-Bal (N-Bal lp rp)) = B-Bal rp
  bal-=>-r-bal (B-Bal (N-Bal lp rp)) = rp






  data _is-almost-balanced : ∀ {n : Nat} → RBT X n → Set₁ where
    balanced :
      ∀ {n} {t : RBT X n}
      → (p : t is-balanced)
      → t is-almost-balanced

    almost-left :
      ∀ {n} {x : X} {l : RBT X n} {r : RBT X n}
      → (lp : l is-R-balanced)
      → (rp : r is-B-balanced)
      → (l < R [ x ]> r) is-almost-balanced

    almost-right :
      ∀ {n} {x : X} {l : RBT X n} {r : RBT X n}
      → (lp : l is-B-balanced)
      → (rp : r is-R-balanced)
      → (l < R [ x ]> r) is-almost-balanced




  mutual
    insert'-on-B :
      ∀ {n} {x : X} {t : RBT X n}
      → (t is-B-balanced)
      → (insert' x t is-balanced)

    insert'-on-B L-Bal = R-Bal (N-Bal L-Bal L-Bal)

    insert'-on-B {_}  {x} {l < B [ v ]> r} (N-Bal lp rp) with compare x v

    insert'-on-B {._} {x} {l < B [ v ]> r} (N-Bal (R-Bal p) rp) | LT =
      balance-lemma (insert'-on-R p) rp
      where
        balance-lemma :
          ∀ {n} {y : X} {l r : RBT X n}
          → (l is-almost-balanced)
          → (r is-balanced)
          → (balance-left (l < B [ y ]> r) is-balanced)

        balance-lemma (balanced (R-Bal (N-Bal L-Bal L-Bal))) rp₂ =
          B-Bal (N-Bal (R-Bal (N-Bal L-Bal L-Bal)) rp₂)
        balance-lemma (balanced (R-Bal (N-Bal L-Bal (N-Bal lp rp₁)))) rp₂ =
          B-Bal (N-Bal (R-Bal (N-Bal L-Bal (N-Bal lp rp₁))) rp₂)
        balance-lemma (balanced (R-Bal (N-Bal (N-Bal lp rp₁) L-Bal))) rp₃ =
          B-Bal (N-Bal (R-Bal (N-Bal (N-Bal lp rp₁) L-Bal)) rp₃)
        balance-lemma (balanced (R-Bal (N-Bal (N-Bal lp rp₁) (N-Bal lp₁ rp₂)))) rp₃ =
          B-Bal (N-Bal (R-Bal (N-Bal (N-Bal lp rp₁) (N-Bal lp₁ rp₂))) rp₃)
        balance-lemma (balanced (B-Bal L-Bal)) rp₁ =
          B-Bal (N-Bal (B-Bal L-Bal) rp₁)
        balance-lemma (balanced (B-Bal (N-Bal lp rp₁))) rp₂ =
          B-Bal (N-Bal (B-Bal (N-Bal lp rp₁)) rp₂)

        balance-lemma (almost-left (N-Bal lp rp₁) L-Bal) rp₃ =
          R-Bal
            (N-Bal
              (N-Bal (B-Bal lp) (B-Bal rp₁))
              (N-Bal (B-Bal L-Bal) rp₃))

        balance-lemma (almost-left (N-Bal lp rp₁) (N-Bal lp₁ rp₂)) rp₃ =
          R-Bal
            (N-Bal
              (N-Bal (B-Bal lp) (B-Bal rp₁))
              (N-Bal (B-Bal (N-Bal lp₁ rp₂)) rp₃))

        balance-lemma (almost-right lp (N-Bal lp₁ rp₁)) rp₂ =
          R-Bal
            (N-Bal (N-Bal (B-Bal lp) (B-Bal lp₁))
                   (N-Bal (B-Bal rp₁) rp₂))

    insert'-on-B {._} {x} {l < B [ v ]> r} (N-Bal (B-Bal p) rp) | LT =
      balance-lemma (insert'-on-B p) rp
      where
        balance-lemma :
          ∀ {n} {y : X} {l r : RBT X n}
          → (l is-balanced)
          → (r is-balanced)
          → (balance-left (l < B [ y ]> r) is-balanced)
        balance-lemma (R-Bal (N-Bal L-Bal L-Bal)) rp₂ =
          B-Bal (N-Bal (R-Bal (N-Bal L-Bal L-Bal)) rp₂)
        balance-lemma (R-Bal (N-Bal L-Bal (N-Bal lp rp₁))) rp₂ =
          B-Bal (N-Bal (R-Bal (N-Bal L-Bal (N-Bal lp rp₁))) rp₂)
        balance-lemma (R-Bal (N-Bal (N-Bal lp rp₁) L-Bal)) rp₃ =
          B-Bal (N-Bal (R-Bal (N-Bal (N-Bal lp rp₁) L-Bal)) rp₃)
        balance-lemma (R-Bal (N-Bal (N-Bal lp rp₁) (N-Bal lp₁ rp₂))) rp₃ =
          B-Bal (N-Bal (R-Bal (N-Bal (N-Bal lp rp₁) (N-Bal lp₁ rp₂))) rp₃)
        balance-lemma (B-Bal L-Bal) rp₁ =
          B-Bal (N-Bal (B-Bal L-Bal) rp₁)
        balance-lemma (B-Bal (N-Bal lp rp₁)) rp₂ =
          B-Bal (N-Bal (B-Bal (N-Bal lp rp₁)) rp₂)

    insert'-on-B {._} {x} {l < B [ v ]> r} (N-Bal lp rp) | EQ = B-Bal (N-Bal lp rp)

    insert'-on-B {._} {x} {l < B [ v ]> r} (N-Bal lp (R-Bal p)) | GT =
      balance-lemma lp (insert'-on-R p)
      where
        balance-lemma :
          ∀ {n} {y : X} {l r : RBT X n}
          → (l is-balanced)
          → (r is-almost-balanced)
          → (balance-right (l < B [ y ]> r) is-balanced)

        balance-lemma lp₁ (balanced (R-Bal (N-Bal L-Bal L-Bal))) =
          B-Bal (N-Bal lp₁ (R-Bal (N-Bal L-Bal L-Bal)))
        balance-lemma lp₁ (balanced (R-Bal (N-Bal L-Bal (N-Bal lp₂ rp)))) =
          B-Bal (N-Bal lp₁ (R-Bal (N-Bal L-Bal (N-Bal lp₂ rp))))
        balance-lemma lp₁ (balanced (R-Bal (N-Bal (N-Bal lp₂ rp) L-Bal))) =
          B-Bal (N-Bal lp₁ (R-Bal (N-Bal (N-Bal lp₂ rp) L-Bal)))
        balance-lemma lp₁ (balanced (R-Bal (N-Bal (N-Bal lp₂ rp) (N-Bal lp₃ rp₁)))) =
          B-Bal (N-Bal lp₁ (R-Bal (N-Bal (N-Bal lp₂ rp) (N-Bal lp₃ rp₁))))
        balance-lemma lp₁ (balanced (B-Bal L-Bal)) =
          B-Bal (N-Bal lp₁ (B-Bal L-Bal))
        balance-lemma lp₁ (balanced (B-Bal (N-Bal lp₂ rp))) =
          B-Bal (N-Bal lp₁ (B-Bal (N-Bal lp₂ rp)))

        balance-lemma lp₁ (almost-left (N-Bal lp₂ rp) L-Bal) =
          R-Bal
            (N-Bal
              (N-Bal lp₁ (B-Bal lp₂))
              (N-Bal (B-Bal rp) (B-Bal L-Bal)))

        balance-lemma lp₁ (almost-left (N-Bal lp₂ rp) (N-Bal lp₃ rp₁)) =
          R-Bal
            (N-Bal
              (N-Bal lp₁ (B-Bal lp₂))
              (N-Bal (B-Bal rp) (B-Bal (N-Bal lp₃ rp₁))))

        balance-lemma lp₁ (almost-right L-Bal (N-Bal lp₂ rp)) =
          R-Bal
            (N-Bal
              (N-Bal lp₁ (B-Bal L-Bal))
              (N-Bal (B-Bal lp₂) (B-Bal rp)))

        balance-lemma lp₁ (almost-right (N-Bal lp₂ rp) (N-Bal lp₃ rp₁)) =
          R-Bal
            (N-Bal
              (N-Bal lp₁ (B-Bal (N-Bal lp₂ rp)))
              (N-Bal (B-Bal lp₃) (B-Bal rp₁)))

    insert'-on-B {._} {x₁} {l < B [ v ]> r} (N-Bal lp (B-Bal p)) | GT =
      balance-lemma lp (insert'-on-B p)
      where
        balance-lemma :
          ∀ {n} {y : X} {l r : RBT X n}
          → (l is-balanced)
          → (r is-balanced)
          → (balance-right (l < B [ y ]> r) is-balanced)

        balance-lemma lp₁ (R-Bal (N-Bal L-Bal L-Bal)) =
          B-Bal (N-Bal lp₁ (R-Bal (N-Bal L-Bal L-Bal)))
        balance-lemma lp₁ (R-Bal (N-Bal (N-Bal lp₂ rp) L-Bal)) =
          B-Bal (N-Bal lp₁ (R-Bal (N-Bal (N-Bal lp₂ rp) L-Bal)))
        balance-lemma lp₁ (R-Bal (N-Bal L-Bal (N-Bal lp₂ rp))) =
          B-Bal (N-Bal lp₁ (R-Bal (N-Bal L-Bal (N-Bal lp₂ rp))))
        balance-lemma lp₁ (R-Bal (N-Bal (N-Bal lp₂ rp) (N-Bal lp₃ rp₁))) =
          B-Bal (N-Bal lp₁ (R-Bal (N-Bal (N-Bal lp₂ rp) (N-Bal lp₃ rp₁))))
        balance-lemma lp₁ (B-Bal L-Bal) =
          B-Bal (N-Bal lp₁ (B-Bal L-Bal))
        balance-lemma lp₁ (B-Bal (N-Bal lp₂ rp)) =
          B-Bal (N-Bal lp₁ (B-Bal (N-Bal lp₂ rp)))

    insert'-on-R :
      ∀ {n} {x : X} {t : RBT X n}
      → (t is-R-balanced)
      → ((insert' x t) is-almost-balanced)

    insert'-on-R {_} {x} {l < R [ v ]> r} (N-Bal lp rp) with compare x v
    insert'-on-R {n} {x} {l < R [ v ]> r} (N-Bal lp rp) | LT with insert'-on-B {n} {x} lp
    insert'-on-R {n} {x} {l < R [ v ]> r} (N-Bal lp rp) | LT | R-Bal p =
      almost-left p rp
    insert'-on-R {n} {x} {l < R [ v ]> r} (N-Bal lp rp) | LT | B-Bal p =
      balanced (R-Bal (N-Bal p rp))
    insert'-on-R {n} {x} {l < R [ v ]> r} (N-Bal lp rp) | EQ =
      balanced (R-Bal (N-Bal lp rp))
    insert'-on-R {n} {x} {l < R [ v ]> r} (N-Bal lp rp) | GT with insert'-on-B {n} {x} rp
    insert'-on-R {n} {x} {l < R [ v ]> r} (N-Bal lp rp) | GT | R-Bal p =
      almost-right lp p
    insert'-on-R {n} {x} {l < R [ v ]> r} (N-Bal lp rp) | GT | B-Bal p =
      balanced (R-Bal (N-Bal lp p))





  blackening-strengthens-balance :
    ∀ {n} {t : RBT X n}
    → (t is-balanced)
    → (snd (blacken-root t) is-B-balanced)

  blackening-strengthens-balance (R-Bal (N-Bal lp rp)) =
    N-Bal (B-Bal lp) (B-Bal rp)
  blackening-strengthens-balance (B-Bal L-Bal) =
    L-Bal
  blackening-strengthens-balance (B-Bal (N-Bal lp rp)) =
    N-Bal lp rp






  insert-preserves-balance :
    ∀ {n} {x : X} {t : RBT X n}
    → (t is-B-balanced)
    → (snd (insert x t) is-B-balanced)

  insert-preserves-balance prf =
    blackening-strengthens-balance
      (insert'-on-B prf)