thodg/libgit2

diff --git a/src/index.c b/src/index.c
index be043a9..8a5bf61 100644
--- a/src/index.c
+++ b/src/index.c
@@ -2114,16 +2114,13 @@ static int write_entries(git_index *index, git_filebuf *file)
 static int write_extension(git_filebuf *file, struct index_extension *header, git_buf *data)
 {
 	struct index_extension ondisk;
-	int error = 0;
 
 	memset(&ondisk, 0x0, sizeof(struct index_extension));
 	memcpy(&ondisk, header, 4);
 	ondisk.extension_size = htonl(header->extension_size);
 
-	if ((error = git_filebuf_write(file, &ondisk, sizeof(struct index_extension))) == 0)
-		error = git_filebuf_write(file, data->ptr, data->size);
-
-	return error;
+	git_filebuf_write(file, &ondisk, sizeof(struct index_extension));
+	return git_filebuf_write(file, data->ptr, data->size);
 }
 
 static int create_name_extension_data(git_buf *name_buf, git_index_name_entry *conflict_name)
@@ -2229,6 +2226,29 @@ done:
 	return error;
 }
 
+static int write_tree_extension(git_index *index, git_filebuf *file)
+{
+	struct index_extension extension;
+	git_buf buf = GIT_BUF_INIT;
+	int error;
+
+	if (index->tree == NULL)
+		return 0;
+
+	if ((error = git_tree_cache_write(&buf, index->tree)) < 0)
+		return error;
+
+	memset(&extension, 0x0, sizeof(struct index_extension));
+	memcpy(&extension.signature, INDEX_EXT_TREECACHE_SIG, 4);
+	extension.extension_size = (uint32_t)buf.size;
+
+	error = write_extension(file, &extension, &buf);
+
+	git_buf_free(&buf);
+
+	return error;
+}
+
 static int write_index(git_index *index, git_filebuf *file)
 {
 	git_oid hash_final;
@@ -2251,7 +2271,9 @@ static int write_index(git_index *index, git_filebuf *file)
 	if (write_entries(index, file) < 0)
 		return -1;
 
-	/* TODO: write tree cache extension */
+	/* write the tree cache extension */
+	if (index->tree != NULL && write_tree_extension(index, file) < 0)
+		return -1;
 
 	/* write the rename conflict extension */
 	if (index->names.length > 0 && write_name_extension(index, file) < 0)
diff --git a/src/tree-cache.c b/src/tree-cache.c
index 1eb046d..48dc46a 100644
--- a/src/tree-cache.c
+++ b/src/tree-cache.c
@@ -31,7 +31,7 @@ void git_tree_cache_invalidate_path(git_tree_cache *tree, const char *path)
 	if (tree == NULL)
 		return;
 
-	tree->entries = -1;
+	tree->entry_count = -1;
 
 	while (ptr != NULL) {
 		end = strchr(ptr, '/');
@@ -43,7 +43,7 @@ void git_tree_cache_invalidate_path(git_tree_cache *tree, const char *path)
 		if (tree == NULL) /* We don't have that tree */
 			return;
 
-		tree->entries = -1;
+		tree->entry_count = -1;
 		ptr = end + 1;
 	}
 }
