thodg/got

diff --git a/lib/got_lib_object_idset.h b/lib/got_lib_object_idset.h
index a20b754..193a271 100644
--- a/lib/got_lib_object_idset.h
+++ b/lib/got_lib_object_idset.h
@@ -29,3 +29,4 @@ int got_object_idset_contains(struct got_object_idset *,
     struct got_object_id *);
 void got_object_idset_for_each(struct got_object_idset *,
     void (*cb)(struct got_object_id *, void *));
+unsigned int got_object_idset_num_elements(struct got_object_idset *);
diff --git a/lib/object_idset.c b/lib/object_idset.c
index 4883efe..0f91478 100644
--- a/lib/object_idset.c
+++ b/lib/object_idset.c
@@ -21,6 +21,7 @@
 #include <sha1.h>
 #include <stdio.h>
 #include <zlib.h>
+#include <limits.h>
 
 #include "got_object.h"
 #include "got_error.h"
@@ -47,6 +48,8 @@ struct got_object_idset {
 	 * which of these lists an object ID is stored in.
 	 */
 	TAILQ_HEAD(, got_object_idset_element) entries[0xff + 1];
+	unsigned int nelem;
+#define GOT_OBJECT_IDSET_MAX_ELEM UINT_MAX
 };
 
 struct got_object_idset *
@@ -89,6 +92,9 @@ got_object_idset_add(struct got_object_idset *set, struct got_object_id *id,
 	struct got_object_idset_element *new, *entry;
 	uint8_t i = id->sha1[0];
 
+	if (set->nelem >= GOT_OBJECT_IDSET_MAX_ELEM)
+		return got_error(GOT_ERR_NO_SPACE);
+
 	new = calloc(1, sizeof(*new));
 	if (new == NULL)
 		return got_error_from_errno();
@@ -98,6 +104,7 @@ got_object_idset_add(struct got_object_idset *set, struct got_object_id *id,
 
 	if (TAILQ_EMPTY(&set->entries[i])) {
 		TAILQ_INSERT_HEAD(&set->entries[i], new, entry);
+		set->nelem++;
 		return NULL;
 	}
 
@@ -114,15 +121,18 @@ got_object_idset_add(struct got_object_idset *set, struct got_object_id *id,
 			return got_error(GOT_ERR_OBJ_EXISTS);
 		} else if (cmp < 0) {
 			TAILQ_INSERT_BEFORE(entry, new, entry);
+			set->nelem++;
 			return NULL;
 		}
 
 		next = TAILQ_NEXT(entry, entry);
 		if (next == NULL) {
 			TAILQ_INSERT_AFTER(&set->entries[i], entry, new, entry);
+			set->nelem++;
 			return NULL;
 		} else if (got_object_id_cmp(&new->id, &next->id) > 0) {
 			TAILQ_INSERT_BEFORE(next, new, entry);
+			set->nelem++;
 			return NULL;
 		}
 	}
@@ -152,9 +162,13 @@ got_object_idset_remove(struct got_object_idset *set,
 	struct got_object_idset_element *entry, *tmp;
 	uint8_t i = id->sha1[0];
 
+	if (set->nelem == 0)
+		return got_error(GOT_ERR_NO_OBJ);
+
 	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);
+			set->nelem--;
 			return NULL;
 		}
 	}
@@ -188,3 +202,9 @@ void got_object_idset_for_each(struct got_object_idset *set,
 			cb(&entry->id, entry->data);
 	}
 }
+
+unsigned int
+got_object_idset_num_elements(struct got_object_idset *set)
+{
+	return set->nelem;
+}
diff --git a/regress/idset/idset_test.c b/regress/idset/idset_test.c
index 80a5654..3490429 100644
--- a/regress/idset/idset_test.c
+++ b/regress/idset/idset_test.c
@@ -77,6 +77,10 @@ idset_add_remove_iter(void)
 		err = got_error_from_errno();
 		goto done;
 	}
+	if (got_object_idset_num_elements(set) != 0) {
+		err = got_error(GOT_ERR_BAD_OBJ_DATA);
+		goto done;
+	}
 
 	if (!got_parse_sha1_digest(id1.sha1, id_str1)) {
 		err = got_error(GOT_ERR_BAD_OBJ_ID_STR);
@@ -94,6 +98,10 @@ idset_add_remove_iter(void)
 	err = got_object_idset_add(set, &id1, (void *)data1);
 	if (err)
 		goto done;
+	if (got_object_idset_num_elements(set) != 1) {
+		err = got_error(GOT_ERR_BAD_OBJ_DATA);
+		goto done;
+	}
 
 	if (!got_object_idset_contains(set, &id1)) {
 		err = got_error(GOT_ERR_BAD_OBJ_DATA);
@@ -104,8 +112,10 @@ idset_add_remove_iter(void)
 	if (err)
 		goto done;
 	err = got_object_idset_add(set, &id2, NULL);
-	if (err == NULL)
-		return 0;
+	if (err == NULL) {
+		err = got_error(GOT_ERR_BAD_OBJ_DATA);
+		goto done;
+	}
 	if (err->code != GOT_ERR_OBJ_EXISTS)
 		goto done;
 	err = NULL;
@@ -118,6 +128,10 @@ idset_add_remove_iter(void)
 		err = got_error(GOT_ERR_BAD_OBJ_DATA);
 		goto done;
 	}
+	if (got_object_idset_num_elements(set) != 2) {
+		err = got_error(GOT_ERR_BAD_OBJ_DATA);
+		goto done;
+	}
 
 	err = got_object_idset_add(set, &id3, (void *)data3);
 	if (err)
@@ -135,10 +149,18 @@ idset_add_remove_iter(void)
 		err = got_error(GOT_ERR_BAD_OBJ_DATA);
 		goto done;
 	}
+	if (got_object_idset_num_elements(set) != 3) {
+		err = got_error(GOT_ERR_BAD_OBJ_DATA);
+		goto done;
+	}
 
 	err = got_object_idset_remove(set, &id2);
 	if (err)
 		goto done;
+	if (got_object_idset_num_elements(set) != 2) {
+		err = got_error(GOT_ERR_BAD_OBJ_DATA);
+		goto done;
+	}
 	if (got_object_idset_contains(set, &id2)) {
 		err = got_error(GOT_ERR_BAD_OBJ_DATA);
 		goto done;
@@ -149,8 +171,8 @@ idset_add_remove_iter(void)
 	}
 
 	got_object_idset_for_each(set, idset_cb);
-	got_object_idset_free(set);
 done:
+	got_object_idset_free(set);
 	return (err == NULL);
 }