1 files changed, 95 insertions, 76 deletions

diff --git a/reftable/merged.c b/reftable/merged.c
index 6adce44f4b..733de07454 100644
--- a/reftable/merged.c
+++ b/reftable/merged.c
@@ -1,10 +1,10 @@
 /*
-Copyright 2020 Google LLC
-
-Use of this source code is governed by a BSD-style
-license that can be found in the LICENSE file or at
-https://developers.google.com/open-source/licenses/bsd
-*/
+ * Copyright 2020 Google LLC
+ *
+ * Use of this source code is governed by a BSD-style
+ * license that can be found in the LICENSE file or at
+ * https://developers.google.com/open-source/licenses/bsd
+ */
 
 #include "merged.h"
 
@@ -12,10 +12,10 @@ https://developers.google.com/open-source/licenses/bsd
 #include "iter.h"
 #include "pq.h"
 #include "record.h"
-#include "generic.h"
 #include "reftable-merged.h"
 #include "reftable-error.h"
 #include "system.h"
+#include "table.h"
 
 struct merged_subiter {
 	struct reftable_iterator iter;
@@ -25,33 +25,17 @@ struct merged_subiter {
 struct merged_iter {
 	struct merged_subiter *subiters;
 	struct merged_iter_pqueue pq;
-	size_t stack_len;
+	size_t subiters_len;
 	int suppress_deletions;
 	ssize_t advance_index;
 };
 
-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;
 
 	merged_iter_pqueue_release(&mi->pq);
-	for (size_t i = 0; i < mi->stack_len; i++) {
+	for (size_t i = 0; i < mi->subiters_len; i++) {
 		reftable_iterator_destroy(&mi->subiters[i].iter);
 		reftable_record_release(&mi->subiters[i].rec);
 	}
@@ -70,7 +54,10 @@ static int merged_iter_advance_subiter(struct merged_iter *mi, size_t idx)
 	if (err)
 		return err;
 
-	merged_iter_pqueue_add(&mi->pq, &e);
+	err = merged_iter_pqueue_add(&mi->pq, &e);
+	if (err)
+		return err;
+
 	return 0;
 }
 
@@ -79,8 +66,13 @@ static int merged_iter_seek(struct merged_iter *mi, struct reftable_record *want
 	int err;
 
 	mi->advance_index = -1;
+	while (!merged_iter_pqueue_is_empty(mi->pq)) {
+		err = merged_iter_pqueue_remove(&mi->pq, NULL);
+		if (err < 0)
+			return err;
+	}
 
-	for (size_t i = 0; i < mi->stack_len; i++) {
+	for (size_t i = 0; i < mi->subiters_len; i++) {
 		err = iterator_seek(&mi->subiters[i].iter, want);
 		if (err < 0)
 			return err;
@@ -131,7 +123,9 @@ static int merged_iter_next_entry(struct merged_iter *mi,
 	if (empty)
 		return 1;
 
-	entry = merged_iter_pqueue_remove(&mi->pq);
+	err = merged_iter_pqueue_remove(&mi->pq, &entry);
+	if (err < 0)
+		return err;
 
 	/*
 	  One can also use reftable as datacenter-local storage, where the ref
@@ -145,18 +139,23 @@ static int merged_iter_next_entry(struct merged_iter *mi,
 		struct pq_entry top = merged_iter_pqueue_top(mi->pq);
 		int cmp;
 
-		cmp = reftable_record_cmp(top.rec, entry.rec);
+		err = reftable_record_cmp(top.rec, entry.rec, &cmp);
+		if (err < 0)
+			return err;
 		if (cmp > 0)
 			break;
 
-		merged_iter_pqueue_remove(&mi->pq);
+		err = merged_iter_pqueue_remove(&mi->pq, NULL);
+		if (err < 0)
+			return err;
+
 		err = merged_iter_advance_subiter(mi, top.index);
 		if (err < 0)
 			return err;
 	}
 
 	mi->advance_index = entry.index;
-	SWAP(*rec, *entry.rec);
+	REFTABLE_SWAP(*rec, *entry.rec);
 	return 0;
 }
 
@@ -192,19 +191,19 @@ static void iterator_from_merged_iter(struct reftable_iterator *it,
 	it->ops = &merged_iter_vtable;
 }
 
-int reftable_new_merged_table(struct reftable_merged_table **dest,
-			      struct reftable_table *stack, size_t n,
-			      uint32_t hash_id)
+int reftable_merged_table_new(struct reftable_merged_table **dest,
+			      struct reftable_table **tables, size_t n,
+			      enum reftable_hash hash_id)
 {
 	struct reftable_merged_table *m = NULL;
 	uint64_t last_max = 0;
 	uint64_t first_min = 0;
 
 	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]);
+		uint64_t min = reftable_table_min_update_index(tables[i]);
+		uint64_t max = reftable_table_max_update_index(tables[i]);
 
-		if (reftable_table_hash_id(&stack[i]) != hash_id) {
+		if (reftable_table_hash_id(tables[i]) != hash_id) {
 			return REFTABLE_FORMAT_ERROR;
 		}
 		if (i == 0 || min < first_min) {
@@ -216,8 +215,11 @@ int reftable_new_merged_table(struct reftable_merged_table **dest,
 	}
 
 	REFTABLE_CALLOC_ARRAY(m, 1);
-	m->stack = stack;
-	m->stack_len = n;
+	if (!m)
+		return REFTABLE_OUT_OF_MEMORY_ERROR;
+
+	m->tables = tables;
+	m->tables_len = n;
 	m->min = first_min;
 	m->max = last_max;
 	m->hash_id = hash_id;
@@ -229,7 +231,6 @@ void reftable_merged_table_free(struct reftable_merged_table *mt)
 {
 	if (!mt)
 		return;
-	FREE_AND_NULL(mt->stack);
 	reftable_free(mt);
 }
 
@@ -245,53 +246,71 @@ reftable_merged_table_min_update_index(struct reftable_merged_table *mt)
 	return mt->min;
 }
 
-void merged_table_init_iter(struct reftable_merged_table *mt,
-			    struct reftable_iterator *it,
-			    uint8_t typ)
+int merged_table_init_iter(struct reftable_merged_table *mt,
+			   struct reftable_iterator *it,
+			   uint8_t typ)
 {
-	struct merged_iter *mi = reftable_malloc(sizeof(*mi));
-	merged_iter_init(mi, mt, typ);
-	iterator_from_merged_iter(it, mi);
-}
+	struct merged_subiter *subiters = NULL;
+	struct merged_iter *mi = NULL;
+	int ret;
+
+	if (mt->tables_len) {
+		REFTABLE_CALLOC_ARRAY(subiters, mt->tables_len);
+		if (!subiters) {
+			ret = REFTABLE_OUT_OF_MEMORY_ERROR;
+			goto out;
+		}
+	}
 
-uint32_t reftable_merged_table_hash_id(struct reftable_merged_table *mt)
-{
-	return mt->hash_id;
-}
+	for (size_t i = 0; i < mt->tables_len; i++) {
+		ret = reftable_record_init(&subiters[i].rec, typ);
+		if (ret < 0)
+			goto out;
 
-static void reftable_merged_table_init_iter_void(void *tab,
-						 struct reftable_iterator *it,
-						 uint8_t typ)
-{
-	merged_table_init_iter(tab, it, typ);
-}
+		ret = table_init_iter(mt->tables[i], &subiters[i].iter, typ);
+		if (ret < 0)
+			goto out;
+	}
 
-static uint32_t reftable_merged_table_hash_id_void(void *tab)
-{
-	return reftable_merged_table_hash_id(tab);
+	REFTABLE_CALLOC_ARRAY(mi, 1);
+	if (!mi) {
+		ret = REFTABLE_OUT_OF_MEMORY_ERROR;
+		goto out;
+	}
+	mi->advance_index = -1;
+	mi->suppress_deletions = mt->suppress_deletions;
+	mi->subiters = subiters;
+	mi->subiters_len = mt->tables_len;
+
+	iterator_from_merged_iter(it, mi);
+	ret = 0;
+
+out:
+	if (ret < 0) {
+		for (size_t i = 0; subiters && i < mt->tables_len; i++) {
+			reftable_iterator_destroy(&subiters[i].iter);
+			reftable_record_release(&subiters[i].rec);
+		}
+		reftable_free(subiters);
+		reftable_free(mi);
+	}
+
+	return ret;
 }
 
-static uint64_t reftable_merged_table_min_update_index_void(void *tab)
+int reftable_merged_table_init_ref_iterator(struct reftable_merged_table *mt,
+					    struct reftable_iterator *it)
 {
-	return reftable_merged_table_min_update_index(tab);
+	return merged_table_init_iter(mt, it, REFTABLE_BLOCK_TYPE_REF);
 }
 
-static uint64_t reftable_merged_table_max_update_index_void(void *tab)
+int reftable_merged_table_init_log_iterator(struct reftable_merged_table *mt,
+					    struct reftable_iterator *it)
 {
-	return reftable_merged_table_max_update_index(tab);
+	return merged_table_init_iter(mt, it, REFTABLE_BLOCK_TYPE_LOG);
 }
 
-static struct reftable_table_vtable merged_table_vtable = {
-	.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,
-};
-
-void reftable_table_from_merged_table(struct reftable_table *tab,
-				      struct reftable_merged_table *merged)
+enum reftable_hash reftable_merged_table_hash_id(struct reftable_merged_table *mt)
 {
-	assert(!tab->ops);
-	tab->ops = &merged_table_vtable;
-	tab->table_arg = merged;
+	return mt->hash_id;
 }