Finish implementation of binary tree

3 files changed, 185 insertions, 9 deletions

diff --git a/src/peaks.c b/src/peaks.c
index d1132a3a..06c5ee35 100644
--- a/src/peaks.c
+++ b/src/peaks.c
@@ -538,7 +538,9 @@ int find_projected_peaks(struct image *image, UnitCell *cell,
 				set_detector_pos(refl, dist, x, y);
 			}
 		} else {
-			add_refl_with_det_pos(reflections, h, k, l, x, y, dist);
+			Reflection *new;
+			new = add_refl(reflections, h, k, l);
+			set_detector_pos(new, dist, x, y);
 			n_reflections++;
 		}
 
diff --git a/src/reflist.c b/src/reflist.c
index 220b32bd..80179f63 100644
--- a/src/reflist.c
+++ b/src/reflist.c
@@ -11,6 +11,7 @@
 
 
 #include <stdlib.h>
+#include <assert.h>
 
 #include "reflist.h"
 
@@ -18,9 +19,12 @@
 struct _reflection {
 
 	/* Listy stuff */
+	RefList *list;
 	unsigned int serial;          /* Unique serial number, key */
 	struct _reflection *child[2]; /* Child nodes */
 	struct _reflection *parent;   /* Parent node */
+	struct _reflection *next;     /* Another reflection with the same
+	                               * indices, or NULL */
 
 	signed int h;
 	signed int k;
@@ -56,6 +60,8 @@ struct _reflist {
 };
 
 
+#define SERIAL(h, k, l) (((h)+256)*512*512 + ((k)+256)*512 + ((l)+256))
+
 /**************************** Creation / deletion *****************************/
 
 /* Create a reflection list */
@@ -87,13 +93,44 @@ void reflist_free(RefList *list)
 
 /********************************** Search ************************************/
 
+static Reflection *recursive_search(Reflection *refl, unsigned int search)
+{
+	int i;
+
+	/* Hit the bottom of the tree? */
+	if ( refl == NULL ) return NULL;
+
+	/* Is this the correct reflection? */
+	if ( refl->serial == search ) return refl;
+
+	for ( i=0; i<2; i++ ) {
+
+		if ( search < refl->serial ) {
+			return recursive_search(refl->child[0], search);
+		}
+
+		if ( search >= refl->serial ) {
+			return recursive_search(refl->child[1], search);
+		}
+
+	}
+
+	return NULL;
+}
+
+
+/* Return the first reflection in 'list' with the given indices, or NULL */
 Reflection *find_refl(RefList *list, INDICES)
 {
+	unsigned int search = SERIAL(h, k, l);
+	return recursive_search(list->head, search);
 }
 
 
+/* Find the next reflection in 'refl's list with the same indices, or NULL */
 Reflection *next_found_refl(Reflection *refl)
 {
+	return refl->next;  /* Well, that was easy... */
 }
 
 
@@ -153,41 +190,127 @@ int get_scalable(Reflection *refl)
 
 void set_detector_pos(Reflection *refl, double exerr, double x, double y)
 {
+	refl->excitation_error = exerr;
+	refl->x = x;
+	refl->y = y;
 }
 
 
 void set_partial(Reflection *refl, double r1, double r2, double p,
                  double clamp_low, double clamp_high)
 {
+	refl->r1 = r1;
+	refl->r2 = r2;
+	refl->p = p;
+	refl->clamp1 = clamp_low;
+	refl->clamp2 = clamp_high;
 }
 
 
-void set_indices(Reflection *refl,
-                        signed int h, signed int k, signed int l)
+void set_indices(Reflection *refl, signed int h, signed int k, signed int l)
 {
+	/* Tricky, because the indices determine the position in the tree */
+	Reflection copy;
+	Reflection *new;
+	Reflection transfer;
+
+	/* Copy all data */
+	copy = *refl;
+
+	/* Delete and re-add with new indices */
+	delete_refl(refl);
+	new = add_refl(copy.list, h, k, l);
+
+	/* Transfer data back */
+	transfer = *new;
+	*new = copy;
+	new->list = transfer.list;
+	new->parent = transfer.parent;
+	new->child[0] = transfer.child[0];
+	new->child[1] = transfer.child[1];
+	new->h = transfer.h;  new->k = transfer.k;  new->l = transfer.l;
+	new->serial = transfer.serial;
 }
 
 
 void set_int(Reflection *refl, double intensity)
 {
+	refl->intensity = intensity;
 }
 
 
 void set_scalable(Reflection *refl, int scalable)
 {
+	refl->scalable = scalable;
 }
 
 
 /********************************* Insertion **********************************/
 
-Reflection *add_refl(RefList *list, INDICES)
+static int recursive_insert(Reflection *refl, Reflection *new)
 {
+	int i;
+
+	if ( refl->serial == new->serial ) {
+
+		/* Found a reflection with identical indices */
+		do {
+			refl = refl->next;
+		} while ( refl != NULL );
+		refl->next = new;
+		new->parent = refl;
+
+		return 1;
+	}
+
+	for ( i=0; i<2; i++ ) {
+
+		if ( new->serial < refl->serial ) {
+			if ( refl->child[0] != NULL ) {
+				return recursive_insert(refl->child[0], new);
+			} else {
+				refl->child[0] = new;
+				new->parent = refl;
+				return 1;
+			}
+		}
+
+		if ( new->serial >= refl->serial ) {
+			if ( refl->child[1] != NULL ) {
+				return recursive_insert(refl->child[1], new);
+			} else {
+				refl->child[1] = new;
+				new->parent = refl;
+				return 1;
+			}
+		}
+
+
+	}
+
+	return 0;
 }
 
 
-Reflection *add_refl_with_det_pos(RefList *refl, INDICES, double exerr,
-                                  double x, double y)
+Reflection *add_refl(RefList *list, INDICES)
 {
+	Reflection *new;
+
+	new = calloc(1, sizeof(struct _reflection));
+	new->h = h;  new->k = k,  new->l = l;
+	new->serial = SERIAL(h, k, l);
+	new->next = NULL;
+	new->child[0] = NULL;
+	new->child[1] = NULL;
+
+	if ( list->head == NULL ) {
+		list->head = new;
+		new->parent = NULL;
+	} else {
+		recursive_insert(list->head, new);
+	}
+
+	return new;
 }
 
 
@@ -195,6 +318,49 @@ Reflection *add_refl_with_det_pos(RefList *refl, INDICES, double exerr,
 
 void delete_refl(Reflection *refl)
 {
+	int i;
+	Reflection **parent_pos = NULL;
+
+	/* Remove parent's reference */
+	for ( i=0; i<2; i++ ) {
+		if ( refl->parent->child[i] == refl ) {
+			parent_pos = &refl->parent->child[i];
+			*parent_pos = NULL;
+		}
+	}
+	assert(parent_pos != NULL);
+
+	/* Two child nodes? */
+	if ( (refl->child[0] != NULL) && (refl->child[1] != NULL ) ) {
+
+		if ( random() > RAND_MAX/2 ) {
+
+			*parent_pos = refl->child[0];
+
+			/* Now sort out the right child */
+			recursive_insert(refl->child[0], refl->child[1]);
+		} else {
+
+			*parent_pos = refl->child[1];
+
+			/* Now sort out the left child */
+			recursive_insert(refl->child[1], refl->child[0]);
+
+		}
+
+	} else if ( refl->child[0] != NULL ) {
+
+		/* One child, left */
+		*parent_pos = refl->child[0];
+
+	} else if (refl->child[1] != NULL ) {
+
+		/* One child, right */
+		*parent_pos = refl->child[1];
+
+	} /* else it was just a leaf node */
+
+	free(refl);
 }
 
 
@@ -202,11 +368,22 @@ void delete_refl(Reflection *refl)
 
 Reflection *first_refl(RefList *list)
 {
+	return list->head;
 }
 
 
 Reflection *next_refl(Reflection *refl)
 {
+	/* Does a left child exist? */
+	if ( refl->child[0] != NULL ) return refl->child[0];
+
+	/* Otherwise move up the tree to find the next right child */
+	while ( refl->child[1] != NULL ) {
+		refl = refl->parent;
+		if ( refl == NULL ) return NULL;
+	}
+
+	return refl->child[1];
 }
 
 
diff --git a/src/reflist.h b/src/reflist.h
index d1f0d53c..c7b2ad42 100644
--- a/src/reflist.h
+++ b/src/reflist.h
@@ -53,9 +53,6 @@ extern void set_scalable(Reflection *refl, int scalable);
 
 /* Insertion */
 extern Reflection *add_refl(RefList *list, INDICES);
-extern Reflection *add_refl_with_det_pos(RefList *list, INDICES, double exerr,
-                                        double x, double y);
-
 
 /* Deletion */
 extern void delete_refl(Reflection *refl);