hgext/states.py
changeset 99 67a3aa020d91
parent 98 8ad5c760c708
child 100 7d0617e584ff
--- a/hgext/states.py	Sun Sep 25 12:46:45 2011 +0200
+++ b/hgext/states.py	Tue Sep 27 11:25:13 2011 +0200
@@ -1032,3 +1032,41 @@
 
     repo.__class__ = statefulrepo
 
+    import mercurial.repair
+    from operator import or_
+    ostrip = mercurial.repair.strip
+
+    def strip(ui, repo, node, backup="all"):
+        cl = repo.changelog
+
+        striprev = cl.rev(node)
+        revstostrip = set(cl.descendants(striprev))
+        revstostrip.add(striprev)
+        tostrip = set(cl.node(rev) for rev in revstostrip) # XXX stay in nodes?
+
+        baserevstostrip = revstostrip - set(cl.descendants(*revstostrip))
+        basetostrip = set(cl.node(rev) for rev in baserevstostrip) # XXX stay in nodes?
+        newstatesheads = reduce(or_,
+                                (set(cl.parents(n)) for n in basetostrip),
+                                set([]))
+
+        statesheads={} # use this cache as repo.nodestate shall work
+        for nd in newstatesheads:
+            state = repo.nodestate(nd)
+            if not state.trackheads:
+                continue
+            statesheads.setdefault(state, set([])).add(nd)
+
+        for state, heads in repo._statesheads.iteritems():
+            if not state.trackheads:
+                continue
+            heads = set(heads)
+            headstostrip = tostrip & heads
+            heads.difference_update(headstostrip)
+            heads.update(statesheads.get(state, set([])))
+            repo._statesheads[state] = list(heads)
+        _writestateshead(repo) # after ostrip ?
+
+
+        return ostrip(ui, repo, node, backup)
+    mercurial.repair.strip = strip