12 files changed, 128 insertions, 54 deletions

diff --git a/src/blob.c b/src/blob.c
index c5a7143f0..f271cc7f6 100644
--- a/src/blob.c
+++ b/src/blob.c
@@ -129,6 +129,10 @@ int git_blob_writefile(git_oid *written_id, git_repository *repo, const char *pa
 		return error;
 
 	git_oid_cpy(written_id, git_object_id((git_object *)blob));
+
+	/* FIXME: maybe we don't want to free this already?
+	 * the user may want to access it again */
+	git_object_close((git_object *)blob);
 	return GIT_SUCCESS;
 }
 
diff --git a/src/commit.c b/src/commit.c
index 3eebb927a..3edc57337 100644
--- a/src/commit.c
+++ b/src/commit.c
@@ -55,6 +55,8 @@ void git_commit__free(git_commit *commit)
 	git_signature_free(commit->author);
 	git_signature_free(commit->committer);
 
+	git_object_close((git_object *)commit->tree);
+
 	free(commit->message);
 	free(commit->message_short);
 	free(commit);
@@ -121,6 +123,7 @@ int commit_parse_buffer(git_commit *commit, void *data, size_t len, unsigned int
 	if ((error = git__parse_oid(&oid, &buffer, buffer_end, "tree ")) < GIT_SUCCESS)
 		return error;
 
+	git_object_close((git_object *)commit->tree);
 	if ((error = git_object_lookup((git_object **)&commit->tree, commit->object.repo, &oid, GIT_OBJ_TREE)) < GIT_SUCCESS)
 		return error;
 
@@ -232,7 +235,21 @@ int git_commit__parse_full(git_commit *commit)
 	if (!commit->object.in_memory && !commit->full_parse)\
 		git_commit__parse_full(commit); 
 
-GIT_COMMIT_GETTER(git_tree *, tree)
+const git_tree *git_commit_tree(git_commit *commit)
+{
+	assert(commit);
+
+	if (!commit->object.in_memory && commit->tree == NULL)
+		git_commit__parse_full(commit);
+
+	if (commit->tree) {
+		GIT_OBJECT_INCREF(commit->tree);
+		return commit->tree;
+	}
+
+	return NULL;
+}
+
 GIT_COMMIT_GETTER(git_signature *, author)
 GIT_COMMIT_GETTER(git_signature *, committer)
 GIT_COMMIT_GETTER(char *, message)
@@ -267,6 +284,9 @@ void git_commit_set_tree(git_commit *commit, git_tree *tree)
 	assert(commit && tree);
 	commit->object.modified = 1;
 	CHECK_FULL_PARSE();
+
+	git_object_close((git_object *)commit->tree);
+	GIT_OBJECT_INCREF(tree);
 	commit->tree = tree;
 }
 
diff --git a/src/git2/object.h b/src/git2/object.h
index 80477d44a..af0f014e3 100644
--- a/src/git2/object.h
+++ b/src/git2/object.h
@@ -128,18 +128,27 @@ GIT_EXTERN(git_otype) git_object_type(const git_object *obj);
 GIT_EXTERN(git_repository *) git_object_owner(const git_object *obj);
 
 /**
- * Free a reference to one of the objects in the repository.
+ * Close an open object
  *
- * Repository objects are managed automatically by the library,
- * but this method can be used to force freeing one of the
- * objects.
+ * This method instructs the library to close an existing
+ * object; note that git_objects are owned by the repository
+ * and are reference counted, so the object may or may not be
+ * freed after this library call, depending on whether any other 
+ * objects still depend on it.
  *
- * Careful: freeing objects in the middle of a repository
- * traversal will most likely cause errors.
+ * IMPORTANT:
+ * It is *not* necessary to call this method when you stop using
+ * an object, since all object memory is automatically reclaimed
+ * by the repository when it is freed.
  *
- * @param object the object to free
+ * Forgetting to call `git_object_close` does not cause memory
+ * leaks, but it's is recommended to close as soon as possible
+ * the biggest objects (e.g. blobs) to prevent wasting memory
+ * space.
+ *
+ * @param object the object to close
  */
