--- a/hgext/obsolete.py	Tue Jun 19 17:05:39 2012 +0200
+++ b/hgext/obsolete.py	Tue Jun 19 18:05:23 2012 +0200
@@ -225,17 +225,23 @@
     rebaseset = repo.revs('%ld - extinct()', rebaseset)
     return orig(repo, dest, rebaseset, *ags, **kws)
 
+def defineparents(orig, repo, rev, target, state, *args, **kwargs):
+    rebasestate = getattr(repo, '_rebasestate', None)
+    if rebasestate is not None:
+        repo._rebasestate = dict(state)
+        repo._rebasetarget = target
+    return orig(repo, rev, target, state, *args, **kwargs)
 
-def concludenode(orig, repo, rev, *args, **kwargs):
+def concludenode(orig, repo, rev, p1, *args, **kwargs):
     """wrapper for rebase 's concludenode that set obsolete relation"""
-    newrev = orig(repo, rev, *args, **kwargs)
-    oldnode = repo[rev].node()
-    if newrev is not None:
-        newnode = repo[newrev].node()
-    else:
-        # Revision was emptied and removed, there is no successor.
-        newnode = nullid
-    repo.addobsolete(newnode, oldnode)
+    newrev = orig(repo, rev, p1, *args, **kwargs)
+    rebasestate = getattr(repo, '_rebasestate', None)
+    if rebasestate is not None:
+        if newrev is not None:
+            nrev = repo[newrev].rev()
+        else:
+            nrev = p1
+        repo._rebasestate[rev] = nrev
     return newrev
 
 def cmdrebase(orig, ui, repo, *args, **kwargs):
@@ -244,8 +250,45 @@
                            'extension'), hint=_("see 'hg help obsolete'"))
     kwargs = dict(kwargs)
     kwargs['keep'] = True
-    return orig(ui, repo, *args, **kwargs)
 
+    # We want to mark rebased revision as obsolete and set their
+    # replacements if any. Doing it in concludenode() prevents
+    # aborting the rebase, and is not called with all relevant
+    # revisions in --collapse case. Instead, we try to track the
+    # rebase state structure by sampling/updating it in
+    # defineparents() and concludenode(). The obsolete markers are
+    # added from this state after a successful call.
+    repo._rebasestate = {}
+    repo._rebasetarget = None
+    maxrev = len(repo) - 1
+    try:
+        res = orig(ui, repo, *args, **kwargs)
+        if not res and not kwargs.get('abort') and repo._rebasetarget:
+            # We have to tell rewritten revisions from removed
+            # ones. When collapsing, removed revisions are considered
+            # to be collapsed onto the final one, while in the normal
+            # case their are marked obsolete without successor.
+            emptynode = nullid
+            if kwargs.get('collapse'):
+                emptynode = repo[max(repo._rebasestate.values())].node()
+            # Rebased revisions are assumed to be descendants of
+            # targetrev. If a source revision is mapped to targetrev
+            # or to another rebased revision, it must have been
+            # removed.
+            targetrev = repo[repo._rebasetarget].rev()
+            newrevs = set([targetrev])
+            for rev, newrev in sorted(repo._rebasestate.items()):
+                oldnode = repo[rev].node()
+                if newrev not in newrevs and newrev >= 0:
+                    newnode = repo[newrev].node()
+                    newrevs.add(newrev)
+                else:
+                    newnode = emptynode
+                repo.addobsolete(newnode, oldnode)
+        return res
+    finally:
+        delattr(repo, '_rebasestate')
+        delattr(repo, '_rebasetarget')
 
 
 def extsetup(ui):
@@ -262,6 +305,7 @@
         rebase = extensions.find('rebase')
         if rebase:
             extensions.wrapfunction(rebase, 'buildstate', buildstate)
+            extensions.wrapfunction(rebase, 'defineparents', defineparents)
             extensions.wrapfunction(rebase, 'concludenode', concludenode)
             extensions.wrapcommand(rebase.cmdtable, "rebase", cmdrebase)
     except KeyError: