--- a/hgext/states.py	Thu Sep 08 17:46:54 2011 +0200
+++ b/hgext/states.py	Thu Sep 08 18:20:01 2011 +0200
@@ -33,6 +33,7 @@
 from mercurial import extensions
 from mercurial import wireproto
 from mercurial import pushkey
+from mercurial import error
 from mercurial.lock import release
 
 
@@ -69,7 +70,12 @@
         assert self.trackheads
         def revsetheads(repo, subset, x):
             args = revset.getargs(x, 0, 0, 'publicheads takes no arguments')
-            heads = map(repo.changelog.rev, repo._statesheads[self])
+            heads = []
+            for h in repo._statesheads[self]:
+                try:
+                    heads.append(repo.changelog.rev(h))
+                except error.LookupError:
+                    pass
             heads.sort()
             return heads
         return revsetheads
@@ -175,8 +181,9 @@
     st = ST0
     w = repo.wlock()
     try:
-        #print 'pushing', key
-        repo.setstate(ST0, [node.bin(key)])
+        newhead = node.bin(key)
+        repo[newhead]
+        repo.setstate(ST0, [newhead])
     finally:
         w.release()
 
@@ -389,9 +396,12 @@
 
         def _pushimmutableheads(self, remote, remoteheads):
             missing = set(self.stateheads(ST0)) - set(remoteheads)
-            for h in missing:
-                #print 'missing', node.short(h)
-                remote.pushkey('immutableheads', node.hex(h), '', '1')
+            while missing:
+                h = missing.pop()
+                try:
+                    remote.pushkey('immutableheads', node.hex(h), '', '1')
+                except error.RepoLookupError:
+                    missing.update(p.node() for p in repo[h].parents())
 
 
         def _pullimmutableheads(self, remote):