(****************************************
 * Description:
 *   Implementiert einen AVL-Baum
 *   in Oberon.
 * Author:
 *   Daniel Hottinger <hodaniel@iiic.ethz.ch>
 *   Department of Computer Science, ETH Zurich
 *   SS 2001
 * Licence: GNU GPL v2 or later
 * Created: Mon May 14 20:16:37 CEST 2001
 * Last update: none
 * Changes:
 ****************************************)

MODULE AVLTree;

IMPORT
  Out;

CONST
  LESS  = -1;
  EQUAL =  0;
  MORE  =  1;

TYPE
  PKey* = POINTER TO TKey;
  TKey* = RECORD
  END;

  PValue* = POINTER TO TValue;
  TValue* = RECORD
  END;

  TCompareFunc* = PROCEDURE (key1, key2: PKey): LONGINT;
  TPrintProc = PROCEDURE(key: PKey);

  PAVLNode = POINTER TO TAVLNode;
  TAVLNode = RECORD
    balance: LONGINT;
    left, right: PAVLNode;
    key: PKey;
    value: PValue;
  END;

  PAVLTree* = POINTER TO TAVLTree;
  TAVLTree* = RECORD
    root: PAVLNode;
    cmp: TCompareFunc;
  END;

(****************************************
 * Speicherverwaltung
 ****************************************)

PROCEDURE NewNode(key: PKey; value: PValue): PAVLNode;
VAR
  node: PAVLNode;
BEGIN
  NEW(node);
  node^.balance := 0;
  node^.left := NIL;
  node^.right := NIL;
  node^.key := key;
  node^.value := value;
  RETURN node;
END NewNode;

