Build & write custom trees in memory

5 files changed, 391 insertions, 17 deletions

diff --git a/include/git2/tree.h b/include/git2/tree.h
index 60150b30..0caf60a4 100644
--- a/include/git2/tree.h
+++ b/include/git2/tree.h
@@ -93,7 +93,7 @@ GIT_EXTERN(size_t) git_tree_entrycount(git_tree *tree);
  * @param filename the filename of the desired entry
  * @return the tree entry; NULL if not found
  */
-GIT_EXTERN(git_tree_entry *) git_tree_entry_byname(git_tree *tree, const char *filename);
+GIT_EXTERN(const git_tree_entry *) git_tree_entry_byname(git_tree *tree, const char *filename);
 
 /**
  * Lookup a tree entry by its position in the tree
@@ -102,7 +102,7 @@ GIT_EXTERN(git_tree_entry *) git_tree_entry_byname(git_tree *tree, const char *f
  * @param idx the position in the entry list
  * @return the tree entry; NULL if not found
  */
-GIT_EXTERN(git_tree_entry *) git_tree_entry_byindex(git_tree *tree, int idx);
+GIT_EXTERN(const git_tree_entry *) git_tree_entry_byindex(git_tree *tree, int idx);
 
 /**
  * Get the UNIX file attributes of a tree entry
@@ -110,7 +110,7 @@ GIT_EXTERN(git_tree_entry *) git_tree_entry_byindex(git_tree *tree, int idx);
  * @param entry a tree entry
  * @return attributes as an integer
  */
-GIT_EXTERN(unsigned int) git_tree_entry_attributes(git_tree_entry *entry);
+GIT_EXTERN(unsigned int) git_tree_entry_attributes(const git_tree_entry *entry);
 
 /**
  * Get the filename of a tree entry
@@ -118,7 +118,7 @@ GIT_EXTERN(unsigned int) git_tree_entry_attributes(git_tree_entry *entry);
  * @param entry a tree entry
  * @return the name of the file
  */
-GIT_EXTERN(const char *) git_tree_entry_name(git_tree_entry *entry);
+GIT_EXTERN(const char *) git_tree_entry_name(const git_tree_entry *entry);
 
 /**
  * Get the id of the object pointed by the entry
@@ -126,7 +126,7 @@ GIT_EXTERN(const char *) git_tree_entry_name(git_tree_entry *entry);
  * @param entry a tree entry
  * @return the oid of the object
  */
-GIT_EXTERN(const git_oid *) git_tree_entry_id(git_tree_entry *entry);
+GIT_EXTERN(const git_oid *) git_tree_entry_id(const git_tree_entry *entry);
 
 /**
  * Convert a tree entry to the git_object it points too.
@@ -136,7 +136,7 @@ GIT_EXTERN(const git_oid *) git_tree_entry_id(git_tree_entry *entry);
  * @param entry a tree entry
  * @return 0 on success; error code otherwise
  */
-GIT_EXTERN(int) git_tree_entry_2object(git_object **object_out, git_repository *repo, git_tree_entry *entry);
+GIT_EXTERN(int) git_tree_entry_2object(git_object **object_out, git_repository *repo, const git_tree_entry *entry);
 
 /**
  * Write a tree to the ODB from the index file
@@ -156,6 +156,111 @@ GIT_EXTERN(int) git_tree_entry_2object(git_object **object_out, git_repository *
  */
 GIT_EXTERN(int) git_tree_create_fromindex(git_oid *oid, git_index *index);
 
