(*------------------------------------------------------------------*)

#define ATS_DYNLOADFLAG 0

#include "share/atspre_staload.hats"

(*------------------------------------------------------------------*)

(*

  Persistent AVL trees.

  References:

    * Niklaus Wirth, 1976. Algorithms + Data Structures =
      Programs. Prentice-Hall, Englewood Cliffs, New Jersey.

    * Niklaus Wirth, 2004. Algorithms and Data Structures. Updated
      by Fyodor Tkachov, 2014.

  (Note: Wirth’s implementations, which are in Pascal and Oberon, are
  for non-persistent trees.)

*)

(*------------------------------------------------------------------*)

(*

  For now, a very simple interface, without much provided in the way
  of proofs.

  You could put all this interface stuff into a .sats file. (You would
  have to remove the word ‘extern’ from the definitions.)

  You might also make avl_t abstract, and put these details in the
  .dats file; you would use ‘assume’ to identify the abstract type
  with an implemented type. That approach would require some name
  changes, and also would make pattern matching on the trees
  impossible outside their implementation. Having users do pattern
  matching on the AVL trees probably is a terrible idea, anyway.

*)

datatype bal_t =
| bal_minus1
| bal_zero
| bal_plus1

datatype avl_t (key_t  : [email protected]+,
                data_t : [email protected]+,
                size   : int) =
| avl_t_nil (key_t, data_t, 0)
| {size_L, size_R : nat}
  avl_t_cons (key_t, data_t, size_L + size_R + 1) of
    (key_t, data_t, bal_t,
     avl_t (key_t, data_t, size_L),
     avl_t (key_t, data_t, size_R))
typedef avl_t (key_t  : [email protected]+,
               data_t : [email protected]+) =
  [size : int] avl_t (key_t, data_t, size)

extern prfun
lemma_avl_t_param :
  {key_t, data_t : [email protected]}
  {size : int}
  avl_t (key_t, data_t, size) -<prf> [0 <= size] void

(* Implement this template, for whichever type of key you are
   using. It should return a negative number if u < v, zero if
   u = v, or a positive number if u > v. *)
extern fun {key_t : [email protected]}
avl_t$compare (u : key_t, v : key_t) :<> int

(* Is the AVL tree empty? *)
extern fun
avl_t_is_empty
          {key_t  : [email protected]}
          {data_t : [email protected]}
          {size   : int}
          (avl    : avl_t (key_t, data_t, size)) :<>
    [b : bool | b == (size == 0)]
    bool b

(* Does the AVL tree contain at least one association? *)
extern fun
avl_t_isnot_empty
          {key_t  : [email protected]}
          {data_t : [email protected]}
          {size   : int}
          (avl    : avl_t (key_t, data_t, size)) :<>
    [b : bool | b == (size <> 0)]
    bool b

(* How many associations are stored in the AVL tree? (Currently we
   have no way to do an avl_t_size that preserves the ‘size’ static
   value. This is the best we can do.) *)
extern fun {key_t  : [email protected]}
           {data_t : [email protected]}
avl_t_size {size : int}
           (avl  : avl_t (key_t, data_t, size)) :<>
    [sz : int | (size == 0 && sz == 0) || (0 < size && 0 < sz)]
    size_t sz

(* Does the AVL tree contain the given key? *)
extern fun {key_t  : [email protected]}
           {data_t : [email protected]}
avl_t_has_key
          {size : int}
          (avl  : avl_t (key_t, data_t, size),
           key  : key_t) :<>
    bool

(* Search for a key. If the key is found, return the data value
   associated with it. Otherwise return the value of ‘dflt’. *)
extern fun {key_t  : [email protected]}
           {data_t : [email protected]}
avl_t_search_dflt
          {size : int}
          (avl  : avl_t (key_t, data_t, size),
           key  : key_t,
           dflt : data_t) :<>
    data_t

(* Search for a key. If the key is found, return
   ‘Some(data)’. Otherwise return ‘None()’. *)
extern fun {key_t  : [email protected]}
           {data_t : [email protected]}
avl_t_search_opt
          {size : int}
          (avl  : avl_t (key_t, data_t, size),
           key  : key_t) :<>
    Option (data_t)

(* Search for a key. If the key is found, set ‘found’ to true, and set
   ‘data’. Otherwise set ‘found’ to false. *)
extern fun {key_t  : [email protected]}
           {data_t : [email protected]}
avl_t_search_ref
          {size  : int}
          (avl   : avl_t (key_t, data_t, size),
           key   : key_t,
           data  : &data_t? >> opt (data_t, found),
           found : &bool? >> bool found) :<!wrt>
    #[found : bool]
    void

(* Overload avl_t_search; these functions are easy for the compiler to
   distinguish. *)
overload avl_t_search with avl_t_search_dflt
overload avl_t_search with avl_t_search_opt
overload avl_t_search with avl_t_search_ref

(* If a key is not present in the AVL tree, insert the key-data
   association; return the new AVL tree. If the key *is* present in
   the AVL tree, then *replace* the key-data association; return the
   new AVL tree. *)
extern fun {key_t  : [email protected]}
           {data_t : [email protected]}
avl_t_insert
          {size : int}
          (avl  : avl_t (key_t, data_t, size),
           key  : key_t,
           data : data_t) :<>
    [sz : pos]
    avl_t (key_t, data_t, sz)

(* If a key is not present in the AVL tree, insert the key-data
   association; return the new AVL tree and ‘true’. If the key *is*
   present in the AVL tree, then *replace* the key-data association;
   return the new AVL tree and ‘false’. *)
extern fun {key_t  : [email protected]}
           {data_t : [email protected]}
avl_t_insert_or_replace
          {size : int}
          (avl  : avl_t (key_t, data_t, size),
           key  : key_t,
           data : data_t) :<>
    [sz : pos]
    (avl_t (key_t, data_t, sz), bool)

(* If a key is present in the AVL tree, delete the key-data
   association; otherwise return the tree as it came. *)
extern fun {key_t  : [email protected]}
           {data_t : [email protected]}
avl_t_delete
          {size : int}
          (avl  : avl_t (key_t, data_t, size),
           key  : key_t) :<>
    [sz : nat]
    avl_t (key_t, data_t, sz)

(* If a key is present in the AVL tree, delete the key-data
   association; otherwise return the tree as it came. Also, return a
   bool to indicate whether or not the key was found; ‘true’ if found,
   ‘false’ if not. *)
extern fun {key_t  : [email protected]}
           {data_t : [email protected]}
avl_t_delete_if_found
          {size : int}
          (avl  : avl_t (key_t, data_t, size),
           key  : key_t) :<>
    [sz : nat]
    (avl_t (key_t, data_t, sz), bool)

(* Return a sorted list of the association pairs in an AVL
   tree. (Currently we have no way to do an avl_t_pairs that preserves
   the ‘size’ static value. This is the best we can do.) *)
extern fun {key_t  : [email protected]}
           {data_t : [email protected]}
avl_t_pairs {size : int}
            (avl  : avl_t (key_t, data_t, size)) :<>
    [sz : int | (size == 0 && sz == 0) || (0 < size && 0 < sz)]
    list ((key_t, data_t), sz)

(* Return a sorted list of the keys in an AVL tree. (Currently we have
   no way to do an avl_t_keys that preserves the ‘size’ static
   value. This is the best we can do.) *)
extern fun {key_t  : [email protected]}
           {data_t : [email protected]}
avl_t_keys {size : int}
           (avl  : avl_t (key_t, data_t, size)) :<>
    [sz : int | (size == 0 && sz == 0) || (0 < size && 0 < sz)]
    list (key_t, sz)

(* Return a list of the data values in an AVL tree, sorted in the
   order of their keys. (Currently we have no way to do an avl_t_data
   that preserves the ‘size’ static value. This is the best we can
   do.) *)
extern fun {key_t  : [email protected]}
           {data_t : [email protected]}
avl_t_data {size : int}
           (avl  : avl_t (key_t, data_t, size)) :<>
    [sz : int | (size == 0 && sz == 0) || (0 < size && 0 < sz)]
    list (data_t, sz)

(* list2avl_t does the reverse of what avl_t_pairs does (although
   they are not inverses of each other).
   Currently we have no way to do a list2avl_t that preserves the
   ‘size’ static value. This is the best we can do. *)
extern fun {key_t  : [email protected]}
           {data_t : [email protected]}
list2avl_t {size : int}
           (lst  : list ((key_t, data_t), size)) :<>
    [sz : int | (size == 0 && sz == 0) || (0 < size && 0 < sz)]
    avl_t (key_t, data_t, sz)

(* Make a closure that returns association pairs in either forwards or
   reverse order. *)
extern fun {key_t  : [email protected]}
           {data_t : [email protected]}
avl_t_make_pairs_generator
          {size      : int}
          (avl       : avl_t (key_t, data_t, size),
           direction : int) :
    () -<cloref1> Option @(key_t, data_t)

(* Make a closure that returns keys in either forwards or reverse
   order. *)
extern fun {key_t  : [email protected]}
           {data_t : [email protected]}
avl_t_make_keys_generator
          {size      : int}
          (avl       : avl_t (key_t, data_t, size),
           direction : int) :
    () -<cloref1> Option key_t

(* Make a closure that returns data values in forwards or reverse
   order of their keys. *)
extern fun {key_t  : [email protected]}
           {data_t : [email protected]}
avl_t_make_data_generator
          {size      : int}
          (avl       : avl_t (key_t, data_t, size),
           direction : int) :
    () -<cloref1> Option data_t

(* Raise an assertion if the AVL condition is not met. This template
   is for testing the code. *)
extern fun {key_t  : [email protected]}
           {data_t : [email protected]}
avl_t_check_avl_condition
          {size : int}
          (avl  : avl_t (key_t, data_t, size)) :
    void

(* Print an AVL tree to standard output, in some useful and perhaps
   even pretty format. *)
extern fun {key_t  : [email protected]}
           {data_t : [email protected]}
avl_t_pretty_print
          {size : int}
          (avl  : avl_t (key_t, data_t, size)) :
    void

(* Implement this template for whichever types of keys and data you
   wish to pretty print. *)
extern fun {key_t  : [email protected]}
           {data_t : [email protected]}
avl_t_pretty_print$key_and_data
          (key  : key_t,
           data : data_t) :
    void

(*------------------------------------------------------------------*)

(*

  What follows is the implementation. It would go into a .dats
  file. Note, however, that the .dats file would have to be staloaded!
  (Preferably anonymously.) This is because the implementation
  contains template functions.

  Notice there are several assertions with ‘$effmask_ntm’ (as opposed
  to proofs) that the routines are terminating. One hopes to remedy
  that problem (with proofs).

  Also there are some ‘$effmask_wrt’, but these effect masks are safe,
  because the writing is to our own stack variables.

*)

#define NIL avl_t_nil ()
#define CONS avl_t_cons
#define LNIL list_nil ()
#define :: list_cons
#define F false
#define T true

typedef fixbal_t = bool

primplement
lemma_avl_t_param avl =
  case+ avl of
  | NIL => ()
  | CONS _ => ()

fn {}
minus_neg_bal (bal : bal_t) :<> bal_t =
  case+ bal of
  | bal_minus1 () => bal_plus1
  | _ => bal_zero ()

fn {}
minus_pos_bal (bal : bal_t) :<> bal_t =
  case+ bal of
  | bal_plus1 () => bal_minus1
  | _ => bal_zero ()

fn {}
bal2int (bal : bal_t) :<> int =
  case+ bal of
  | bal_minus1 () => ~1
  | bal_zero () => 0
  | bal_plus1 () => 1

implement
avl_t_is_empty avl =
  case+ avl of
  | NIL => T
  | CONS _ => F

