--- a/hgext/states.py	Thu Sep 08 18:20:01 2011 +0200
+++ b/hgext/states.py	Thu Sep 08 19:03:42 2011 +0200
@@ -93,6 +93,7 @@
 ST0 = state('published', next=ST1)
 
 STATES = (ST0, ST1, ST2)
+STATESMAP =dict([(st.name, st) for st in STATES])
 
 @util.cachefunc
 def laststatewithout(prop):
@@ -177,20 +178,28 @@
 # Pushkey mechanism for mutable
 #########################################
 
-def pushimmutableheads(repo, key, old, new):
-    st = ST0
+def pushstatesheads(repo, key, old, new):
+    st = STATESMAP[new]
     w = repo.wlock()
     try:
         newhead = node.bin(key)
         repo[newhead]
-        repo.setstate(ST0, [newhead])
+        repo.setstate(st, [newhead])
     finally:
         w.release()
 
-def listimmutableheads(repo):
-    return dict.fromkeys(map(node.hex, repo.stateheads(ST0)), '1')
+def liststatesheads(repo):
+    keys = {}
+    for state in [st for st in STATES if st.trackheads]:
+        for head in repo.stateheads(state):
+            head = node.hex(head)
+            if head in keys:
+                keys[head] += ',' + state.name
+            else:
+                keys[head] = state.name
+    return keys
 
-pushkey.register('immutableheads', pushimmutableheads, listimmutableheads)
+pushkey.register('states-heads', pushstatesheads, liststatesheads)
 
 
 
@@ -311,7 +320,7 @@
             """change state of targets changeset and it's ancestors.
 
             Simplify the list of head."""
-            assert not isinstance(nodes, basestring)
+            assert not isinstance(nodes, basestring), repr(nodes)
             heads = self._statesheads[state]
             olds = heads[:]
             heads.extend(nodes)
@@ -378,42 +387,45 @@
 
         def pull(self, remote, *args, **kwargs):
             result = opull(remote, *args, **kwargs)
-            remoteheads = self._pullimmutableheads(remote)
+            remoteheads = self._pullstatesheads(remote)
             #print [node.short(h) for h in remoteheads]
-            self.setstate(ST0, remoteheads)
+            for st, heads in remoteheads.iteritems():
+                self.setstate(st, heads)
             return result
 
         def push(self, remote, *args, **opts):
             result = opush(remote, *args, **opts)
-            remoteheads = self._pullimmutableheads(remote)
-            self.setstate(ST0, remoteheads)
-            if remoteheads != self.stateheads(ST0):
-                #print 'stuff to push'
-                #print 'remote', [node.short(h) for h in remoteheads]
-                #print 'local',  [node.short(h) for h in self._statesheads[ST0]]
-                self._pushimmutableheads(remote, remoteheads)
+            remoteheads = self._pullstatesheads(remote)
+            for st, heads in remoteheads.iteritems():
+                self.setstate(st, heads)
+                if heads != self.stateheads(st):
+                    self._pushstatesheads(remote, st,  heads)
             return result
 
-        def _pushimmutableheads(self, remote, remoteheads):
-            missing = set(self.stateheads(ST0)) - set(remoteheads)
+        def _pushstatesheads(self, remote, state, remoteheads):
+            local = set(self.stateheads(state))
+            missing = local - set(remoteheads)
             while missing:
                 h = missing.pop()
                 try:
-                    remote.pushkey('immutableheads', node.hex(h), '', '1')
+                    remote.pushkey('states-heads', node.hex(h), '', state.name)
                 except error.RepoLookupError:
                     missing.update(p.node() for p in repo[h].parents())
 
 
-        def _pullimmutableheads(self, remote):
-            self.ui.debug('checking for immutableheadshg on server')
-            if 'immutableheads' not in remote.listkeys('namespaces'):
-                self.ui.debug('immutableheads not enabled on the remote server, '
-                              'marking everything as frozen')
-                remote = remote.heads()
+        def _pullstatesheads(self, remote):
+            remoteheads = {}
+            self.ui.debug('checking for states-heads on remote server')
+            if 'states-heads' not in remote.listkeys('namespaces'):
+                self.ui.debug('states-heads not enabled on the remote server, '
+                              'marking everything as published')
+                remoteheads[ST0] = remote.heads()
             else:
-                self.ui.debug('server has immutableheads enabled, merging lists')
-                remote = map(node.bin, remote.listkeys('immutableheads'))
-            return remote
+                self.ui.debug('server has states-heads enabled, merging lists')
+                for hex, statenames in  remote.listkeys('states-heads').iteritems():
+                    for stn in statenames.split(','):
+                        remoteheads.setdefault(STATESMAP[stn], []).append(node.bin(hex))
+            return remoteheads
 
         ### Tag support