+/**
+ * Create a new tree builder.
+ *
+ * The tree builder can be used to create or modify
+ * trees in memory and write them as tree objects to the
+ * database.
+ *
+ * If the `source` parameter is not NULL, the tree builder
+ * will be initialized with the entries of the given tree.
+ * 
+ * If the `source` parameter is NULL, the tree builder will
+ * have no entries and will have to be filled manually.
+ *
+ * @param builder_p Pointer where to store the tree builder
+ * @param source Source tree to initialize the builder (optional)
+ * @return 0 on sucess; error code otherwise
+ */
+GIT_EXTERN(int) git_treebuilder_create(git_treebuilder **builder_p, const git_tree *source);
+
+/**
+ * Clear all the entires in the builder
+ *
+ * @param bld Builder to clear
+ */
+GIT_EXTERN(void) git_treebuilder_clear(git_treebuilder *bld);
+
+/**
+ * Free a tree builder
+ *
+ * This will clear all the entries and free to builder.
+ * Failing to free the builder after you're done using it
+ * will result in a memory leak
+ *
+ * @param bld Builder to free
+ */
+GIT_EXTERN(void) git_treebuilder_free(git_treebuilder *bld);
+
+/**
+ * Get an entry from the builder from its filename
+ *
+ * The returned entry is owned by the builder and should
+ * not be freed manually.
+ *
+ * @param bld Tree builder
+ * @param filename Name of the entry
+ * @return pointer to the entry; NULL if not found
+ */
+GIT_EXTERN(const git_tree_entry *) git_treebuilder_get(git_treebuilder *bld, const char *filename);
+
+/**
+ * Add or update an entry to the builder
+ *
+ * Insert a new entry for `filename` in the builder with the
+ * given attributes.
+ *
+ * if an entry named `filename` already exists, its attributes
+ * will be updated with the given ones.
+ *
+ * The optional pointer `entry_out` can be used to retrieve a
+ * pointer to the newly created/updated entry.
+ *
+ * @param entry_out Pointer to store the entry (optional)
+ * @param bld Tree builder
+ * @param filename Filename of the entry
+ * @param id SHA1 oid of the entry
+ * @param attributes Folder attributes of the entry
+ * @return 0 on success; error code otherwise
+ */
+GIT_EXTERN(int) git_treebuilder_insert(git_tree_entry **entry_out, git_treebuilder *bld, const char *filename, const git_oid *id, unsigned int attributes);
+
+/**
+ * Remove an entry from the builder by its filename
+ *
+ * @param bld Tree builder
+ * @param filename Filename of the entry to remove
+ */
+GIT_EXTERN(int) git_treebuilder_remove(git_treebuilder *bld, const char *filename);
+
+/**
+ * Filter the entries in the tree
+ *
+ * The `filter` callback will be called for each entry
+ * in the tree with a pointer to the entry and the
+ * provided `payload`: if the callback returns 1, the
+ * entry will be filtered (removed from the builder).
+ *
+ * @param bld Tree builder
+ * @param filter Callback to filter entries
+ */
+GIT_EXTERN(void) git_treebuilder_filter(git_treebuilder *bld, int (*filter)(const git_tree_entry *, void *), void *payload);
+
+/**
+ * Write the contents of the tree builder as a tree object
+ *
+ * The tree builder will be written to the given `repo`, and
+ * it's identifying SHA1 hash will be stored in the `oid`
+ * pointer.
+ *
+ * @param oid Pointer where to store the written OID
+ * @param repo Repository where to store the object
+ * @param bld Tree builder to write
+ * @return 0 on success; error code otherwise
+ */
+GIT_EXTERN(int) git_treebuilder_write(git_oid *oid, git_repository *repo, git_treebuilder *bld);
+
 /** @} */
 GIT_END_DECL
 #endif
diff --git a/include/git2/types.h b/include/git2/types.h
index 88f6b7d5..6123abc8 100644
--- a/include/git2/types.h
+++ b/include/git2/types.h
@@ -124,6 +124,9 @@ typedef struct git_tree_entry git_tree_entry;
 /** Representation of a tree object. */
 typedef struct git_tree git_tree;
 
+/** Constructor for in-memory trees */
+typedef struct git_treebuilder git_treebuilder;
+
 /** Memory representation of an index file. */
 typedef struct git_index git_index;
 
diff --git a/src/tree.c b/src/tree.c
index 588b5e9a..693ed1ca 100644
--- a/src/tree.c
+++ b/src/tree.c
@@ -33,6 +33,10 @@
 #define MAX_FILEMODE 0777777
 #define MAX_FILEMODE_BYTES 6
 
+static int valid_attributes(const int attributes) {
+	return attributes >= 0 && attributes <= MAX_FILEMODE; 
+}
+
 int entry_search_cmp(const void *key, const void *array_member)
 {
 	const char *filename = (const char *)key;
@@ -79,30 +83,30 @@ const git_oid *git_tree_id(git_tree *c)
 	return git_object_id((git_object *)c);
 }
 
-unsigned int git_tree_entry_attributes(git_tree_entry *entry)
+unsigned int git_tree_entry_attributes(const git_tree_entry *entry)
 {
 	return entry->attr;
 }
 