implement
avl_t_isnot_empty avl =
  ~avl_t_is_empty avl

implement {key_t} {data_t}
avl_t_size {siz} avl =
  let
    fun
    traverse {size : int}
             (p    : avl_t (key_t, data_t, size)) :<!ntm>
        [sz : int | (size == 0 && sz == 0) ||
                    (0 < size && 0 < sz)]
        size_t sz =
      case+ p of
      | NIL => i2sz 0
      | CONS (_, _, _, left, right) =>
        let
          val [sz_L : int] sz_L = traverse left
          val [sz_R : int] sz_R = traverse right
          prval _ = prop_verify {0 <= sz_L} ()
          prval _ = prop_verify {0 <= sz_R} ()
        in
          succ (sz_L + sz_R)
        end

    val [sz : int] sz = $effmask_ntm (traverse {siz} avl)
    prval _ = prop_verify {(siz == 0 && sz == 0) ||
                           (0 < siz && 0 < sz)} ()
  in
    sz
  end

implement {key_t} {data_t}
avl_t_has_key (avl, key) =
  let
    fun
    search {size : int}
           (p    : avl_t (key_t, data_t, size)) :<!ntm>
        bool =
      case+ p of
      | NIL => F
      | CONS (k, _, _, left, right) =>
        begin
          case+ avl_t$compare<key_t> (key, k) of
          | cmp when cmp < 0 => search left
          | cmp when cmp > 0 => search right
          | _ => T
        end
  in
    $effmask_ntm search avl
  end

implement {key_t} {data_t}
avl_t_search_dflt (avl, key, dflt) =
  let
    var data : data_t?
    var found : bool?
    val _ = $effmask_wrt avl_t_search_ref (avl, key, data, found)
  in
    if found then
      let
        prval _ = opt_unsome data
      in
        data
      end
    else
      let
        prval _ = opt_unnone data
      in
        dflt
      end
  end

implement {key_t} {data_t}
avl_t_search_opt (avl, key) =
  let
    var data : data_t?
    var found : bool?
    val _ = $effmask_wrt avl_t_search_ref (avl, key, data, found)
  in
    if found then
      let
        prval _ = opt_unsome data
      in
        Some {data_t} data
      end
    else
      let
        prval _ = opt_unnone data
      in
        None {data_t} ()
      end
  end

implement {key_t} {data_t}
avl_t_search_ref (avl, key, data, found) =
  let
    fun
    search (p     : avl_t (key_t, data_t),
            data  : &data_t? >> opt (data_t, found),
            found : &bool? >> bool found) :<!wrt,!ntm>
        #[found : bool] void =
      case+ p of
      | NIL =>
        {
          prval _ = opt_none {data_t} data
          val _ = found := F
        }
      | CONS (k, d, _, left, right) =>
        begin
          case+ avl_t$compare<key_t> (key, k) of
          | cmp when cmp < 0 => search (left, data, found)
          | cmp when cmp > 0 => search (right, data, found)
          | _ =>
            {
              val _ = data := d
              prval _ = opt_some {data_t} data
              val _ = found := T
            }
        end
  in
    $effmask_ntm search (avl, data, found)
  end

implement {key_t} {data_t}
avl_t_insert (avl, key, data) =
  let
    val (avl, _) =
      avl_t_insert_or_replace<key_t><data_t> (avl, key, data)
  in
    avl
  end

implement {key_t} {data_t}
avl_t_insert_or_replace (avl, key, data) =
  let
    fun
    search {size   : nat}
           (p      : avl_t (key_t, data_t, size),
            fixbal : fixbal_t,
            found  : bool) :<!ntm>
        [sz : pos]
        (avl_t (key_t, data_t, sz), fixbal_t, bool) =
      case+ p of
      | NIL =>
        (* The key was not found. Insert a new node. The tree will
            need rebalancing. *)
        (CONS (key, data, bal_zero, NIL, NIL), T, F)
      | CONS (k, d, bal, left, right) =>
        case+ avl_t$compare<key_t> (key, k) of
        | cmp when cmp < 0 =>
          let
            val (p1, fixbal, found) = search (left, fixbal, found)
          in
            (* If fixbal is T, then a node has been inserted
               on the left, and rebalancing may be necessary. *)
            case+ (fixbal, bal) of
            | (F, _) =>
              (* No rebalancing is necessary. *)
              (CONS (k, d, bal, p1, right), F, found)
            | (T, bal_plus1 ()) =>
              (* No rebalancing is necessary. *)
              (CONS (k, d, bal_zero (), p1, right), F, found)
            | (T, bal_zero ()) =>
              (* Rebalancing might still be necessary. *)
              (CONS (k, d, bal_minus1 (), p1, right), fixbal, found)
            | (T, bal_minus1 ()) =>
              (* Rebalancing is necessary. *)
              let
                val+ CONS (k1, d1, bal1, left1, right1) = p1
              in
                case+ bal1 of
                | bal_minus1 () =>
                  (* A single LL rotation. *)
                  let
                    val q = CONS (k, d, bal_zero (), right1, right)
                    val q1 = CONS (k1, d1, bal_zero (), left1, q)
                  in
                    (q1, F, found)
                  end
                | _ =>
                  (* A double LR rotation. *)
                  let
                    val p2 = right1
                    val- CONS (k2, d2, bal2, left2, right2) = p2
                    val q = CONS (k, d, minus_neg_bal bal2,
                                  right2, right)
                    val q1 = CONS (k1, d1, minus_pos_bal bal2,
                                   left1, left2)
                    val q2 = CONS (k2, d2, bal_zero (), q1, q)
                  in
                    (q2, F, found)
                  end
              end
          end
        | cmp when cmp > 0 =>
          let
            val (p1, fixbal, found) = search (right, fixbal, found)
          in
            (* If fixbal is T, then a node has been inserted
               on the right, and rebalancing may be necessary. *)
            case+ (fixbal, bal) of
            | (F, _) =>
              (* No rebalancing is necessary. *)
              (CONS (k, d, bal, left, p1), F, found)
            | (T, bal_minus1 ()) =>
              (* No rebalancing is necessary. *)
              (CONS (k, d, bal_zero (), left, p1), F, found)
            | (T, bal_zero ()) =>
              (* Rebalancing might still be necessary. *)
              (CONS (k, d, bal_plus1 (), left, p1), fixbal, found)
            | (T, bal_plus1 ()) =>
              (* Rebalancing is necessary. *)
              let
                val+ CONS (k1, d1, bal1, left1, right1) = p1
              in
                case+ bal1 of
                | bal_plus1 () =>
                  (* A single RR rotation. *)
                  let
                    val q = CONS (k, d, bal_zero (), left, left1)
                    val q1 = CONS (k1, d1, bal_zero (), q, right1)
                  in
                    (q1, F, found)
                  end
                | _ =>
                  (* A double RL rotation. *)
                  let
                    val p2 = left1
                    val- CONS (k2, d2, bal2, left2, right2) = p2
                    val q = CONS (k, d, minus_pos_bal bal2,
                                  left, left2)
                    val q1 = CONS (k1, d1, minus_neg_bal bal2,
                                   right2, right1)
                    val q2 = CONS (k2, d2, bal_zero (), q, q1)
                  in
                    (q2, F, found)
                  end
              end
          end
        | _ =>
          (* The key was found; p is an existing node. Replace
             it. The tree needs no rebalancing. *)
          (CONS (key, data, bal, left, right), F, T)
  in
    if avl_t_is_empty avl then
      (* Start a new tree. *)
      (CONS (key, data, bal_zero, NIL, NIL), F)
    else
      let
        prval _ = lemma_avl_t_param avl
        val (avl, _, found) = $effmask_ntm search (avl, F, F)
      in
        (avl, found)
      end
  end

fn {key_t  : [email protected]}
   {data_t : [email protected]}
balance_for_shrunken_left
          {size : pos}
          (p    : avl_t (key_t, data_t, size)) :<>
    (* Returns a new avl_t, and a ‘fixbal’ flag. *)
    [sz : pos]
    (avl_t (key_t, data_t, sz), fixbal_t) =
  let
    val+ CONS (k, d, bal, left, right) = p
  in
    case+ bal of
    | bal_minus1 () => (CONS (k, d, bal_zero, left, right), T)
    | bal_zero () => (CONS (k, d, bal_plus1, left, right), F)
    | bal_plus1 () =>
      (* Rebalance. *)
      let
        val p1 = right
        val- CONS (k1, d1, bal1, left1, right1) = p1
      in
        case+ bal1 of
        | bal_zero () =>
          (* A single RR rotation. *)
          let
            val q = CONS (k, d, bal_plus1, left, left1)
            val q1 = CONS (k1, d1, bal_minus1, q, right1)
          in
            (q1, F)
          end
        | bal_plus1 () =>
          (* A single RR rotation. *)
          let
            val q = CONS (k, d, bal_zero, left, left1)
            val q1 = CONS (k1, d1, bal_zero, q, right1)
          in
            (q1, T)
          end
        | bal_minus1 () =>
          (* A double RL rotation. *)
          let
            val p2 = left1
            val- CONS (k2, d2, bal2, left2, right2) = p2
            val q = CONS (k, d, minus_pos_bal bal2, left, left2)
            val q1 = CONS (k1, d1, minus_neg_bal bal2, right2, right1)
            val q2 = CONS (k2, d2, bal_zero, q, q1)
          in
            (q2, T)
          end
      end
  end

fn {key_t  : [email protected]}
   {data_t : [email protected]}
balance_for_shrunken_right
          {size : pos}
          (p    : avl_t (key_t, data_t, size)) :<>
    (* Returns a new avl_t, and a ‘fixbal’ flag. *)
    [sz : pos]
    (avl_t (key_t, data_t, sz), fixbal_t) =
  let
    val+ CONS (k, d, bal, left, right) = p
  in
    case+ bal of
    | bal_plus1 () => (CONS (k, d, bal_zero, left, right), T)
    | bal_zero () => (CONS (k, d, bal_minus1, left, right), F)
    | bal_minus1 () =>
      (* Rebalance. *)
      let
        val p1 = left
        val- CONS (k1, d1, bal1, left1, right1) = p1
      in
        case+ bal1 of
        | bal_zero () =>
          (* A single LL rotation. *)
          let
            val q = CONS (k, d, bal_minus1, right1, right)
            val q1 = CONS (k1, d1, bal_plus1, left1, q)
          in
            (q1, F)
          end
        | bal_minus1 () =>
          (* A single LL rotation. *)
          let
            val q = CONS (k, d, bal_zero, right1, right)
            val q1 = CONS (k1, d1, bal_zero, left1, q)
          in
            (q1, T)
          end
        | bal_plus1 () =>
          (* A double LR rotation. *)
          let
            val p2 = right1
            val- CONS (k2, d2, bal2, left2, right2) = p2
            val q = CONS (k, d, minus_neg_bal bal2, right2, right)
            val q1 = CONS (k1, d1, minus_pos_bal bal2, left1, left2)
            val q2 = CONS (k2, d2, bal_zero, q1, q)
          in
            (q2, T)
          end
      end
  end

implement {key_t} {data_t}
avl_t_delete (avl, key) =
  (avl_t_delete_if_found (avl, key)).0