(* One day we'll reuse these nodes, *)
(* so oberon can do less wrong      *)
PROCEDURE DestroyNode(VAR node: PAVLNode);
BEGIN
  IF node # NIL THEN
    DestroyNode(node^.left);
    DestroyNode(node^.right);
  END;
  (* Hope the oberon garbage collector works *)
  node := NIL; 
END DestroyNode;

PROCEDURE NewTree*(cmp: TCompareFunc): PAVLTree;
VAR
  tree: PAVLTree;
BEGIN
  NEW(tree);
  tree^.root := NIL;
  tree^.cmp := cmp;
  RETURN tree;
END NewTree;

PROCEDURE DestroyTree*(VAR tree: PAVLTree);
BEGIN
  IF tree # NIL THEN
    DestroyNode(tree^.root);
    tree := NIL;
  END;
END DestroyTree;

(****************************************
 * Rotationen
 ****************************************)

(****************************************
 * n(l,r(rl,rr)) =>
 * r(n(l,rl),rr)
 * 
 *   n                r
 *  / `_           _.' \
 * l    r    =>   n     rr
 *     / \       / \
 *    rl  rr    l   rl
 ****************************************)
PROCEDURE RotLeft(node: PAVLNode): PAVLNode;
VAR
  left, right: PAVLNode;
  abal, bbal: LONGINT;
BEGIN
  left := node^.left;
  right := node^.right;

  node^.right := right^.left;
  right^.left := node;

  abal := node^.balance;
  bbal := right^.balance;

  IF bbal <= EQUAL THEN
    IF abal >= MORE THEN
      right^.balance := bbal - 1;
    ELSE
      right^.balance := abal + bbal - 2;
    END;
    node^.balance := abal - 1;
  ELSE
    IF abal <= bbal THEN
      right^.balance := abal - 2;
    ELSE
      right^.balance := bbal - 1;
    END;
    node^.balance := abal - bbal - 1;
  END;
  RETURN right;
END RotLeft;

(****************************************
 * n(l(ll,lr),r) =>
 * l(ll,n(lr,r))
 *
 *       n        l
 *    _.' \      / `_
 *   l     r => ll   n
 *  / \             / \
 * ll  lr          lr  r
 ****************************************)
PROCEDURE RotRight(node: PAVLNode): PAVLNode;
VAR
  left, right: PAVLNode;
  abal, bbal: LONGINT;
BEGIN
  left := node^.left;
  right := node^.right;

  node^.left := left^.right;
  left^.right := node;

  abal := node^.balance;
  bbal := left^.balance;

  IF bbal <= EQUAL THEN
    IF abal < bbal THEN
      left^.balance := bbal + 1;
    ELSE
      left^.balance := abal + 2;
    END;
    node^.balance := abal - bbal + 1;
  ELSE
    IF abal <= LESS THEN
      left^.balance := bbal + 1;
    ELSE
      left^.balance := abal + bbal + 2;
    END;
    node^.balance := abal + 1;
  END;
  RETURN left;
END RotRight;

(****************************************
 * Ausbalancieren
 ****************************************)

(****************************************
 * Beispiel: (Doppelrotation)
 * 2(1,6(5(4),7)) =>
 * 2(1,5(4,6(,7))) =>
 * 5(2(1,4),6(,7))
 * 
 *   2             2
 *  / `-._        / `_            5
 * 1      6      1    5          / \
 *       / \  =>     / \    =>  2   6
 *      5   7       4   6       ^    \
 *     /                 \     1 4    7
 *    4                   7
 ****************************************)
PROCEDURE BalanceNode(node: PAVLNode): PAVLNode;
BEGIN
  IF node^.balance < LESS THEN
    IF node^.left^.balance > EQUAL THEN
      node^.left := RotLeft(node^.left);
    END;
    node := RotRight(node);
  ELSIF node^.balance > MORE THEN
    IF node^.right^.balance < EQUAL THEN
      node^.right := RotRight(node^.right); (* siehe Beispiel *)
    END;
    node := RotLeft(node);
  END;
  RETURN node;
END BalanceNode;

(* Called after a node of node^.left was removed *)
PROCEDURE RestoreLeftBalance(node: PAVLNode; oldbalance: LONGINT): PAVLNode;
BEGIN
  IF node^.left = NIL THEN
    INC(node^.balance);
  ELSIF (node^.left^.balance # oldbalance) & (node^.left^.balance = 0) THEN
    (* left tree shrunk *)
    INC(node^.balance);
  END;
  IF node^.balance > MORE THEN
    RETURN BalanceNode(node);
  ELSE
    RETURN node;
  END;
END RestoreLeftBalance;

(* Called after a node of node^.right was removed *)
PROCEDURE RestoreRightBalance(node: PAVLNode; oldbalance: LONGINT): PAVLNode;
BEGIN
  IF node^.right = NIL THEN
    DEC(node^.balance);
  ELSIF (node^.right^.balance # oldbalance) & (node^.right^.balance = 0) THEN
    (* right tree shrunk *)
    DEC(node^.balance);
  END;
  IF node^.balance > LESS THEN
    RETURN BalanceNode(node);
  ELSE
    RETURN node;
  END;
END RestoreRightBalance;

PROCEDURE RemoveNodeMostLeft(node: PAVLNode; VAR leftmost: PAVLNode): PAVLNode;
VAR
  oldbalance: LONGINT;
BEGIN
  IF node^.left = NIL THEN
    leftmost := node;
    RETURN node^.right;
  END;

  oldbalance := node^.left^.balance;
  node^.left := RemoveNodeMostLeft(node^.left, leftmost);
  RETURN RestoreLeftBalance(node, oldbalance);
END RemoveNodeMostLeft;

(****************************************
 * grundlegende Operationen
 ****************************************)

PROCEDURE InsertNode(node: PAVLNode; cmp: TCompareFunc; key: PKey; 
  value: PValue; VAR inserted: BOOLEAN): PAVLNode;
VAR
  relation: LONGINT;
  oldbalance: LONGINT;
BEGIN
  IF node = NIL THEN
    inserted := TRUE;
    RETURN NewNode(key, value);
  END;

  relation := cmp(key, node^.key);
  IF relation = EQUAL THEN
    (* Don't insert dublicate key/value *)
    inserted := FALSE;
    RETURN node;
  ELSIF relation = LESS THEN
    IF node^.left # NIL THEN
      oldbalance := node^.left^.balance;
      node^.left := InsertNode(node^.left, cmp, key, value, inserted);
      IF (oldbalance # node^.left^.balance) & (node^.left^.balance # 0) THEN
	(* Tree has grown *)
	DEC(node^.balance);
      END;
    ELSE
      inserted := TRUE;
      node^.left := NewNode(key, value);
      DEC(node^.balance);
    END;
  ELSIF relation = MORE THEN
    IF node^.right # NIL THEN
      oldbalance := node^.right^.balance;
      node^.right := InsertNode(node^.right, cmp, key, value, inserted);
      IF (oldbalance # node^.right^.balance) & (node^.right^.balance # 0) THEN
	(* Tree has grown *)
	INC(node^.balance);
      END;
    ELSE
      inserted := TRUE;
      node^.right := NewNode(key, value);
      INC(node^.balance);
    END;
  END;

  IF inserted THEN
    IF ABS(node^.balance) > 1 THEN
      node := BalanceNode(node);
    END;
  END;

  RETURN node;

END InsertNode;

PROCEDURE Insert*(tree: PAVLTree; key: PKey; value: PValue);
VAR
  inserted: BOOLEAN;
BEGIN
  IF tree # NIL THEN
    inserted := FALSE;
    tree^.root := InsertNode(tree^.root, tree^.cmp, key, value, inserted);
  END;
END Insert;

(****************************************
 * Beispiel:
 * n(l(:,:),r(rl(rll,rlr),rr)) =>
 * n(l(:,:),r(rl(,rlr),rr)) =>
 * rll(l(:,:),r(rl(,rlr),rr))
 *    n                    n                 rll
 *   / `--...___          / `-..__          /   `-..__
 *  l           r        l        r        l          r
 *  ^      __.-' \   =>  ^   __.-' \   =>  ^     __.-' \
 * : :    rl      rr    : : rl      rr    : :   rl      rr
 *      _'  \                 \                   \
 *     rll   rlr               rlr                 rlr
 ****************************************)
PROCEDURE RemoveNode(node: PAVLNode; cmp: TCompareFunc; key: PKey): PAVLNode;
VAR
  relation, oldbalance: LONGINT;
  garbage, newroot: PAVLNode;
BEGIN
  IF node = NIL THEN
    RETURN NIL;
  END;

  relation := cmp(key, node^.key);
  IF relation = EQUAL THEN
    garbage := node;
    IF node^.right = NIL THEN
      node := node^.left;
    ELSE
      oldbalance := node^.right^.balance;
      (* new right node is the leftmost of the right tree *)
      (* Beispiel *)
      node^.right := RemoveNodeMostLeft(node^.right, newroot);
      newroot^.left := node^.left;
      newroot^.right := node^.right;
      newroot^.balance := node^.balance;
      node := RestoreRightBalance(newroot, oldbalance);
    END;
    (* free *only* the removed node *)
    garbage^.right := NIL;
    garbage^.left := NIL;
    DestroyNode(garbage);
  ELSIF relation = LESS THEN
    IF node^.left # NIL THEN
      oldbalance := node^.left^.balance;
      node^.left := RemoveNode(node^.left, cmp, key);
      node := RestoreLeftBalance(node, oldbalance);
    END;
  ELSIF relation = MORE THEN
    IF node^.right # NIL THEN
      oldbalance := node^.right^.balance;
      node^.right := RemoveNode(node^.right, cmp, key);
      node := RestoreRightBalance(node, oldbalance);
    END;
  END;

  RETURN node;

END RemoveNode;

PROCEDURE Remove*(tree: PAVLTree; key: PKey);
BEGIN
  IF tree # NIL THEN
    tree^.root := RemoveNode(tree^.root, tree^.cmp, key);
  END;
END Remove;

(****************************************
 * Debug
 ****************************************)
PROCEDURE Traverse(node: PAVLNode; print: TPrintProc);
BEGIN
  IF node # NIL THEN
    print(node^.key);
    IF (node^.left # NIL) OR (node^.right # NIL) THEN
      Out.String("(");
      Traverse(node^.left, print);
      IF node^.right # NIL  THEN
	Out.String(",");
      END;
      Traverse(node^.right, print);
      Out.String(")");
    END;
  END;
END Traverse;

PROCEDURE Dump*(tree: PAVLTree; print: TPrintProc);
BEGIN
  Traverse(tree^.root, print);
END Dump;

END AVLTree.