-const char *git_tree_entry_name(git_tree_entry *entry)
+const char *git_tree_entry_name(const git_tree_entry *entry)
 {
 	assert(entry);
 	return entry->filename;
 }
 
-const git_oid *git_tree_entry_id(git_tree_entry *entry)
+const git_oid *git_tree_entry_id(const git_tree_entry *entry)
 {
 	assert(entry);
 	return &entry->oid;
 }
 
-int git_tree_entry_2object(git_object **object_out, git_repository *repo, git_tree_entry *entry)
+int git_tree_entry_2object(git_object **object_out, git_repository *repo, const git_tree_entry *entry)
 {
 	assert(entry && object_out);
 	return git_object_lookup(object_out, repo, &entry->oid, GIT_OBJ_ANY);
 }
 
-git_tree_entry *git_tree_entry_byname(git_tree *tree, const char *filename)
+const git_tree_entry *git_tree_entry_byname(git_tree *tree, const char *filename)
 {
 	int idx;
 
@@ -115,7 +119,7 @@ git_tree_entry *git_tree_entry_byname(git_tree *tree, const char *filename)
 	return git_vector_get(&tree->entries, idx);
 }
 
-git_tree_entry *git_tree_entry_byindex(git_tree *tree, int idx)
+const git_tree_entry *git_tree_entry_byindex(git_tree *tree, int idx)
 {
 	assert(tree);
 	return git_vector_get(&tree->entries, (unsigned int)idx);
@@ -159,6 +163,7 @@ static int tree_parse_buffer(git_tree *tree, const char *buffer, const char *buf
 		}
 
 		entry->filename = git__strdup(buffer);
+		entry->filename_len = strlen(buffer);
 
 		while (buffer < buffer_end && *buffer != 0)
 			buffer++;
@@ -178,7 +183,7 @@ int git_tree__parse(git_tree *tree, git_odb_object *obj)
 	return tree_parse_buffer(tree, (char *)obj->raw.data, (char *)obj->raw.data + obj->raw.len);
 }
 
-static int write_entry(char *buffer, int mode, const char *path, size_t path_len, const git_oid *oid)
+static int write_index_entry(char *buffer, int mode, const char *path, size_t path_len, const git_oid *oid)
 {
 	int written;
 	written = sprintf(buffer, "%o %.*s%c", mode, (int)path_len, path, 0);
@@ -254,7 +259,7 @@ static int write_index(git_oid *oid, git_index *index, const char *base, int bas
 				return GIT_ENOMEM;
 		}
 
-		offset += write_entry(buffer + offset, write_mode, filename, entrylen, write_oid);
+		offset += write_index_entry(buffer + offset, write_mode, filename, entrylen, write_oid);
 	}
 	
 	error = git_odb_write(oid, index->repository->db, buffer, offset, GIT_OBJ_TREE);
@@ -273,3 +278,209 @@ int git_tree_create_fromindex(git_oid *oid, git_index *index)
 	error = write_index(oid, index, "", 0, 0, git_index_entrycount(index));
 	return (error < GIT_SUCCESS) ? error : GIT_SUCCESS;
 }