implement {key_t} {data_t}
avl_t_delete_if_found (avl, key) =
  let
    fn
    balance_L__ {size : pos}
                (p    : avl_t (key_t, data_t, size)) :<>
        [sz : pos]
        (avl_t (key_t, data_t, sz), fixbal_t) =
      balance_for_shrunken_left<key_t><data_t> p
    fn
    balance_R__ {size : pos}
                (p    : avl_t (key_t, data_t, size)) :<>
        [sz : pos]
        (avl_t (key_t, data_t, sz), fixbal_t) =
      balance_for_shrunken_right<key_t><data_t> p

    fn {}
    balance_L {size   : pos}
              (p      : avl_t (key_t, data_t, size),
               fixbal : fixbal_t) :<>
        [sz : pos]
        (avl_t (key_t, data_t, sz), fixbal_t) =
      if fixbal then
        balance_L__ p
      else
        (p, F)

    fn {}
    balance_R {size   : pos}
              (p      : avl_t (key_t, data_t, size),
               fixbal : fixbal_t) :<>
        [sz : pos]
        (avl_t (key_t, data_t, sz), fixbal_t) =
      if fixbal then
        balance_R__ p
      else
        (p, F)

    fun
    del {size   : pos}
        (r      : avl_t (key_t, data_t, size),
         fixbal : fixbal_t) :<!ntm>
        (* Returns a new avl_t, a new fixbal, and key and data to be
           ‘moved up the tree’. *)
        [sz : nat]
        (avl_t (key_t, data_t, sz), fixbal_t, key_t, data_t) =
      case+ r of
      | CONS (k, d, bal, left, right) =>
        begin
          case+ right of
          | CONS _ =>
            let
              val (q, fixbalq, kq, dq) = del (right, fixbal)
              val q1 = CONS (k, d, bal, left, q)
              val (q1bal, fixbal) = balance_R (q1, fixbalq)
            in
              (q1bal, fixbal, kq, dq)
            end
          | NIL => (left, T, k, d)
        end

    fun
    search {size   : nat}
           (p      : avl_t (key_t, data_t, size),
            fixbal : fixbal_t) :<!ntm>
        (* Return three values: a new avl_t, a new fixbal, and
           whether the key was found. *)
        [sz : nat]
        (avl_t (key_t, data_t, sz), fixbal_t, bool) =
      case+ p of
      | NIL => (p, F, F)
      | CONS (k, d, bal, left, right) =>
        case+ avl_t$compare<key_t> (key, k) of
        | cmp when cmp < 0 =>
          (* Recursive search down the left branch. *)
          let
            val (q, fixbal, found) = search (left, fixbal)
            val (q1, fixbal) =
              balance_L (CONS (k, d, bal, q, right), fixbal)
          in
            (q1, fixbal, found)
          end
        | cmp when cmp > 0 =>
          (* Recursive search down the right branch. *)
          let
            val (q, fixbal, found) = search (right, fixbal)
            val (q1, fixbal) =
              balance_R (CONS (k, d, bal, left, q), fixbal)
          in
            (q1, fixbal, found)
          end
        | _ =>
          if avl_t_is_empty right then
            (* Delete p, replace it with its left branch, then
               rebalance. *)
            (left, T, T)
          else if avl_t_is_empty left then
            (* Delete p, replace it with its right branch, then
               rebalance. *)
            (right, T, T)
          else
            (* Delete p, but it has both left and right branches, and
               therefore may have complicated branch structure. *)
            let
              val (q, fixbal, k1, d1) = del (left, fixbal)
              val (q1, fixbal) =
                balance_L (CONS (k1, d1, bal, q, right), fixbal)
            in
              (q1, fixbal, T)
            end
  in
    if avl_t_is_empty avl then
      (avl, F)
    else
      let
        prval _ = lemma_avl_t_param avl
        val (avl1, _, found) = $effmask_ntm search (avl, F)
      in
        (avl1, found)
      end
  end

implement {key_t} {data_t}
avl_t_pairs (avl) =
  let
    fun
    traverse {size : pos}
             {n    : nat}
             (p    : avl_t (key_t, data_t, size),
              lst  : list ((key_t, data_t), n)) :<!ntm>
        [sz : pos] list ((key_t, data_t), sz) =
      (* Reverse in-order traversal, to make an in-order list by
         consing. *)
      case+ p of
      | CONS (k, d, _, left, right) =>
        if avl_t_is_empty left then
          begin
            if avl_t_is_empty right then
              (k, d) :: lst
            else
              (k, d) :: traverse (right, lst)
          end
        else
          begin
            if avl_t_is_empty right then
              traverse (left, (k, d) :: lst)
            else
              traverse (left, (k, d) :: traverse (right, lst))
          end
  in
    case+ avl of
    | NIL => LNIL
    | CONS _ => $effmask_ntm traverse (avl, LNIL)
  end

implement {key_t} {data_t}
avl_t_keys (avl) =
  let
    fun
    traverse {size : pos}
             {n    : nat}
             (p    : avl_t (key_t, data_t, size),
              lst  : list (key_t, n)) :<!ntm>
        [sz : pos] list (key_t, sz) =
      (* Reverse in-order traversal, to make an in-order list by
         consing. *)
      case+ p of
      | CONS (k, _, _, left, right) =>
        if avl_t_is_empty left then
          begin
            if avl_t_is_empty right then
              k :: lst
            else
              k :: traverse (right, lst)
          end
        else
          begin
            if avl_t_is_empty right then
              traverse (left, k :: lst)
            else
              traverse (left, k :: traverse (right, lst))
          end
  in
    case+ avl of
    | NIL => LNIL
    | CONS _ => $effmask_ntm traverse (avl, LNIL)
  end

implement {key_t} {data_t}
avl_t_data (avl) =
  let
    fun
    traverse {size : pos}
             {n    : nat}
             (p    : avl_t (key_t, data_t, size),
              lst  : list (data_t, n)) :<!ntm>
        [sz : pos] list (data_t, sz) =
      (* Reverse in-order traversal, to make an in-order list by
         consing. *)
      case+ p of
      | CONS (_, d, _, left, right) =>
        if avl_t_is_empty left then
          begin
            if avl_t_is_empty right then
              d :: lst
            else
              d :: traverse (right, lst)
          end
        else
          begin
            if avl_t_is_empty right then
              traverse (left, d :: lst)
            else
              traverse (left, d :: traverse (right, lst))
          end
  in
    case+ avl of
    | NIL => LNIL
    | CONS _ => $effmask_ntm traverse (avl, LNIL)
  end

implement {key_t} {data_t}
list2avl_t lst =
  let
    fun
    traverse {n    : pos}
             {size : nat} .<n>.
             (lst  : list ((key_t, data_t), n),
              p    : avl_t (key_t, data_t, size)) :<>
        [sz : pos] avl_t (key_t, data_t, sz) =
      case+ lst of
      | (k, d) :: LNIL => avl_t_insert<key_t><data_t> (p, k, d)
      | (k, d) :: (_ :: _) =>
        let
          val+ _ :: tail = lst
        in
          traverse (tail, avl_t_insert<key_t><data_t> (p, k, d))
        end
  in
    case+ lst of
    | LNIL => NIL
    | (_ :: _) => traverse (lst, NIL)
  end

fun {key_t  : [email protected]}
    {data_t : [email protected]}
push_all_the_way_left (stack : List (avl_t (key_t, data_t)),
                       p     : avl_t (key_t, data_t)) :
    List0 (avl_t (key_t, data_t)) =
  let
    prval _ = lemma_list_param stack
  in
    case+ p of
    | NIL => stack
    | CONS (_, _, _, left, _) =>
      push_all_the_way_left (p :: stack, left)
  end

fun {key_t  : [email protected]}
    {data_t : [email protected]}
push_all_the_way_right (stack : List (avl_t (key_t, data_t)),
                        p     : avl_t (key_t, data_t)) :
    List0 (avl_t (key_t, data_t)) =
  let
    prval _ = lemma_list_param stack
  in
    case+ p of
    | NIL => stack
    | CONS (_, _, _, _, right) =>
      push_all_the_way_right (p :: stack, right)
  end

fun {key_t  : [email protected]}
    {data_t : [email protected]}
push_all_the_way (stack     : List (avl_t (key_t, data_t)),
                  p         : avl_t (key_t, data_t),
                  direction : int) :
    List0 (avl_t (key_t, data_t)) =
  if direction < 0 then
    push_all_the_way_right<key_t><data_t> (stack, p)
  else
    push_all_the_way_left<key_t><data_t> (stack, p)

fun {key_t  : [email protected]}
    {data_t : [email protected]}
update_generator_stack (stack     : List (avl_t (key_t, data_t)),
                        left      : avl_t (key_t, data_t),
                        right     : avl_t (key_t, data_t),
                        direction : int) :
    List0 (avl_t (key_t, data_t)) =
  let
    prval _ = lemma_list_param stack
  in
    if direction < 0 then
      begin
        if avl_t_is_empty left then
          stack
        else
          push_all_the_way_right<key_t><data_t> (stack, left)
      end
    else
      begin
        if avl_t_is_empty right then
          stack
        else
          push_all_the_way_left<key_t><data_t> (stack, right)
      end
  end

implement {key_t} {data_t}
avl_t_make_pairs_generator (avl, direction) =
  let
    typedef avl_t = avl_t (key_t, data_t)

    val stack = push_all_the_way (LNIL, avl, direction)
    val stack_ref = ref stack

    (* Cast stack_ref to its (otherwise untyped) pointer, so it can be
       enclosed within ‘generate’. *)
    val p_stack_ref = $UNSAFE.castvwtp0{ptr} stack_ref

    fun
    generate () :<cloref1> Option @(key_t, data_t) =
      let
        (* Restore the type information for stack_ref. *)
        val stack_ref =
          $UNSAFE.castvwtp0{ref (List avl_t)} p_stack_ref

        var stack : List0 avl_t = !stack_ref
        var retval : Option @(key_t, data_t)
      in
        begin
          case+ stack of
          | LNIL => retval := None ()
          | p :: tail =>
            let
              val- CONS (k, d, _, left, right) = p
            in
              retval := Some @(k, d);
              stack :=
                update_generator_stack<key_t><data_t>
                  (tail, left, right, direction)
            end
        end;
        !stack_ref := stack;
        retval
      end
  in
    generate
  end

implement {key_t} {data_t}
avl_t_make_keys_generator (avl, direction) =
  let
    typedef avl_t = avl_t (key_t, data_t)

    val stack = push_all_the_way (LNIL, avl, direction)
    val stack_ref = ref stack

    (* Cast stack_ref to its (otherwise untyped) pointer, so it can be
       enclosed within ‘generate’. *)
    val p_stack_ref = $UNSAFE.castvwtp0{ptr} stack_ref

    fun
    generate () :<cloref1> Option key_t =
      let
        (* Restore the type information for stack_ref. *)
        val stack_ref =
          $UNSAFE.castvwtp0{ref (List avl_t)} p_stack_ref

        var stack : List0 avl_t = !stack_ref
        var retval : Option key_t
      in
        begin
          case+ stack of
          | LNIL => retval := None ()
          | p :: tail =>
            let
              val- CONS (k, _, _, left, right) = p
            in
              retval := Some k;
              stack :=
                update_generator_stack<key_t><data_t>
                  (tail, left, right, direction)
            end
        end;
        !stack_ref := stack;
        retval
      end
  in
    generate
  end

implement {key_t} {data_t}
avl_t_make_data_generator (avl, direction) =
  let
    typedef avl_t = avl_t (key_t, data_t)

    val stack = push_all_the_way (LNIL, avl, direction)
    val stack_ref = ref stack

    (* Cast stack_ref to its (otherwise untyped) pointer, so it can be
       enclosed within ‘generate’. *)
    val p_stack_ref = $UNSAFE.castvwtp0{ptr} stack_ref

    fun
    generate () :<cloref1> Option data_t =
      let
        (* Restore the type information for stack_ref. *)
        val stack_ref =
          $UNSAFE.castvwtp0{ref (List avl_t)} p_stack_ref

        var stack : List0 avl_t = !stack_ref
        var retval : Option data_t
      in
        begin
          case+ stack of
          | LNIL => retval := None ()
          | p :: tail =>
            let
              val- CONS (_, d, _, left, right) = p
            in
              retval := Some d;
              stack :=
                update_generator_stack<key_t><data_t>
                  (tail, left, right, direction)
            end
        end;
        !stack_ref := stack;
        retval
      end
  in
    generate
  end

