path: root/rb.c

/*
Generic C code for red-black trees.
Copyright (C) 2000 James S. Plank

This library is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public
License as published by the Free Software Foundation; either
version 2.1 of the License, or (at your option) any later version.

This library is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public
License along with this library; if not, write to the Free Software
Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */
/* Revision 1.2.  Jim Plank */

/* Original code by Jim Plank (plank@cs.utk.edu) */
/* modified for THINK C 6.0 for Macintosh by Chris Bartley */

/* Modified by Russell King to support embedded keys only */

#include <assert.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <ctype.h>
#include "rb.h"

static void mk_new_int(Rb_node l, Rb_node r, Rb_node p, int il);
static Rb_node lprev(Rb_node n);
static Rb_node rprev(Rb_node n);
static void recolor(Rb_node n);
static void single_rotate(Rb_node y, int l);
#ifdef DEBUG
static void rb_print_tree(Rb_node t, int level);
static void rb_iprint_tree(Rb_node t, int level);
#endif


#define isred(n) (n->s.red)
#define isblack(n) (!isred(n))
#define isleft(n) (n->s.left)
#define isright(n) (!isleft(n))
#define isint(n) (n->s.internal)
#define isext(n) (!isint(n))
#define ishead(n) (n->s.head)
#define isroot(n) (n->s.root)
#define setred(n) n->s.red = 1
#define setblack(n) n->s.red = 0
#define setleft(n) n->s.left = 1
#define setright(n) n->s.left = 0
#define sethead(n) n->s.head = 1
#define setroot(n) n->s.root = 1
#define setint(n) n->s.internal = 1
#define setext(n) n->s.internal = 0
#define setnormal(n) { n->s.root = 0; n ->s.head = 0; }
#define sibling(n) ((isleft(n)) ? n->p.parent->c.child.right \
                                : n->p.parent->c.child.left)

static void insert(Rb_node item, Rb_node list)	/* Inserts to the end of a list */
{
  Rb_node last_node;

  last_node = list->c.list.blink;

  list->c.list.blink = item;
  last_node->c.list.flink = item;
  item->c.list.blink = last_node;
  item->c.list.flink = list;
}

static void delete_item(Rb_node item)		/* Deletes an arbitrary iterm */
{
  item->c.list.flink->c.list.blink = item->c.list.blink;
  item->c.list.blink->c.list.flink = item->c.list.flink;
}

#define mk_new_ext(new, vvval) {\
  new = (Rb_node) malloc(sizeof(struct rb_node));\
  new->v.val = vvval;\
  setext(new);\
  setblack(new);\
  setnormal(new);\
}

static void mk_new_int(Rb_node l, Rb_node r, Rb_node p, int il)
{
  Rb_node newnode;

  newnode = (Rb_node) malloc(sizeof(struct rb_node));
  setint(newnode);
  setred(newnode);
  setnormal(newnode);
  newnode->c.child.left = l;
  newnode->c.child.right = r;
  newnode->p.parent = p;
  newnode->k.lext = l;
  newnode->v.rext = r;
  l->p.parent = newnode;
  r->p.parent = newnode;
  setleft(l);
  setright(r);
  if (ishead(p)) {
    p->p.root = newnode;
    setroot(newnode);
  } else if (il) {
    setleft(newnode);
    p->c.child.left = newnode;
  } else {
    setright(newnode);
    p->c.child.right = newnode;
  }
  recolor(newnode);
}


Rb_node lprev(Rb_node n)
{
  if (ishead(n)) return n;
  while (!isroot(n)) {
    if (isright(n)) return n->p.parent;
    n = n->p.parent;
  }
  return n->p.parent;
}

Rb_node rprev(Rb_node n)
{
  if (ishead(n)) return n;
  while (!isroot(n)) {
    if (isleft(n)) return n->p.parent;
    n = n->p.parent;
  }
  return n->p.parent;
}

Rb_node make_rb(void)
{
  Rb_node head;

  head = (Rb_node) malloc (sizeof(struct rb_node));
  head->c.list.flink = head;
  head->c.list.blink = head;
  head->p.root = head;
  sethead(head);
  return head;
}

Rb_node rb_find_key_n(Rb_node n, const void *key, rb_key_cmp_fn cmp, int *fnd)
{
  int rc;
  *fnd = 0;
#ifdef DEBUG
  if (!ishead(n)) {
    fprintf(stderr, "%s called on non-head %p\n", __FUNCTION__, n);
    exit(1);
  }
#else
  assert(ishead(n));
#endif
  if (n->p.root == n) return n;
  rc = cmp(key, n->c.list.blink->v.val);
  if (rc == 0) {
    *fnd = 1;
    return n->c.list.blink;
  }
  if (rc > 0) return n;
  else n = n->p.root;
  while (1) {
    if (isext(n)) return n;
    rc = cmp(key, n->k.lext->v.val);
    if (rc == 0) {
      *fnd = 1;
      return n->k.lext;
    }
    n = rc < 0 ? n->c.child.left : n->c.child.right;
  }
}

