module rbt_t where

open import Data.Nat hiding (compare)
open import Data.Nat.Properties as NatProp -- <-cmp
open import Relation.Binary
open import Data.List

-- !$\rightarrow$! t を適用するために必要だった
open import Level renaming ( suc to succ ; zero to Zero )
open import Level


data col : Set where
  black : col
  red : col

record node (A B : Set) : Set where
  field
    coler : A
    number : B
open node

record tree (A B C : Set) : Set where
  field
    key : A
    ltree : B
    rtree : C
open tree

data rbt : Set where
  bt-empty : rbt
  bt-node  : (node : tree (node col !$\mathbb{N}$!) rbt rbt ) !$\rightarrow$! rbt

record Env : Set (succ Zero) where
  field
    vart : rbt
    varn : !$\mathbb{N}$!
    varl : List rbt
open Env


whileTest-next : {l : Level} {t : Set l}  !$\rightarrow$! (c10 : rbt) !$\rightarrow$! (n : !$\mathbb{N}$!) !$\rightarrow$! (list : List rbt) !$\rightarrow$! (next : Env !$\rightarrow$! t) !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t
whileTest-next c10 n list next exit = next (record {vart = c10 ; varn = n ; varl = list} )

merge-tree :  {le : Level} {t : Set le} !$\rightarrow$! Env !$\rightarrow$! (next : Env !$\rightarrow$! t) !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t

-- 全てmerge-treeへ
split : {le : Level} {t : Set le} !$\rightarrow$! Env !$\rightarrow$! (next : Env !$\rightarrow$! t) !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t
split node@record { vart = bt-empty ; varn = varn ; varl = varl } next exit = exit node
split node@record { vart = (bt-node record { key = key ; ltree = bt-empty ; rtree = rtree }) ; varn = varn ; varl = varl } next exit = exit node
split node@record { vart = (bt-node record { key = key ; ltree = (bt-node node!$\text{1}$!) ; rtree = bt-empty }) ; varn = varn ; varl = varl } next exit = exit node
split record { vart = (bt-node record { key = record { coler = coler!$\text{1}$! ; number = number!$\text{1}$! }
  ; ltree = (bt-node record { key = record { coler = ly ; number = ln } ; ltree = ll ; rtree = lr })
  ; rtree = (bt-node record { key = record { coler = ry ; number = rn } ; ltree = rl ; rtree = rr }) })
  ; varn = varn ; varl = varl } next exit
  = next record { vart = (bt-node record { key = record { coler = red ; number = number!$\text{1}$! }
    ; ltree = (bt-node record { key = record { coler = black ; number = ln } ; ltree = ll ; rtree = lr })
    ; rtree = (bt-node record { key = record { coler = black ; number = rn } ; ltree = rl ; rtree = rr }) })
    ; varn = varn ; varl = varl }


-- 右回転
-- 実行時splitへ，それ以外はmerge-treeへ
merge-rotate-right : {le : Level} {t : Set le} !$\rightarrow$! Env !$\rightarrow$! (next : Env !$\rightarrow$! t) !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t
merge-rotate-right node@record { vart = bt-empty ; varn = varn ; varl = varl } next exit = exit node
merge-rotate-right node@record { vart = bt-node record { key = record { coler = coler ; number = number } ; ltree = bt-empty ; rtree = r } ; varn = varn ; varl = varl } next exit = exit node
merge-rotate-right record { vart = bt-node record { key = record { coler = y ; number = x }
  ; ltree = (bt-node record { key =  record { coler = ly ; number = lx } ; ltree = ll ; rtree = lr })
  ; rtree = r }
  ; varn = varn ; varl = varl } next exit
    = next record { vart = bt-node record { key = record { coler = y ; number = lx }
      ; ltree = ll
      ; rtree = (bt-node record { key = record { coler = red ; number = x } ; ltree = lr; rtree = r }) }
      ; varn = varn ; varl = varl }


