2 files changed, 74 insertions, 14 deletions

diff --git a/src/checkout.c b/src/checkout.c
index a10507aaf..261dee112 100644
--- a/src/checkout.c
+++ b/src/checkout.c
@@ -418,6 +418,8 @@ static int checkout_action_with_wd_dir(
 	return checkout_action_common(data, action, delta, wd);
 }
 
+#define EXPAND_DIRS_FOR_STRATEGY (GIT_CHECKOUT_FORCE | GIT_CHECKOUT_REMOVE_UNTRACKED | GIT_CHECKOUT_REMOVE_IGNORED)
+
 static int checkout_action(
 	checkout_data *data,
 	git_diff_delta *delta,
@@ -429,6 +431,7 @@ static int checkout_action(
 	int cmp = -1, act;
 	int (*strcomp)(const char *, const char *) = data->diff->strcomp;
 	int (*pfxcomp)(const char *str, const char *pfx) = data->diff->pfxcomp;
+	bool expand_dirs = (data->strategy & EXPAND_DIRS_FOR_STRATEGY) != 0;
 
 	/* move workdir iterator to follow along with deltas */
 
@@ -449,14 +452,14 @@ static int checkout_action(
 		if (cmp < 0) {
 			cmp = pfxcomp(delta->old_file.path, wd->path);
 
-			if (cmp == 0) {
-				if (wd->mode == GIT_FILEMODE_TREE) {
-					/* case 2 - descend in wd */
-					if (git_iterator_advance_into_directory(workdir, &wd) < 0)
-						goto fail;
-					continue;
-				}
+			if (wd->mode == GIT_FILEMODE_TREE && (cmp == 0 || expand_dirs)) {
+				/* case 2 or untracked wd item that might need removal */
+				if (git_iterator_advance_into_directory(workdir, &wd) < 0)
+					goto fail;
+				continue;
+			}
 
+			if (cmp == 0) {
 				/* case 3 -  wd contains non-dir where dir expected */
 				act = checkout_action_with_wd_blocker(data, delta, wd);
 				*wditem_ptr = git_iterator_advance(workdir, &wd) ? NULL : wd;
@@ -519,6 +522,26 @@ fail:
 	return -1;
 }
 
+static int checkout_remaining_wd_items(
+	checkout_data *data,
+	git_iterator *workdir,
+	const git_index_entry *wd,
+	git_vector *spec)
+{
+	int error = 0;
+	bool expand_dirs = (data->strategy & EXPAND_DIRS_FOR_STRATEGY) != 0;
+
+	while (wd && !error) {
+		if (wd->mode == GIT_FILEMODE_TREE && expand_dirs)
+			error = git_iterator_advance_into_directory(workdir, &wd);
+
+		else if (!(error = checkout_action_wd_only(data, workdir, wd, spec)))
+			error = git_iterator_advance(workdir, &wd);
+	}
+
+	return error;
+}
+
 static int checkout_get_actions(
 	uint32_t **actions_ptr,
 	size_t **counts_ptr,
@@ -570,13 +593,9 @@ static int checkout_get_actions(
 			counts[CHECKOUT_ACTION__CONFLICT]++;
 	}
 
-	while (wditem != NULL) {
-		error = checkout_action_wd_only(data, workdir, wditem, &pathspec);
-		if (!error)
-			error = git_iterator_advance(workdir, &wditem);
-		if (error < 0)
-			goto fail;
-	}
+	error = checkout_remaining_wd_items(data, workdir, wditem, &pathspec);
+	if (error < 0)
+		goto fail;
 
 	counts[CHECKOUT_ACTION__REMOVE] += data->removes.length;
 
diff --git a/tests-clar/checkout/head.c b/tests-clar/checkout/head.c
index aed203a06..8b3099303 100644
--- a/tests-clar/checkout/head.c
+++ b/tests-clar/checkout/head.c
@@ -1,6 +1,8 @@
 #include "clar_libgit2.h"
 #include "refs.h"
 #include "repo/repo_helpers.h"
+#include "path.h"
+#include "fileops.h"
 
 static git_repository *g_repo;
 
@@ -20,3 +22,42 @@ void test_checkout_head__orphaned_head_returns_GIT_EORPHANEDHEAD(void)
 
 	cl_assert_equal_i(GIT_EORPHANEDHEAD, git_checkout_head(g_repo, NULL));
 }
+
+void test_checkout_head__with_index_only_tree(void)
+{
+	git_checkout_opts opts = GIT_CHECKOUT_OPTS_INIT;
+	git_index *index;
+
+	/* let's start by getting things into a known state */
+
+	opts.checkout_strategy = GIT_CHECKOUT_FORCE;
+	cl_git_pass(git_checkout_head(g_repo, &opts));
+
+	/* now let's stage some new stuff including a new directory */
+
+	cl_git_pass(git_repository_index(&index, g_repo));
+
+	p_mkdir("testrepo/newdir", 0777);
+    cl_git_mkfile("testrepo/newdir/newfile.txt", "new file\n");
+
+	cl_git_pass(git_index_add_from_workdir(index, "newdir/newfile.txt"));
+	cl_git_pass(git_index_write(index));
+
+	cl_assert(git_path_isfile("testrepo/newdir/newfile.txt"));
+	cl_assert(git_index_get_bypath(index, "newdir/newfile.txt", 0) != NULL);
+
+	git_index_free(index);
+
+	/* okay, so now we have staged this new file; let's see if we can remove */
+
+	opts.checkout_strategy = GIT_CHECKOUT_FORCE | GIT_CHECKOUT_REMOVE_UNTRACKED;
+	cl_git_pass(git_checkout_head(g_repo, &opts));
+
+	cl_git_pass(git_repository_index(&index, g_repo));
+	cl_git_pass(git_index_read(index)); /* reload if needed */
+
+	cl_assert(!git_path_isfile("testrepo/newdir/newfile.txt"));
+	cl_assert(git_index_get_bypath(index, "newdir/newfile.txt", 0) == NULL);
+
+	git_index_free(index);
+}