-GIT_EXTERN(void) git_object_free(git_object *object);
+GIT_EXTERN(void) git_object_close(git_object *object);
 
 /**
  * Convert an object type to it's string representation.
diff --git a/src/object.c b/src/object.c
index c9809c52a..e23271442 100644
--- a/src/object.c
+++ b/src/object.c
@@ -166,8 +166,12 @@ static int write_back(git_object *object)
 	if ((error = git_odb_write(&new_id, object->repo->db, &object->source.raw)) < GIT_SUCCESS)
 		return error;
 
-	if (!object->in_memory)
+	if (object->in_memory) {
+		int idx = git_vector_search(&object->repo->memory_objects, object);
+		git_vector_remove(&object->repo->memory_objects, idx);
+	} else {
 		git_hashtable_remove(object->repo->objects, &object->id);
+	}
 
 	git_oid_cpy(&object->id, &new_id);
 	git_hashtable_insert(object->repo->objects, &object->id, object);
@@ -257,6 +261,7 @@ int git_object_new(git_object **object_out, git_repository *repo, git_otype type
 
 	object->source.raw.type = type;
 
+	object->refcount++;
 	*object_out = object;
 	return GIT_SUCCESS;
 }
@@ -317,13 +322,14 @@ int git_object_lookup(git_object **object_out, git_repository *repo, const git_o
 	}
 
 	if (error < GIT_SUCCESS) {
-		git_object_free(object);
+		git_object__free(object);
 		return error;
 	}
 
 	git_object__source_close(object);
 	git_hashtable_insert(repo->objects, &object->id, object);
 
+	object->refcount++;
 	*object_out = object;
 	return GIT_SUCCESS;
 }
@@ -371,13 +377,18 @@ int git_object_write(git_object *object)
 	return write_back(object);
 }
 
-void git_object_free(git_object *object)
+void git_object__free(git_object *object)
 {
-	if (object == NULL)
-		return;
+	assert(object);
 
 	git_object__source_close(object);
-	git_hashtable_remove(object->repo->objects, &object->id);
+
+	if (object->in_memory) {
+		int idx = git_vector_search(&object->repo->memory_objects, object);
+		git_vector_remove(&object->repo->memory_objects, idx);
+	} else {
+		git_hashtable_remove(object->repo->objects, &object->id);
+	}
 
 	switch (object->source.raw.type) {
 	case GIT_OBJ_COMMIT:
@@ -402,6 +413,15 @@ void git_object_free(git_object *object)
 	}
 }
 
+void git_object_close(git_object *object)
+{
+	if (object == NULL)
+		return;
+
+	if (--object->refcount <= 0)
+		git_object__free(object);
+}
+
 const git_oid *git_object_id(const git_object *obj)
 {
 	assert(obj);
diff --git a/src/repository.c b/src/repository.c
index 77ce0c0f4..9f27a38e8 100644
--- a/src/repository.c
+++ b/src/repository.c
@@ -209,6 +209,13 @@ static git_repository *repository_alloc()
 		return NULL;
 	}
 
+	if (git_vector_init(&repo->memory_objects, 16, NULL) < GIT_SUCCESS) {
+		git_hashtable_free(repo->objects);
+		git_repository__refcache_free(&repo->references);
+		free(repo);
+		return NULL;
+	}
+
 	return repo;
 }
 
@@ -328,7 +335,8 @@ int git_repository_open(git_repository **repo_out, const char *path)
 void git_repository_free(git_repository *repo)
 {
 	git_object *object;
-	const git_oid *oid;
+	const void *_unused;
+	unsigned int i;
 
 	if (repo == NULL)
 		return;
@@ -338,11 +346,24 @@ void git_repository_free(git_repository *repo)
 	free(repo->path_repository);
 	free(repo->path_odb);
 
-	GIT_HASHTABLE_FOREACH(repo->objects, oid, object, {
-		git_object_free(object);
-	});
+	/* Increment the refcount of all the objects in the repository
+	 * to prevent freeing dependencies */
+	GIT_HASHTABLE_FOREACH(repo->objects, _unused, object,
+		GIT_OBJECT_INCREF(object);
+	);
+
+	/* force free all the objects */
+	GIT_HASHTABLE_FOREACH(repo->objects, _unused, object,
+		git_object__free(object);
+	);
+
+	for (i = 0; i < repo->memory_objects.length; ++i) {
+		object = git_vector_get(&repo->memory_objects, i);
+		git_object__free(object);
+	}
 
 	git_hashtable_free(repo->objects);
+	git_vector_free(&repo->memory_objects);
 
 	git_repository__refcache_free(&repo->references);
 
diff --git a/src/repository.h b/src/repository.h
index c8e4186b3..2e8d187b0 100644
--- a/src/repository.h
+++ b/src/repository.h
@@ -15,6 +15,8 @@
 #define GIT_OBJECTS_DIR "objects/"
 #define GIT_INDEX_FILE "index"
 