implement {key_t} {data_t}
avl_t_check_avl_condition (avl) =
  (* If any of the assertions here is triggered, there is a bug. *)
  let
    fun
    get_heights (p : avl_t (key_t, data_t)) : (int, int) =
      case+ p of
      | NIL => (0, 0)
      | CONS (k, d, bal, left, right) =>
        let
          val (height_LL, height_LR) = get_heights left
          val (height_RL, height_RR) = get_heights right
        in
          assertloc (abs (height_LL - height_LR) <= 1);
          assertloc (abs (height_RL - height_RR) <= 1);
          (height_LL + height_LR, height_RL + height_RR)
        end
  in
    if avl_t_isnot_empty avl then
      let
        val (height_L, height_R) = get_heights avl
      in
        assertloc (abs (height_L - height_R) <= 1)
      end
  end

implement {key_t} {data_t}
avl_t_pretty_print (avl) =
  let
    fun
    pad {depth : nat} .<depth>.
        (depth : int depth) : void =
      if depth <> 0 then
        begin
          print! ("  ");
          pad (pred depth)
        end

    fun
    traverse {size  : nat}
             {depth : nat}
             (p     : avl_t (key_t, data_t, size),
              depth : int depth) : void =
      if avl_t_isnot_empty p then
        let
          val+ CONS (k, d, bal, left, right) = p
        in
          traverse (left, succ depth);
          pad depth;
          avl_t_pretty_print$key_and_data<key_t><data_t> (k, d);
          println! ("\t\tdepth = ", depth, " bal = ", bal2int bal);
          traverse (right, succ depth)
        end
  in
    if avl_t_isnot_empty avl then
      let
        val+ CONS (k, d, bal, left, right) = avl
      in
        traverse (left, 1);
        avl_t_pretty_print$key_and_data<key_t><data_t> (k, d);
        println! ("\t\tdepth = 0  bal = ", bal2int bal);
        traverse (right, 1)
      end
  end

(*------------------------------------------------------------------*)

(*

  Here is a little demonstration program.

  Assuming you are using Boehm GC, compile this source file with

      patscc -O2 -DATS_MEMALLOC_GCBDW avl_trees-postiats.dats -lgc

  and run it with

      ./a.out

*)

%{^
#include <time.h>

ATSinline() atstype_uint64
get_the_time (void)
{
  return (atstype_uint64) time (NULL);
}
%}

(* An implementation of avl_t$compare for keys of type ‘int’. *)
implement
avl_t$compare<int> (u, v) =
  if u < v then
    ~1
  else if u > v then
    1
  else
    0

(* An implementation of avl_t_pretty_print$key_and_data for keys of
   type ‘int’ and values of type ‘double’. *)
implement
avl_t_pretty_print$key_and_data<int><double> (key, data) =
  print! ("(", key, ", ", data, ")")

implement
main0 () =
  let
    (* A linear congruential random number generator attributed
       to Donald Knuth. *)
    fn
    next_random (seed : &uint64) : uint64 =
      let
        val a : uint64 = $UNSAFE.cast 6364136223846793005ULL
        val c : uint64 = $UNSAFE.cast 1442695040888963407ULL
        val retval = seed
      in
        seed := (a * seed) + c;
        retval
      end

    fn {t : [email protected]}
    fisher_yates_shuffle
              {n    : nat}
              (a    : &(@[t][n]),
               n    : size_t n,
               seed : &uint64) : void =
      let
        var i : [i : nat | i <= n] size_t i
      in
        for (i := i2sz 0; i < n; i := succ i)
          let
            val randnum = $UNSAFE.cast{Size_t} (next_random seed)
            val j = randnum mod n (* This is good enough for us. *)
            val xi = a[i]
            val xj = a[j]
          in
            a[i] := xj;
            a[j] := xi
          end
      end

    var seed : uint64 = $extfcall (uint64, "get_the_time")

    #define N 20
    var keys : @[int][N] = @[int][N] (0)

    var a : avl_t (int, double)
    var a_saved : avl_t (int, double)
    var a1 : (avl_t (int, double), bool)

    var i : [i : nat] int i

    val dflt = ~99999999.0
    val not_dflt = 123456789.0
  in
    println! ("----------------------------------------------------");
    print! ("\n");

    (* Initialize a shuffled array of keys. *)
    for (i := 0; i < N; i := succ i)
      keys[i] := succ i;
    fisher_yates_shuffle<int> {N} (keys, i2sz N, seed);

    print! ("The keys\n ");
    for (i := 0; i < N; i := succ i)
      print! (" ", keys[i]);
    print! ("\n");

    print! ("\nRunning some tests... ");

    (* Insert key-data pairs in the shuffled order, checking aspects
       of the implementation while doing so. *)
    a := avl_t_nil ();
    for (i := 0; i < N; i := succ i)
      let
        var j : [j : nat] int j
      in
        a := avl_t_insert<int> (a, keys[i], g0i2f keys[i]);
        avl_t_check_avl_condition (a);
        assertloc (avl_t_size a = succ i);
        assertloc (avl_t_is_empty a = iseqz (avl_t_size a));
        assertloc (avl_t_isnot_empty a = isneqz (avl_t_size a));
        a := avl_t_insert<int> (a, keys[i], not_dflt);
        avl_t_check_avl_condition (a);
        assertloc (avl_t_search<int><double> (a, keys[i], dflt)
                      = not_dflt);
        assertloc (avl_t_size a = succ i);
        assertloc (avl_t_is_empty a = iseqz (avl_t_size a));
        assertloc (avl_t_isnot_empty a = isneqz (avl_t_size a));
        a := avl_t_insert<int> (a, keys[i], g0i2f keys[i]);
        avl_t_check_avl_condition (a);
        assertloc (avl_t_size a = succ i);
        assertloc (avl_t_is_empty a = iseqz (avl_t_size a));
        assertloc (avl_t_isnot_empty a = isneqz (avl_t_size a));
        for (j := 0; j < N; j := succ j)
          let
            val k = keys[j]
            val has_key = avl_t_has_key<int> (a, k)
            val data_opt = avl_t_search<int><double> (a, k)
            val data_dflt = avl_t_search<int><double> (a, k, dflt)
          in
            assertloc (has_key = (j <= i));
            assertloc (option_is_some data_opt = (j <= i));
            if (j <= i) then
              let
                val- Some data = data_opt
              in
                assertloc (data = g0i2f k);
                assertloc (data_dflt = g0i2f k);
              end
            else
              let
                val- None () = data_opt
              in
                assertloc (data_dflt = dflt);
              end
          end
      end;

    (* Do it again, but using avl_t_insert_or_replace and checking its
       second return value. *)
    a1 := (avl_t_nil (), false);
    for (i := 0; i < N; i := succ i)
      let
        var j : [j : nat] int j
      in
        a1 :=
          avl_t_insert_or_replace<int> (a1.0, keys[i], g0i2f keys[i]);
        avl_t_check_avl_condition (a1.0);
        assertloc (~(a1.1));
        assertloc (avl_t_size (a1.0) = succ i);
        assertloc (avl_t_is_empty a1.0 = iseqz (avl_t_size a1.0));
        assertloc (avl_t_isnot_empty a1.0 = isneqz (avl_t_size a1.0));
        a1 := avl_t_insert_or_replace<int> (a1.0, keys[i], not_dflt);
        avl_t_check_avl_condition (a1.0);
        assertloc (avl_t_search<int><double> (a1.0, keys[i], dflt)
                      = not_dflt);
        assertloc (avl_t_size (a1.0) = succ i);
        assertloc (avl_t_is_empty a1.0 = iseqz (avl_t_size a1.0));
        assertloc (avl_t_isnot_empty a1.0 = isneqz (avl_t_size a1.0));
        a1 :=
          avl_t_insert_or_replace<int> (a1.0, keys[i], g0i2f keys[i]);
        avl_t_check_avl_condition (a1.0);
        assertloc (a1.1);
        assertloc (avl_t_size (a1.0) = succ i);
        assertloc (avl_t_is_empty a1.0 = iseqz (avl_t_size a1.0));
        assertloc (avl_t_isnot_empty a1.0 = isneqz (avl_t_size a1.0));
        for (j := 0; j < N; j := succ j)
          let
            val k = keys[j]
            val has_key = avl_t_has_key<int> (a1.0, k)
            val data_opt = avl_t_search<int><double> (a1.0, k)
            val data_dflt = avl_t_search<int><double> (a1.0, k, dflt)
          in
            assertloc (has_key = (j <= i));
            assertloc (option_is_some data_opt = (j <= i));
            if (j <= i) then
              let
                val- Some data = data_opt
              in
                assertloc (data = g0i2f k);
                assertloc (data_dflt = g0i2f k);
              end
            else
              let
                val- None () = data_opt
              in
                assertloc (data_dflt = dflt);
              end
          end
      end;
    a := a1.0;

    (* The trees are PERSISTENT, so SAVE THE CURRENT VALUE! *)
    a_saved := a;

    (* Reshuffle the keys, and test deletion, using the reshuffled
       keys. *)
    fisher_yates_shuffle<int> {N} (keys, i2sz N, seed);
    for (i := 0; i < N; i := succ i)
      let
        val ix = keys[i]
        var j : [j : nat] int j
      in
        a := avl_t_delete<int> (a, ix);
        avl_t_check_avl_condition (a);
        assertloc (avl_t_size a = N - succ i);
        assertloc (avl_t_is_empty a = iseqz (avl_t_size a));
        assertloc (avl_t_isnot_empty a = isneqz (avl_t_size a));
        a := avl_t_delete<int> (a, ix);
        avl_t_check_avl_condition (a);
        assertloc (avl_t_size a = N - succ i);
        assertloc (avl_t_is_empty a = iseqz (avl_t_size a));
        assertloc (avl_t_isnot_empty a = isneqz (avl_t_size a));
        for (j := 0; j < N; j := succ j)
          let
            val k = keys[j]
            val has_key = avl_t_has_key<int> (a, k)
            val data_opt = avl_t_search<int><double> (a, k)
            val data_dflt = avl_t_search<int><double> (a, k, dflt)
          in
            assertloc (has_key = (i < j));
            assertloc (option_is_some data_opt = (i < j));
            if (i < j) then
              let
                val- Some data = data_opt
              in
                assertloc (data = g0i2f k);
                assertloc (data_dflt = g0i2f k);
              end
            else
              let
                val- None () = data_opt
              in
                assertloc (data_dflt = dflt);
              end
          end
      end;

    (* Get back the PERSISTENT VALUE from before the deletions. *)
    a := a_saved;

    (* Reshuffle the keys, and test deletion again, this time using
       avl_t_delete_if_found. *)
    fisher_yates_shuffle<int> {N} (keys, i2sz N, seed);
    for (i := 0; i < N; i := succ i)
      let
        val ix = keys[i]
        var j : [j : nat] int j
      in
        a1 := avl_t_delete_if_found<int> (a, ix);
        a := a1.0;
        avl_t_check_avl_condition (a);
        assertloc (a1.1);
        assertloc (avl_t_size a = N - succ i);
        assertloc (avl_t_is_empty a = iseqz (avl_t_size a));
        assertloc (avl_t_isnot_empty a = isneqz (avl_t_size a));
        a1 := avl_t_delete_if_found<int> (a, ix);
        a := a1.0;
        avl_t_check_avl_condition (a);
        assertloc (~(a1.1));
        assertloc (avl_t_size a = N - succ i);
        assertloc (avl_t_is_empty a = iseqz (avl_t_size a));
        assertloc (avl_t_isnot_empty a = isneqz (avl_t_size a));
        for (j := 0; j < N; j := succ j)
          let
            val k = keys[j]
            val has_key = avl_t_has_key<int> (a, k)
            val data_opt = avl_t_search<int><double> (a, k)
            val data_dflt = avl_t_search<int><double> (a, k, dflt)
          in
            assertloc (has_key = (i < j));
            assertloc (option_is_some data_opt = (i < j));
            if (i < j) then
              let
                val- Some data = data_opt
              in
                assertloc (data = g0i2f k);
                assertloc (data_dflt = g0i2f k);
              end
            else
              let
                val- None () = data_opt
              in
                assertloc (data_dflt = dflt);
              end
          end
      end;

    print! ("passed\n");

    (* Get back the PERSISTENT VALUE from before the deletions. *)
    a := a_saved;

    print! ("\n");
    println! ("----------------------------------------------------");
    print! ("\n");
    print! ("*** PRETTY-PRINTING OF THE TREE ***\n\n");

    avl_t_pretty_print<int><double> a;

    print! ("\n");
    println! ("----------------------------------------------------");
    print! ("\n");
    print! ("*** GENERATORS ***\n\n");

    let
      val gen = avl_t_make_pairs_generator (a, 1)
      var x : Option @(int, double)
    in
      print! ("Association pairs in order\n ");
      for (x := gen (); option_is_some (x); x := gen ())
        let
          val @(k, d) = option_unsome x
        in
          print! (" (", k : int, ", ", d : double, ")")
        end
    end;

    print! ("\n\n");

    let
      val gen = avl_t_make_pairs_generator (a, ~1)
      var x : Option @(int, double)
    in
      print! ("Association pairs in reverse order\n ");
      for (x := gen (); option_is_some (x); x := gen ())
        let
          val @(k, d) = option_unsome x
        in
          print! (" (", k : int, ", ", d : double, ")")
        end
    end;

    print! ("\n\n");

    let
      val gen = avl_t_make_keys_generator (a, 1)
      var x : Option int
    in
      print! ("Keys in order\n ");
      for (x := gen (); option_is_some (x); x := gen ())
        print! (" ", (option_unsome x) : int)
    end;

    print! ("\n\n");

    let
      val gen = avl_t_make_keys_generator (a, ~1)
      var x : Option int
    in
      print! ("Keys in reverse order\n ");
      for (x := gen (); option_is_some (x); x := gen ())
        print! (" ", (option_unsome x) : int)
    end;

    print! ("\n\n");

    let
      val gen = avl_t_make_data_generator (a, 1)
      var x : Option double
    in
      print! ("Data values in order of their keys\n ");
      for (x := gen (); option_is_some (x); x := gen ())
        print! (" ", (option_unsome x) : double)
    end;

    print! ("\n\n");

    let
      val gen = avl_t_make_data_generator (a, ~1)
      var x : Option double
    in
      print! ("Data values in reverse order of their keys\n ");
      for (x := gen (); option_is_some (x); x := gen ())
        print! (" ", (option_unsome x) : double)
    end;

    print! ("\n");

    print! ("\n");
    println! ("----------------------------------------------------");
    print! ("\n");
    print! ("*** AVL TREES TO LISTS ***\n\n");

    print! ("Association pairs in order\n  ");
    print! (avl_t_pairs<int><double> a);

    print! ("\n\n");

    print! ("Keys in order\n  ");
    print! (avl_t_keys<int> a);
    print! ("\n\n");

    print! ("Data values in order of their keys\n  ");
    print! (avl_t_data<int><double> a);
    print! ("\n");

    print! ("\n");
    println! ("----------------------------------------------------");
    print! ("\n");
    print! ("*** LISTS TO AVL TREES ***\n\n");

    let
      val lst = (3, 3.0) :: (1, 1.0) :: (4, 4.0) :: (2, 2.0) :: LNIL
      val avl = list2avl_t<int><double> lst
    in
      print! (lst : List @(int, double));
      print! ("\n\n  =>\n\n");
      avl_t_pretty_print<int><double> avl
    end;

    print! ("\n");
    println! ("----------------------------------------------------")
  end