Rb_node rb_find_key(Rb_node n, const void *key, rb_key_cmp_fn cmp)
{
  int fnd;
  return rb_find_key_n(n, key, cmp, &fnd);
}

Rb_node rb_insert_b(Rb_node n, void *val)
{
  Rb_node newleft, newright, newnode, p;

  if (ishead(n)) {
    if (n->p.root == n) {         /* Tree is empty */
      mk_new_ext(newnode, val);
      insert(newnode, n);
      n->p.root = newnode;
      newnode->p.parent = n;
      setroot(newnode);
      return newnode;
    } else {
      mk_new_ext(newright, val);
      insert(newright, n);
      newleft = newright->c.list.blink;
      setnormal(newleft);
      mk_new_int(newleft, newright, newleft->p.parent, isleft(newleft));
      p = rprev(newright);
      if (!ishead(p)) p->k.lext = newright;
      return newright;
    }
  } else {
    mk_new_ext(newleft, val);
    insert(newleft, n);
    setnormal(n);
    mk_new_int(newleft, n, n->p.parent, isleft(n));
    p = lprev(newleft);
    if (!ishead(p)) p->v.rext = newleft;
    return newleft;
  }
}

static void recolor(Rb_node n)
{
  Rb_node p, gp, s;
  int done = 0;

  while(!done) {
    if (isroot(n)) {
      setblack(n);
      return;
    }

    p = n->p.parent;

    if (isblack(p)) return;

    if (isroot(p)) {
      setblack(p);
      return;
    }

    gp = p->p.parent;
    s = sibling(p);
    if (isred(s)) {
      setblack(p);
      setred(gp);
      setblack(s);
      n = gp;
    } else {
      done = 1;
    }
  }
  /* p's sibling is black, p is red, gp is black */

  if ((isleft(n) == 0) == (isleft(p) == 0)) {
    single_rotate(gp, isleft(n));
    setblack(p);
    setred(gp);
  } else {
    single_rotate(p, isleft(n));
    single_rotate(gp, isleft(n));
    setblack(n);
    setred(gp);
  }
}

static void single_rotate(Rb_node y, int l)
{
  int rl, ir;
  Rb_node x, yp;

  ir = isroot(y);
  yp = y->p.parent;
  if (!ir) {
    rl = isleft(y);
  }

  if (l) {
    x = y->c.child.left;
    y->c.child.left = x->c.child.right;
    setleft(y->c.child.left);
    y->c.child.left->p.parent = y;
    x->c.child.right = y;
    setright(y);
  } else {
    x = y->c.child.right;
    y->c.child.right = x->c.child.left;
    setright(y->c.child.right);
    y->c.child.right->p.parent = y;
    x->c.child.left = y;
    setleft(y);
  }

  x->p.parent = yp;
  y->p.parent = x;
  if (ir) {
    yp->p.root = x;
    setnormal(y);
    setroot(x);
  } else {
    if (rl) {
      yp->c.child.left = x;
      setleft(x);
    } else {
      yp->c.child.right = x;
      setright(x);
    }
  }
}

