1 files changed, 117 insertions, 181 deletions
diff --git a/reftable/merged.c b/reftable/merged.c
index 574394092d..6adce44f4b 100644
--- a/reftable/merged.c
+++ b/reftable/merged.c
@@ -17,85 +17,121 @@ https://developers.google.com/open-source/licenses/bsd
 #include "reftable-error.h"
 #include "system.h"
 
-static int merged_iter_init(struct merged_iter *mi)
-{
-	int i = 0;
-	for (i = 0; i < mi->stack_len; i++) {
-		struct reftable_record rec = reftable_new_record(mi->typ);
-		int err = iterator_next(&mi->stack[i], &rec);
-		if (err < 0) {
-			return err;
-		}
+struct merged_subiter {
+	struct reftable_iterator iter;
+	struct reftable_record rec;
+};
 
-		if (err > 0) {
-			reftable_iterator_destroy(&mi->stack[i]);
-			reftable_record_release(&rec);
-		} else {
-			struct pq_entry e = {
-				.rec = rec,
-				.index = i,
-			};
-			merged_iter_pqueue_add(&mi->pq, &e);
-		}
-	}
+struct merged_iter {
+	struct merged_subiter *subiters;
+	struct merged_iter_pqueue pq;
+	size_t stack_len;
+	int suppress_deletions;
+	ssize_t advance_index;
+};
 
-	return 0;
+static void merged_iter_init(struct merged_iter *mi,
+			     struct reftable_merged_table *mt,
+			     uint8_t typ)
+{
+	memset(mi, 0, sizeof(*mi));
+	mi->advance_index = -1;
+	mi->suppress_deletions = mt->suppress_deletions;
+
+	REFTABLE_CALLOC_ARRAY(mi->subiters, mt->stack_len);
+	for (size_t i = 0; i < mt->stack_len; i++) {
+		reftable_record_init(&mi->subiters[i].rec, typ);
+		table_init_iter(&mt->stack[i], &mi->subiters[i].iter, typ);
+	}
+	mi->stack_len = mt->stack_len;
 }
 
 static void merged_iter_close(void *p)
 {
 	struct merged_iter *mi = p;
-	int i = 0;
+
 	merged_iter_pqueue_release(&mi->pq);
-	for (i = 0; i < mi->stack_len; i++) {
-		reftable_iterator_destroy(&mi->stack[i]);
+	for (size_t i = 0; i < mi->stack_len; i++) {
+		reftable_iterator_destroy(&mi->subiters[i].iter);
+		reftable_record_release(&mi->subiters[i].rec);
 	}
-	reftable_free(mi->stack);
-	strbuf_release(&mi->key);
-	strbuf_release(&mi->entry_key);
+	reftable_free(mi->subiters);
 }
 
-static int merged_iter_advance_nonnull_subiter(struct merged_iter *mi,
-					       size_t idx)
+static int merged_iter_advance_subiter(struct merged_iter *mi, size_t idx)
 {
 	struct pq_entry e = {
-		.rec = reftable_new_record(mi->typ),
 		.index = idx,
+		.rec = &mi->subiters[idx].rec,
 	};
-	int err = iterator_next(&mi->stack[idx], &e.rec);
-	if (err < 0)
-		return err;
+	int err;
 
-	if (err > 0) {
-		reftable_iterator_destroy(&mi->stack[idx]);
-		reftable_record_release(&e.rec);
-		return 0;
-	}
+	err = iterator_next(&mi->subiters[idx].iter, &mi->subiters[idx].rec);
+	if (err)
+		return err;
 
 	merged_iter_pqueue_add(&mi->pq, &e);
 	return 0;
 }
 
-static int merged_iter_advance_subiter(struct merged_iter *mi, size_t idx)
+static int merged_iter_seek(struct merged_iter *mi, struct reftable_record *want)
 {
-	if (iterator_is_null(&mi->stack[idx]))
-		return 0;
-	return merged_iter_advance_nonnull_subiter(mi, idx);
+	int err;
+
+	mi->advance_index = -1;
+
+	for (size_t i = 0; i < mi->stack_len; i++) {
+		err = iterator_seek(&mi->subiters[i].iter, want);
+		if (err < 0)
+			return err;
+		if (err > 0)
+			continue;
+
+		err = merged_iter_advance_subiter(mi, i);
+		if (err < 0)
+			return err;
+	}
+
+	return 0;
 }
 
 static int merged_iter_next_entry(struct merged_iter *mi,
 				  struct reftable_record *rec)
 {
 	struct pq_entry entry = { 0 };
-	int err = 0;
+	int err = 0, empty;
+
+	empty = merged_iter_pqueue_is_empty(mi->pq);
+
+	if (mi->advance_index >= 0) {
+		/*
+		 * When there are no pqueue entries then we only have a single
+		 * subiter left. There is no need to use the pqueue in that
+		 * case anymore as we know that the subiter will return entries
+		 * in the correct order already.
+		 *
+		 * While this may sound like a very specific edge case, it may
+		 * happen more frequently than you think. Most repositories
+		 * will end up having a single large base table that contains
+		 * most of the refs. It's thus likely that we exhaust all
+		 * subiters but the one from that base ref.
+		 */
+		if (empty)
+			return iterator_next(&mi->subiters[mi->advance_index].iter,
+					     rec);
+
+		err = merged_iter_advance_subiter(mi, mi->advance_index);
+		if (err < 0)
+			return err;
+		if (!err)
+			empty = 0;
+		mi->advance_index = -1;
+	}
 
-	if (merged_iter_pqueue_is_empty(mi->pq))
+	if (empty)
 		return 1;
 
 	entry = merged_iter_pqueue_remove(&mi->pq);
-	err = merged_iter_advance_subiter(mi, entry.index);
-	if (err < 0)
-		return err;
 
 	/*
 	  One can also use reftable as datacenter-local storage, where the ref
@@ -105,56 +141,45 @@ static int merged_iter_next_entry(struct merged_iter *mi,
 	  such a deployment, the loop below must be changed to collect all
 	  entries for the same key, and return new the newest one.
 	*/
-	reftable_record_key(&entry.rec, &mi->entry_key);
 	while (!merged_iter_pqueue_is_empty(mi->pq)) {
 		struct pq_entry top = merged_iter_pqueue_top(mi->pq);
-		int cmp = 0;
-
-		reftable_record_key(&top.rec, &mi->key);
+		int cmp;
 
-		cmp = strbuf_cmp(&mi->key, &mi->entry_key);
+		cmp = reftable_record_cmp(top.rec, entry.rec);
 		if (cmp > 0)
 			break;
 
 		merged_iter_pqueue_remove(&mi->pq);
 		err = merged_iter_advance_subiter(mi, top.index);
 		if (err < 0)
-			goto done;
-		reftable_record_release(&top.rec);
+			return err;
 	}
 
-	reftable_record_copy_from(rec, &entry.rec, hash_size(mi->hash_id));
-
-done:
-	reftable_record_release(&entry.rec);
-	strbuf_release(&mi->entry_key);
-	strbuf_release(&mi->key);
-	return err;
+	mi->advance_index = entry.index;
+	SWAP(*rec, *entry.rec);
+	return 0;
 }
 
-static int merged_iter_next(struct merged_iter *mi, struct reftable_record *rec)
+static int merged_iter_seek_void(void *it, struct reftable_record *want)
 {
-	while (1) {
-		int err = merged_iter_next_entry(mi, rec);
-		if (err == 0 && mi->suppress_deletions &&
-		    reftable_record_is_deletion(rec)) {
-			continue;
-		}
-
-		return err;
-	}
+	return merged_iter_seek(it, want);
 }
 
 static int merged_iter_next_void(void *p, struct reftable_record *rec)
 {
 	struct merged_iter *mi = p;
-	if (merged_iter_pqueue_is_empty(mi->pq))
-		return 1;
-
-	return merged_iter_next(mi, rec);
+	while (1) {
+		int err = merged_iter_next_entry(mi, rec);
+		if (err)
+			return err;
+		if (mi->suppress_deletions && reftable_record_is_deletion(rec))
+			continue;
+		return 0;
+	}
 }
 
 static struct reftable_iterator_vtable merged_iter_vtable = {
+	.seek = merged_iter_seek_void,
 	.next = &merged_iter_next_void,
 	.close = &merged_iter_close,
 };
@@ -168,14 +193,14 @@ static void iterator_from_merged_iter(struct reftable_iterator *it,
 }
 
 int reftable_new_merged_table(struct reftable_merged_table **dest,
-			      struct reftable_table *stack, int n,
+			      struct reftable_table *stack, size_t n,
 			      uint32_t hash_id)
 {
 	struct reftable_merged_table *m = NULL;
 	uint64_t last_max = 0;
 	uint64_t first_min = 0;
-	int i = 0;
-	for (i = 0; i < n; i++) {
+
+	for (size_t i = 0; i < n; i++) {
 		uint64_t min = reftable_table_min_update_index(&stack[i]);
 		uint64_t max = reftable_table_max_update_index(&stack[i]);
 
@@ -190,7 +215,7 @@ int reftable_new_merged_table(struct reftable_merged_table **dest,
 		}
 	}
 
-	m = reftable_calloc(sizeof(struct reftable_merged_table));
+	REFTABLE_CALLOC_ARRAY(m, 1);
 	m->stack = stack;
 	m->stack_len = n;
 	m->min = first_min;
@@ -200,19 +225,11 @@ int reftable_new_merged_table(struct reftable_merged_table **dest,
 	return 0;
 }
 
-/* clears the list of subtable, without affecting the readers themselves. */
-void merged_table_release(struct reftable_merged_table *mt)
-{
-	FREE_AND_NULL(mt->stack);
-	mt->stack_len = 0;
-}
-
 void reftable_merged_table_free(struct reftable_merged_table *mt)
 {
-	if (!mt) {
+	if (!mt)
 		return;
-	}
-	merged_table_release(mt);
+	FREE_AND_NULL(mt->stack);
 	reftable_free(mt);
 }
 
@@ -228,94 +245,13 @@ reftable_merged_table_min_update_index(struct reftable_merged_table *mt)
 	return mt->min;
 }
 
-static int reftable_table_seek_record(struct reftable_table *tab,
-				      struct reftable_iterator *it,
-				      struct reftable_record *rec)
-{
-	return tab->ops->seek_record(tab->table_arg, it, rec);
-}
-
-static int merged_table_seek_record(struct reftable_merged_table *mt,
-				    struct reftable_iterator *it,
-				    struct reftable_record *rec)
-{
-	struct reftable_iterator *iters = reftable_calloc(
-		sizeof(struct reftable_iterator) * mt->stack_len);
-	struct merged_iter merged = {
-		.stack = iters,
-		.typ = reftable_record_type(rec),
-		.hash_id = mt->hash_id,
-		.suppress_deletions = mt->suppress_deletions,
-		.key = STRBUF_INIT,
-		.entry_key = STRBUF_INIT,
-	};
-	int n = 0;
-	int err = 0;
-	int i = 0;
-	for (i = 0; i < mt->stack_len && err == 0; i++) {
-		int e = reftable_table_seek_record(&mt->stack[i], &iters[n],
-						   rec);
-		if (e < 0) {
-			err = e;
-		}
-		if (e == 0) {
-			n++;
-		}
-	}
-	if (err < 0) {
-		int i = 0;
-		for (i = 0; i < n; i++) {
-			reftable_iterator_destroy(&iters[i]);
-		}
-		reftable_free(iters);
-		return err;
-	}
-
-	merged.stack_len = n;
-	err = merged_iter_init(&merged);
-	if (err < 0) {
-		merged_iter_close(&merged);
-		return err;
-	} else {
-		struct merged_iter *p =
-			reftable_malloc(sizeof(struct merged_iter));
-		*p = merged;
-		iterator_from_merged_iter(it, p);
-	}
-	return 0;
-}
-
-int reftable_merged_table_seek_ref(struct reftable_merged_table *mt,
-				   struct reftable_iterator *it,
-				   const char *name)
-{
-	struct reftable_record rec = {
-		.type = BLOCK_TYPE_REF,
-		.u.ref = {
-			.refname = (char *)name,
-		},
-	};
-	return merged_table_seek_record(mt, it, &rec);
-}
-
-int reftable_merged_table_seek_log_at(struct reftable_merged_table *mt,
-				      struct reftable_iterator *it,
-				      const char *name, uint64_t update_index)
-{
-	struct reftable_record rec = { .type = BLOCK_TYPE_LOG,
-				       .u.log = {
-					       .refname = (char *)name,
-					       .update_index = update_index,
-				       } };
-	return merged_table_seek_record(mt, it, &rec);
-}
-
-int reftable_merged_table_seek_log(struct reftable_merged_table *mt,
-				   struct reftable_iterator *it,
-				   const char *name)
+void merged_table_init_iter(struct reftable_merged_table *mt,
+			    struct reftable_iterator *it,
+			    uint8_t typ)
 {
-	uint64_t max = ~((uint64_t)0);
-	return reftable_merged_table_seek_log_at(mt, it, name, max);
+	struct merged_iter *mi = reftable_malloc(sizeof(*mi));
+	merged_iter_init(mi, mt, typ);
+	iterator_from_merged_iter(it, mi);
 }
 
 uint32_t reftable_merged_table_hash_id(struct reftable_merged_table *mt)
@@ -323,11 +259,11 @@ uint32_t reftable_merged_table_hash_id(struct reftable_merged_table *mt)
 	return mt->hash_id;
 }
 
-static int reftable_merged_table_seek_void(void *tab,
-					   struct reftable_iterator *it,
-					   struct reftable_record *rec)
+static void reftable_merged_table_init_iter_void(void *tab,
+						 struct reftable_iterator *it,
+						 uint8_t typ)
 {
-	return merged_table_seek_record(tab, it, rec);
+	merged_table_init_iter(tab, it, typ);
 }
 
 static uint32_t reftable_merged_table_hash_id_void(void *tab)
@@ -346,7 +282,7 @@ static uint64_t reftable_merged_table_max_update_index_void(void *tab)
 }
 
 static struct reftable_table_vtable merged_table_vtable = {
-	.seek_record = reftable_merged_table_seek_void,
+	.init_iter = reftable_merged_table_init_iter_void,
 	.hash_id = reftable_merged_table_hash_id_void,
 	.min_update_index = reftable_merged_table_min_update_index_void,
 	.max_update_index = reftable_merged_table_max_update_index_void,