thodg/got

diff --git a/lib/commit_graph.c b/lib/commit_graph.c
index 958b66e..4023b5d 100644
--- a/lib/commit_graph.c
+++ b/lib/commit_graph.c
@@ -230,7 +230,7 @@ advance_open_branches(struct got_commit_graph *graph,
 	const struct got_error *err;
 	struct got_object_qid *qid;
 
-	err = got_object_idset_remove(graph->open_branches, commit_id);
+	err = got_object_idset_remove(NULL, graph->open_branches, commit_id);
 	if (err && err->code != GOT_ERR_NO_OBJ)
 		return err;
 
diff --git a/lib/got_lib_object_idset.h b/lib/got_lib_object_idset.h
index f5c7fb0..7d98296 100644
--- a/lib/got_lib_object_idset.h
+++ b/lib/got_lib_object_idset.h
@@ -22,8 +22,8 @@ void got_object_idset_free(struct got_object_idset *);
 const struct got_error *got_object_idset_add(void **,
     struct got_object_idset *, struct got_object_id *, void *);
 void *got_object_idset_get(struct got_object_idset *, struct got_object_id *);
-const struct got_error *got_object_idset_remove(struct got_object_idset *,
-    struct got_object_id *);
+const struct got_error *got_object_idset_remove(void **,
+    struct got_object_idset *, struct got_object_id *);
 const struct got_error *got_object_idset_remove_random(void **,
     struct got_object_idset *);
 int got_object_idset_contains(struct got_object_idset *,
diff --git a/lib/object_idset.c b/lib/object_idset.c
index 38f091d..cce68c1 100644
--- a/lib/object_idset.c
+++ b/lib/object_idset.c
@@ -161,7 +161,7 @@ got_object_idset_get(struct got_object_idset *set, struct got_object_id *id)
 }
 
 const struct got_error *
-got_object_idset_remove(struct got_object_idset *set,
+got_object_idset_remove(void **data, struct got_object_idset *set,
     struct got_object_id *id)
 {
 	struct got_object_idset_element *entry, *tmp;
@@ -173,6 +173,8 @@ got_object_idset_remove(struct got_object_idset *set,
 	TAILQ_FOREACH_SAFE(entry, &set->entries[i], entry, tmp) {
 		if (got_object_id_cmp(&entry->id, id) == 0) {
 			TAILQ_REMOVE(&set->entries[i], entry, entry);
+			if (data)
+				*data = entry->data;
 			free(entry);
 			set->nelem--;
 			return NULL;