1 files changed, 47 insertions, 25 deletions

diff --git a/plugins/hg4idea/src/org/zmlx/hg4idea/provider/update/HgRegularUpdater.java b/plugins/hg4idea/src/org/zmlx/hg4idea/provider/update/HgRegularUpdater.java
index 355405f33e7f..98f6cad13ba3 100644
--- a/plugins/hg4idea/src/org/zmlx/hg4idea/provider/update/HgRegularUpdater.java
+++ b/plugins/hg4idea/src/org/zmlx/hg4idea/provider/update/HgRegularUpdater.java
@@ -23,9 +23,12 @@ import com.intellij.openapi.vfs.VirtualFile;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.zmlx.hg4idea.*;
+import org.zmlx.hg4idea.action.HgCommandResultNotifier;
 import org.zmlx.hg4idea.command.*;
 import org.zmlx.hg4idea.execution.HgCommandException;
 import org.zmlx.hg4idea.execution.HgCommandResult;
+import org.zmlx.hg4idea.repo.HgRepository;
+import org.zmlx.hg4idea.util.HgErrorUtil;
 import org.zmlx.hg4idea.util.HgUtil;
 
 import java.util.ArrayList;
@@ -47,12 +50,8 @@ public class HgRegularUpdater implements HgUpdater {
     this.updateConfiguration = configuration;
   }
 
-  private boolean shouldPull() {
-    return updateConfiguration.shouldPull();
-  }
-
-  private boolean shouldUpdate() {
-    return updateConfiguration.shouldUpdate();
+  private boolean shouldRebase() {
+    return updateConfiguration.shouldRebase();
   }
 
   private boolean shouldMerge() {
@@ -88,27 +87,28 @@ public class HgRegularUpdater implements HgUpdater {
 //      throw new VcsException("working dir not at branch tip (use \"Update to...\" to check out branch tip)");
 //    }
 
-    if (shouldPull()) {
-      boolean pullResult = pull(repoRoot, indicator);
-      if (!pullResult) {
-        return false;
-      }
+    HgRevisionNumber parentBeforeUpdate = new HgWorkingCopyRevisionsCommand(project).firstParent(repoRoot);
+    HgCommandExitCode pullResult = pull(repoRoot, indicator, shouldRebase());
+    if (pullResult == HgCommandExitCode.ERROR) {
+      return false;
     }
 
-    if (shouldUpdate()) {
+    if (pullResult == HgCommandExitCode.SUCCESS) {
+      HgRevisionNumber parentAfterUpdate = new HgWorkingCopyRevisionsCommand(project).firstParent(repoRoot);
+      addUpdatedFiles(repoRoot, updatedFiles, parentBeforeUpdate, parentAfterUpdate);
+      return true;
+    }
 
-      List<HgRevisionNumber> parentsBeforeUpdate = new HgWorkingCopyRevisionsCommand(project).parents(repoRoot);
-      if (parentsBeforeUpdate.size() > 1) {
-        throw new VcsException(HgVcsMessages.message("hg4idea.update.error.uncommittedMerge", repoRoot.getPath()));
-      }
+    if (shouldMerge()) {
 
       indicator.setText2(HgVcsMessages.message("hg4idea.progress.countingHeads"));
 
       List<HgRevisionNumber> branchHeadsAfterPull = new HgHeadsCommand(project, repoRoot).execute();
       List<HgRevisionNumber> pulledBranchHeads = determinePulledBranchHeads(branchHeadsBeforePull, branchHeadsAfterPull);
-      List<HgRevisionNumber> remainingOriginalBranchHeads = determingRemainingOriginalBranchHeads(branchHeadsBeforePull, branchHeadsAfterPull);
+      List<HgRevisionNumber> remainingOriginalBranchHeads =
+        determingRemainingOriginalBranchHeads(branchHeadsBeforePull, branchHeadsAfterPull);
 
-      if (branchHeadsAfterPull.size() > 1 && shouldMerge()) {
+      if (branchHeadsAfterPull.size() > 1) {
         abortOnLocalChanges();
         abortOnMultiplePulledHeads(pulledBranchHeads);
         abortOnMultipleLocalHeads(remainingOriginalBranchHeads);
@@ -118,14 +118,14 @@ public class HgRegularUpdater implements HgUpdater {
         if (shouldCommitAfterMerge()) {
           commitOrWarnAboutConflicts(warnings, mergeResult);
         }
-      } else {
-        //in case of multiple heads the update will report the appropriate error
-        update(repoRoot, indicator, updatedFiles, warnings);
       }
-
       //any kind of update could have resulted in merges and merge conflicts, so run the resolver
       resolvePossibleConflicts(updatedFiles);
     }
+    else {
+      processRebase(updatedFiles);
+    }
+
     return true;
   }
 
@@ -207,6 +207,24 @@ public class HgRegularUpdater implements HgUpdater {
     return new HgHeadMerger(project, mergeCommand).merge(repoRoot);
   }
 
+  private void processRebase(final UpdatedFiles updatedFiles) throws VcsException {
+    HgRepository repository = HgUtil.getRepositoryManager(project).getRepositoryForRoot(repoRoot);
+    HgCommandResult result;
+    do {
+      resolvePossibleConflicts(updatedFiles);
+      if (repository == null || !HgConflictResolver.findConflicts(project, repoRoot).isEmpty()) {
+        return;
+      }
+      HgRebaseCommand rebaseCommand = new HgRebaseCommand(project, repository);
+      result = rebaseCommand.continueRebase();
+      if (HgErrorUtil.isAbort(result)) {
+        new HgCommandResultNotifier(project).notifyError(result, "Hg Error", "Couldn't continue rebasing");
+        return;
+      }
+    }
+    while (result.getExitValue() == 1);
+  }
+
   private void abortOnLocalChanges() throws VcsException {
     if (getLocalChanges().size() != 0) {
       throw new VcsException(HgVcsMessages.message("hg4idea.update.error.localchanges", repoRoot.getPath()));
@@ -222,14 +240,18 @@ public class HgRegularUpdater implements HgUpdater {
     return statusCommand.execute(repoRoot);
   }
 
-  private boolean pull(VirtualFile repo, ProgressIndicator indicator)
+  private HgCommandExitCode pull(VirtualFile repo, ProgressIndicator indicator, boolean isRebase)
     throws VcsException {
     indicator.setText2(HgVcsMessages.message("hg4idea.progress.pull.with.update"));
     HgPullCommand hgPullCommand = new HgPullCommand(project, repo);
     final String defaultPath = HgUtil.getRepositoryDefaultPath(project, repo);
     hgPullCommand.setSource(defaultPath);
-    hgPullCommand.setUpdate(false);
-    hgPullCommand.setRebase(false);
+    if (isRebase) {
+      hgPullCommand.setRebase(true);
+    }
+    else {
+      hgPullCommand.setUpdate(true);
+    }
     return hgPullCommand.execute();
   }