@@ -93,7 +93,7 @@ static int read_tree_internal(git_tree_cache **out,
 	if (git__strtol32(&count, buffer, &buffer, 10) < 0)
 		goto corrupted;
 
-	tree->entries = count;
+	tree->entry_count = count;
 
 	if (*buffer != ' ' || ++buffer >= buffer_end)
 		goto corrupted;
@@ -108,7 +108,7 @@ static int read_tree_internal(git_tree_cache **out,
 		goto corrupted;
 
 	/* The SHA1 is only there if it's not invalidated */
-	if (tree->entries >= 0) {
+	if (tree->entry_count >= 0) {
 		/* 160-bit SHA-1 for this tree and it's children */
 		if (buffer + GIT_OID_RAWSZ > buffer_end)
 			goto corrupted;
@@ -244,3 +244,55 @@ int git_tree_cache_new(git_tree_cache **out, const char *name, git_pool *pool)
 	*out = tree;
 	return 0;
 }
+
+/**
+ * Recursively recalculate the total entry count, which we need to
+ * write out to the index
+ */
+static void recount_entries(git_tree_cache *tree)
+{
+	size_t i;
+	ssize_t entry_count;
+	git_tree_cache *child;
+
+	for (i = 0; i < tree->children_count; i++)
+		recount_entries(tree->children[i]);
+
+	if (tree->entry_count == -1)
+		return;
+
+	entry_count = 0;
+	for (i = 0; i < tree->children_count; i++) {
+		child = tree->children[i];
+
+		if (child->entry_count == -1)
+			continue;
+
+		entry_count += tree->children[i]->children_count;
+	}
+
+	tree->entry_count = entry_count;
+}
+
+static void write_tree(git_buf *out, git_tree_cache *tree)
+{
+	size_t i;
+
+	git_buf_printf(out, "%s%c%zd %"PRIuZ"\n", tree->name, 0, tree->entry_count, tree->children_count);
+
+	if (tree->entry_count == -1)
+		return;
+
+	git_buf_put(out, (const char *) &tree->oid, GIT_OID_RAWSZ);
+
+	for (i = 0; i < tree->children_count; i++)
+		write_tree(out, tree->children[i]);
+}
+
+int git_tree_cache_write(git_buf *out, git_tree_cache *tree)
+{
+	recount_entries(tree);
+	write_tree(out, tree);
+
+	return git_buf_oom(out) ? -1 : 0;
+}
diff --git a/src/tree-cache.h b/src/tree-cache.h
index b92a401..c44ca7c 100644
--- a/src/tree-cache.h
+++ b/src/tree-cache.h
@@ -10,18 +10,20 @@
 
 #include "common.h"
 #include "pool.h"
+#include "buffer.h"
 #include "git2/oid.h"
 
 typedef struct git_tree_cache {
 	struct git_tree_cache **children;
 	size_t children_count;
 
-	ssize_t entries;
+	ssize_t entry_count;
 	git_oid oid;
 	size_t namelen;
 	char name[GIT_FLEX_ARRAY];
 } git_tree_cache;
 
+int git_tree_cache_write(git_buf *out, git_tree_cache *tree);
 int git_tree_cache_read(git_tree_cache **tree, const char *buffer, size_t buffer_size, git_pool *pool);
 void git_tree_cache_invalidate_path(git_tree_cache *tree, const char *path);
 const git_tree_cache *git_tree_cache_get(const git_tree_cache *tree, const char *path);
diff --git a/src/tree.c b/src/tree.c
index 28190d6..adc6fe5 100644
--- a/src/tree.c
+++ b/src/tree.c
@@ -488,7 +488,7 @@ static int write_tree(
 	const git_tree_cache *cache;
 
 	cache = git_tree_cache_get(index->tree, dirname);
-	if (cache != NULL && cache->entries >= 0){
+	if (cache != NULL && cache->entry_count >= 0){
 		git_oid_cpy(oid, &cache->oid);
 		return (int)find_next_dir(dirname, index, start);
 	}
@@ -589,7 +589,7 @@ int git_tree__write_index(
 		return GIT_EUNMERGED;
 	}
 
-	if (index->tree != NULL && index->tree->entries >= 0) {
+	if (index->tree != NULL && index->tree->entry_count >= 0) {
 		git_oid_cpy(oid, &index->tree->oid);
 		return 0;
 	}
diff --git a/tests/index/cache.c b/tests/index/cache.c
index 79b431b..ba2d51f 100644
--- a/tests/index/cache.c
+++ b/tests/index/cache.c
@@ -16,6 +16,79 @@ void test_index_cache__cleanup(void)
 	g_repo = NULL;
 }
 
+void test_index_cache__write_extension_at_root(void)
+{
+	git_index *index;
+	git_tree *tree;
+	git_oid id;
+	const char *tree_id_str = "45dd856fdd4d89b884c340ba0e047752d9b085d6";
+	const char *index_file = "index-tree";
+
+	cl_git_pass(git_index_open(&index, index_file));
+	cl_assert(index->tree == NULL);
+	cl_git_pass(git_oid_fromstr(&id, tree_id_str));
+	cl_git_pass(git_tree_lookup(&tree, g_repo, &id));
+	cl_git_pass(git_index_read_tree(index, tree));
+	git_tree_free(tree);
+
+	cl_assert(index->tree);
+	cl_git_pass(git_index_write(index));
+	git_index_free(index);
+
+	cl_git_pass(git_index_open(&index, index_file));
+	cl_assert(index->tree);
+
+	cl_assert_equal_i(0, index->tree->entry_count);
+	cl_assert_equal_i(0, index->tree->children_count);
+
+	cl_assert(git_oid_equal(&id, &index->tree->oid));
+
+	cl_git_pass(p_unlink(index_file));
+	git_index_free(index);
+}
+
+void test_index_cache__write_extension_invalidated_root(void)
+{
+	git_index *index;
+	git_tree *tree;
+	git_oid id;
+	const char *tree_id_str = "45dd856fdd4d89b884c340ba0e047752d9b085d6";
+	const char *index_file = "index-tree-invalidated";
+	git_index_entry entry;
+
+	cl_git_pass(git_index_open(&index, index_file));
+	cl_assert(index->tree == NULL);
+	cl_git_pass(git_oid_fromstr(&id, tree_id_str));
+	cl_git_pass(git_tree_lookup(&tree, g_repo, &id));
+	cl_git_pass(git_index_read_tree(index, tree));
+	git_tree_free(tree);
+
+	cl_assert(index->tree);
+
+	memset(&entry, 0x0, sizeof(git_index_entry));
+	git_oid_cpy(&entry.id, &git_index_get_byindex(index, 0)->id);
+	entry.mode = GIT_FILEMODE_BLOB;
+	entry.path = "some-new-file.txt";
+
+	cl_git_pass(git_index_add(index, &entry));
+
+	cl_assert_equal_i(-1, index->tree->entry_count);
+
+	cl_git_pass(git_index_write(index));
+	git_index_free(index);
+
+	cl_git_pass(git_index_open(&index, index_file));
+	cl_assert(index->tree);
+
+	cl_assert_equal_i(-1, index->tree->entry_count);
+	cl_assert_equal_i(0, index->tree->children_count);
+
+	cl_assert(git_oid_cmp(&id, &index->tree->oid));
+
+	cl_git_pass(p_unlink(index_file));
+	git_index_free(index);
+}
+
 void test_index_cache__read_tree_no_children(void)
 {
 	git_index *index;
@@ -33,7 +106,7 @@ void test_index_cache__read_tree_no_children(void)
 	cl_assert(index->tree);
 	cl_assert(git_oid_equal(&id, &index->tree->oid));
 	cl_assert_equal_i(0, index->tree->children_count);
-	cl_assert_equal_i(0, index->tree->entries); /* 0 is a placeholder here */
+	cl_assert_equal_i(0, index->tree->entry_count); /* 0 is a placeholder here */
 
 	memset(&entry, 0x0, sizeof(git_index_entry));
 	entry.path = "new.txt";
@@ -41,7 +114,7 @@ void test_index_cache__read_tree_no_children(void)
 	git_oid_fromstr(&entry.id, "45b983be36b73c0788dc9cbcb76cbb80fc7bb057");
 
 	cl_git_pass(git_index_add(index, &entry));
-	cl_assert_equal_i(-1, index->tree->entries);
+	cl_assert_equal_i(-1, index->tree->entry_count);
 
 	git_index_free(index);
 }
@@ -92,19 +165,19 @@ void test_index_cache__read_tree_children(void)
 	git_oid_fromstr(&entry.id, "45b983be36b73c0788dc9cbcb76cbb80fc7bb058");
 	cl_git_pass(git_index_add(index, &entry));
 
-	cl_assert_equal_i(-1, index->tree->entries);
+	cl_assert_equal_i(-1, index->tree->entry_count);
 
 	cache = git_tree_cache_get(index->tree, "subdir");
 	cl_assert(cache);
-	cl_assert_equal_i(-1, cache->entries);
+	cl_assert_equal_i(-1, cache->entry_count);
 
 	cache = git_tree_cache_get(index->tree, "subdir/even-deeper");
 	cl_assert(cache);
-	cl_assert_equal_i(0, cache->entries);
+	cl_assert_equal_i(0, cache->entry_count);
 
 	cache = git_tree_cache_get(index->tree, "subdir2");
 	cl_assert(cache);
-	cl_assert_equal_i(0, cache->entries);
+	cl_assert_equal_i(0, cache->entry_count);
 
 	git_index_free(index);
 }