+
+static void sort_entries(git_treebuilder *bld)
+{
+	git_vector_sort(&bld->entries);
+}
+
+int git_treebuilder_create(git_treebuilder **builder_p, const git_tree *source)
+{
+	git_treebuilder *bld;
+	size_t i, source_entries = DEFAULT_TREE_SIZE;
+
+	assert(builder_p);
+
+	bld = git__calloc(1, sizeof(git_treebuilder));
+	if (bld == NULL)
+		return GIT_ENOMEM;
+
+	if (source != NULL)
+		source_entries = source->entries.length;
+
+	if (git_vector_init(&bld->entries, source_entries, entry_sort_cmp) < GIT_SUCCESS) {
+		free(bld);
+		return GIT_ENOMEM;
+	}
+
+	if (source != NULL) {
+		bld->entry_count = source_entries;
+		for (i = 0; i < source->entries.length; ++i) {
+			git_tree_entry *entry_src = source->entries.contents[i];
+			git_tree_entry *entry = git__calloc(1, sizeof(git_tree_entry));
+
+			if (entry == NULL) {
+				git_treebuilder_free(bld);
+				return GIT_ENOMEM;
+			}
+
+			entry->filename = git__strdup(entry_src->filename);
+
+			if (entry->filename == NULL) {
+				free(entry);
+				git_treebuilder_free(bld);
+				return GIT_ENOMEM;
+			}
+
+			entry->filename_len = entry_src->filename_len;
+			git_oid_cpy(&entry->oid, &entry_src->oid);
+			entry->attr = entry_src->attr;
+
+			git_vector_insert(&bld->entries, entry);
+		}
+	}
+
+	*builder_p = bld;
+	return GIT_SUCCESS;
+}
+
+int git_treebuilder_insert(git_tree_entry **entry_out, git_treebuilder *bld, const char *filename, const git_oid *id, unsigned int attributes)
+{
+	git_tree_entry *entry;
+	int pos;
+
+	assert(bld && id && filename);
+
+	if (!valid_attributes(attributes))
+		return GIT_ERROR;
+
+	if ((pos = git_vector_bsearch2(&bld->entries, entry_search_cmp, filename)) != GIT_ENOTFOUND) {
+		entry = git_vector_get(&bld->entries, pos);
+		if (entry->removed) {
+			entry->removed = 0;
+			bld->entry_count++;
+		}
+	} else {
+		if ((entry = git__malloc(sizeof(git_tree_entry))) == NULL)
+			return GIT_ENOMEM;
+
+		memset(entry, 0x0, sizeof(git_tree_entry));
+		entry->filename = git__strdup(filename);
+		entry->filename_len = strlen(entry->filename);
+
+		bld->entry_count++;
+	}
+
+	git_oid_cpy(&entry->oid, id);
+	entry->attr = attributes;
+
+	if (pos != GIT_ENOTFOUND) {
+		if (git_vector_insert(&bld->entries, entry) < 0)
+			return GIT_ENOMEM;
+	}
+
+	if (entry_out != NULL)
+		*entry_out = entry;
+
+	return GIT_SUCCESS;
+}
+
+const git_tree_entry *git_treebuilder_get(git_treebuilder *bld, const char *filename)
+{
+	int idx;
+	git_tree_entry *entry;
+
+	assert(bld && filename);
+
+	sort_entries(bld);
+	idx = git_vector_bsearch2(&bld->entries, entry_search_cmp, filename);
+	if (idx == GIT_ENOTFOUND)
+		return NULL;
+
+	entry = git_vector_get(&bld->entries, idx);
+	if (entry->removed)
+		return NULL;
+
+	return entry;
+}
+
+int git_treebuilder_remove(git_treebuilder *bld, const char *filename)
+{
+	git_tree_entry *remove_ptr = (git_tree_entry *)git_treebuilder_get(bld, filename);
+
+	if (remove_ptr == NULL || remove_ptr->removed)
+		return GIT_ENOTFOUND;
+
+	remove_ptr->removed = 1;
+	bld->entry_count--;
+	return GIT_SUCCESS;
+}
+
+int git_treebuilder_write(git_oid *oid, git_repository *repo, git_treebuilder *bld)
+{
+	size_t i, size = 0;
+	char filemode[MAX_FILEMODE_BYTES + 1 + 1];
+	git_odb_stream *stream;
+	int error;
+
+	assert(bld);
+
+	sort_entries(bld);
+
+	for (i = 0; i < bld->entries.length; ++i) {
+		git_tree_entry *entry = bld->entries.contents[i];
+
+		if (entry->removed)
+			continue;
+
+		size += (entry->attr > 0x7FF) ? 7 : 6;
+		size += entry->filename_len + 1;
+		size += GIT_OID_RAWSZ;
+	}
+
+	if ((error = git_odb_open_wstream(&stream, git_repository_database(repo), size, GIT_OBJ_TREE)) < GIT_SUCCESS)
+		return error;
+
+	for (i = 0; i < bld->entries.length; ++i) {
+		git_tree_entry *entry = bld->entries.contents[i];
+
+		if (entry->removed)
+			continue;
+
+		snprintf(filemode, sizeof(filemode), "%o ", entry->attr);
+		stream->write(stream, filemode, strlen(filemode));
+		stream->write(stream, entry->filename, entry->filename_len + 1);
+		stream->write(stream, (char *)entry->oid.id, GIT_OID_RAWSZ);
+	} 
+
+	error = stream->finalize_write(oid, stream);
+	stream->free(stream);
+
+	return error;
+}
+
+void git_treebuilder_filter(git_treebuilder *bld, int (*filter)(const git_tree_entry *, void *), void *payload)
+{
+	size_t i;
+
+	assert(bld && filter);
+
+	for (i = 0; i < bld->entries.length; ++i) {
+		git_tree_entry *entry = bld->entries.contents[i];
+		if (!entry->removed && filter(entry, payload))
+			entry->removed = 1;
+	}
+}
+
+void git_treebuilder_clear(git_treebuilder *bld)
+{
+	size_t i;
+	assert(bld);
+
+	for (i = 0; i < bld->entries.length; ++i) {
+		git_tree_entry *e = bld->entries.contents[i];
+		free(e->filename);
+		free(e);
+	}
+
+	git_vector_clear(&bld->entries);
+}
+
+void git_treebuilder_free(git_treebuilder *bld)
+{
+	git_treebuilder_clear(bld);
+	git_vector_free(&bld->entries);
+	free(bld);
+}
+
+
diff --git a/src/tree.h b/src/tree.h
index b4e910a9..bff3f8ed 100644
--- a/src/tree.h
+++ b/src/tree.h
@@ -10,6 +10,8 @@ struct git_tree_entry {
 	unsigned int attr;
 	char *filename;
 	git_oid oid;
+	size_t filename_len;
+	int removed;
 };
 
 struct git_tree {
@@ -17,6 +19,12 @@ struct git_tree {
 	git_vector entries;
 };
 
+struct git_treebuilder {
+	git_vector entries;
+	size_t entry_count;
+};
+
+
 void git_tree__free(git_tree *tree);
 int git_tree__parse(git_tree *tree, git_odb_object *obj);
 
diff --git a/tests/t09-tree.c b/tests/t09-tree.c
index 6c1b2e64..e4252dbc 100644
--- a/tests/t09-tree.c
+++ b/tests/t09-tree.c
@@ -29,6 +29,36 @@
 
 static const char *tree_oid = "1810dff58d8a660512d4832e740f692884338ccd";
 
+#if 0
+static int print_tree(git_repository *repo, const git_oid *tree_oid, int depth)
+{
+	static const char *indent = "                              ";
+	git_tree *tree;
+	unsigned int i;
+
+	if (git_tree_lookup(&tree, repo, tree_oid) < GIT_SUCCESS)
+		return GIT_ERROR;
+
+	for (i = 0; i < git_tree_entrycount(tree); ++i) {
+		const git_tree_entry *entry = git_tree_entry_byindex(tree, i);
+		char entry_oid[40];
+
+		git_oid_fmt(entry_oid, &entry->oid);
+		printf("%.*s%o [%.*s] %s\n", depth*2, indent, entry->attr, 40, entry_oid, entry->filename);
+
+		if (entry->attr == S_IFDIR) {
+			if (print_tree(repo, &entry->oid, depth + 1) < GIT_SUCCESS) {
+				git_tree_close(tree);
+				return GIT_ERROR;
+			}
+		}
+	}
+
+	git_tree_close(tree);
+	return GIT_SUCCESS;
+}
+#endif
+
 BEGIN_TEST(read0, "acces randomly the entries on a loaded tree")
 	git_oid id;
 	git_repository *repo;
@@ -55,7 +85,7 @@ BEGIN_TEST(read1, "read a tree from the repository")
 	git_oid id;
 	git_repository *repo;
 	git_tree *tree;
-	git_tree_entry *entry;
+	const git_tree_entry *entry;
 	git_object *obj;
 
 	must_pass(git_repository_open(&repo, REPOSITORY_FOLDER));
@@ -80,10 +110,27 @@ BEGIN_TEST(read1, "read a tree from the repository")
 	git_repository_free(repo);
 END_TEST
 
+#if 0
+BEGIN_TEST(write0, "write a tree from an index")
+	git_repository *repo;
+	git_index *index;
+	git_oid tree_oid;
+
+	must_pass(git_repository_open(&repo, "/tmp/redtmp/.git"));
+	must_pass(git_repository_index(&index, repo));
+
+	must_pass(git_tree_create_fromindex(&tree_oid, index));
+	must_pass(print_tree(repo, &tree_oid, 0));
+
+	git_repository_free(repo);
+END_TEST
+#endif
+
 BEGIN_SUITE(tree)
+	//ADD_TEST(print0);
 	ADD_TEST(read0);
 	ADD_TEST(read1);
-//	ADD_TEST(write0); /* TODO THREADSAFE */
-//	ADD_TEST(write1);
+	//ADD_TEST(write0);
+	//ADD_TEST(write1);
 END_SUITE