Allow to include untracked files in stash operations.

Unstashed changes are saved in a commit which is added as an additional
parent to the stash commit.
This behaviour is fully compatible with C Git stashing of untracked
files.

Bug: 434411
Change-Id: I2af784deb0c2320bb57bc4fd472a8daad8674e7d
Signed-off-by: Andreas Hermann <a.v.hermann@gmail.com>

4 files changed, 162 insertions, 14 deletions

diff --git a/org.eclipse.jgit/resources/org/eclipse/jgit/internal/JGitText.properties b/org.eclipse.jgit/resources/org/eclipse/jgit/internal/JGitText.properties
index 7b073217a5..fd5801e6a5 100644
--- a/org.eclipse.jgit/resources/org/eclipse/jgit/internal/JGitText.properties
+++ b/org.eclipse.jgit/resources/org/eclipse/jgit/internal/JGitText.properties
@@ -476,7 +476,7 @@ stashApplyConflictInIndex=Applying stashed index changes resulted in a conflict.
 stashApplyFailed=Applying stashed changes did not successfully complete
 stashApplyOnUnsafeRepository=Cannot apply stashed commit on a repository with state: {0}
 stashApplyWithoutHead=Cannot apply stashed commit in an empty repository or onto an unborn branch
-stashCommitMissingTwoParents=Stashed commit ''{0}'' does not have two parent commits
+stashCommitIncorrectNumberOfParents=Stashed commit ''{0}'' does have {1} parent commits instead of 2 or 3.
 stashDropDeleteRefFailed=Deleting stash reference failed with result: {0}
 stashDropFailed=Dropping stashed commit failed
 stashDropMissingReflog=Stash reflog does not contain entry ''{0}''
diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/api/StashApplyCommand.java b/org.eclipse.jgit/src/org/eclipse/jgit/api/StashApplyCommand.java
index f73ce831fe..d935857cec 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/api/StashApplyCommand.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/api/StashApplyCommand.java
@@ -42,6 +42,7 @@
  */
 package org.eclipse.jgit.api;
 
+import java.io.File;
 import java.io.IOException;
 import java.text.MessageFormat;
 
@@ -56,6 +57,7 @@ import org.eclipse.jgit.dircache.DirCacheBuilder;
 import org.eclipse.jgit.dircache.DirCacheCheckout;
 import org.eclipse.jgit.dircache.DirCacheEntry;
 import org.eclipse.jgit.dircache.DirCacheIterator;
+import org.eclipse.jgit.errors.CheckoutConflictException;
 import org.eclipse.jgit.internal.JGitText;
 import org.eclipse.jgit.lib.Constants;
 import org.eclipse.jgit.lib.ObjectId;
@@ -90,6 +92,8 @@ public class StashApplyCommand extends GitCommand<ObjectId> {
 
 	private boolean applyIndex = true;
 
+	private boolean applyUntracked = true;
+
 	private boolean ignoreRepositoryState;
 
 	private MergeStrategy strategy = MergeStrategy.RECURSIVE;
@@ -173,15 +177,20 @@ public class StashApplyCommand extends GitCommand<ObjectId> {
 
 			final ObjectId stashId = getStashId();
 			RevCommit stashCommit = revWalk.parseCommit(stashId);
-			if (stashCommit.getParentCount() != 2)
+			if (stashCommit.getParentCount() < 2
+					|| stashCommit.getParentCount() > 3)
 				throw new JGitInternalException(MessageFormat.format(
-						JGitText.get().stashCommitMissingTwoParents,
-						stashId.name()));
+						JGitText.get().stashCommitIncorrectNumberOfParents,
+						stashId.name(),
+						Integer.valueOf(stashCommit.getParentCount())));
 
 			ObjectId headTree = repo.resolve(Constants.HEAD + "^{tree}"); //$NON-NLS-1$
 			ObjectId stashIndexCommit = revWalk.parseCommit(stashCommit
 					.getParent(1));
 			ObjectId stashHeadCommit = stashCommit.getParent(0);