(*------------------------------------------------------------------*)


#include <algorithm>
#include <iostream>

/* AVL node */
template <class T>
class AVLnode {
public:
    T key;
    int balance;
    AVLnode *left, *right, *parent;

    AVLnode(T k, AVLnode *p) : key(k), balance(0), parent(p),
                        left(NULL), right(NULL) {}

    ~AVLnode() {
        delete left;
        delete right;
    }
};

/* AVL tree */
template <class T>
class AVLtree {
public:
    AVLtree(void);
    ~AVLtree(void);
    bool insert(T key);
    void deleteKey(const T key);
    void printBalance();

private:
    AVLnode<T> *root;

    AVLnode<T>* rotateLeft          ( AVLnode<T> *a );
    AVLnode<T>* rotateRight         ( AVLnode<T> *a );
    AVLnode<T>* rotateLeftThenRight ( AVLnode<T> *n );
    AVLnode<T>* rotateRightThenLeft ( AVLnode<T> *n );
    void rebalance                  ( AVLnode<T> *n );
    int height                      ( AVLnode<T> *n );
    void setBalance                 ( AVLnode<T> *n );
    void printBalance               ( AVLnode<T> *n );
};

/* AVL class definition */
template <class T>
void AVLtree<T>::rebalance(AVLnode<T> *n) {
    setBalance(n);

    if (n->balance == -2) {
        if (height(n->left->left) >= height(n->left->right))
            n = rotateRight(n);
        else
            n = rotateLeftThenRight(n);
    }
    else if (n->balance == 2) {
        if (height(n->right->right) >= height(n->right->left))
            n = rotateLeft(n);
        else
            n = rotateRightThenLeft(n);
    }

    if (n->parent != NULL) {
        rebalance(n->parent);
    }
    else {
        root = n;
    }
}

template <class T>
AVLnode<T>* AVLtree<T>::rotateLeft(AVLnode<T> *a) {
    AVLnode<T> *b = a->right;
    b->parent = a->parent;
    a->right = b->left;

    if (a->right != NULL)
        a->right->parent = a;

    b->left = a;
    a->parent = b;

    if (b->parent != NULL) {
        if (b->parent->right == a) {
            b->parent->right = b;
        }
        else {
            b->parent->left = b;
        }
    }

    setBalance(a);
    setBalance(b);
    return b;
}

template <class T>
AVLnode<T>* AVLtree<T>::rotateRight(AVLnode<T> *a) {
    AVLnode<T> *b = a->left;
    b->parent = a->parent;
    a->left = b->right;

    if (a->left != NULL)
        a->left->parent = a;

    b->right = a;
    a->parent = b;

    if (b->parent != NULL) {
        if (b->parent->right == a) {
            b->parent->right = b;
        }
        else {
            b->parent->left = b;
        }
    }

    setBalance(a);
    setBalance(b);
    return b;
}

template <class T>
AVLnode<T>* AVLtree<T>::rotateLeftThenRight(AVLnode<T> *n) {
    n->left = rotateLeft(n->left);
    return rotateRight(n);
}

template <class T>
AVLnode<T>* AVLtree<T>::rotateRightThenLeft(AVLnode<T> *n) {
    n->right = rotateRight(n->right);
    return rotateLeft(n);
}

template <class T>
int AVLtree<T>::height(AVLnode<T> *n) {
    if (n == NULL)
        return -1;
    return 1 + std::max(height(n->left), height(n->right));
}

template <class T>
void AVLtree<T>::setBalance(AVLnode<T> *n) {
    n->balance = height(n->right) - height(n->left);
}

template <class T>
void AVLtree<T>::printBalance(AVLnode<T> *n) {
    if (n != NULL) {
        printBalance(n->left);
        std::cout << n->balance << " ";
        printBalance(n->right);
    }
}

template <class T>
AVLtree<T>::AVLtree(void) : root(NULL) {}

template <class T>
AVLtree<T>::~AVLtree(void) {
    delete root;
}

template <class T>
bool AVLtree<T>::insert(T key) {
    if (root == NULL) {
        root = new AVLnode<T>(key, NULL);
    }
    else {
        AVLnode<T>
            *n = root,
            *parent;

        while (true) {
            if (n->key == key)
                return false;

            parent = n;

            bool goLeft = n->key > key;
            n = goLeft ? n->left : n->right;

            if (n == NULL) {
                if (goLeft) {
                    parent->left = new AVLnode<T>(key, parent);
                }
                else {
                    parent->right = new AVLnode<T>(key, parent);
                }

                rebalance(parent);
                break;
            }
        }
    }

    return true;
}

template <class T>
void AVLtree<T>::deleteKey(const T delKey) {
    if (root == NULL)
        return;

    AVLnode<T>
        *n       = root,
        *parent  = root,
        *delNode = NULL,
        *child   = root;

    while (child != NULL) {
        parent = n;
        n = child;
        child = delKey >= n->key ? n->right : n->left;
        if (delKey == n->key)
            delNode = n;
    }

    if (delNode != NULL) {
        delNode->key = n->key;

        child = n->left != NULL ? n->left : n->right;

        if (root->key == delKey) {
            root = child;
        }
        else {
            if (parent->left == n) {
                parent->left = child;
            }
            else {
                parent->right = child;
            }

            rebalance(parent);
        }
    }
}

template <class T>
void AVLtree<T>::printBalance() {
    printBalance(root);
    std::cout << std::endl;
}

int main(void)
{
    AVLtree<int> t;

    std::cout << "Inserting integer values 1 to 10" << std::endl;
    for (int i = 1; i <= 10; ++i)
        t.insert(i);

    std::cout << "Printing balance: ";
    t.printBalance();
}


(defpackage :avl-tree
  (:use :cl)
  (:export
   :avl-tree
   :make-avl-tree
   :avl-tree-count
   :avl-tree-p
   :avl-tree-key<=
   :gettree
   :remtree
   :clrtree
   :dfs-maptree
   :bfs-maptree))

(in-package :avl-tree)

(defstruct %tree
  key
  value
  (height 0 :type fixnum)
  left
  right)

(defstruct (avl-tree (:constructor %make-avl-tree))
  key<=
  tree
  (count 0 :type fixnum))

(defun make-avl-tree (key<=)
  "Create a new AVL tree using the given comparison function KEY<=
for emplacing keys into the tree."
  (%make-avl-tree :key<= key<=))

(declaim (inline
          recalc-height
          height balance
          swap-kv
          right-right-rotate
          right-left-rotate
          left-right-rotate
          left-left-rotate
          rotate))

(defun recalc-height (tree)
  "Calculate the new height of the tree from the heights of the children."
  (when tree
    (setf (%tree-height tree)
          (1+ (the fixnum (max (height (%tree-right tree))
                               (height (%tree-left tree))))))))

(declaim (ftype (function (t) fixnum) height balance))
(defun height (tree)
  (if tree (%tree-height tree) 0))

(defun balance (tree)
  (if tree
      (- (height (%tree-right tree))
         (height (%tree-left tree)))
      0))

