CheckoutCommand: Support checking out ours and theirs

The checkoutPaths body is split into two implementations, depending on
whether we are checking out the index or a branch. This improves
readability, as in the index case we now also need to have access to
DirCacheIterator.

Bug: 390147
Change-Id: I99fd599b25b2ace9bdd84535a56565286a3cb7f1
Signed-off-by: Chris Aniszczyk <zx@twitter.com>

org.eclipse.jgit.test/tst/org/eclipse/jgit/api/PathCheckoutCommandTest.java

@@ -47,6 +47,7 @@ import static org.junit.Assert.assertEquals;
 import java.io.File;
 import java.io.IOException;
 
+import org.eclipse.jgit.api.CheckoutCommand.Stage;
 import org.eclipse.jgit.api.errors.JGitInternalException;
 import org.eclipse.jgit.dircache.DirCache;
 import org.eclipse.jgit.dircache.DirCacheEntry;
@@ -246,10 +247,41 @@ public class PathCheckoutCommandTest extends RepositoryTestCase {
 		assertEquals("a", read(test2));
 	}
 
+
 	@Test(expected = JGitInternalException.class)
 	public void testCheckoutOfConflictingFileShouldThrow()
 			throws Exception {
-		// Setup
+		setupConflictingState();
+
+		git.checkout().addPath(FILE1).call();
+	}
+
+	@Test
+	public void testCheckoutOurs() throws Exception {
+		setupConflictingState();
+
+		git.checkout().setStage(Stage.OURS).addPath(FILE1).call();
+
+		assertEquals("3", read(FILE1));
+		assertStageOneToThree(FILE1);
+	}
+
+	@Test
+	public void testCheckoutTheirs() throws Exception {
+		setupConflictingState();
+
+		git.checkout().setStage(Stage.THEIRS).addPath(FILE1).call();
+
+		assertEquals("Conflicting", read(FILE1));
+		assertStageOneToThree(FILE1);
+	}
+
+	@Test(expected = IllegalStateException.class)
+	public void testStageNotPossibleWithBranch() throws Exception {
+		git.checkout().setStage(Stage.OURS).setStartPoint("master").call();
+	}
+
+	private void setupConflictingState() throws Exception {
 		git.checkout().setCreateBranch(true).setName("conflict")
 				.setStartPoint(initialCommit).call();
 		writeTrashFile(FILE1, "Conflicting");
@@ -260,8 +292,18 @@ public class PathCheckoutCommandTest extends RepositoryTestCase {
 
 		git.merge().include(conflict).call();
 		assertEquals(RepositoryState.MERGING, db.getRepositoryState());
+		assertStageOneToThree(FILE1);
+	}
 
-		// Now check out the conflicting path
+	private void assertStageOneToThree(String name) throws Exception {
-		git.checkout().addPath(FILE1).call();
+		DirCache cache = DirCache.read(db.getIndexFile(), db.getFS());
+		int i = cache.findEntry(name);
+		DirCacheEntry stage1 = cache.getEntry(i);
+		DirCacheEntry stage2 = cache.getEntry(i + 1);
+		DirCacheEntry stage3 = cache.getEntry(i + 2);
+
+		assertEquals(DirCacheEntry.STAGE_1, stage1.getStage());
+		assertEquals(DirCacheEntry.STAGE_2, stage2.getStage());
+		assertEquals(DirCacheEntry.STAGE_3, stage3.getStage());
 	}
 }

org.eclipse.jgit/src/org/eclipse/jgit/api/CheckoutCommand.java