+			ObjectId untrackedCommit = null;
+			if (applyUntracked && stashCommit.getParentCount() == 3)
+				untrackedCommit = revWalk.parseCommit(stashCommit.getParent(2));
 
 			ResolveMerger merger = (ResolveMerger) strategy.newMerger(repo);
 			merger.setCommitNames(new String[] { "stashed HEAD", "HEAD",
@@ -209,6 +218,29 @@ public class StashApplyCommand extends GitCommand<ObjectId> {
 								JGitText.get().stashApplyConflict);
 					}
 				}
+
+				if (untrackedCommit != null) {
+					ResolveMerger untrackedMerger = (ResolveMerger) strategy
+							.newMerger(repo, true);
+					untrackedMerger.setCommitNames(new String[] {
+							"stashed HEAD", "HEAD", "untracked files" }); //$NON-NLS-1$//$NON-NLS-2$//$NON-NLS-3$
+					untrackedMerger.setBase(stashHeadCommit);
+					boolean ok = untrackedMerger.merge(headCommit,
+							untrackedCommit);
+					if (ok)
+						try {
+							RevTree untrackedTree = revWalk
+									.parseTree(untrackedMerger
+											.getResultTreeId());
+							resetUntracked(untrackedTree);
+						} catch (CheckoutConflictException e) {
+							throw new StashApplyFailureException(
+									JGitText.get().stashApplyConflict);
+						}
+					else
+						throw new StashApplyFailureException(
+								JGitText.get().stashApplyConflict);
+				}
 			} else {
 				throw new StashApplyFailureException(
 						JGitText.get().stashApplyConflict);
@@ -244,6 +276,15 @@ public class StashApplyCommand extends GitCommand<ObjectId> {
 		return this;
 	}
 
+	/**
+	 * @param applyUntracked
+	 *            true (default) if the command should restore untracked files
+	 * @since 3.4
+	 */
+	public void setApplyUntracked(boolean applyUntracked) {
+		this.applyUntracked = applyUntracked;
+	}
+
 	private void resetIndex(RevTree tree) throws IOException {
 		DirCache dc = repo.lockDirCache();
 		TreeWalk walk = null;
@@ -285,4 +326,55 @@ public class StashApplyCommand extends GitCommand<ObjectId> {
 				walk.release();
 		}
 	}
+
+	private void resetUntracked(RevTree tree) throws CheckoutConflictException,
+			IOException {
+		TreeWalk walk = null;
+		try {
+			walk = new TreeWalk(repo); // maybe NameConflictTreeWalk?
+			walk.addTree(tree);
+			walk.addTree(new FileTreeIterator(repo));
+			walk.setRecursive(true);
+
+			final ObjectReader reader = walk.getObjectReader();
+
+			while (walk.next()) {
+				final AbstractTreeIterator cIter = walk.getTree(0,
+						AbstractTreeIterator.class);
+				if (cIter == null)
+					// Not in commit, don't create untracked
+					continue;
+
+				final DirCacheEntry entry = new DirCacheEntry(walk.getRawPath());
+				entry.setFileMode(cIter.getEntryFileMode());
+				entry.setObjectIdFromRaw(cIter.idBuffer(), cIter.idOffset());
+
+				FileTreeIterator fIter = walk
+						.getTree(1, FileTreeIterator.class);
+				if (fIter != null) {
+					if (fIter.isModified(entry, true, reader)) {
+						// file exists and is dirty
+						throw new CheckoutConflictException(
+								entry.getPathString());
+					}
+				}
+
+				checkoutPath(entry, reader);
+			}
+		} finally {
+			if (walk != null)
+				walk.release();
+		}
+	}
+
+	private void checkoutPath(DirCacheEntry entry, ObjectReader reader) {
+		try {
+			File file = new File(repo.getWorkTree(), entry.getPathString());
+			DirCacheCheckout.checkoutEntry(repo, file, entry, reader);
+		} catch (IOException e) {
+			throw new JGitInternalException(MessageFormat.format(
+					JGitText.get().checkoutConflictWithFile,
+					entry.getPathString()), e);
+		}
+	}
 }
diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/api/StashCreateCommand.java b/org.eclipse.jgit/src/org/eclipse/jgit/api/StashCreateCommand.java
index cf0b6d1d9c..af35f772ce 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/api/StashCreateCommand.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/api/StashCreateCommand.java
@@ -42,6 +42,7 @@
  */
 package org.eclipse.jgit.api;
 
+import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
 import java.text.MessageFormat;
@@ -54,6 +55,7 @@ import org.eclipse.jgit.api.errors.JGitInternalException;
 import org.eclipse.jgit.api.errors.NoHeadException;
 import org.eclipse.jgit.api.errors.UnmergedPathsException;
 import org.eclipse.jgit.dircache.DirCache;
+import org.eclipse.jgit.dircache.DirCacheBuilder;
 import org.eclipse.jgit.dircache.DirCacheEditor;
 import org.eclipse.jgit.dircache.DirCacheEditor.DeletePath;
 import org.eclipse.jgit.dircache.DirCacheEditor.PathEdit;
@@ -80,6 +82,7 @@ import org.eclipse.jgit.treewalk.WorkingTreeIterator;
 import org.eclipse.jgit.treewalk.filter.AndTreeFilter;
 import org.eclipse.jgit.treewalk.filter.IndexDiffFilter;
 import org.eclipse.jgit.treewalk.filter.SkipWorkTreeFilter;