(defmacro swap (place-a place-b)
  "Swap the values of two places."
  (let ((tmp (gensym)))
    `(let ((,tmp ,place-a))
       (setf ,place-a ,place-b)
       (setf ,place-b ,tmp))))

(defun swap-kv (tree-a tree-b)
  "Swap the keys and values of two trees."
  (swap (%tree-value tree-a) (%tree-value tree-b))
  (swap (%tree-key tree-a) (%tree-key tree-b)))

;; We should really use gensyms for the variables in here.
(defmacro slash-rotate (tree right left)
  "Rotate nodes in a slash `/` imbalance."
  `(let* ((a ,tree)
          (b (,right a))
          (c (,right b))
          (a-left (,left a))
          (b-left (,left b)))
     (setf (,right a) c)
     (setf (,left a) b)
     (setf (,left b) a-left)
     (setf (,right b) b-left)
     (swap-kv a b)
     (recalc-height b)
     (recalc-height a)))

(defmacro angle-rotate (tree right left)
  "Rotate nodes in an angle bracket `<` imbalance."
  `(let* ((a ,tree)
          (b (,right a))
          (c (,left b))
          (a-left (,left a))
          (c-left (,left c))
          (c-right (,right c)))
     (setf (,left a) c)
     (setf (,left c) a-left)
     (setf (,right c) c-left)
     (setf (,left b) c-right)
     (swap-kv a c)
     (recalc-height c)
     (recalc-height b)
     (recalc-height a)))

(defun right-right-rotate (tree)
  (slash-rotate tree %tree-right %tree-left))

(defun left-left-rotate (tree)
  (slash-rotate tree %tree-left %tree-right))

(defun right-left-rotate (tree)
  (angle-rotate tree %tree-right %tree-left))

(defun left-right-rotate (tree)
  (angle-rotate tree %tree-left %tree-right))

(defun rotate (tree)
  (declare (type %tree tree))
  "Perform a rotation on the given TREE if it is imbalanced."
  (recalc-height tree)
  (with-slots (left right) tree
    (let ((balance (balance tree)))
      (cond ((< 1 balance) ;; Right imbalanced tree
             (if (<= 0 (balance right))
                 (right-right-rotate tree)
                 (right-left-rotate tree)))
            ((> -1 balance) ;; Left imbalanced tree
             (if (<= 0 (balance left))
                 (left-right-rotate tree)
                 (left-left-rotate tree)))))))

(defun gettree (key avl-tree &optional default)
  "Finds an entry in AVL-TREE whos key is KEY and returns the
associated value and T as multiple values, or returns DEFAULT and NIL
if there was no such entry. Entries can be added using SETF."
  (with-slots (key<= tree) avl-tree
    (labels
        ((rec (tree)
           (if tree
               (with-slots ((t-key key) left right value) tree
                 (if (funcall key<= t-key key)
                     (if (funcall key<= key t-key)
                         (values value t)
                         (rec right))
                     (rec left)))
               (values default nil))))
      (rec tree))))

(defun puttree (value key avl-tree)
  ;;(declare (optimize speed))
  (declare (type avl-tree avl-tree))
  "Emplace the the VALUE with the given KEY into the AVL-TREE, or
overwrite the value if the given key already exists."
  (let ((node (make-%tree :key key :value value)))
    (with-slots (key<= tree count) avl-tree
      (cond (tree
             (labels
                 ((rec (tree)
                    (with-slots ((t-key key) left right) tree
                      (if (funcall key<= t-key key)
                          (if (funcall key<= key t-key)
                              (setf (%tree-value tree) value)
                              (cond (right (rec right))
                                    (t (setf right node)
                                       (incf count))))
                          (cond (left (rec left))
                                (t (setf left node)
                                   (incf count))))
                      (rotate tree))))
               (rec tree)))
            (t (setf tree node)
               (incf count))))
    value))

(defun (setf gettree) (value key avl-tree &optional default)
  (declare (ignore default))
  (puttree value key avl-tree))

(defun remtree (key avl-tree)
  (declare (type avl-tree avl-tree))
  "Remove the entry in AVL-TREE associated with KEY. Return T if
there was such an entry, or NIL if not."
  (with-slots (key<= tree count) avl-tree
    (labels
        ((find-left (tree)
           (with-slots ((t-key key) left right) tree
             (if left
                 (find-left left)
                 tree)))
         (rec (tree &optional parent type)
           (when tree
             (prog1
                 (with-slots ((t-key key) left right) tree
                   (if (funcall key<= t-key key)
                       (cond
                         ((funcall key<= key t-key)
                          (cond
                            ((and left right)
                             (let ((sub-left (find-left right)))
                               (swap-kv sub-left tree)
                               (rec right tree :right)))
                            (t
                             (let ((sub (or left right)))
                               (case type
                                 (:right (setf (%tree-right parent) sub))
                                 (:left (setf (%tree-left parent) sub))
                                 (nil (setf (avl-tree-tree avl-tree) sub))))
                             (decf count)))
                          t)
                         (t (rec right tree :right)))
                       (rec left tree :left)))
               (when parent (rotate parent))))))
      (rec tree))))

(defun clrtree (avl-tree)
  "This removes all the entries from AVL-TREE and returns the tree itself."
  (setf (avl-tree-tree avl-tree) nil)
  (setf (avl-tree-count avl-tree) 0)
  avl-tree)

(defun dfs-maptree (function avl-tree)
  "For each entry in AVL-TREE call the two-argument FUNCTION on
the key and value of each entry in depth-first order from left to right.
Consequences are undefined if AVL-TREE is modified during this call."
  (with-slots (key<= tree) avl-tree
    (labels
        ((rec (tree)
           (when tree
             (with-slots ((t-key key) left right key value) tree
               (rec left)
               (funcall function key value)
               (rec right)))))
      (rec tree))))

(defun bfs-maptree (function avl-tree)
  "For each entry in AVL-TREE call the two-argument FUNCTION on
the key and value of each entry in breadth-first order from left to right.
Consequences are undefined if AVL-TREE is modified during this call."
  (with-slots (key<= tree) avl-tree
    (let* ((queue (cons nil nil))
           (end queue))
      (labels ((pushend (value)
                 (when value
                   (setf (cdr end) (cons value nil))
                   (setf end (cdr end))))
               (empty-p () (eq nil (cdr queue)))
               (popfront ()
                 (prog1 (pop (cdr queue))
                   (when (empty-p) (setf end queue)))))
        (when tree
          (pushend tree)
          (loop until (empty-p)
             do (let ((current (popfront)))
                  (with-slots (key value left right) current
                    (funcall function key value)
                    (pushend left)
                    (pushend right)))))))))

(defun test ()
  (let ((tree (make-avl-tree #'<=))
        (printer (lambda (k v) (print (list k v)))))
    (loop for key in '(0 8 6 4 2 3 7 9 1 5 5)
       for value in '(a b c d e f g h i j k)
       do (setf (gettree key tree) value))
    (loop for key in '(0 1 2 3 4 10)
       do (print (multiple-value-list (gettree key tree))))
    (terpri)
    (print tree)
    (terpri)
    (dfs-maptree printer tree)
    (terpri)
    (bfs-maptree printer tree)
    (terpri)
    (loop for key in '(0 1 2 3 10 7)
       do (print (remtree key tree)))
    (terpri)
    (print tree)
    (terpri)
    (clrtree tree)
    (print tree))
  (values))

(defun profile-test ()
  (let ((tree (make-avl-tree #'<=))
        (randoms (loop repeat 1000000 collect (random 100.0))))
    (loop for key in randoms do (setf (gettree key tree) key))))


import std.stdio, std.algorithm;

class AVLtree {
    private Node* root;

    private static struct Node {
        private int key, balance;
        private Node* left, right, parent;

        this(in int k, Node* p) pure nothrow @safe @nogc {
            key = k;
            parent = p;
        }
    }

    public bool insert(in int key) pure nothrow @safe {
        if (root is null)
            root = new Node(key, null);
        else {
            Node* n = root;
            Node* parent;
            while (true) {
                if (n.key == key)
                    return false;

                parent = n;

                bool goLeft = n.key > key;
                n = goLeft ? n.left : n.right;

                if (n is null) {
                    if (goLeft) {
                        parent.left = new Node(key, parent);
                    } else {
                        parent.right = new Node(key, parent);
                    }
                    rebalance(parent);
                    break;
                }
            }
        }
        return true;
    }

    public void deleteKey(in int delKey) pure nothrow @safe @nogc {
        if (root is null)
            return;
        Node* n = root;
        Node* parent = root;
        Node* delNode = null;
        Node* child = root;

        while (child !is null) {
            parent = n;
            n = child;
            child = delKey >= n.key ? n.right : n.left;
            if (delKey == n.key)
                delNode = n;
        }

        if (delNode !is null) {
            delNode.key = n.key;

            child = n.left !is null ? n.left : n.right;

            if (root.key == delKey) {
                root = child;
            } else {
                if (parent.left is n) {
                    parent.left = child;
                } else {
                    parent.right = child;
                }
                rebalance(parent);
            }
        }
    }

    private void rebalance(Node* n) pure nothrow @safe @nogc {
        setBalance(n);

        if (n.balance == -2) {
            if (height(n.left.left) >= height(n.left.right))
                n = rotateRight(n);
            else
                n = rotateLeftThenRight(n);

        } else if (n.balance == 2) {
            if (height(n.right.right) >= height(n.right.left))
                n = rotateLeft(n);
            else
                n = rotateRightThenLeft(n);
        }

        if (n.parent !is null) {
            rebalance(n.parent);
        } else {
            root = n;
        }
    }

    private Node* rotateLeft(Node* a) pure nothrow @safe @nogc {
        Node* b = a.right;
        b.parent = a.parent;

        a.right = b.left;

        if (a.right !is null)
            a.right.parent = a;

        b.left = a;
        a.parent = b;

        if (b.parent !is null) {
            if (b.parent.right is a) {
                b.parent.right = b;
            } else {
                b.parent.left = b;
            }
        }

        setBalance(a, b);

        return b;
    }

    private Node* rotateRight(Node* a) pure nothrow @safe @nogc {
        Node* b = a.left;
        b.parent = a.parent;

        a.left = b.right;

        if (a.left !is null)
            a.left.parent = a;

        b.right = a;
        a.parent = b;

        if (b.parent !is null) {
            if (b.parent.right is a) {
                b.parent.right = b;
            } else {
                b.parent.left = b;
            }
        }

        setBalance(a, b);

        return b;
    }

    private Node* rotateLeftThenRight(Node* n) pure nothrow @safe @nogc {
        n.left = rotateLeft(n.left);
        return rotateRight(n);
    }

    private Node* rotateRightThenLeft(Node* n) pure nothrow @safe @nogc {
        n.right = rotateRight(n.right);
        return rotateLeft(n);
    }

    private int height(in Node* n) const pure nothrow @safe @nogc {
        if (n is null)
            return -1;
        return 1 + max(height(n.left), height(n.right));
    }

    private void setBalance(Node*[] nodes...) pure nothrow @safe @nogc {
        foreach (n; nodes)
            n.balance = height(n.right) - height(n.left);
    }

    public void printBalance() const @safe {
        printBalance(root);
    }

    private void printBalance(in Node* n) const @safe {
        if (n !is null) {
            printBalance(n.left);
            write(n.balance, ' ');
            printBalance(n.right);
        }
    }
}

void main() @safe {
    auto tree = new AVLtree();

    writeln("Inserting values 1 to 10");
    foreach (immutable i; 1 .. 11)
        tree.insert(i);

    write("Printing balance: ");
    tree.printBalance;
}


module avl_trees
  !
  ! References:
  !
  !   * Niklaus Wirth, 1976. Algorithms + Data Structures =
  !     Programs. Prentice-Hall, Englewood Cliffs, New Jersey.
  !
  !   * Niklaus Wirth, 2004. Algorithms and Data Structures. Updated
  !     by Fyodor Tkachov, 2014.
  !

  implicit none
  private

  ! The type for an AVL tree.
  public :: avl_tree_t

  ! The type for a pair of pointers to key and data within the tree.
  ! (Be careful with these!)
  public :: avl_pointer_pair_t

  ! Insertion, replacement, modification, etc.
  public :: avl_insert_or_modify

  ! Insert or replace.
  public :: avl_insert

  ! Is the key in the tree?
  public :: avl_contains

  ! Retrieve data from a tree.
  public :: avl_retrieve

  ! Delete data from a tree. This is a generic function.
  public :: avl_delete

  ! Implementations of avl_delete.
  public :: avl_delete_with_found
  public :: avl_delete_without_found

  ! How many nodes are there in the tree?
  public :: avl_size

  ! Return a list of avl_pointer_pair_t for the elements in the
  ! tree. The list will be in order.
  public :: avl_pointer_pairs

  ! Print a representation of the tree to an output unit.
  public :: avl_write

  ! Check the AVL condition (that the heights of the two branches from
  ! a node should differ by zero or one). ERROR STOP if the condition
  ! is not met.
  public :: avl_check

  ! Procedure types.
  public :: avl_less_than_t
  public :: avl_insertion_t
  public :: avl_key_data_writer_t

  type :: avl_node_t
     class(*), allocatable :: key, data
     type(avl_node_t), pointer :: left
     type(avl_node_t), pointer :: right
     integer :: bal             ! bal == -1, 0, 1
  end type avl_node_t

  type :: avl_tree_t
     type(avl_node_t), pointer :: p => null ()
   contains
     final :: avl_tree_t_final
  end type avl_tree_t

  type :: avl_pointer_pair_t
     class(*), pointer :: p_key, p_data
     class(avl_pointer_pair_t), pointer :: next => null ()
   contains
     final :: avl_pointer_pair_t_final
  end type avl_pointer_pair_t

  interface avl_delete
     module procedure avl_delete_with_found
     module procedure avl_delete_without_found
  end interface avl_delete

  interface
     function avl_less_than_t (key1, key2) result (key1_lt_key2)
       !
       ! The ordering predicate (‘<’).
       !
       ! Two keys a,b are considered equivalent if neither a<b nor
       ! b<a.
       !
       class(*), intent(in) :: key1, key2
       logical key1_lt_key2
     end function avl_less_than_t

     subroutine avl_insertion_t (key, data, p_is_new, p)
       !
       ! Insertion or modification of a found node.
       !
       import avl_node_t
       class(*), intent(in) :: key, data
       logical, intent(in) :: p_is_new
       type(avl_node_t), pointer, intent(inout) :: p
     end subroutine avl_insertion_t

     subroutine avl_key_data_writer_t (unit, key, data)
       !
       ! Printing the key and data of a node.
       !
       integer, intent(in) :: unit
       class(*), intent(in) :: key, data
     end subroutine avl_key_data_writer_t
  end interface

contains

  subroutine avl_tree_t_final (tree)
    type(avl_tree_t), intent(inout) :: tree

    type(avl_node_t), pointer :: p

    p => tree%p
    call free_the_nodes (p)

  contains

    recursive subroutine free_the_nodes (p)
      type(avl_node_t), pointer, intent(inout) :: p

      if (associated (p)) then
         call free_the_nodes (p%left)
         call free_the_nodes (p%right)
         deallocate (p)
      end if
    end subroutine free_the_nodes

  end subroutine avl_tree_t_final

  recursive subroutine avl_pointer_pair_t_final (node)
    type(avl_pointer_pair_t), intent(inout) :: node

    if (associated (node%next)) deallocate (node%next)
  end subroutine avl_pointer_pair_t_final

  function avl_contains (less_than, key, tree) result (found)
    procedure(avl_less_than_t) :: less_than
    class(*), intent(in) :: key
    class(avl_tree_t), intent(in) :: tree
    logical :: found

    found = avl_contains_recursion (less_than, key, tree%p)
  end function avl_contains

  recursive function avl_contains_recursion (less_than, key, p) result (found)
    procedure(avl_less_than_t) :: less_than
    class(*), intent(in) :: key
    type(avl_node_t), pointer, intent(in) :: p
    logical :: found

    if (.not. associated (p)) then
       found = .false.
    else if (less_than (key, p%key)) then
       found = avl_contains_recursion (less_than, key, p%left)
    else if (less_than (p%key, key)) then
       found = avl_contains_recursion (less_than, key, p%right)
    else
       found = .true.
    end if
  end function avl_contains_recursion

  subroutine avl_retrieve (less_than, key, tree, found, data)
    procedure(avl_less_than_t) :: less_than
    class(*), intent(in) :: key
    class(avl_tree_t), intent(in) :: tree
    logical, intent(out) :: found
    class(*), allocatable, intent(inout) :: data

    call avl_retrieve_recursion (less_than, key, tree%p, found, data)
  end subroutine avl_retrieve

  recursive subroutine avl_retrieve_recursion (less_than, key, p, found, data)
    procedure(avl_less_than_t) :: less_than
    class(*), intent(in) :: key
    type(avl_node_t), pointer, intent(in) :: p
    logical, intent(out) :: found
    class(*), allocatable, intent(inout) :: data

    if (.not. associated (p)) then
       found = .false.
    else if (less_than (key, p%key)) then
       call avl_retrieve_recursion (less_than, key, p%left, found, data)
    else if (less_than (p%key, key)) then
       call avl_retrieve_recursion (less_than, key, p%right, found, data)
    else
       found = .true.
       data = p%data
    end if
  end subroutine avl_retrieve_recursion

  subroutine avl_insert (less_than, key, data, tree)
    procedure(avl_less_than_t) :: less_than
    class(*), intent(in) :: key, data
    class(avl_tree_t), intent(inout) :: tree

    call avl_insert_or_modify (less_than, insert_or_replace, key, data, tree)
  end subroutine avl_insert

  subroutine insert_or_replace (key, data, p_is_new, p)
    class(*), intent(in) :: key, data
    logical, intent(in) :: p_is_new
    type(avl_node_t), pointer, intent(inout) :: p

    p%data = data
  end subroutine insert_or_replace

  subroutine avl_insert_or_modify (less_than, insertion, key, data, tree)
    procedure(avl_less_than_t) :: less_than
    procedure(avl_insertion_t) :: insertion ! Or modification in place.
    class(*), intent(in) :: key, data
    class(avl_tree_t), intent(inout) :: tree

    logical :: fix_balance

    fix_balance = .false.
    call insertion_search (less_than, insertion, key, data, tree%p, fix_balance)
  end subroutine avl_insert_or_modify

  recursive subroutine insertion_search (less_than, insertion, key, data, p, fix_balance)
    procedure(avl_less_than_t) :: less_than
    procedure(avl_insertion_t) :: insertion
    class(*), intent(in) :: key, data
    type(avl_node_t), pointer, intent(inout) :: p
    logical, intent(inout) :: fix_balance

    type(avl_node_t), pointer :: p1, p2

    if (.not. associated (p)) then
       ! The key was not found. Make a new node.
       allocate (p)
       p%key = key
       p%left => null ()
       p%right => null ()
       p%bal = 0
       call insertion (key, data, .true., p)
       fix_balance = .true.
    else if (less_than (key, p%key)) then
       ! Continue searching.
       call insertion_search (less_than, insertion, key, data, p%left, fix_balance)
       if (fix_balance) then
          ! A new node has been inserted on the left side.
          select case (p%bal)
          case (1)
             p%bal = 0
             fix_balance = .false.
          case (0)
             p%bal = -1
          case (-1)
             ! Rebalance.
             p1 => p%left
             select case (p1%bal)
             case (-1)
                ! A single LL rotation.
                p%left => p1%right
                p1%right => p
                p%bal = 0
                p => p1
                p%bal = 0
                fix_balance = .false.
             case (0, 1)
                ! A double LR rotation.
                p2 => p1%right
                p1%right => p2%left
                p2%left => p1
                p%left => p2%right
                p2%right => p
                p%bal = -(min (p2%bal, 0))
                p1%bal = -(max (p2%bal, 0))
                p => p2
                p%bal = 0
                fix_balance = .false.
             case default
                error stop
             end select
          case default
             error stop
          end select
       end if
    else if (less_than (p%key, key)) then
       call insertion_search (less_than, insertion, key, data, p%right, fix_balance)
       if (fix_balance) then
          ! A new node has been inserted on the right side.
          select case (p%bal)
          case (-1)
             p%bal = 0
             fix_balance = .false.
          case (0)
             p%bal = 1
          case (1)
             ! Rebalance.
             p1 => p%right
             select case (p1%bal)
             case (1)
                ! A single RR rotation.
                p%right => p1%left
                p1%left => p
                p%bal = 0
                p => p1
                p%bal = 0
                fix_balance = .false.
             case (-1, 0)
                ! A double RL rotation.
                p2 => p1%left
                p1%left => p2%right
                p2%right => p1
                p%right => p2%left
                p2%left => p
                p%bal = -(max (p2%bal, 0))
                p1%bal = -(min (p2%bal, 0))
                p => p2
                p%bal = 0
                fix_balance = .false.
             case default
                error stop
             end select
          case default
             error stop
          end select
       end if
    else
       ! The key was found. The pointer p points to an existing node.
       call insertion (key, data, .false., p)
    end if
  end subroutine insertion_search

  subroutine avl_delete_with_found (less_than, key, tree, found)
    procedure(avl_less_than_t) :: less_than
    class(*), intent(in) :: key
    class(avl_tree_t), intent(inout) :: tree
    logical, intent(out) :: found

    logical :: fix_balance

    fix_balance = .false.
    call deletion_search (less_than, key, tree%p, fix_balance, found)
  end subroutine avl_delete_with_found

  subroutine avl_delete_without_found (less_than, key, tree)
    procedure(avl_less_than_t) :: less_than
    class(*), intent(in) :: key
    class(avl_tree_t), intent(inout) :: tree

    logical :: found

    call avl_delete_with_found (less_than, key, tree, found)
  end subroutine avl_delete_without_found

  recursive subroutine deletion_search (less_than, key, p, fix_balance, found)
    procedure(avl_less_than_t) :: less_than
    class(*), intent(in) :: key
    type(avl_node_t), pointer, intent(inout) :: p
    logical, intent(inout) :: fix_balance
    logical, intent(out) :: found

    type(avl_node_t), pointer :: q

    if (.not. associated (p)) then
       ! The key is not in the tree.
       found = .false.
    else if (less_than (key, p%key)) then
       call deletion_search (less_than, key, p%left, fix_balance, found)
       if (fix_balance) call balance_for_shrunken_left (p, fix_balance)
    else if (less_than (p%key, key)) then
       call deletion_search (less_than, key, p%right, fix_balance, found)
       if (fix_balance) call balance_for_shrunken_right (p, fix_balance)
    else
       q => p
       if (.not. associated (q%right)) then
          p => q%left
          fix_balance = .true.
       else if (.not. associated (q%left)) then
          p => q%right
          fix_balance = .true.
       else
          call del (q%left, q, fix_balance)
          if (fix_balance) call balance_for_shrunken_left (p, fix_balance)
       end if
       deallocate (q)
       found = .true.
    end if
  end subroutine deletion_search

  recursive subroutine del (r, q, fix_balance)
    type(avl_node_t), pointer, intent(inout) :: r, q
    logical, intent(inout) :: fix_balance

    if (associated (r%right)) then
       call del (r%right, q, fix_balance)
       if (fix_balance) call balance_for_shrunken_right (r, fix_balance)
    else
       q%key = r%key
       q%data = r%data
       q => r
       r => r%left
       fix_balance = .true.
    end if
  end subroutine del

  subroutine balance_for_shrunken_left (p, fix_balance)
    type(avl_node_t), pointer, intent(inout) :: p
    logical, intent(inout) :: fix_balance

    ! The left side has lost a node.

    type(avl_node_t), pointer :: p1, p2

    if (.not. fix_balance) error stop

    select case (p%bal)
    case (-1)
       p%bal = 0
    case (0)
       p%bal = 1
       fix_balance = .false.
    case (1)
       ! Rebalance.
       p1 => p%right
       select case (p1%bal)
       case (0)
          ! A single RR rotation.
          p%right => p1%left
          p1%left => p
          p%bal = 1
          p1%bal = -1
          p => p1
          fix_balance = .false.
       case (1)
          ! A single RR rotation.
          p%right => p1%left
          p1%left => p
          p%bal = 0
          p1%bal = 0
          p => p1
          fix_balance = .true.
       case (-1)
          ! A double RL rotation.
          p2 => p1%left
          p1%left => p2%right
          p2%right => p1
          p%right => p2%left
          p2%left => p
          p%bal = -(max (p2%bal, 0))
          p1%bal = -(min (p2%bal, 0))
          p => p2
          p2%bal = 0
       case default
          error stop
       end select
    case default
       error stop
    end select
  end subroutine balance_for_shrunken_left

  subroutine balance_for_shrunken_right (p, fix_balance)
    type(avl_node_t), pointer, intent(inout) :: p
    logical, intent(inout) :: fix_balance

    ! The right side has lost a node.

    type(avl_node_t), pointer :: p1, p2

    if (.not. fix_balance) error stop

    select case (p%bal)
    case (1)
       p%bal = 0
    case (0)
       p%bal = -1
       fix_balance = .false.
    case (-1)
       ! Rebalance.
       p1 => p%left
       select case (p1%bal)
       case (0)
          ! A single LL rotation.
          p%left => p1%right
          p1%right => p
          p1%bal = 1
          p%bal = -1
          p => p1
          fix_balance = .false.
       case (-1)
          ! A single LL rotation.
          p%left => p1%right
          p1%right => p
          p1%bal = 0
          p%bal = 0
          p => p1
          fix_balance = .true.
       case (1)
          ! A double LR rotation.
          p2 => p1%right
          p1%right => p2%left
          p2%left => p1
          p%left => p2%right
          p2%right => p
          p%bal = -(min (p2%bal, 0))
          p1%bal = -(max (p2%bal, 0))
          p => p2
          p2%bal = 0
       case default
          error stop
       end select
    case default
       error stop
    end select
  end subroutine balance_for_shrunken_right

  function avl_size (tree) result (size)
    class(avl_tree_t), intent(in) :: tree
    integer :: size

    size = traverse (tree%p)

  contains

    recursive function traverse (p) result (size)
      type(avl_node_t), pointer, intent(in) :: p
      integer :: size

      if (associated (p)) then
         ! The order of traversal is arbitrary.
         size = 1 + traverse (p%left) + traverse (p%right)
      else
         size = 0
      end if
    end function traverse

  end function avl_size

  function avl_pointer_pairs (tree) result (lst)
    class(avl_tree_t), intent(in) :: tree
    type(avl_pointer_pair_t), pointer :: lst

    ! Reverse in-order traversal of the tree, to produce a CONS-list
    ! of pointers to the contents.

    lst => null ()
    if (associated (tree%p)) lst => traverse (tree%p, lst)

  contains

    recursive function traverse (p, lst1) result (lst2)
      type(avl_node_t), pointer, intent(in) :: p
      type(avl_pointer_pair_t), pointer, intent(in) :: lst1
      type(avl_pointer_pair_t), pointer :: lst2

      type(avl_pointer_pair_t), pointer :: new_entry

      lst2 => lst1
      if (associated (p%right)) lst2 => traverse (p%right, lst2)
      allocate (new_entry)
      new_entry%p_key => p%key
      new_entry%p_data => p%data
      new_entry%next => lst2
      lst2 => new_entry
      if (associated (p%left)) lst2 => traverse (p%left, lst2)
    end function traverse

  end function avl_pointer_pairs

  subroutine avl_write (write_key_data, unit, tree)
    procedure(avl_key_data_writer_t) :: write_key_data
    integer, intent(in) :: unit
    class(avl_tree_t), intent(in) :: tree

    character(len = *), parameter :: tab = achar (9)

    type(avl_node_t), pointer :: p

    p => tree%p
    if (.not. associated (p)) then
       continue
    else
       call traverse (p%left, 1, .true.)
       call write_key_data (unit, p%key, p%data)
       write (unit, '(2A, "depth = ", I0, "  bal = ", I0)') tab, tab, 0, p%bal
       call traverse (p%right, 1, .false.)
    end if

  contains

    recursive subroutine traverse (p, depth, left)
      type(avl_node_t), pointer, intent(in) :: p
      integer, value :: depth
      logical, value :: left

      if (.not. associated (p)) then
         continue
      else
         call traverse (p%left, depth + 1, .true.)
         call pad (depth, left)
         call write_key_data (unit, p%key, p%data)
         write (unit, '(2A, "depth = ", I0, "  bal = ", I0)') tab, tab, depth, p%bal
         call traverse (p%right, depth + 1, .false.)
      end if
    end subroutine traverse

    subroutine pad (depth, left)
      integer, value :: depth
      logical, value :: left

      integer :: i

      do i = 1, depth
         write (unit, '(2X)', advance = 'no')
      end do
    end subroutine pad

  end subroutine avl_write

  subroutine avl_check (tree)
    use, intrinsic :: iso_fortran_env, only: error_unit
    class(avl_tree_t), intent(in) :: tree

    type(avl_node_t), pointer :: p
    integer :: height_L, height_R

    p => tree%p
    call get_heights (p, height_L, height_R)
    call check_heights (height_L, height_R)

  contains

    recursive subroutine get_heights (p, height_L, height_R)
      type(avl_node_t), pointer, intent(in) :: p
      integer, intent(out) :: height_L, height_R

      integer :: height_LL, height_LR
      integer :: height_RL, height_RR

      height_L = 0
      height_R = 0
      if (associated (p)) then
         call get_heights (p%left, height_LL, height_LR)
         call check_heights (height_LL, height_LR)
         height_L = height_LL + height_LR
         call get_heights (p%right, height_RL, height_RR)
         call check_heights (height_RL, height_RR)
         height_R = height_RL + height_RR
      end if
    end subroutine get_heights

    subroutine check_heights (height_L, height_R)
      integer, value :: height_L, height_R

      if (2 <= abs (height_L - height_R)) then
         write (error_unit, '("*** AVL condition violated ***")')
         error stop
      end if
    end subroutine check_heights

  end subroutine avl_check

end module avl_trees

program avl_trees_demo
  use, intrinsic :: iso_fortran_env, only: output_unit
  use, non_intrinsic :: avl_trees

  implicit none

  integer, parameter :: keys_count = 20

  type(avl_tree_t) :: tree
  logical :: found
  class(*), allocatable :: retval
  integer :: the_keys(1:keys_count)
  integer :: i, j

  do i = 1, keys_count
     the_keys(i) = i
  end do
  call fisher_yates_shuffle (the_keys, keys_count)

  call avl_check (tree)
  do i = 1, keys_count
     call avl_insert (lt, the_keys(i), real (the_keys(i)), tree)
     call avl_check (tree)
     if (avl_size (tree) /= i) error stop
     do j = 1, keys_count
        if (avl_contains (lt, the_keys(j), tree) .neqv. (j <= i)) error stop
     end do
     do j = 1, keys_count
        call avl_retrieve (lt, the_keys(j), tree, found, retval)
        if (found .neqv. (j <= i)) error stop
        if (found) then
           ! This crazy way to write ‘/=’ is to quell those tiresome
           ! warnings about using ‘==’ or ‘/=’ with floating point
           ! numbers. Floating point numbers can represent integers
           ! *exactly*.
           if (0 < abs (real_cast (retval) - real (the_keys(j)))) error stop
        end if
        if (found) then
           block
             character(len = 1), parameter :: ch = '*'
             !
             ! Try replacing the data with a character and then
             ! restoring the number.
             !
             call avl_insert (lt, the_keys(j), ch, tree)
             call avl_retrieve (lt, the_keys(j), tree, found, retval)
             if (.not. found) error stop
             if (char_cast (retval) /= ch) error stop
             call avl_insert (lt, the_keys(j), real (the_keys(j)), tree)
             call avl_retrieve (lt, the_keys(j), tree, found, retval)
             if (.not. found) error stop
             if (0 < abs (real_cast (retval) - real (the_keys(j)))) error stop
           end block
        end if
     end do
  end do

  write (output_unit, '(70("-"))')
  call avl_write (int_real_writer, output_unit, tree)
  write (output_unit, '(70("-"))')
  call print_contents (output_unit, tree)
  write (output_unit, '(70("-"))')

  call fisher_yates_shuffle (the_keys, keys_count)
  do i = 1, keys_count
     call avl_delete (lt, the_keys(i), tree)
     call avl_check (tree)
     if (avl_size (tree) /= keys_count - i) error stop
     ! Try deleting a second time.
     call avl_delete (lt, the_keys(i), tree)
     call avl_check (tree)
     if (avl_size (tree) /= keys_count - i) error stop
     do j = 1, keys_count
        if (avl_contains (lt, the_keys(j), tree) .neqv. (i < j)) error stop
     end do
     do j = 1, keys_count
        call avl_retrieve (lt, the_keys(j), tree, found, retval)
        if (found .neqv. (i < j)) error stop
        if (found) then
           if (0 < abs (real_cast (retval) - real (the_keys(j)))) error stop
        end if
     end do
  end do

contains

  subroutine fisher_yates_shuffle (keys, n)
    integer, intent(inout) :: keys(*)
    integer, intent(in) :: n

    integer :: i, j
    real :: randnum
    integer :: tmp

    do i = 1, n - 1
       call random_number (randnum)
       j = i + floor (randnum * (n - i + 1))
       tmp = keys(i)
       keys(i) = keys(j)
       keys(j) = tmp
    end do
  end subroutine fisher_yates_shuffle

  function int_cast (u) result (v)
    class(*), intent(in) :: u
    integer :: v

    select type (u)
    type is (integer)
       v = u
    class default
       ! This case is not handled.
       error stop
    end select
  end function int_cast

  function real_cast (u) result (v)
    class(*), intent(in) :: u
    real :: v

    select type (u)
    type is (real)
       v = u
    class default
       ! This case is not handled.
       error stop
    end select
  end function real_cast

  function char_cast (u) result (v)
    class(*), intent(in) :: u
    character(len = 1) :: v

    select type (u)
    type is (character(*))
       v = u
    class default
       ! This case is not handled.
       error stop
    end select
  end function char_cast

  function lt (u, v) result (u_lt_v)
    class(*), intent(in) :: u, v
    logical :: u_lt_v

    select type (u)
    type is (integer)
       select type (v)
       type is (integer)
          u_lt_v = (u < v)
       class default
          ! This case is not handled.
          error stop
       end select
    class default
       ! This case is not handled.
       error stop
    end select
  end function lt

  subroutine int_real_writer (unit, key, data)
    integer, intent(in) :: unit
    class(*), intent(in) :: key, data

    write (unit, '("(", I0, ", ", F0.1, ")")', advance = 'no') &
         & int_cast(key), real_cast(data)
  end subroutine int_real_writer

  subroutine print_contents (unit, tree)
    integer, intent(in) :: unit
    class(avl_tree_t), intent(in) :: tree

    type(avl_pointer_pair_t), pointer :: ppairs, pp

    write (unit, '("tree size = ", I0)') avl_size (tree)
    ppairs => avl_pointer_pairs (tree)
    pp => ppairs
    do while (associated (pp))
       write (unit, '("(", I0, ", ", F0.1, ")")') &
            & int_cast (pp%p_key), real_cast (pp%p_data)
       pp => pp%next
    end do
    if (associated (ppairs)) deallocate (ppairs)
  end subroutine print_contents

end program avl_trees_demo