--- a/states.py	Wed May 25 02:04:12 2011 +0200
+++ b/states.py	Wed May 25 02:13:02 2011 +0200
@@ -153,7 +153,7 @@
     # Write protocols
     ####################
     def heads(repo, proto):
-        h = repo._readyheads
+        h = repo.stateheads(ST1)
         return wireproto.encodelist(h) + "\n"
 
     def _reducehead(wirerepo, heads):
@@ -173,27 +173,21 @@
 
         def nodestate(self, node):
             rev = self.changelog.rev(node)
-            for head in self._readyheads:
-                revhead = self.changelog.rev(head)
-                if self.changelog.descendant(revhead, rev):
-                    return STATES[2]
-            for head in self._publishedheads:
-                revhead = self.changelog.rev(head)
-                if self.changelog.descendant(revhead, rev):
-                    return STATES[1]
-            return STATES[0]
+            for state in STATES[::-1]:
+                # XXX avoid for untracked heads
+                if state.next is not None:
+                    for head in self.stateheads(state):
+                        revhead = self.changelog.rev(head)
+                        if self.changelog.descendant(revhead, rev):
+                            return state.next
+            return state
 
 
-        @property
-        def _readyheads(self):
-            if self.ui.configbool('states', ST1.next.name, False):
-                return self._statesheads[ST1]
-            return self.heads()
 
-        @property
-        def _publishedheads(self):
-            if self.ui.configbool('states', ST0.next.name, False):
-                return self._statesheads[ST0]
+        def stateheads(self, state):
+            if state.trackheads:
+                if self.ui.configbool('states', state.next.name, False):
+                    return self._statesheads[state]
             return self.heads()
 
         @util.propertycache
@@ -258,7 +252,7 @@
             for candidate in candidates:
                 rev = self.changelog.rev(candidate)
                 ok = True
-                for h in self._readyheads:
+                for h in self.stateheads(ST1):
                     revh = self.changelog.rev(h)
                     if self.changelog.descendant(revh, rev):
                         ok = False
@@ -268,7 +262,7 @@
             return sorted(selected)
 
         def cancopy(self):
-            return o_cancopy() and (self._readyheads == self.heads())
+            return o_cancopy() and (self.stateheads(ST1) == self.heads())
 
     repo.__class__ = statefulrepo