[PATCH RESEND] rebase: reset bookmarks (issue2265 and issue2873)

[Stefano Tortarolo]

# HG changeset patch
# User Stefano Tortarolo <stefano.tortarolo at gmail.com>
# Date 1310753239 -7200
# Node ID 7e5c7364e1af8bb3e0792cdb1259038c3a1ceac0
# Parent  647071c6dfcfc19a4dedb37c1e4c0b60378d953b
rebase: reset bookmarks (issue2265 and issue2873)

diff --git a/hgext/rebase.py b/hgext/rebase.py
--- a/hgext/rebase.py
+++ b/hgext/rebase.py
@@ -14,7 +14,7 @@
 http://mercurial.selenic.com/wiki/RebaseExtension
 '''
 
-from mercurial import hg, util, repair, merge, cmdutil, commands
+from mercurial import hg, util, repair, merge, cmdutil, commands, bookmarks
 from mercurial import extensions, copies, patch
 from mercurial.commands import templateopts
 from mercurial.node import nullrev
@@ -181,6 +181,9 @@
             targetancestors = set(repo.changelog.ancestors(target))
             targetancestors.add(target)
 
+        # Keep track of the current bookmarks in order to reset them later
+        currentbookmarks = repo._bookmarks.copy()
+
         sortedstate = sorted(state)
         total = len(sortedstate)
         pos = 0
@@ -241,6 +244,13 @@
         if 'qtip' in repo.tags():
             updatemq(repo, state, skipped, **opts)
 
+        if currentbookmarks:
+            # Nodeids are needed to reset bookmarks
+            nstate = {}
+            for k, v in state.iteritems():
+                if v != nullmerge:
+                    nstate[repo[k].node()] = repo[v].node()
+
         if not keepf:
             # Remove no more useful revisions
             rebased = [rev for rev in state if state[rev] != nullmerge]
@@ -252,6 +262,9 @@
                     # backup the old csets by default
                     repair.strip(ui, repo, repo[min(rebased)].node(), "all")
 
+        if currentbookmarks:
+            updatebookmarks(repo, nstate, currentbookmarks, **opts)
+
         clearstatus(repo)
         ui.note(_("rebase completed\n"))
         if os.path.exists(repo.sjoin('undo')):
@@ -401,6 +414,18 @@
         mq.series_dirty = True
         mq.savedirty()
 
+def updatebookmarks(repo, nstate, originalbookmarks, **opts):
+    'Move bookmarks to their correct changesets'
+    current = repo._bookmarkcurrent
+    for k, v in originalbookmarks.iteritems():
+        if v in nstate:
+            if nstate[v] != nullmerge:
+                # reset the pointer if the bookmark was moved incorrectly
+                if k != current:
+                    repo._bookmarks[k] = nstate[v]
+
+    bookmarks.write(repo)
+
 def storestatus(repo, originalwd, target, state, collapse, keep, keepbranches,
                                                                 external):
     'Store the current status to allow recovery'
diff --git a/tests/test-rebase-bookmarks.t b/tests/test-rebase-bookmarks.t
new file mode 100644
--- /dev/null
+++ b/tests/test-rebase-bookmarks.t
@@ -0,0 +1,85 @@
+  $ cat >> $HGRCPATH <<EOF
+  > [extensions]
+  > graphlog=
+  > rebase=
+  > 
+  > [alias]
+  > tglog = log -G --template "{rev}: '{desc}' bookmarks: {bookmarks}\n"
+  > EOF
+
+Create a repo with several bookmarks
+  $ hg init a
+  $ cd a
+
+  $ echo a > a
+  $ hg ci -Am A
+  adding a
+
+  $ echo b > b
+  $ hg ci -Am B
+  adding b
+  $ hg book 'X'
+  $ hg book 'Y'
+ 
+  $ echo c > c
+  $ hg ci -Am C
+  adding c
+  $ hg book 'Z'
+
+  $ hg up -q 0
+
+  $ echo d > d
+  $ hg ci -Am D
+  adding d
+  created new head
+
+  $ hg tglog 
+  @  3: 'D' bookmarks:
+  |
+  | o  2: 'C' bookmarks: Y Z
+  | |
+  | o  1: 'B' bookmarks: X
+  |/
+  o  0: 'A' bookmarks:
+  
+ 
+Move only rebased bookmarks
+
+  $ cd ..
+  $ hg clone -q a a1
+
+  $ cd a1
+  $ hg up -q Z
+
+  $ hg rebase --detach -s Y -d 3
+  saved backup bundle to $TESTTMP/a1/.hg/strip-backup/*-backup.hg (glob)
+
+  $ hg tglog 
+  @  3: 'C' bookmarks: Y Z
+  |
+  o  2: 'D' bookmarks:
+  |
+  | o  1: 'B' bookmarks: X
+  |/
+  o  0: 'A' bookmarks:
+  
+Keep bookmarks to the correct rebased changeset
+
+  $ cd ..
+  $ hg clone -q a a2
+
+  $ cd a2
+  $ hg up -q Z
+
+  $ hg rebase -s 1 -d 3
+  saved backup bundle to $TESTTMP/a2/.hg/strip-backup/*-backup.hg (glob)
+
+  $ hg tglog 
+  @  3: 'C' bookmarks: Y Z
+  |
+  o  2: 'B' bookmarks: X
+  |
+  o  1: 'D' bookmarks:
+  |
+  o  0: 'A' bookmarks:
+