Add index updating to checkout

Make checkout update entries in the index for all files that are
updated and/or removed, unless flag GIT_CHECKOUT_DONT_UPDATE_INDEX
is given.  To do this, iterators were extended to allow a little
more introspection into the index being iterated over, etc.

5 files changed, 187 insertions, 62 deletions

diff --git a/include/git2/checkout.h b/include/git2/checkout.h
index 196962bb9..884ea27f6 100644
--- a/include/git2/checkout.h
+++ b/include/git2/checkout.h
@@ -140,8 +140,11 @@ typedef enum {
 	/** Only update existing files, don't create new ones */
 	GIT_CHECKOUT_UPDATE_ONLY = (1u << 7),
 
+	/** Normally checkout updates index entries as it goes; this stops that */
+	GIT_CHECKOUT_DONT_UPDATE_INDEX = (1u << 8),
+
 	/** Don't refresh index/config/etc before doing checkout */
-	GIT_CHECKOUT_NO_REFRESH = (1u << 8),
+	GIT_CHECKOUT_NO_REFRESH = (1u << 9),
 
 	/**
 	 * THE FOLLOWING OPTIONS ARE NOT YET IMPLEMENTED
diff --git a/src/checkout.c b/src/checkout.c
index 5aeb0624c..a62df5efd 100644
--- a/src/checkout.c
+++ b/src/checkout.c
@@ -194,6 +194,7 @@ typedef struct {
 	bool opts_free_baseline;
 	char *pfx;
 	git_iterator *baseline;
+	git_index *index;
 	git_pool pool;
 	git_vector removes;
 	git_buf path;
@@ -261,16 +262,20 @@ static int checkout_notify(
 
 static bool checkout_is_workdir_modified(
 	checkout_data *data,
-	const git_diff_file *wditem,
-	const git_index_entry *baseitem)
+	const git_diff_file *baseitem,
+	const git_index_entry *wditem)
 {
 	git_oid oid;
 
-	if (wditem->size != baseitem->file_size)
+	/* depending on where base is coming from, we may or may not know
+	 * the actual size of the data, so we can't rely on this shortcut.
+	 */
+	if (baseitem->size && wditem->file_size != baseitem->size)
 		return true;
 
 	if (git_diff__oid_for_file(
-			data->repo, wditem->path, wditem->mode, wditem->size, &oid) < 0)
+			data->repo, wditem->path, wditem->mode,
+			wditem->file_size, &oid) < 0)
 		return false;
 
 	return (git_oid_cmp(&baseitem->oid, &oid) != 0);