void rb_delete_node(Rb_node n)
{
  Rb_node s, p, gp;
  char ir;

#ifdef DEBUG
  if (isint(n)) {
    fprintf(stderr, "Cannot delete an internal node: %p\n", n);
    exit(1);
  }
  if (ishead(n)) {
    fprintf(stderr, "Cannot delete the head of an rb_tree: %p\n", n);
    exit(1);
  }
#else
  assert(!isint(n));
  assert(!ishead(n));
#endif
  delete_item(n); /* Delete it from the list */
  p = n->p.parent;  /* The only node */
  if (isroot(n)) {
    p->p.root = p;
    free(n);
    return;
  }
  s = sibling(n);    /* The only node after deletion */
  if (isroot(p)) {
    s->p.parent = p->p.parent;
    s->p.parent->p.root = s;
    setroot(s);
    free(p);
    free(n);
    return;
  }
  gp = p->p.parent;  /* Set parent to sibling */
  s->p.parent = gp;
  if (isleft(p)) {
    gp->c.child.left = s;
    setleft(s);
  } else {
    gp->c.child.right = s;
    setright(s);
  }
  ir = isred(p);
  free(p);
  free(n);

  if (isext(s)) {      /* Update proper rext and lext values */
    p = lprev(s);
    if (!ishead(p)) p->v.rext = s;
    p = rprev(s);
    if (!ishead(p)) p->k.lext = s;
  } else if (isblack(s)) {
#ifdef DEBUG
    fprintf(stderr, "DELETION PROB -- sib is black, internal\n");
    exit(1);
#else
    assert(1);
#endif
  } else {
    p = lprev(s);
    if (!ishead(p)) p->v.rext = s->c.child.left;
    p = rprev(s);
    if (!ishead(p)) p->k.lext = s->c.child.right;
    setblack(s);
    return;
  }

  if (ir) return;

  /* Recolor */

  n = s;
  p = n->p.parent;
  s = sibling(n);
  while(isblack(p) && isblack(s) && isint(s) &&
        isblack(s->c.child.left) && isblack(s->c.child.right)) {
    setred(s);
    n = p;
    if (isroot(n)) return;
    p = n->p.parent;
    s = sibling(n);
  }

  if (isblack(p) && isred(s)) {  /* Rotation 2.3b */
    single_rotate(p, isright(n));
    setred(p);
    setblack(s);
    s = sibling(n);
  }

  { Rb_node x, z; char il;
#ifdef DEBUG
    if (isext(s)) {
      fprintf(stderr, "DELETION ERROR: sibling not internal\n");
      exit(1);
    }
#else
    assert(!isext(s));
#endif

    il = isleft(n);
    x = il ? s->c.child.left : s->c.child.right;
    z = sibling(x);

    if (isred(z)) {  /* Rotation 2.3f */
      single_rotate(p, !il);
      setblack(z);
      if (isred(p)) setred(s); else setblack(s);
      setblack(p);
    } else if (isblack(x)) {   /* Recoloring only (2.3c) */
#ifdef DEBUG
      if (isred(s) || isblack(p)) {
        fprintf(stderr, "DELETION ERROR: 2.3c not quite right\n");
        exit(1);
      }
#else
      assert(!isred(s) && !isblack(p));
#endif
      setblack(p);
      setred(s);
      return;
    } else if (isred(p)) { /* 2.3d */
      single_rotate(s, il);
      single_rotate(p, !il);
      setblack(x);
      setred(s);
      return;
    } else {  /* 2.3e */
      single_rotate(s, il);
      single_rotate(p, !il);
      setblack(x);
      return;
    }
  }
}


#ifdef DEBUG
void rb_print_tree(Rb_node t, int level)
{
  int i;
  if (ishead(t) && t->p.parent == t) {
    printf("tree %p is empty\n", t);
  } else if (ishead(t)) {
    printf("Head: %p.  Root = %p\n", t, t->p.root);
    rb_print_tree(t->p.root, 0);
  } else {
    if (isext(t)) {
      for (i = 0; i < level; i++) putchar(' ');
      printf("Ext node %p: %c,%c: p=0x%x, k=%s\n",
              t, isred(t)?'R':'B', isleft(t)?'l':'r', t->p.parent, t->k.key);
    } else {
      rb_print_tree(t->c.child.left, level+2);
      rb_print_tree(t->c.child.right, level+2);
      for (i = 0; i < level; i++) putchar(' ');
      printf("Int node %p: %c,%c: l=0x%x, r=0x%x, p=0x%x, lr=(%s,%s)\n",
              t, isred(t)?'R':'B', isleft(t)?'l':'r', t->c.child.left,
              t->c.child.right,
              t->p.parent, t->k.lext->k.key, t->v.rext->k.key);
    }
  }
}
#endif

int rb_nblack(Rb_node n)
{
  int nb;
#ifdef DEBUG
  if (ishead(n) || isint(n)) {
    fprintf(stderr, "ERROR: rb_nblack called on a non-external node %p\n",
            n);
    exit(1);
  }
#else
  assert(!ishead(n));
  assert(!isint(n));
#endif
  nb = 0;
  while(!ishead(n)) {
    if (isblack(n)) nb++;
    n = n->p.parent;
  }
  return nb;
}

int rb_plength(Rb_node n)
{
  int pl;
#ifdef DEBUG
  if (ishead(n) || isint(n)) {
    fprintf(stderr, "ERROR: rb_plength called on a non-external node %p\n",
            n);
    exit(1);
  }
#else
  assert(!ishead(n));
  assert(!isint(n));
#endif
  pl = 0;
  while(!ishead(n)) {
    pl++;
    n = n->p.parent;
  }
  return pl;
}

void rb_free_tree(Rb_node n)
{
#ifdef DEBUG
  if (!ishead(n)) {
    fprintf(stderr, "ERROR: Rb_free_tree called on a non-head node\n");
    exit(1);
  }
#else
  assert(ishead(n));
#endif

  while(rb_first(n) != rb_nil(n)) {
    rb_delete_node(rb_first(n));
  }
  free(n);
}

void *rb_val(Rb_node n)
{
  return n->v.val;
}

Rb_node rb_insert_a(Rb_node nd, void *val)
{
  return rb_insert_b(nd->c.list.flink, val);
}

Rb_node rb_insert(Rb_node tree, const void *key, rb_key_cmp_fn cmp, void *val)
{
  return rb_insert_b(rb_find_key(tree, key, cmp), val);
}