-- split
split-branch : {le : Level} {t : Set le} !$\rightarrow$! Env !$\rightarrow$! (next : Env !$\rightarrow$! t) !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t
split-branch node@record{ vart = bt-empty ; varn = varn ; varl = varl } next exit = exit node
split-branch node@record{ vart = bt-node record { key = key ; ltree = bt-empty ; rtree = rtree } ; varn = varn ; varl = varl } next exit = exit node
split-branch node@record{ vart = bt-node record { key = key!$\text{1}$! ; ltree = (bt-node record { key = record { coler = lc ; number = number } ; ltree = bt-empty ; rtree = lr }) ; rtree = rtree } ; varn = varn ; varl = varl } next exit = exit node
split-branch node@record{ vart = bt-node record { key = key!$\text{1}$! ; ltree = (bt-node record { key = record { coler = black ; number = number } ; ltree = (bt-node node!$\text{1}$!) ; rtree = lr }) ; rtree = rtree } ; varn = varn ; varl = varl } next exit = exit node
split-branch node@record{ vart = bt-node record { key = key!$\text{1}$! ; ltree = (bt-node record { key = record { coler = red ; number = number!$\text{1}$! } ; ltree = (bt-node record { key = record { coler = black ; number = number } ; ltree = ltree ; rtree = rtree!$\text{1}$! }) ; rtree = lr }) ; rtree = rtree } ; varn = varn ; varl = varl } next exit = exit node
split-branch node@record{ vart = bt-node record { key = key!$\text{1}$!
  ; ltree = (bt-node record { key = record { coler = red ; number = number!$\text{1}$! }
    ; ltree = (bt-node record { key = record { coler = red ; number = number } ; ltree = ltree ; rtree = rtree!$\text{1}$! })
    ; rtree = lr })
  ; rtree = rtree }
  ; varn = varn ; varl = varl } next exit
  = next node


-- 左回転，exitはsplit_branchへ　nextもsplit_branchへ
merge-rotate-left : {le : Level} {t : Set le} !$\rightarrow$! Env !$\rightarrow$! (next : Env !$\rightarrow$! t) !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t
merge-rotate-left node@record { vart = bt-empty ; varn = varn ; varl = varl }  next exit = exit node
merge-rotate-left node@record { vart = bt-node record { key = record { coler = coler ; number = number } ; ltree = ltree ; rtree = bt-empty } ; varn = varn ; varl = varl } next exit = exit node
merge-rotate-left record { vart = bt-node record { key = record { coler = y ; number = x } ; ltree = l ; rtree = (bt-node record { key = record { coler = ry ; number = rx } ; ltree = rl ; rtree = rr }) } ; varn = varn ; varl = varl } next exit
  = next record { vart = bt-node record { key = record { coler = y ; number = rx }
  　　; ltree = (bt-node record { key = record { coler = red ; number = x } ; ltree = l ; rtree = rl })
  　　; rtree = rr}
  　　; varn = varn ; varl = varl }


-- skew exitがsplitへ　nextが左回転
skew-bt : {le : Level} {t : Set le} !$\rightarrow$! Env !$\rightarrow$! (next : Env !$\rightarrow$! t) !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t
skew-bt node@record { vart = bt-empty ; varn = varn ; varl = varl } next exit =  exit node
skew-bt node@record { vart = (bt-node record { key = key ; ltree = ltree ; rtree = bt-empty }) ; varn = varn ; varl = varl } next exit = exit node
skew-bt node@record { vart = (bt-node record { key = key!$\text{1}$!
  ; ltree = ltree!$\text{1}$!
  ; rtree = (bt-node record { key = record { coler = black ; number = number }
  ; ltree = ltree ; rtree = rtree }) }) ; varn = varn ; varl = varl } next exit
    = exit node
skew-bt node@record { vart = (bt-node record { key = key!$\text{1}$!
  ; ltree = ltree!$\text{1}$!
  ; rtree = (bt-node record { key = record { coler = red ; number = number } ; ltree = ltree ; rtree = rtree }) })
  ; varn = varn ; varl = varl } next exit
    = next node


set-node-coler : Env !$\rightarrow$! Env
set-node-coler node@record { vart = bt-empty ; varn = varn ; varl = varl } = node
set-node-coler record { vart = (bt-node record { key = record { coler = coler ; number = number } ; ltree = ltree ; rtree = rtree }) ; varn = varn ; varl = varl }
  = record { vart = (bt-node record { key = record { coler = black ; number = number } ; ltree = ltree ; rtree = rtree }) ; varn = varn ; varl = varl }

init-node-coler : {le : Level} {t : Set le} !$\rightarrow$! Env !$\rightarrow$! (next : Env !$\rightarrow$! t) !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t
init-node-coler node@record { vart = bt-empty ; varn = varn ; varl = varl } next exit = exit node
init-node-coler record { vart = (bt-node record { key = record { coler = coler ; number = number } ; ltree = ltree ; rtree = rtree }) ; varn = varn ; varl = varl } next exit
  = exit record { vart = (bt-node record { key = record { coler = black ; number = number } ; ltree = ltree ; rtree = rtree }) ; varn = varn ; varl = varl }