@@ -279,33 +284,6 @@ static bool checkout_is_workdir_modified(
 #define CHECKOUT_ACTION_IF(FLAG,YES,NO) \
 	((data->strategy & GIT_CHECKOUT_##FLAG) ? CHECKOUT_ACTION__##YES : CHECKOUT_ACTION__##NO)
 
-static const char *checkout_action_name_debug(int act)
-{
-	if (act & CHECKOUT_ACTION__CONFLICT)
-		return "CONFLICT";
-
-	if (act & CHECKOUT_ACTION__REMOVE) {
-		if (act & CHECKOUT_ACTION__UPDATE_BLOB)
-			return "REMOVE+UPDATE";
-		if (act & CHECKOUT_ACTION__UPDATE_SUBMODULE)
-			return "REMOVE+UPDATE SUB";
-		return "REMOVE";
-	}
-	if (act & CHECKOUT_ACTION__DEFER_REMOVE) {
-		if (act & CHECKOUT_ACTION__UPDATE_BLOB)
-			return "UPDATE (WITH REMOVE)";
-		if (act & CHECKOUT_ACTION__UPDATE_SUBMODULE)
-			return "UPDATE SUB (WITH REMOVE)";
-		return "DEFERRED REMOVE";
-	}
-	if (act & CHECKOUT_ACTION__UPDATE_BLOB)
-		return "UPDATE";
-	if (act & CHECKOUT_ACTION__UPDATE_SUBMODULE)
-		return "UPDATE SUB";
-	assert(act == 0);
-	return "NONE";
-}
-
 static int checkout_action_common(
 	checkout_data *data,
 	int action,
@@ -372,19 +350,26 @@ static int checkout_action_wd_only(
 	const git_index_entry *wd,
 	git_vector *pathspec)
 {
-	bool ignored, remove;
+	bool remove = false;
 	git_checkout_notify_t notify = GIT_CHECKOUT_NOTIFY_NONE;
+	const git_index_entry *entry;
 
 	if (!git_pathspec_match_path(
 			pathspec, wd->path, false, workdir->ignore_case))
 		return 0;
 
-	ignored = git_iterator_current_is_ignored(workdir);
-
-	if (ignored) {
+	/* check if item is tracked in the index but not in the checkout diff */
+	if (data->index != NULL &&
+		(entry = git_index_get_bypath(data->index, wd->path, 0)) != NULL)
+	{
+		notify = GIT_CHECKOUT_NOTIFY_DIRTY;
+		remove = ((data->strategy & GIT_CHECKOUT_FORCE) != 0);
+	}
+	else if (git_iterator_current_is_ignored(workdir)) {
 		notify = GIT_CHECKOUT_NOTIFY_IGNORED;
 		remove = ((data->strategy & GIT_CHECKOUT_REMOVE_IGNORED) != 0);
-	} else {
+	}
+	else {
 		notify = GIT_CHECKOUT_NOTIFY_UNTRACKED;
 		remove = ((data->strategy & GIT_CHECKOUT_REMOVE_UNTRACKED) != 0);
 	}
@@ -697,6 +682,7 @@ fail:
 }
 
 static int buffer_to_file(
+	struct stat *st,
 	git_buf *buffer,
 	const char *path,
 	mode_t dir_mode,
@@ -717,6 +703,9 @@ static int buffer_to_file(
 		giterr_set(GITERR_OS, "Could not write to '%s'", path);
 		(void)p_close(fd);
 	} else {
+		if ((error = p_fstat(fd, st)) < 0)
+			giterr_set(GITERR_OS, "Error while statting '%s'", path);
+
 		if ((error = p_close(fd)) < 0)
 			giterr_set(GITERR_OS, "Error while closing '%s'", path);
 	}
@@ -730,6 +719,7 @@ static int buffer_to_file(
 }
 
 static int blob_content_to_file(
+	struct stat *st,
 	git_blob *blob,
 	const char *path,
 	mode_t entry_filemode,
@@ -772,7 +762,12 @@ static int blob_content_to_file(
 		file_mode = entry_filemode;
 
 	error = buffer_to_file(
-		&filtered, path, opts->dir_mode, opts->file_open_flags, file_mode);
+		st, &filtered, path, opts->dir_mode, opts->file_open_flags, file_mode);
+
+	if (!error) {
+		st->st_size = blob->odb_object->raw.len;
+		st->st_mode = entry_filemode;
+	}
 
 cleanup:
 	git_filters_free(&filters);
@@ -784,7 +779,7 @@ cleanup:
 }
 
 static int blob_content_to_link(
-	git_blob *blob, const char *path, int can_symlink)
+	struct stat *st, git_blob *blob, const char *path, int can_symlink)
 {
 	git_buf linktarget = GIT_BUF_INIT;
 	int error;
@@ -792,28 +787,57 @@ static int blob_content_to_link(
 	if ((error = git_blob__getbuf(&linktarget, blob)) < 0)
 		return error;
 
-	if (can_symlink)
-		error = p_symlink(git_buf_cstr(&linktarget), path);
-	else
+	if (can_symlink) {
+		if ((error = p_symlink(git_buf_cstr(&linktarget), path)) < 0)
+			giterr_set(GITERR_CHECKOUT, "Could not create symlink %s\n", path);
+	} else {
 		error = git_futils_fake_symlink(git_buf_cstr(&linktarget), path);
+	}
+
+	if (!error) {
+		if ((error = p_lstat(path, st)) < 0)
+			giterr_set(GITERR_CHECKOUT, "Could not stat symlink %s", path);
+
+		st->st_mode = GIT_FILEMODE_LINK;
+	}
 
 	git_buf_free(&linktarget);
 
 	return error;
 }
 
+static int checkout_update_index(
+	checkout_data *data,
+	const git_diff_file *file,
+	struct stat *st)
+{
+	git_index_entry entry;
+
+	if (!data->index)
+		return 0;
+
+	memset(&entry, 0, sizeof(entry));
+	entry.path = (char *)file->path; /* cast to prevent warning */
+	git_index_entry__init_from_stat(&entry, st);
+	git_oid_cpy(&entry.oid, &file->oid);
+
+	return git_index_add(data->index, &entry);
+}
+
 static int checkout_submodule(
 	checkout_data *data,
 	const git_diff_file *file)
 {
+	int error = 0;
+
 	/* Until submodules are supported, UPDATE_ONLY means do nothing here */
 	if ((data->strategy & GIT_CHECKOUT_UPDATE_ONLY) != 0)
 		return 0;
 
-	if (git_futils_mkdir(
+	if ((error = git_futils_mkdir(
 			file->path, git_repository_workdir(data->repo),
-			data->opts.dir_mode, GIT_MKDIR_PATH) < 0)
-		return -1;
+			data->opts.dir_mode, GIT_MKDIR_PATH)) < 0)
+		return error;
 
 	/* TODO: Support checkout_strategy options.  Two circumstances:
 	 * 1 - submodule already checked out, but we need to move the HEAD
@@ -824,7 +848,26 @@ static int checkout_submodule(
 	 * command should probably be able to.  Do we need a submodule callback?
 	 */
 
-	return 0;
+	/* update the index unless prevented */
+	if ((data->strategy & GIT_CHECKOUT_DONT_UPDATE_INDEX) == 0) {
+		struct stat st;
+
+		git_buf_truncate(&data->path, data->workdir_len);
+		if (git_buf_puts(&data->path, file->path) < 0)
+			return -1;
+
+		if ((error = p_stat(git_buf_cstr(&data->path), &st)) < 0) {
+			giterr_set(
+				GITERR_CHECKOUT, "Could not stat submodule %s\n", file->path);
+			return error;
+		}
+
+		st.st_mode = GIT_FILEMODE_COMMIT;
+
+		error = checkout_update_index(data, file, &st);
+	}
+
+	return error;
 }
 
 static void report_progress(
@@ -843,6 +886,7 @@ static int checkout_blob(
 {
 	int error = 0;
 	git_blob *blob;
+	struct stat st;
 
 	git_buf_truncate(&data->path, data->workdir_len);
 	if (git_buf_puts(&data->path, file->path) < 0)
@@ -853,10 +897,10 @@ static int checkout_blob(
 
 	if (S_ISLNK(file->mode))
 		error = blob_content_to_link(
-			blob, git_buf_cstr(&data->path), data->can_symlink);
+			&st, blob, git_buf_cstr(&data->path), data->can_symlink);
 	else
 		error = blob_content_to_file(
-			blob, git_buf_cstr(&data->path), file->mode, &data->opts);
+			&st, blob, git_buf_cstr(&data->path), file->mode, &data->opts);
 
 	git_blob_free(blob);
 
@@ -871,6 +915,10 @@ static int checkout_blob(
 		error = 0;
 	}
 
+	/* update the index unless prevented */
+	if (!error && (data->strategy & GIT_CHECKOUT_DONT_UPDATE_INDEX) == 0)
+		error = checkout_update_index(data, file, &st);
+
 	return error;
 }
 
@@ -896,6 +944,13 @@ static int checkout_remove_the_old(
 
 			data->completed_steps++;
 			report_progress(data, delta->old_file.path);
+
+			if ((actions[i] & CHECKOUT_ACTION__UPDATE_BLOB) == 0 &&
+				(data->strategy & GIT_CHECKOUT_DONT_UPDATE_INDEX) == 0 &&
+				data->index != NULL)
+			{
+				(void)git_index_remove(data->index, delta->old_file.path, 0);
+			}
 		}
 	}
 
@@ -906,6 +961,12 @@ static int checkout_remove_the_old(
 
 		data->completed_steps++;
 		report_progress(data, str);
+
+		if ((data->strategy & GIT_CHECKOUT_DONT_UPDATE_INDEX) == 0 &&
+			data->index != NULL)
+		{
+			(void)git_index_remove(data->index, str, 0);
+		}
 	}
 
 	return 0;
@@ -926,6 +987,7 @@ static int checkout_deferred_remove(git_repository *repo, const char *path)
 #else
 	GIT_UNUSED(repo);
 	GIT_UNUSED(path);
+	assert(false);
 	return 0;
 #endif
 }
@@ -1021,15 +1083,19 @@ static void checkout_data_clear(checkout_data *data)
 	data->pfx = NULL;
 
 	git_buf_free(&data->path);
+
+	git_index_free(data->index);
+	data->index = NULL;
 }
 
 static int checkout_data_init(
 	checkout_data *data,
-	git_repository *repo,
+	git_iterator *target,
 	git_checkout_opts *proposed)
 {
 	int error = 0;
 	git_config *cfg;
+	git_repository *repo = git_iterator_owner(target);
 
 	memset(data, 0, sizeof(*data));
 
@@ -1054,8 +1120,24 @@ static int checkout_data_init(
 	else
 		memmove(&data->opts, proposed, sizeof(git_checkout_opts));
 
-	/* if you are forcing, definitely allow safe updates */
+	/* refresh config and index content unless NO_REFRESH is given */
+	if ((data->opts.checkout_strategy & GIT_CHECKOUT_NO_REFRESH) == 0) {
+		if ((error = git_config_refresh(cfg)) < 0)
+			goto cleanup;
+
+		if (git_iterator_inner_type(target) == GIT_ITERATOR_INDEX) {
+			/* if we are iterating over the index, don't reload */
+			data->index = git_iterator_index_get_index(target);
+			GIT_REFCOUNT_INC(data->index);
+		} else {
+			/* otherwise, grab and reload the index */
+			if ((error = git_repository_index(&data->index, data->repo)) < 0 ||
+				(error = git_index_read(data->index)) < 0)
+				goto cleanup;
+		}
+	}
 
+	/* if you are forcing, definitely allow safe updates */
 	if ((data->opts.checkout_strategy & GIT_CHECKOUT_FORCE) != 0)
 		data->opts.checkout_strategy |= GIT_CHECKOUT_SAFE_CREATE;
 	if ((data->opts.checkout_strategy & GIT_CHECKOUT_SAFE_CREATE) != 0)
@@ -1116,7 +1198,7 @@ int git_checkout_iterator(
 	size_t *counts = NULL;
 
 	/* initialize structures and options */
-	error = checkout_data_init(&data, git_iterator_owner(target), opts);
+	error = checkout_data_init(&data, target, opts);
 	if (error < 0)
 		return error;
 
@@ -1204,6 +1286,10 @@ cleanup:
 	if (error == GIT_EUSER)
 		giterr_clear();
 
+	if (!error && data.index != NULL &&
+		(data.strategy & GIT_CHECKOUT_DONT_UPDATE_INDEX) == 0)
+		error = git_index_write(data.index);
+
 	git_diff_list_free(data.diff);
 	git_iterator_free(workdir);
 	git_iterator_free(data.baseline);
diff --git a/src/iterator.c b/src/iterator.c
index 28fccce0e..b15453400 100644
--- a/src/iterator.c
+++ b/src/iterator.c
@@ -988,6 +988,33 @@ fail:
 	return -1;
 }
 
+git_index *git_iterator_index_get_index(git_iterator *iter)
+{
+	if (iter->type == GIT_ITERATOR_SPOOLANDSORT)
+		iter = ((spoolandsort_iterator *)iter)->wrapped;
+
+	if (iter->type == GIT_ITERATOR_INDEX)
+		return ((index_iterator *)iter)->index;
+
+	return NULL;
+}
+
+git_iterator_type_t git_iterator_inner_type(git_iterator *iter)
+{
+	if (iter->type == GIT_ITERATOR_SPOOLANDSORT)
+		iter = ((spoolandsort_iterator *)iter)->wrapped;
+
+	return iter->type;
+}
+
+git_iterator *git_iterator_spoolandsort_inner_iterator(git_iterator *iter)
+{
+	if (iter->type == GIT_ITERATOR_SPOOLANDSORT)
+		return ((spoolandsort_iterator *)iter)->wrapped;
+
+	return NULL;
+}
+
 int git_iterator_current_tree_entry(
 	git_iterator *iter, const git_tree_entry **tree_entry)
 {
diff --git a/src/iterator.h b/src/iterator.h
index 8bcb6fb0c..ccdab4d94 100644
--- a/src/iterator.h
+++ b/src/iterator.h
@@ -193,4 +193,12 @@ extern int git_iterator_cmp(
 extern int git_iterator_current_workdir_path(
 	git_iterator *iter, git_buf **path);
 
+
+extern git_index *git_iterator_index_get_index(git_iterator *iter);
+
+extern git_iterator_type_t git_iterator_inner_type(git_iterator *iter);
+
+extern git_iterator *git_iterator_spoolandsort_inner_iterator(
+	git_iterator *iter);
+
 #endif
diff --git a/tests-clar/checkout/tree.c b/tests-clar/checkout/tree.c
index ab502dd30..6e7175a87 100644
--- a/tests-clar/checkout/tree.c
+++ b/tests-clar/checkout/tree.c
@@ -50,28 +50,29 @@ void test_checkout_tree__can_checkout_a_subdirectory_from_a_commit(void)
 
 void test_checkout_tree__can_checkout_and_remove_directory(void)
 {
-	git_reference *head;
 	cl_assert_equal_i(false, git_path_isdir("./testrepo/ab/"));
 
-	// Checkout brach "subtrees" and update HEAD, so that HEAD matches the current working tree
+	/* Checkout brach "subtrees" and update HEAD, so that HEAD matches the
+	 * current working tree
+	 */
 	cl_git_pass(git_revparse_single(&g_object, g_repo, "subtrees"));
 	cl_git_pass(git_checkout_tree(g_repo, g_object, &g_opts));
-	cl_git_pass(git_reference_lookup(&head, g_repo, "HEAD"));
-	cl_git_pass(git_reference_symbolic_set_target(head, "refs/heads/subtrees"));
-	git_reference_free(head);
-	
+
+	cl_git_pass(git_repository_set_head(g_repo, "refs/heads/subtrees"));
+
 	cl_assert_equal_i(true, git_path_isdir("./testrepo/ab/"));
 	cl_assert_equal_i(true, git_path_isfile("./testrepo/ab/de/2.txt"));
 	cl_assert_equal_i(true, git_path_isfile("./testrepo/ab/de/fgh/1.txt"));
 
-	// Checkout brach "master" and update HEAD, so that HEAD matches the current working tree
+	/* Checkout brach "master" and update HEAD, so that HEAD matches the
+	 * current working tree
+	 */
 	cl_git_pass(git_revparse_single(&g_object, g_repo, "master"));
 	cl_git_pass(git_checkout_tree(g_repo, g_object, &g_opts));
-	cl_git_pass(git_reference_lookup(&head, g_repo, "HEAD"));
-	cl_git_pass(git_reference_symbolic_set_target(head, "refs/heads/master"));
-	git_reference_free(head);
 
-	// This directory should no longer exist
+	cl_git_pass(git_repository_set_head(g_repo, "refs/heads/master"));
+
+	/* This directory should no longer exist */
 	cl_assert_equal_i(false, git_path_isdir("./testrepo/ab/"));
 }