+#define GIT_OBJECT_INCREF(ob) ++(((git_object *)(ob))->refcount)
+
 typedef struct {
 	git_rawobj raw;
 	void *write_ptr;
@@ -26,13 +28,16 @@ struct git_object {
 	git_oid id;
 	git_repository *repo;
 	git_odb_source source;
-	int in_memory:1, modified:1;
+	unsigned short refcount;
+	short in_memory:1, modified:1;
 };
 
 struct git_repository {
 	git_odb *db;
 	git_index *index;
+
 	git_hashtable *objects;
+	git_vector memory_objects;
 
 	git_refcache references;
 
@@ -47,6 +52,10 @@ struct git_repository {
 int git_object__source_open(git_object *object);
 void git_object__source_close(git_object *object);
 
+/* fully free the object; internal method, do not
+ * export */
+void git_object__free(git_object *object);
+
 int git__source_printf(git_odb_source *source, const char *format, ...);
 int git__source_write(git_odb_source *source, const void *bytes, size_t len);
 
diff --git a/src/revwalk.c b/src/revwalk.c
index c073be13f..872cdbc43 100644
--- a/src/revwalk.c
+++ b/src/revwalk.c
@@ -114,6 +114,7 @@ static git_revwalk_commit *commit_to_walkcommit(git_revwalk *walk, git_commit *c
 	memset(commit, 0x0, sizeof(git_revwalk_commit));
 
 	commit->commit_object = commit_object;
+	GIT_OBJECT_INCREF(commit_object);
 
 	git_hashtable_insert(walk->commits, commit_object, commit);
 
@@ -229,6 +230,7 @@ int git_revwalk_next(git_commit **commit, git_revwalk *walk)
 	while ((next = walk->next(&walk->iterator)) != NULL) {
 		if (!next->uninteresting) {
 			*commit = next->commit_object;
+			GIT_OBJECT_INCREF(*commit);
 			return GIT_SUCCESS;
 		}
 	}
@@ -246,6 +248,7 @@ void git_revwalk_reset(git_revwalk *walk)
 	assert(walk);
 
 	GIT_HASHTABLE_FOREACH(walk->commits, _unused, commit, {
+		git_object_close((git_object *)commit->commit_object);
 		git_revwalk_list_clear(&commit->parents);
 		free(commit);
 	});
diff --git a/src/tag.c b/src/tag.c
index 7d57ee075..01cc0dc8f 100644
--- a/src/tag.c
+++ b/src/tag.c
@@ -35,6 +35,7 @@
 void git_tag__free(git_tag *tag)
 {
 	git_signature_free(tag->tagger);
+	git_object_close(tag->target);
 	free(tag->message);
 	free(tag->tag_name);
 	free(tag);
@@ -48,6 +49,7 @@ const git_oid *git_tag_id(git_tag *c)
 const git_object *git_tag_target(git_tag *t)
 {
 	assert(t);
+	GIT_OBJECT_INCREF(t->target);
 	return t->target;
 }
 
@@ -55,6 +57,9 @@ void git_tag_set_target(git_tag *tag, git_object *target)
 {
 	assert(tag && target);
 
+	git_object_close(tag->target);
+	GIT_OBJECT_INCREF(target);
+
 	tag->object.modified = 1;
 	tag->target = target;
 	tag->type = git_object_type(target);
@@ -66,14 +71,6 @@ git_otype git_tag_type(git_tag *t)
 	return t->type;
 }
 
-void git_tag_set_type(git_tag *tag, git_otype type)
-{
-	assert(tag);
-
-	tag->object.modified = 1;
-	tag->type = type;
-}
-
 const char *git_tag_name(git_tag *t)
 {
 	assert(t);
@@ -164,6 +161,7 @@ static int parse_tag_buffer(git_tag *tag, char *buffer, const char *buffer_end)
 	if (tag->type == GIT_OBJ_BAD)
 		return GIT_EOBJCORRUPTED;
 
+	git_object_close(tag->target);
 	error = git_object_lookup(&tag->target, tag->object.repo, &target_oid, tag->type);
 	if (error < 0)
 		return error;
diff --git a/src/tree.c b/src/tree.c
index 5ea062286..30938f258 100644
--- a/src/tree.c
+++ b/src/tree.c
@@ -67,9 +67,7 @@ void git_tree_clear_entries(git_tree *tree)
 	}
 
 	git_vector_clear(&tree->entries);
-
 	tree->object.modified = 1;
-	tree->sorted = 1;
 }
 
 
@@ -88,7 +86,6 @@ git_tree *git_tree__new(void)
 		return NULL;
 	}
 
-	tree->sorted = 1;
 	return tree;
 }
 
@@ -155,10 +152,7 @@ int git_tree_entry_2object(git_object **object_out, git_tree_entry *entry)
 
 static void sort_entries(git_tree *tree)
 {
-	if (tree->sorted == 0) {
-        git_vector_sort(&tree->entries);
-		tree->sorted = 1;
-	}
+	git_vector_sort(&tree->entries);
 }
 
 git_tree_entry *git_tree_entry_byname(git_tree *tree, const char *filename)
@@ -167,8 +161,7 @@ git_tree_entry *git_tree_entry_byname(git_tree *tree, const char *filename)
 
 	assert(tree && filename);
 
-	if (!tree->sorted)
-		sort_entries(tree);
+	sort_entries(tree);
 
 	idx = git_vector_bsearch2(&tree->entries, entry_search_cmp, filename);
 	if (idx == GIT_ENOTFOUND)
@@ -181,8 +174,7 @@ git_tree_entry *git_tree_entry_byindex(git_tree *tree, int idx)
 {
 	assert(tree);
 
-	if (!tree->sorted)
-		sort_entries(tree);
+	sort_entries(tree);
 
 	return git_vector_get(&tree->entries, (unsigned int)idx);
 }
@@ -216,7 +208,6 @@ int git_tree_add_entry(git_tree_entry **entry_out, git_tree *tree, const git_oid
 		*entry_out = entry;
 
 	tree->object.modified = 1;
-	tree->sorted = 0;
 	return GIT_SUCCESS;
 }
 
@@ -226,8 +217,7 @@ int git_tree_remove_entry_byindex(git_tree *tree, int idx)
 
 	assert(tree);
 
-	if (!tree->sorted)
-		sort_entries(tree);
+	sort_entries(tree);
 
 	remove_ptr = git_vector_get(&tree->entries, (unsigned int)idx);
 	if (remove_ptr == NULL)
@@ -247,8 +237,7 @@ int git_tree_remove_entry_byname(git_tree *tree, const char *filename)
 
 	assert(tree && filename);
 
-	if (!tree->sorted)
-		sort_entries(tree);
+	sort_entries(tree);
 
 	idx = git_vector_bsearch2(&tree->entries, entry_search_cmp, filename);
 	if (idx == GIT_ENOTFOUND)
@@ -267,8 +256,7 @@ int git_tree__writeback(git_tree *tree, git_odb_source *src)
 	if (tree->entries.length == 0)
 		return GIT_EMISSINGOBJDATA;
 
-	if (!tree->sorted)
-		sort_entries(tree);
+	sort_entries(tree);
 
 	for (i = 0; i < tree->entries.length; ++i) {
 		git_tree_entry *entry;
diff --git a/src/tree.h b/src/tree.h
index 796c5b950..78500c471 100644
--- a/src/tree.h
+++ b/src/tree.h
@@ -16,7 +16,6 @@ struct git_tree_entry {
 struct git_tree {
 	git_object object;
 	git_vector entries;
-	unsigned sorted:1;
 };
 
 void git_tree__free(git_tree *tree);
diff --git a/src/vector.c b/src/vector.c
index 126aec1bb..631364031 100644
--- a/src/vector.c
+++ b/src/vector.c
@@ -104,6 +104,10 @@ int git_vector_bsearch2(git_vector *v, git_vector_cmp key_lookup, const void *ke
 
 	assert(v && key && key_lookup);
 
+	/* need comparison function to sort the vector */
+	if (v->_cmp == NULL)
+		return GIT_ENOTFOUND;
+
 	git_vector_sort(v);
 
 	find = bsearch(key, v->contents, v->length, sizeof(void *), key_lookup);
@@ -127,12 +131,14 @@ int git_vector_search2(git_vector *v, git_vector_cmp key_lookup, const void *key
 	return GIT_ENOTFOUND;
 }
 
-int git_vector_search(git_vector *v, const void *key)
+static int strict_comparison(const void *a, const void *b)
 {
-	if (v->_cmp == NULL)
-		return GIT_ENOTFOUND;
+	return a - b;
+}
 
-	return git_vector_search2(v, v->_cmp, key);
+int git_vector_search(git_vector *v, const void *entry)
+{
+	return git_vector_search2(v, v->_cmp ? v->_cmp : strict_comparison, entry);
 }
 
 int git_vector_bsearch(git_vector *v, const void *key)
diff --git a/tests/t09-tree.c b/tests/t09-tree.c
index cb127556b..6bc2a84bd 100644
--- a/tests/t09-tree.c
+++ b/tests/t09-tree.c
@@ -101,8 +101,6 @@ BEGIN_TEST(write0, "add a new entry to a tree and write it back to disk")
 	must_pass(git_object_write((git_object *)tree));
 	must_pass(remove_loose_object(REPOSITORY_FOLDER, (git_object *)tree));
 
-	git_object_free((git_object *)tree);
-
 	git_repository_free(repo);
 END_TEST
 
@@ -152,7 +150,6 @@ BEGIN_TEST(write1, "add several entries in-memory and validate that they exist;
 */
 
 	must_pass(remove_loose_object(REPOSITORY_FOLDER, (git_object *)tree));
-	git_object_free((git_object *)tree);
 	git_repository_free(repo);
 END_TEST