--修正前
merge-tree node@record { vart = vart ; varn = varn ; varl = [] } next exit = exit node
merge-tree node@record { vart = vart ; varn = varn ; varl = (bt-empty !$\text{::}$! varl!$\text{1}$!) } next exit = exit node

merge-tree record { vart = vart ; varn = varn ; varl = (bt-node record { key = record { coler = coler ; number = number } ; ltree = ltree ; rtree = rtree } !$\text{::}$! varl!$\text{1}$!) } next exit with <-cmp varn number

merge-tree record { vart = vart ; varn = varn ; varl = (bt-node record { key = record { coler = coler ; number = number } ; ltree = ltree ; rtree = rtree } !$\text{::}$! varl!$\text{1}$!) } next exit | tri≈ !$\neg$!a b !$\neg$!c
  = next record { vart = (bt-node record { key = record { coler = coler ; number = number } ; ltree = ltree ; rtree = rtree }) ; varn = number ; varl = varl!$\text{1}$! }

merge-tree record { vart = vart ; varn = varn ; varl = (bt-node record { key = record { coler = coler ; number = number } ; ltree = ltree ; rtree = rtree } !$\text{::}$! varl!$\text{1}$!) } next exit | tri> !$\neg$!a !$\neg$!b c
  = next record { vart = (bt-node record { key = record { coler = coler ; number = number } ; ltree = ltree ; rtree = vart }) ; varn = number ; varl = varl!$\text{1}$! }

merge-tree record { vart = vart ; varn = varn ; varl = (bt-node record { key = record { coler = coler ; number = number } ; ltree = ltree ; rtree = rtree } !$\text{::}$! varl!$\text{1}$!) } next exit | tri< a !$\neg$!b !$\neg$!c
  = next record { vart = (bt-node record { key = record { coler = coler ; number = number } ; ltree = vart ; rtree = rtree }) ; varn = number ; varl = varl!$\text{1}$! }

-- do marge-tree
-- next merge-tree-c
-- exit fin
merge-tree-c :  {le : Level} {t : Set le} !$\rightarrow$! Env !$\rightarrow$! (next : Env !$\rightarrow$! t) !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t
merge-tree-c env next exit with varl env
... | [] = exit env
... | bt-empty !$\text{::}$! nl = exit env
... | bt-node xtree !$\text{::}$! varl with <-cmp (varn env) (number ( key xtree ))
... | tri≈ !$\neg$!a b !$\neg$!c = next (record env { vart = bt-node xtree ; varn = number (key xtree) ; varl = varl })
... | tri> !$\neg$!a !$\neg$!b c = next (record env { vart = (bt-node record xtree{rtree = Env.vart env}) ; varn = number (key xtree) ; varl = varl })
... | tri< a !$\neg$!b !$\neg$!c = next (record env { vart = (bt-node record xtree{ltree = Env.vart env}) ; varn = number (key xtree) ; varl = varl })

-- do brandh
-- next insert-c
-- exit merge-tree
insert-c : {le : Level} {t : Set le} !$\rightarrow$! Env !$\rightarrow$! (next : Env !$\rightarrow$! t) !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t
insert-c env next exit with Env.vart env
... | bt-empty = exit record env{vart = (bt-node (record { key = record { coler = red ; number = Env.varn env }; ltree = bt-empty; rtree = bt-empty }))}
... | bt-node node with <-cmp (Env.varn env) (node.number (tree.key node))
... | tri≈ !$\neg$!a b !$\neg$!c = exit env
... | tri> !$\neg$!a !$\neg$!b c = next record env{vart = (tree.ltree node) ; varl = (bt-node record node{ltree = bt-empty}) !$\text{::}$! (Env.varl env) }
... | tri< a !$\neg$!b !$\neg$!c = next record env{vart = (tree.rtree node) ; varl = (bt-node record node{rtree = bt-empty}) !$\text{::}$! (Env.varl env) }

-- insert
bt-insert : {le : Level} {t : Set le} !$\rightarrow$! Env !$\rightarrow$! (next : Env !$\rightarrow$! t) !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t
-- mergeへ遷移する
bt-insert (record { vart = bt-empty ; varn = varn ; varl = varl }) next exit
  = exit (record { vart = (bt-node (record { key = record { coler = red ; number = varn }; ltree = bt-empty; rtree = bt-empty })) ; varn = varn ; varl = varl })