@@ -125,6 +125,33 @@ import org.eclipse.jgit.util.FileUtils;
  *      >Git documentation about Checkout</a>
  */
 public class CheckoutCommand extends GitCommand<Ref> {
+
+	/**
+	 * Stage to check out, see {@link CheckoutCommand#setStage(Stage)}.
+	 */
+	public static enum Stage {
+		/**
+		 * Base stage (#1)
+		 */
+		BASE(DirCacheEntry.STAGE_1),
+
+		/**
+		 * Ours stage (#2)
+		 */
+		OURS(DirCacheEntry.STAGE_2),
+
+		/**
+		 * Theirs stage (#3)
+		 */
+		THEIRS(DirCacheEntry.STAGE_3);
+
+		private final int number;
+
+		private Stage(int number) {
+			this.number = number;
+		}
+	}
+
 	private String name;
 
 	private boolean force = false;
@@ -137,6 +164,8 @@ public class CheckoutCommand extends GitCommand<Ref> {
 
 	private RevCommit startCommit;
 
+	private Stage checkoutStage = null;
+
 	private CheckoutResult status;
 
 	private List<String> paths;
@@ -327,58 +356,94 @@ public class CheckoutCommand extends GitCommand<Ref> {
 		RevWalk revWalk = new RevWalk(repo);
 		DirCache dc = repo.lockDirCache();
 		try {
-			DirCacheEditor editor = dc.editor();
+			TreeWalk treeWalk = new TreeWalk(revWalk.getObjectReader());
-			TreeWalk startWalk = new TreeWalk(revWalk.getObjectReader());
+			treeWalk.setRecursive(true);
-			startWalk.setRecursive(true);
 			if (!checkoutAllPaths)
-				startWalk.setFilter(PathFilterGroup.createFromStrings(paths));
+				treeWalk.setFilter(PathFilterGroup.createFromStrings(paths));
-			final boolean checkoutIndex = startCommit == null
-					&& startPoint == null;
-			if (!checkoutIndex)
-				startWalk.addTree(revWalk.parseCommit(getStartPoint())
-						.getTree());
-			else
-				startWalk.addTree(new DirCacheIterator(dc));
-
-			final File workTree = repo.getWorkTree();
-			final ObjectReader r = repo.getObjectDatabase().newReader();
 			try {
-				while (startWalk.next()) {
+				if (isCheckoutIndex())
-					final ObjectId blobId = startWalk.getObjectId(0);
+					checkoutPathsFromIndex(treeWalk, dc);
-					final FileMode mode = startWalk.getFileMode(0);
+				else {
-					editor.add(new PathEdit(startWalk.getPathString()) {
+					RevCommit commit = revWalk.parseCommit(getStartPoint());
+					checkoutPathsFromCommit(treeWalk, dc, commit);
+				}
+			} finally {
+				treeWalk.release();
+			}
+		} finally {
+			dc.unlock();
+			revWalk.release();
+		}
+		return this;
+	}
+
+	private void checkoutPathsFromIndex(TreeWalk treeWalk, DirCache dc)
+			throws IOException {
+		DirCacheIterator dci = new DirCacheIterator(dc);
+		treeWalk.addTree(dci);
+
+		final ObjectReader r = treeWalk.getObjectReader();
+		DirCacheEditor editor = dc.editor();
+		while (treeWalk.next()) {
+			DirCacheEntry entry = dci.getDirCacheEntry();
+			// Only add one edit per path
+			if (entry != null && entry.getStage() > DirCacheEntry.STAGE_1)
+				continue;
+			editor.add(new PathEdit(treeWalk.getPathString()) {
 				public void apply(DirCacheEntry ent) {
-							if (checkoutIndex
+					int stage = ent.getStage();
-									&& ent.getStage() > DirCacheEntry.STAGE_0) {
+					if (stage > DirCacheEntry.STAGE_0) {
-								UnmergedPathException e = new UnmergedPathException(ent);
+						if (checkoutStage != null) {
+							if (stage == checkoutStage.number)
+								checkoutPath(ent, r);
+						} else {
+							UnmergedPathException e = new UnmergedPathException(
+									ent);
 							throw new JGitInternalException(e.getMessage(), e);
 						}
+					} else {
+						checkoutPath(ent, r);
+					}
+				}
+			});
+		}
+		editor.commit();
+	}
+
+	private void checkoutPathsFromCommit(TreeWalk treeWalk, DirCache dc,
+			RevCommit commit) throws IOException {
+		treeWalk.addTree(commit.getTree());
+		final ObjectReader r = treeWalk.getObjectReader();
+		DirCacheEditor editor = dc.editor();
+		while (treeWalk.next()) {
+			final ObjectId blobId = treeWalk.getObjectId(0);
+			final FileMode mode = treeWalk.getFileMode(0);
+			editor.add(new PathEdit(treeWalk.getPathString()) {
+				public void apply(DirCacheEntry ent) {
 					ent.setObjectId(blobId);
 					ent.setFileMode(mode);
-							File file = new File(workTree, ent.getPathString());
+					checkoutPath(ent, r);
+				}
+			});
+		}
+		editor.commit();
+	}
+
+	private void checkoutPath(DirCacheEntry entry, ObjectReader reader) {
+		File file = new File(repo.getWorkTree(), entry.getPathString());
 		File parentDir = file.getParentFile();
 		try {
 			FileUtils.mkdirs(parentDir, true);
-								DirCacheCheckout.checkoutEntry(repo, file, ent, r);
+			DirCacheCheckout.checkoutEntry(repo, file, entry, reader);
 		} catch (IOException e) {
-								throw new JGitInternalException(
+			throw new JGitInternalException(MessageFormat.format(
-										MessageFormat.format(
 					JGitText.get().checkoutConflictWithFile,
-												ent.getPathString()), e);
+					entry.getPathString()), e);
-							}
 		}
-					});
-				}
-				editor.commit();
-			} finally {
-				startWalk.release();
-				r.release();
 	}
-		} finally {
+
-			dc.unlock();
+	private boolean isCheckoutIndex() {
-			revWalk.release();
+		return startCommit == null && startPoint == null;
-		}
-		return this;
 	}
 
 	private ObjectId getStartPoint() throws AmbiguousObjectException,
@@ -483,6 +548,7 @@ public class CheckoutCommand extends GitCommand<Ref> {
 		checkCallable();
 		this.startPoint = startPoint;
 		this.startCommit = null;
+		checkOptions();
 		return this;
 	}
 
@@ -503,6 +569,7 @@ public class CheckoutCommand extends GitCommand<Ref> {
 		checkCallable();
 		this.startCommit = startCommit;
 		this.startPoint = null;
+		checkOptions();
 		return this;
 	}
 
@@ -522,6 +589,24 @@ public class CheckoutCommand extends GitCommand<Ref> {
 		return this;
 	}
 
+	/**
+	 * When checking out the index, check out the specified stage (ours or
+	 * theirs) for unmerged paths.
+	 * <p>
+	 * This can not be used when checking out a branch, only when checking out
+	 * the index.
+	 *
+	 * @param stage
+	 *            the stage to check out
+	 * @return this
+	 */
+	public CheckoutCommand setStage(Stage stage) {
+		checkCallable();
+		this.checkoutStage = stage;
+		checkOptions();
+		return this;
+	}
+
 	/**
 	 * @return the result, never <code>null</code>
 	 */
@@ -530,4 +615,11 @@ public class CheckoutCommand extends GitCommand<Ref> {
 			return CheckoutResult.NOT_TRIED_RESULT;
 		return status;
 	}
+
+	private void checkOptions() {
+		if (checkoutStage != null && !isCheckoutIndex())
+			throw new IllegalStateException(
+					"Checking out ours/theirs is only possible when checking out index, "
+							+ "not when switching branches.");
+	}
 }