Store xattr writer block description in a red-black tree

By storing the blocks in a tree, the de-duplication can lookup
existing blocks in logartihmic instead of linear time.

The linked list is still maintained, because we need to iterate
over the blocks in creation order during serialization.

Signed-off-by: David Oberhollenzer <david.oberhollenzer@sigma-star.at>

5 files changed, 72 insertions, 63 deletions

diff --git a/lib/sqfs/Makemodule.am b/lib/sqfs/Makemodule.am
index fedb7a8..7dafd8e 100644
--- a/lib/sqfs/Makemodule.am
+++ b/lib/sqfs/Makemodule.am
@@ -41,6 +41,7 @@ libsquashfs_la_LIBADD += $(ZSTD_LIBS) $(PTHREAD_LIBS)
 libsquashfs_la_SOURCES += lib/util/str_table.c lib/util/alloc.c
 libsquashfs_la_SOURCES += lib/util/xxhash.c
 libsquashfs_la_SOURCES += lib/util/hash_table.c include/hash_table.h
+libsquashfs_la_SOURCES += lib/util/rbtree.c include/rbtree.h
 
 if WINDOWS
 libsquashfs_la_SOURCES += lib/sqfs/win32/io_file.c
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;