+import org.eclipse.jgit.util.FileUtils;
 
 /**
  * Command class to stash changes in the working directory and index in a
@@ -93,6 +96,8 @@ public class StashCreateCommand extends GitCommand<RevCommit> {
 
 	private static final String MSG_INDEX = "index on {0}: {1} {2}";
 
+	private static final String MSG_UNTRACKED = "untracked files on {0}: {1} {2}";
+
 	private static final String MSG_WORKING_DIR = "WIP on {0}: {1} {2}";
 
 	private String indexMessage = MSG_INDEX;
@@ -103,6 +108,8 @@ public class StashCreateCommand extends GitCommand<RevCommit> {
 
 	private PersonIdent person;
 
+	private boolean includeUntracked;
+
 	/**
 	 * Create a command to stash changes in the working directory and index
 	 *
@@ -166,6 +173,18 @@ public class StashCreateCommand extends GitCommand<RevCommit> {
 		return this;
 	}
 
+	/**
+	 * Whether to include untracked files in the stash.
+	 *
+	 * @param includeUntracked
+	 * @return {@code this}
+	 * @since 3.4
+	 */
+	public StashCreateCommand setIncludeUntracked(boolean includeUntracked) {
+		this.includeUntracked = includeUntracked;
+		return this;
+	}
+
 	private RevCommit parseCommit(final ObjectReader reader,
 			final ObjectId headId) throws IOException {
 		final RevWalk walk = new RevWalk(reader);
@@ -173,14 +192,13 @@ public class StashCreateCommand extends GitCommand<RevCommit> {
 		return walk.parseCommit(headId);
 	}
 
-	private CommitBuilder createBuilder(ObjectId headId) {
+	private CommitBuilder createBuilder() {
 		CommitBuilder builder = new CommitBuilder();
 		PersonIdent author = person;
 		if (author == null)
 			author = new PersonIdent(repo);
 		builder.setAuthor(author);
 		builder.setCommitter(author);
-		builder.setParentId(headId);
 		return builder;
 	}
 
@@ -244,6 +262,7 @@ public class StashCreateCommand extends GitCommand<RevCommit> {
 				MutableObjectId id = new MutableObjectId();
 				List<PathEdit> wtEdits = new ArrayList<PathEdit>();
 				List<String> wtDeletes = new ArrayList<String>();
+				List<DirCacheEntry> untracked = new ArrayList<DirCacheEntry>();
 				boolean hasChanges = false;
 				do {
 					AbstractTreeIterator headIter = treeWalk.getTree(0,
@@ -258,7 +277,8 @@ public class StashCreateCommand extends GitCommand<RevCommit> {
 								new UnmergedPathException(
 										indexIter.getDirCacheEntry()));
 					if (wtIter != null) {
-						if (indexIter == null && headIter == null)
+						if (indexIter == null && headIter == null
+								&& !includeUntracked)
 							continue;
 						hasChanges = true;
 						if (indexIter != null && wtIter.idEqual(indexIter))
@@ -279,11 +299,15 @@ public class StashCreateCommand extends GitCommand<RevCommit> {
 						} finally {
 							in.close();
 						}
-						wtEdits.add(new PathEdit(entry) {
-							public void apply(DirCacheEntry ent) {
-								ent.copyMetaData(entry);
-							}
-						});
+
+						if (indexIter == null && headIter == null)
+							untracked.add(entry);
+						else
+							wtEdits.add(new PathEdit(entry) {
+								public void apply(DirCacheEntry ent) {
+									ent.copyMetaData(entry);
+								}
+							});
 					}
 					hasChanges = true;
 					if (wtIter == null && headIter != null)
@@ -297,13 +321,32 @@ public class StashCreateCommand extends GitCommand<RevCommit> {
 						.getName());
 
 				// Commit index changes
-				CommitBuilder builder = createBuilder(headCommit);
+				CommitBuilder builder = createBuilder();
+				builder.setParentId(headCommit);
 				builder.setTreeId(cache.writeTree(inserter));
 				builder.setMessage(MessageFormat.format(indexMessage, branch,
 						headCommit.abbreviate(7).name(),
 						headCommit.getShortMessage()));
 				ObjectId indexCommit = inserter.insert(builder);
 
+				// Commit untracked changes
+				ObjectId untrackedCommit = null;
+				if (!untracked.isEmpty()) {
+					DirCache untrackedDirCache = DirCache.newInCore();
+					DirCacheBuilder untrackedBuilder = untrackedDirCache
+							.builder();
+					for (DirCacheEntry entry : untracked)
+						untrackedBuilder.add(entry);
+					untrackedBuilder.finish();
+
+					builder.setParentIds(new ObjectId[0]);
+					builder.setTreeId(untrackedDirCache.writeTree(inserter));
+					builder.setMessage(MessageFormat.format(MSG_UNTRACKED,
+							branch, headCommit.abbreviate(7).name(),
+							headCommit.getShortMessage()));
+					untrackedCommit = inserter.insert(builder);
+				}
+
 				// Commit working tree changes
 				if (!wtEdits.isEmpty() || !wtDeletes.isEmpty()) {
 					DirCacheEditor editor = cache.editor();
@@ -313,7 +356,10 @@ public class StashCreateCommand extends GitCommand<RevCommit> {
 						editor.add(new DeletePath(path));
 					editor.finish();
 				}
+				builder.setParentId(headCommit);
 				builder.addParentId(indexCommit);
+				if (untrackedCommit != null)
+					builder.addParentId(untrackedCommit);
 				builder.setMessage(MessageFormat.format(
 						workingDirectoryMessage, branch,
 						headCommit.abbreviate(7).name(),
@@ -324,6 +370,16 @@ public class StashCreateCommand extends GitCommand<RevCommit> {
 
 				updateStashRef(commitId, builder.getAuthor(),
 						builder.getMessage());
+
+				// Remove untracked files
+				if (includeUntracked) {
+					for (DirCacheEntry entry : untracked) {
+						File file = new File(repo.getWorkTree(),
+								entry.getPathString());
+						FileUtils.delete(file);
+					}
+				}
+
 			} finally {
 				inserter.release();
 				cache.unlock();
diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/internal/JGitText.java b/org.eclipse.jgit/src/org/eclipse/jgit/internal/JGitText.java
index 39f203ca8b..8acfb54b85 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/internal/JGitText.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/internal/JGitText.java
@@ -538,7 +538,7 @@ public class JGitText extends TranslationBundle {
 	/***/ public String stashApplyFailed;
 	/***/ public String stashApplyWithoutHead;
 	/***/ public String stashApplyOnUnsafeRepository;
-	/***/ public String stashCommitMissingTwoParents;
+	/***/ public String stashCommitIncorrectNumberOfParents;
 	/***/ public String stashDropDeleteRefFailed;
 	/***/ public String stashDropFailed;
 	/***/ public String stashDropMissingReflog;