avlset.dats

datatype avlt(x: t@ype, n: int) =
    Empty(x, 0)
  | {l, m: nat |
    (l + 2 == n && m + 1 == n) ||
    (l + 1 == n && m + 1 == n) ||
    (l + 1 == n && m + 2 == n)
} Node of (avlt(x, l), x, avlt(x, m), int(n))




fun {x:t@ype}height{n: nat}(tree: avlt(x, n)): int(n) = 
  case+ tree of
  | Empty () => 0
  | Node (_, _, _, n) => n


(* 
 Output type is not acculate. According to type signature, it can happen where
 l + 1 == m and n == m + 1.
*)  
fun {x:t@ype}create{
  l, m: nat |
  l + 1 == m ||
  l == m + 1 ||
  l == m
  } (l: avlt(x, l), v: x, r: avlt(x, m)): [n: nat | n == m + 1 || n == l + 1] avlt(x, n) = let
  val hl = height l
  val hr = height r
in
  if hl >= hr
  then Node(l, v, r, hl + 1)
  else Node(l, v, r, hr + 1)
end

fun {x:t@ype}rotate_right{hl, hr: nat |
  hl == hr + 2
}(l: avlt(x, hl), v: x, r: avlt(x, hr)): [n: nat| n == hl || n == hl + 1] avlt(x, n) = let
  val hl = height l
  val hr = height r
  val+ Node(ll, lv, lr, _) = l
  val hll = height ll
  val hlr = height lr
in
  if  hll >= hlr
  then create(ll, lv, create(lr, v, r))
  else let
       val+ Node(lrl, lrv, lrr, _) =  lr
       val hlrl = height lrl
       val hlrr = height lrr
  in
      if  hlrl = hlrr
      then create(create(ll, lv, lrl), lrv, create(lrr, v, r))
        (* because `create` is a bit fuzzy, you need to write height manually *)  
      else Node(create(ll, lv, lrl), lrv, create(lrr, v, r), hr+2)
  end
end


fun {x:t@ype}rotate_left{hl, hr: nat |
  hl + 2 == hr
}(l: avlt(x, hl), v: x, r: avlt(x, hr)): [n: nat| n == hr || n == hr + 1] avlt(x, n) = let
  val hl = height l
  val hr = height r
  val+ Node(rl, rv, rr, _) = r
  val hrl = height rl
  val hrr = height rr
in
  if  hrr >= hrl
  then create(create(l, v, rl), rv, rr)
  else let
    val+ Node(rll, rlv, rlr, _) =  rl
    val hrll = height rll
    val hrlr = height rlr
  in
      if  hrll = hrlr
      then create(create(l, v, rll), rlv, create(rlr, rv, rr))
        (* because `create` is a bit fuzzy, you need to write height manually *)  
      else Node(create(l, v, rll), rlv, create(rlr, rv, rr), hl+2)
  end
end




extern fun {x:t@ype}cmp(x: x, y: x): int

fun {x:t@ype}empty(): avlt(x, 0) = Empty()
fun {x:t@ype}singleton(x: x): avlt(x, 1) = Node(Empty, x, Empty, 1)


fun {x:t@ype}insert{m: nat}(x: x, tree: avlt(x, m)): [n: nat | n == m || n == m + 1]avlt(x, n) =
  case+ tree of
  | Empty () => singleton(x)
  | t as Node(l, v, r, _) => let
    val c = cmp(x, v)
  in
    if c = 0 then t
    else if c < 0
    then let
       val tmp = insert(x, l)
       val hr = height r
       val ht = height tmp
     in
      if ht = hr + 2
      then rotate_right(tmp, v, r)
      else if ht = hr + 1
      then Node(tmp, v, r, ht + 1)
      else Node(tmp, v, r, hr + 1)
     end
    else let
       val tmp = insert(x, r)
       val hl = height l
       val ht = height tmp
     in
      if ht = hl + 2
      then rotate_left(l, v, tmp)
      else if ht = hl + 1
      then Node(l, v, tmp, ht + 1)
      else Node(l, v, tmp, hl + 1)
     end
  end


fun {x:t@ype}find{m: nat}(x: x, tree: avlt(x, m)): Option(x) =
  case+ tree of
  | Empty () => option_none()
  | Node(l, v, r, _) => let
    val c = cmp(x, v)
  in
    if c = 0
    then option_some(v)
    else if c < 0
    then find(x, r)
    else find(x, l)
  end
  
  
implement
{x: string}
cmp(x, y) = strcmp<>(x, y)
  
val tree = Empty
val tree = insert("a", tree)