4 files changed, 71 insertions, 63 deletions

diff --git a/lib/sqfs/xattr/xattr_writer.c b/lib/sqfs/xattr/xattr_writer.c
index c04bd4a..d23e063 100644
--- a/lib/sqfs/xattr/xattr_writer.c
+++ b/lib/sqfs/xattr/xattr_writer.c
@@ -9,8 +9,8 @@
 static sqfs_object_t *xattr_writer_copy(const sqfs_object_t *obj)
 {
 	const sqfs_xattr_writer_t *xwr = (const sqfs_xattr_writer_t *)obj;
-	kv_block_desc_t *blk, *it, **next;
 	sqfs_xattr_writer_t *copy;
+	kv_block_desc_t *it;
 
 	copy = calloc(1, sizeof(*copy));
 	if (copy == NULL)
@@ -34,27 +34,26 @@ static sqfs_object_t *xattr_writer_copy(const sqfs_object_t *obj)
 	memcpy(copy->kv_pairs, xwr->kv_pairs,
 	       sizeof(copy->kv_pairs[0]) * xwr->num_pairs);
 
-	next = &(copy->kv_blocks);
+	if (rbtree_copy(&xwr->kv_block_tree, &copy->kv_block_tree) != 0)
+		goto fail_tree;
 
-	for (it = xwr->kv_blocks; it != NULL; it = it->next) {
-		blk = malloc(sizeof(*blk));
-		if (blk == NULL)
-			goto fail_blk;
+	for (it = xwr->kv_block_first; it != NULL; it = it->next) {
+		rbtree_node_t *n = rbtree_lookup(&copy->kv_block_tree, it);
 
-		memcpy(blk, it, sizeof(*it));
-		blk->next = NULL;
+		if (copy->kv_block_last == NULL) {
+			copy->kv_block_first = rbtree_node_key(n);
+			copy->kv_block_last = copy->kv_block_first;
+		} else {
+			copy->kv_block_last->next = rbtree_node_key(n);
+			copy->kv_block_last = copy->kv_block_last->next;
+		}
 
-		*next = blk;
-		next = &(blk->next);
+		copy->kv_block_last->next = NULL;
 	}
 
 	return (sqfs_object_t *)copy;
-fail_blk:
-	while (copy->kv_blocks != NULL) {
-		blk = copy->kv_blocks;
-		copy->kv_blocks = copy->kv_blocks->next;
-		free(blk);
-	}
+fail_tree:
+	free(copy->kv_pairs);
 fail_pairs:
 	str_table_cleanup(&copy->values);
 fail_values:
@@ -67,20 +66,31 @@ fail_keys:
 static void xattr_writer_destroy(sqfs_object_t *obj)
 {
 	sqfs_xattr_writer_t *xwr = (sqfs_xattr_writer_t *)obj;
-	kv_block_desc_t *blk;
-
-	while (xwr->kv_blocks != NULL) {
-		blk = xwr->kv_blocks;
-		xwr->kv_blocks = xwr->kv_blocks->next;
-		free(blk);
-	}
 
+	rbtree_cleanup(&xwr->kv_block_tree);
 	free(xwr->kv_pairs);
 	str_table_cleanup(&xwr->values);
 	str_table_cleanup(&xwr->keys);
 	free(xwr);
 }
 
+static int block_compare(const void *context,
+			 const void *lhs, const void *rhs)
+{
+	const sqfs_xattr_writer_t *xwr = context;
+	const kv_block_desc_t *l = lhs, *r = rhs;
+
+	if (l->count != r->count)
+		return l->count < r->count ? -1 : 1;
+
+	if (l->start == r->start)
+		return 0;
+
+	return memcmp(xwr->kv_pairs + l->start,
+		      xwr->kv_pairs + r->start,
+		      l->count * sizeof(xwr->kv_pairs[0]));
+}
+
 sqfs_xattr_writer_t *sqfs_xattr_writer_create(sqfs_u32 flags)
 {
 	sqfs_xattr_writer_t *xwr;
@@ -104,9 +114,18 @@ sqfs_xattr_writer_t *sqfs_xattr_writer_create(sqfs_u32 flags)
 	if (xwr->kv_pairs == NULL)
 		goto fail_pairs;
 
+	if (rbtree_init(&xwr->kv_block_tree, sizeof(kv_block_desc_t),
+			sizeof(sqfs_u32), block_compare)) {
+		goto fail_tree;
+	}
+
+	xwr->kv_block_tree.key_context = xwr;
+
 	((sqfs_object_t *)xwr)->copy = xattr_writer_copy;
 	((sqfs_object_t *)xwr)->destroy = xattr_writer_destroy;
 	return xwr;
+fail_tree:
+	free(xwr->kv_pairs);
 fail_pairs:
 	str_table_cleanup(&xwr->values);
 fail_values:
diff --git a/lib/sqfs/xattr/xattr_writer.h b/lib/sqfs/xattr/xattr_writer.h
index 826b6f6..8ff9691 100644
--- a/lib/sqfs/xattr/xattr_writer.h
+++ b/lib/sqfs/xattr/xattr_writer.h
@@ -19,6 +19,7 @@
 #include "sqfs/io.h"
 
 #include "str_table.h"
+#include "rbtree.h"
 #include "util.h"
 
 #include <stdlib.h>
@@ -56,7 +57,9 @@ struct sqfs_xattr_writer_t {
 
 	size_t kv_start;
 
-	kv_block_desc_t *kv_blocks;
+	rbtree_t kv_block_tree;
+	kv_block_desc_t *kv_block_first;
+	kv_block_desc_t *kv_block_last;
 	size_t num_blocks;
 };
 
diff --git a/lib/sqfs/xattr/xattr_writer_flush.c b/lib/sqfs/xattr/xattr_writer_flush.c
index 53c3187..92f9be3 100644
--- a/lib/sqfs/xattr/xattr_writer_flush.c
+++ b/lib/sqfs/xattr/xattr_writer_flush.c
@@ -207,7 +207,7 @@ static int write_kv_pairs(const sqfs_xattr_writer_t *xwr,
 	for (i = 0; i < xwr->values.num_strings; ++i)
 		ool_locations[i] = 0xFFFFFFFFFFFFFFFFUL;
 
-	for (blk = xwr->kv_blocks; blk != NULL; blk = blk->next) {
+	for (blk = xwr->kv_block_first; blk != NULL; blk = blk->next) {
 		sqfs_meta_writer_get_position(mw, &block, &offset);
 		blk->start_ref = (block << 16) | (offset & 0xFFFF);
 
@@ -237,7 +237,7 @@ static int write_id_table(const sqfs_xattr_writer_t *xwr,
 
 	locations[i++] = 0;
 
-	for (blk = xwr->kv_blocks; blk != NULL; blk = blk->next) {
+	for (blk = xwr->kv_block_first; blk != NULL; blk = blk->next) {
 		memset(&id_ent, 0, sizeof(id_ent));
 		id_ent.xattr = htole64(blk->start_ref);
 		id_ent.count = htole32(blk->count);
diff --git a/lib/sqfs/xattr/xattr_writer_record.c b/lib/sqfs/xattr/xattr_writer_record.c
index aa87246..510ca17 100644
--- a/lib/sqfs/xattr/xattr_writer_record.c
+++ b/lib/sqfs/xattr/xattr_writer_record.c
@@ -114,59 +114,45 @@ int sqfs_xattr_writer_add(sqfs_xattr_writer_t *xwr, const char *key,
 
 int sqfs_xattr_writer_end(sqfs_xattr_writer_t *xwr, sqfs_u32 *out)
 {
-	kv_block_desc_t *blk, *blk_prev;
+	kv_block_desc_t blk;
+	rbtree_node_t *n;
 	sqfs_u32 index;
-	size_t count;
 	int ret;
 
-	count = xwr->num_pairs - xwr->kv_start;
-	if (count == 0) {
+	memset(&blk, 0, sizeof(blk));
+	blk.start = xwr->kv_start;
+	blk.count = xwr->num_pairs - xwr->kv_start;
+
+	if (blk.count == 0) {
 		*out = 0xFFFFFFFF;
 		return 0;
 	}
 
-	qsort(xwr->kv_pairs + xwr->kv_start, count,
+	qsort(xwr->kv_pairs + blk.start, blk.count,
 	      sizeof(xwr->kv_pairs[0]), compare_u64);
 
-	blk_prev = NULL;
-	blk = xwr->kv_blocks;
-	index = 0;
-
-	while (blk != NULL) {
-		if (blk->count == count) {
-			ret = memcmp(xwr->kv_pairs + blk->start,
-				     xwr->kv_pairs + xwr->kv_start,
-				     sizeof(xwr->kv_pairs[0]) * count);
-
-			if (ret == 0)
-				break;
-		}
+	n = rbtree_lookup(&xwr->kv_block_tree, &blk);
 
-		if (index == 0xFFFFFFFF)
-			return SQFS_ERROR_OVERFLOW;
-
-		++index;
-		blk_prev = blk;
-		blk = blk->next;
-	}
-
-	if (blk != NULL) {
+	if (n != NULL) {
+		index = *((sqfs_u32 *)rbtree_node_value(n));
 		xwr->num_pairs = xwr->kv_start;
 	} else {
-		blk = calloc(1, sizeof(*blk));
-		if (blk == NULL)
-			return SQFS_ERROR_ALLOC;
+		index = xwr->num_blocks;
+
+		ret = rbtree_insert(&xwr->kv_block_tree, &blk, &index);
+		if (ret != 0)
+			return ret;
 
-		blk->start = xwr->kv_start;
-		blk->count = count;
+		xwr->num_blocks += 1;
+		n = rbtree_lookup(&xwr->kv_block_tree, &blk);
 
-		if (blk_prev == NULL) {
-			xwr->kv_blocks = blk;
+		if (xwr->kv_block_last == NULL) {
+			xwr->kv_block_first = rbtree_node_key(n);
+			xwr->kv_block_last = xwr->kv_block_first;
 		} else {
-			blk_prev->next = blk;
+			xwr->kv_block_last->next = rbtree_node_key(n);
+			xwr->kv_block_last = xwr->kv_block_last->next;
 		}
-
-		xwr->num_blocks += 1;
 	}
 
 	*out = index;