--- a/mercurial/localrepo.py	Fri Jun 01 08:56:17 2012 -0700
+++ b/mercurial/localrepo.py	Fri May 18 12:45:47 2012 -0700
@@ -550,6 +550,9 @@
                     continue
                 node, label = l.split(" ", 1)
                 label = encoding.tolocal(label.strip())
+                if not node in self:
+                    raise ValueError('invalidating branch cache because node '+
+                                     '%s does not exist' % node)
                 partial.setdefault(label, []).append(bin(node))
         except KeyboardInterrupt:
             raise
@@ -571,22 +574,46 @@
             pass
 
     def _updatebranchcache(self, partial, ctxgen):
+        """Given a branchhead cache, partial, that may have extra nodes or be
+        missing heads, and a generator of nodes that are at least a superset of
+        heads missing, this function updates partial to be correct.
+        """
         # collect new branch entries
         newbranches = {}
         for c in ctxgen:
-            newbranches.setdefault(c.branch(), []).append(c.rev())
+            newbranches.setdefault(c.branch(), []).append(c.node())
         # if older branchheads are reachable from new ones, they aren't
         # really branchheads. Note checking parents is insufficient:
         # 1 (branch a) -> 2 (branch b) -> 3 (branch a)
-        for branch, newrevs in newbranches.iteritems():
-            bheadrevs = [self.changelog.rev(node) for node in
-                         partial.setdefault(branch, [])]
-            bheadrevs.extend(newrevs)
-            bheadrevs.sort()
-            # starting from tip means fewer passes over ancestors
-            newrevs.sort()
-            while newrevs:
-                latest = newrevs.pop()
+        for branch, newnodes in newbranches.iteritems():
+            bheads = partial.setdefault(branch, [])
+            # Remove candidate heads that no longer are in the repo (e.g., as
+            # the result of a strip that just happened).  Avoid using 'node in
+            # self' here because that dives down into branchcache code somewhat
+            # recrusively.
+            bheadrevs = [self.changelog.rev(node) for node in bheads
+                         if self.changelog.hasnode(node)]
+            newheadrevs = [self.changelog.rev(node) for node in newnodes
+                           if self.changelog.hasnode(node)]
+            ctxisnew = bheadrevs and min(newheadrevs) > max(bheadrevs)
+            # Remove duplicates - nodes that are in newheadrevs and are already
+            # in bheadrevs.  This can happen if you strip a node whose parent
+            # was already a head (because they're on different branches).
+            bheadrevs = sorted(set(bheadrevs).union(newheadrevs))
+
+            # Starting from tip means fewer passes over reachable.  If we know
+            # the new candidates are not ancestors of existing heads, we don't
+            # have to examine ancestors of existing heads
+            if ctxisnew:
+                iterrevs = sorted(newheadrevs)
+            else:
+                iterrevs = list(bheadrevs)
+
+            # This loop prunes out two kinds of heads - heads that are
+            # superceded by a head in newheadrevs, and newheadrevs that are not
+            # heads because an existing head is their descendant.
+            while iterrevs:
+                latest = iterrevs.pop()
                 if latest not in bheadrevs:
                     continue
                 ancestors = set(self.changelog.ancestors([latest],
@@ -595,6 +622,17 @@
                     bheadrevs = [b for b in bheadrevs if b not in ancestors]
             partial[branch] = [self.changelog.node(rev) for rev in bheadrevs]
 
+        # There may be branches that cease to exist when the last commit in the
+        # branch was stripped.  This code filters them out.  Note that the
+        # branch that ceased to exist may not be in newbranches because
+        # newbranches is the set of candidate heads, which when you strip the
+        # last commit in a branch will be the parent branch.
+        for branch in partial:
+            nodes = [head for head in partial[branch]
+                     if self.changelog.hasnode(head)]
+            if not nodes:
+                del partial[branch]
+
     def lookup(self, key):
         return self[key].node()
 
@@ -857,6 +895,9 @@
             else:
                 ui.status(_('working directory now based on '
                             'revision %d\n') % parents)
+        # TODO: if we know which new heads may result from this rollback, pass
+        # them to destroy(), which will prevent the branchhead cache from being
+        # invalidated.
         self.destroyed()
         return 0
 
@@ -1303,12 +1344,27 @@
                 tr.release()
             lock.release()
 
-    def destroyed(self):
+    def destroyed(self, newheadnodes=None):
         '''Inform the repository that nodes have been destroyed.
         Intended for use by strip and rollback, so there's a common
-        place for anything that has to be done after destroying history.'''
-        # XXX it might be nice if we could take the list of destroyed
-        # nodes, but I don't see an easy way for rollback() to do that
+        place for anything that has to be done after destroying history.
+
+        If you know the branchheadcache was uptodate before nodes were removed
+        and you also know the set of candidate new heads that may have resulted
+        from the destruction, you can set newheadnodes.  This will enable the
+        code to update the branchheads cache, rather than having future code
+        decide it's invalid and regenrating it from scratch.
+        '''
+        # If we have info, newheadnodes, on how to update the branch cache, do
+        # it, Otherwise, since nodes were destroyed, the cache is stale and this
+        # will be caught the next time it is read.
+        if newheadnodes:
+            tiprev = len(self) - 1
+            ctxgen = (self[node] for node in newheadnodes
+                      if self.changelog.hasnode(node))
+            self._updatebranchcache(self._branchcache, ctxgen)
+            self._writebranchcache(self._branchcache, self.changelog.tip(),
+                                   tiprev)
 
         # Ensure the persistent tag cache is updated.  Doing it now
         # means that the tag cache only has to worry about destroyed

--- a/mercurial/repair.py	Fri Jun 01 08:56:17 2012 -0700
+++ b/mercurial/repair.py	Fri May 18 12:45:47 2012 -0700
@@ -56,6 +56,11 @@
     return s
 
 def strip(ui, repo, nodelist, backup="all", topic='backup'):
+    # It simplifies the logic around updating the branchheads cache if we only
+    # have to consider the effect of the stripped revisions and not revisions
+    # missing because the cache is out-of-date.
+    repo.updatebranchcache()
+
     cl = repo.changelog
     # TODO handle undo of merge sets
     if isinstance(nodelist, str):
@@ -63,6 +68,17 @@
     striplist = [cl.rev(node) for node in nodelist]
     striprev = min(striplist)
 
+    # Generate set of branches who will have nodes stripped.
+    striprevs = repo.revs("%ld::", striplist)
+    stripbranches = set([repo[rev].branch() for rev in striprevs])
+
+    # Set of potential new heads resulting from the strip.  The parents of any
+    # node removed could be a new head because the node to be removed could have
+    # been the only child of the parent.
+    newheadrevs = repo.revs("parents(%ld::) - %ld::", striprevs, striprevs)
+    newheadnodes = set([cl.node(rev) for rev in newheadrevs])
+    newheadbranches = set([repo[rev].branch() for rev in newheadrevs])
+
     keeppartialbundle = backup == 'strip'
 
     # Some revisions with rev > striprev may not be descendants of striprev.
@@ -169,4 +185,11 @@
                     % chgrpfile)
         raise
 
-    repo.destroyed()
+    if len(stripbranches) == 1 and len(newheadbranches) == 1 \
+            and stripbranches == newheadbranches:
+        repo.destroyed(newheadnodes)
+    else:
+        # Multiple branches involved in strip. Will allow branchcache to become
+        # invalid and later on rebuilt from scratch
+        repo.destroyed()
+

--- a/tests/test-keyword.t	Fri Jun 01 08:56:17 2012 -0700
+++ b/tests/test-keyword.t	Fri May 18 12:45:47 2012 -0700
@@ -502,7 +502,6 @@
   $ hg -q commit -d '1 14' -m 'prepare amend'
 
   $ hg --debug commit --amend -d '1 15' -m 'amend without changes' | grep keywords
-  invalidating branch cache (tip differs)
   overwriting a expanding keywords
   $ hg -q id
   a71343332ea9

--- a/tests/test-mq-caches.t	Fri Jun 01 08:56:17 2012 -0700
+++ b/tests/test-mq-caches.t	Fri May 18 12:45:47 2012 -0700
@@ -36,7 +36,8 @@
   $ hg qrefresh -m 'patch 1'
   $ show_branch_cache
   tip: 0
-  No branch cache
+  d986d5caac23a7d44a46efc0ddaf5eb9665844cf 0
+  d986d5caac23a7d44a46efc0ddaf5eb9665844cf default
 
 some regular revisions
 
@@ -75,8 +76,8 @@
   $ hg qrefresh -m 'patch 2'
   $ show_branch_cache 1
   tip: 3
-  c229711f16da3d7591f89b1b8d963b79bda22714 1
-  c229711f16da3d7591f89b1b8d963b79bda22714 bar
+  982611f6955f9c48d3365decea203217c945ef0d 2
+  982611f6955f9c48d3365decea203217c945ef0d bar
   dc25e3827021582e979f600811852e36cbe57341 foo
   branch foo: 3
   branch bar: 2

--- a/tests/test-rebase-cache.t	Fri Jun 01 08:56:17 2012 -0700
+++ b/tests/test-rebase-cache.t	Fri May 18 12:45:47 2012 -0700
@@ -2,6 +2,7 @@
   > [extensions]
   > graphlog=
   > rebase=
+  > mq=
   > 
   > [phases]
   > publish=False
@@ -262,4 +263,125 @@
   
   $ hg verify -q
 
+Stripping multiple branches in one go bypasses the fast-case code to
+update the branch cache.
+
+  $ hg strip 2
+  0 files updated, 0 files merged, 4 files removed, 0 files unresolved
+  saved backup bundle to $TESTTMP/a3/.hg/strip-backup/*-backup.hg (glob)
+
+  $ hg tglog
+  o  3: 'C' branch2
+  |
+  o  2: 'branch2' branch2
+  |
+  | @  1: 'branch1' branch1
+  |/
+  o  0: 'A'
+  
+
+  $ hg branches
+  branch2                        3:e4fdb121d036
+  branch1                        1:63379ac49655
+  default                        0:1994f17a630e (inactive)
+
+  $ hg theads
+  3: 'C' branch2
+  1: 'branch1' branch1
+  0: 'A' 
+
+Fast path branchcache code should not be invoked if branches stripped is not
+the same as branches remaining.
+
+  $ hg init b
+  $ cd b
+
+  $ hg branch branch1
+  marked working directory as branch branch1
+  (branches are permanent and global, did you want a bookmark?)
+  $ hg ci -m 'branch1'
+
+  $ hg branch branch2
+  marked working directory as branch branch2
+  (branches are permanent and global, did you want a bookmark?)
+  $ hg ci -m 'branch2'
+
+  $ hg branch -f branch1
+  marked working directory as branch branch1
+  (branches are permanent and global, did you want a bookmark?)
+
+  $ echo a > A
+  $ hg ci -Am A
+  adding A
+  created new head
+
+  $ hg tglog
+  @  2: 'A' branch1
+  |
+  o  1: 'branch2' branch2
+  |
+  o  0: 'branch1' branch1
+  
+
+  $ hg theads
+  2: 'A' branch1
+  1: 'branch2' branch2
+
+  $ hg strip 2
+  0 files updated, 0 files merged, 1 files removed, 0 files unresolved
+  saved backup bundle to $TESTTMP/a3/b/.hg/strip-backup/*-backup.hg (glob)
+
+  $ hg theads
+  1: 'branch2' branch2
+  0: 'branch1' branch1
+
+
+Make sure requesting to strip a revision already stripped does not confuse things.
+Try both orders.
+
   $ cd ..
+
+  $ hg init c
+  $ cd c
+
+  $ echo a > a
+  $ hg ci -Am A
+  adding a
+  $ echo b > b
+  $ hg ci -Am B
+  adding b
+  $ echo c > c
+  $ hg ci -Am C
+  adding c
+  $ echo d > d
+  $ hg ci -Am D
+  adding d
+  $ echo e > e
+  $ hg ci -Am E
+  adding e
+
+  $ hg tglog
+  @  4: 'E'
+  |
+  o  3: 'D'
+  |
+  o  2: 'C'
+  |
+  o  1: 'B'
+  |
+  o  0: 'A'
+  
+
+  $ hg strip 3 4
+  0 files updated, 0 files merged, 2 files removed, 0 files unresolved
+  saved backup bundle to $TESTTMP/a3/c/.hg/strip-backup/*-backup.hg (glob)
+
+  $ hg theads
+  2: 'C' 
+
+  $ hg strip 2 1
+  0 files updated, 0 files merged, 2 files removed, 0 files unresolved
+  saved backup bundle to $TESTTMP/a3/c/.hg/strip-backup/*-backup.hg (glob)
+
+  $ hg theads
+  0: 'A' 

--- a/tests/test-rebase-collapse.t	Fri Jun 01 08:56:17 2012 -0700
+++ b/tests/test-rebase-collapse.t	Fri May 18 12:45:47 2012 -0700
@@ -2,6 +2,7 @@
   > [extensions]
   > graphlog=
   > rebase=
+  > mq=
   > 
   > [phases]
   > publish=False
@@ -259,6 +260,44 @@
   $ cd ..
 
 
+
+
+Test that branchheads cache is updated correctly when doing a strip in which
+the parent of the ancestor node to be stripped does not become a head and
+also, the parent of a node that is a child of the node stripped becomes a head
+(node 3).
+
+  $ hg clone -q -u . b b2
+  $ cd b2
+
+  $ hg heads --template="{rev}:{node} {branch}\n"
+  7:c65502d4178782309ce0574c5ae6ee9485a9bafa default
+  6:c772a8b2dc17629cec88a19d09c926c4814b12c7 default
+
+  $ cat $TESTTMP/b2/.hg/cache/branchheads
+  c65502d4178782309ce0574c5ae6ee9485a9bafa 7
+  c772a8b2dc17629cec88a19d09c926c4814b12c7 default
+  c65502d4178782309ce0574c5ae6ee9485a9bafa default
+
+  $ hg strip 4
+  saved backup bundle to $TESTTMP/b2/.hg/strip-backup/8a5212ebc852-backup.hg
+
+  $ cat $TESTTMP/b2/.hg/cache/branchheads
+  c65502d4178782309ce0574c5ae6ee9485a9bafa 4
+  2870ad076e541e714f3c2bc32826b5c6a6e5b040 default
+  c65502d4178782309ce0574c5ae6ee9485a9bafa default
+
+  $ hg heads --template="{rev}:{node} {branch}\n"
+  4:c65502d4178782309ce0574c5ae6ee9485a9bafa default
+  3:2870ad076e541e714f3c2bc32826b5c6a6e5b040 default
+
+  $ cd ..
+
+
+
+
+
+
 Create repo c:
 
   $ hg init c
@@ -628,3 +667,56 @@
   b
   b
   $ cd ..
+
+
+Test stripping a revision with another child
+
+  $ hg init f
+  $ cd f
+
+  $ echo A > A
+  $ hg ci -Am A
+  adding A
+  $ echo B > B
+  $ hg ci -Am B
+  adding B
+
+  $ hg up -q 0
+
+  $ echo C > C
+  $ hg ci -Am C
+  adding C
+  created new head
+
+  $ hg tglog
+  @  2: 'C'
+  |
+  | o  1: 'B'
+  |/
+  o  0: 'A'
+  
+
+
+  $ hg heads --template="{rev}:{node} {branch}: {desc}\n"
+  2:c5cefa58fd557f84b72b87f970135984337acbc5 default: C
+  1:27547f69f25460a52fff66ad004e58da7ad3fb56 default: B
+
+  $ hg strip 2
+  0 files updated, 0 files merged, 1 files removed, 0 files unresolved
+  saved backup bundle to $TESTTMP/f/.hg/strip-backup/*-backup.hg (glob)
+
+  $ hg tglog
+  o  1: 'B'
+  |
+  @  0: 'A'
+  
+
+
+  $ hg heads --template="{rev}:{node} {branch}: {desc}\n"
+  1:27547f69f25460a52fff66ad004e58da7ad3fb56 default: B
+
+  $ cd ..
+
+
+
+