-- 場合分けを行う
bt-insert record {vart = (bt-node record { key = record { coler = y; number = x } ; ltree = ltree ; rtree = rtree }) ; varn = n ; varl = varl } next exit with <-cmp x n
bt-insert node@(record { vart = (bt-node record { key = key ; ltree = ltree ; rtree = rtree }) ; varn = varn ; varl = varl } ) next exit | tri≈ !$\neg$!a b !$\neg$!c
  = exit node

bt-insert record {vart = (bt-node record { key = key ; ltree = ltree ; rtree = rtree }) ; varn = varn ; varl = varl } next exit | tri> !$\neg$!a !$\neg$!b c
  = next (record {vart = ltree ; varn = varn ; varl = (bt-node record { key = key ; ltree = bt-empty ; rtree = rtree }) !$\text{::}$! varl })

bt-insert record {vart = (bt-node record { key = key ; ltree = ltree ; rtree = rtree }) ; varn = varn ; varl = varl } next exit | tri< a !$\neg$!b !$\neg$!c
  = next (record {vart = rtree ; varn = varn ; varl = (bt-node record { key = key ; ltree = ltree ; rtree = bt-empty }) !$\text{::}$! varl })

-- normal loop without termination
{-# TERMINATING #-}
insert : {l : Level} {t : Set l} !$\rightarrow$! Env !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t
insert env exit = bt-insert env (!$\lambda$! env !$\rightarrow$! insert env exit ) exit

init-col : {l : Level} {t : Set l} !$\rightarrow$! Env !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t
init-col env exit = init-node-coler env exit exit

{-
{-# TERMINATING #-}
merge : {l : Level} {t : Set l} !$\rightarrow$! Env !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t

{-# TERMINATING #-}
split-p : {l : Level} {t : Set l} !$\rightarrow$! Env !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t
split-p env exit = split env (!$\lambda$! env !$\rightarrow$! merge env (!$\lambda$! env !$\rightarrow$! merge env exit ) ) exit

{-# TERMINATING #-}
rotate_right : {l : Level} {t : Set l} !$\rightarrow$! Env !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t
rotate_right env exit = merge-rotate-right env (!$\lambda$! env !$\rightarrow$! split-p env (!$\lambda$! env !$\rightarrow$! merge env exit ) ) exit

{-# TERMINATING #-}
split-b : {l : Level} {t : Set l} !$\rightarrow$! Env !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t
split-b env exit = split-branch env (!$\lambda$! env !$\rightarrow$! rotate_right env exit ) !$\lambda$! env !$\rightarrow$! merge env exit

{-# TERMINATING #-}
rotate_left : {l : Level} {t : Set l} !$\rightarrow$! Env !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t
rotate_left env exit = merge-rotate-left env (!$\lambda$! env !$\rightarrow$! split-b env exit ) exit

{-# TERMINATING #-}
skew : {l : Level} {t : Set l} !$\rightarrow$! Env !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t
skew env exit = skew-bt env (!$\lambda$! env !$\rightarrow$! rotate_left env (!$\lambda$! env !$\rightarrow$! split-b env exit ) ) exit

merge env exit = merge-tree env (!$\lambda$! env !$\rightarrow$! skew env exit ) exit
-}
-- skewはnextがrotate-right．exitはsplitとなる
-- skewの中にrotate-rightが内包され，実行後は必ずsplitに遷移する
-- それはskewのexitと等しいので同時に記述してやる．
skew!$\prime$! : {l : Level} {t : Set l} !$\rightarrow$! Env !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t
skew!$\prime$! env exit = skew-bt env (!$\lambda$! env !$\rightarrow$! merge-rotate-left env exit exit ) exit

split!$\prime$! : {l : Level} {t : Set l} !$\rightarrow$! Env !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t
split!$\prime$! env exit = split-branch env (!$\lambda$! env !$\rightarrow$! merge-rotate-right env (!$\lambda$! env !$\rightarrow$! split env exit exit ) (!$\lambda$! env !$\rightarrow$! split env exit exit ) ) exit


{-
merge!$\prime$! : {l : Level} {t : Set l} !$\rightarrow$! Env !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t
merge!$\prime$! node@record { vart = vart ; varn = varn ; varl = [] } exit = exit node
merge!$\prime$! node@record { vart = vart ; varn = varn ; varl = (x !$\text{::}$! varl!$\text{1}$!) } exit = merge!$\prime$! (merge-tree node (!$\lambda$! env !$\rightarrow$! skew!$\prime$! env (!$\lambda$! env !$\rightarrow$! split!$\prime$! env (!$\lambda$! env !$\rightarrow$! env) ) ) (!$\lambda$! env !$\rightarrow$! env) ) exit
-}

whileTestP : {l : Level} {t : Set l} !$\rightarrow$! (tree : rbt) !$\rightarrow$! (n : !$\mathbb{N}$!) !$\rightarrow$! (Code : Env !$\rightarrow$! t) !$\rightarrow$! t
whileTestP tree n next = next (record {vart = tree ; varn = n ; varl = [] } )



{-# TERMINATING #-}
mergeP : {l : Level} {t : Set l} !$\rightarrow$! Env !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t
mergeP env exit = merge-tree env (!$\lambda$! env !$\rightarrow$! skew!$\prime$! env (!$\lambda$! env !$\rightarrow$! split!$\prime$! env (!$\lambda$! env !$\rightarrow$! mergeP env exit)) ) exit

{-
mergeP1 : {l : Level} {t : Set l} !$\rightarrow$! Env !$\rightarrow$! (exit : Env !$\rightarrow$! t) !$\rightarrow$! t
mergeP1 record { vart = vart ; varn = varn ; varl = [] } exit = {!!}
mergeP1 record { vart = vart ; varn = varn ; varl = (x !$\text{::}$! varl!$\text{1}$!) } exit = {!!}
-}

{-
merge-tree env
  (!$\lambda$! env !$\rightarrow$! skew-bt env
    (!$\lambda$! env !$\rightarrow$! merge-rotate-left env
      (!$\lambda$! env !$\rightarrow$! split-branch env
        (!$\lambda$! env !$\rightarrow$!  merge-rotate-right env
          (!$\lambda$! env !$\rightarrow$! split env (!$\lambda$! env !$\rightarrow$! mergeP env exit ) (!$\lambda$! env !$\rightarrow$! mergeP env exit ) ) exit)
      (!$\lambda$! env !$\rightarrow$! mergeP env exit ) )
    exit )
  exit ) exit
-}

--whileTestP : {l : Level} {t : Set l} !$\rightarrow$! (c10 : !$\mathbb{N}$!) !$\rightarrow$! (c11 : !$\mathbb{N}$!) !$\rightarrow$! (Code : Envc !$\rightarrow$! t) !$\rightarrow$! t
--whileTestP c10 n next = next (record {varn = c10 ; vari = 0 ; c10 = c10 } )

--whileTestPCall : (tree : rbt) !$\rightarrow$! (n : !$\mathbb{N}$!)  !$\rightarrow$! Env
--whileTestPCall tree n = whileTestP {_} {_} tree n (!$\lambda$! env !$\rightarrow$! insert env (!$\lambda$! env !$\rightarrow$! merge env (!$\lambda$! env !$\rightarrow$! init-col env (!$\lambda$! env !$\rightarrow$! env ) ) ) )

whileTestPCall!$\prime$! : (tree : rbt) !$\rightarrow$! (n : !$\mathbb{N}$!)  !$\rightarrow$! Env
whileTestPCall!$\prime$! tree n = whileTestP {_} {_} tree n (!$\lambda$! env !$\rightarrow$! insert env (!$\lambda$! env !$\rightarrow$! mergeP env (!$\lambda$! env !$\rightarrow$! init-col env (!$\lambda$! env !$\rightarrow$! env ) ) ) )

-- 以下test部分
test1 = whileTestPCall!$\prime$! bt-empty 8
test2 = whileTestPCall!$\prime$! (vart test1) 10
test3 = whileTestPCall!$\prime$! (vart test2) 24
test4 = whileTestPCall!$\prime$! (vart test3) 31
test5 = whileTestPCall!$\prime$! (vart test4) 41
test6 = whileTestPCall!$\prime$! (vart test5) 45
test7 = whileTestPCall!$\prime$! (vart test6) 50
test8 = whileTestPCall!$\prime$! (vart test7) 59
test9 = whileTestPCall!$\prime$! (vart test8) 73
test10 = whileTestPCall!$\prime$! (vart test9) 74
test11 = whileTestPCall!$\prime$! (vart test10) 79