1 files changed, 36 insertions, 24 deletions

diff --git a/builtin/checkout.c b/builtin/checkout.c
index 011c0561d9..f1984d9933 100644
--- a/builtin/checkout.c
+++ b/builtin/checkout.c
@@ -115,16 +115,21 @@ static int check_stage(int stage, struct cache_entry *ce, int pos)
 		return error(_("path '%s' does not have their version"), ce->name);
 }
 
-static int check_all_stages(struct cache_entry *ce, int pos)
+static int check_stages(unsigned stages, struct cache_entry *ce, int pos)
 {
-	if (ce_stage(ce) != 1 ||
-	    active_nr <= pos + 2 ||
-	    strcmp(active_cache[pos+1]->name, ce->name) ||
-	    ce_stage(active_cache[pos+1]) != 2 ||
-	    strcmp(active_cache[pos+2]->name, ce->name) ||
-	    ce_stage(active_cache[pos+2]) != 3)
-		return error(_("path '%s' does not have all three versions"),
-			     ce->name);
+	unsigned seen = 0;
+	const char *name = ce->name;
+
+	while (pos < active_nr) {
+		ce = active_cache[pos];
+		if (strcmp(name, ce->name))
+			break;
+		seen |= (1 << ce_stage(ce));
+		pos++;
+	}
+	if ((stages & seen) != stages)
+		return error(_("path '%s' does not have all necessary versions"),
+			     name);
 	return 0;
 }
 
@@ -151,18 +156,27 @@ static int checkout_merged(int pos, struct checkout *state)
 	int status;
 	unsigned char sha1[20];
 	mmbuffer_t result_buf;
+	unsigned char threeway[3][20];
+	unsigned mode = 0;
+
+	memset(threeway, 0, sizeof(threeway));
+	while (pos < active_nr) {
+		int stage;
+		stage = ce_stage(ce);
+		if (!stage || strcmp(path, ce->name))
+			break;
+		hashcpy(threeway[stage - 1], ce->sha1);
+		if (stage == 2)
+			mode = create_ce_mode(ce->ce_mode);
+		pos++;
+		ce = active_cache[pos];
+	}
+	if (is_null_sha1(threeway[1]) || is_null_sha1(threeway[2]))
+		return error(_("path '%s' does not have necessary versions"), path);
 
-	if (ce_stage(ce) != 1 ||
-	    active_nr <= pos + 2 ||
-	    strcmp(active_cache[pos+1]->name, path) ||
-	    ce_stage(active_cache[pos+1]) != 2 ||
-	    strcmp(active_cache[pos+2]->name, path) ||
-	    ce_stage(active_cache[pos+2]) != 3)
-		return error(_("path '%s' does not have all 3 versions"), path);
-
-	read_mmblob(&ancestor, active_cache[pos]->sha1);
-	read_mmblob(&ours, active_cache[pos+1]->sha1);
-	read_mmblob(&theirs, active_cache[pos+2]->sha1);
+	read_mmblob(&ancestor, threeway[0]);
+	read_mmblob(&ours, threeway[1]);
+	read_mmblob(&theirs, threeway[2]);
 
 	/*
 	 * NEEDSWORK: re-create conflicts from merges with
@@ -193,9 +207,7 @@ static int checkout_merged(int pos, struct checkout *state)
 	if (write_sha1_file(result_buf.ptr, result_buf.size,
 			    blob_type, sha1))
 		die(_("Unable to add merge result for '%s'"), path);
-	ce = make_cache_entry(create_ce_mode(active_cache[pos+1]->ce_mode),
-			      sha1,
-			      path, 2, 0);
+	ce = make_cache_entry(mode, sha1, path, 2, 0);
 	if (!ce)
 		die(_("make_cache_entry failed for path '%s'"), path);
 	status = checkout_entry(ce, state, NULL);
@@ -253,7 +265,7 @@ static int checkout_paths(struct tree *source_tree, const char **pathspec,
 			} else if (stage) {
 				errs |= check_stage(stage, ce, pos);
 			} else if (opts->merge) {
-				errs |= check_all_stages(ce, pos);
+				errs |= check_stages((1<<2) | (1<<3), ce, pos);
 			} else {
 				errs = 1;
 				error(_("path '%s' is unmerged"), ce->name);