diff --git a/modules/core/src/main/java/org/apache/ignite/internal/util/GridPartitionStateMap.java b/modules/core/src/main/java/org/apache/ignite/internal/util/GridPartitionStateMap.java
index f5893c02f77b4..b256da9edcb1e 100644
--- a/modules/core/src/main/java/org/apache/ignite/internal/util/GridPartitionStateMap.java
+++ b/modules/core/src/main/java/org/apache/ignite/internal/util/GridPartitionStateMap.java
@@ -42,7 +42,19 @@ public class GridPartitionStateMap extends AbstractMap<Integer, GridDhtPartition
     private static final int BITS = Integer.SIZE -
         Integer.numberOfLeadingZeros(GridDhtPartitionState.values().length + 1);
 
-    /** */
+    /**
+     * Contains partition map.
+     * For a map containing 3 partitions with a size of 3 bits storage will be done this way:
+     * <pre>
+     *     +-----------+-----------+-----------+
+     *     |  p0 - 100 |  p1 - 011 |  p2 - 001 |
+     *     +---+---+---+---+---+---+---+---+---+
+     *     | 0 | 0 | 1 | 1 | 1 | 0 | 1 | 0 | 0 |
+     *     +---+---+---+---+---+---+---+---+---+
+     * </pre>
+     * The first element takes the first {@link GridPartitionStateMap#BITS} bits in reverse order,
+     * the second element next {@link GridPartitionStateMap#BITS} bits in reverse order, etc.
+     */
     private final BitSet states;
 
     /** */
@@ -52,25 +64,36 @@ public class GridPartitionStateMap extends AbstractMap<Integer, GridDhtPartition
     @Override public Set<Entry<Integer, GridDhtPartitionState>> entrySet() {
         return new AbstractSet<Entry<Integer, GridDhtPartitionState>>() {
             @Override public Iterator<Entry<Integer, GridDhtPartitionState>> iterator() {
-                final int size = states.isEmpty() ? 0 : (states.length() - 1) / BITS + 1;
-
                 return new Iterator<Entry<Integer, GridDhtPartitionState>>() {
-                    private int next;
+                    /** Current {@link GridPartitionStateMap#states} index. */
+                    private int idx;
+
+                    /** Current key value. */
                     private int cur;
 
                     @Override public boolean hasNext() {
-                        while(state(next) == null && next < size)
-                            next++;
+                        idx = states.nextSetBit(idx);
 
-                        return next < size;
+                        return idx != -1;
                     }
 
                     @Override public Entry<Integer, GridDhtPartitionState> next() {
                         if (!hasNext())
                             throw new NoSuchElementException();
 
-                        cur = next;
-                        next++;
+                        cur = idx / BITS;
+
+                        int bitN = idx % BITS;
+
+                        // Get value from BitSet like in GridPartitionStateMap#state, but don't process known bits.
+                        int st = 1 << bitN;
+
+                        for (int i = 1; i < BITS - bitN; i++)
+                            st |= (states.get(idx + i) ? 1 : 0) << i + bitN;
+
+                        final int ordinal = st - 1;
+
+                        idx += (BITS - bitN);
 
                         return new Entry<Integer, GridDhtPartitionState>() {
                             int p = cur;
@@ -80,7 +103,7 @@ public class GridPartitionStateMap extends AbstractMap<Integer, GridDhtPartition
                             }
 
                             @Override public GridDhtPartitionState getValue() {
-                                return state(p);
+                                return GridDhtPartitionState.fromOrdinal(ordinal);
                             }
 
                             @Override public GridDhtPartitionState setValue(GridDhtPartitionState val) {
@@ -219,4 +242,4 @@ private GridDhtPartitionState state(int part) {
     @Override public int hashCode() {
         return 31 * states.hashCode() + size;
     }
-}
\ No newline at end of file
+}
diff --git a/modules/core/src/test/java/org/apache/ignite/util/GridPartitionMapSelfTest.java b/modules/core/src/test/java/org/apache/ignite/util/GridPartitionMapSelfTest.java
index ebb4cd73b31e9..b099cf93e8633 100644
--- a/modules/core/src/test/java/org/apache/ignite/util/GridPartitionMapSelfTest.java
+++ b/modules/core/src/test/java/org/apache/ignite/util/GridPartitionMapSelfTest.java
@@ -19,8 +19,11 @@
 
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.Iterator;
 import java.util.Map;
 import java.util.Set;
+import java.util.TreeMap;
+import java.util.concurrent.ThreadLocalRandom;
 import org.apache.ignite.internal.processors.cache.distributed.dht.GridDhtPartitionState;
 import org.apache.ignite.internal.util.GridPartitionStateMap;
 import org.apache.ignite.testframework.junits.common.GridCommonAbstractTest;
@@ -144,6 +147,94 @@ public void testCopyNoActive() {
         assertEquals(1, cp2.size());
     }
 
+    /**
+     * Tests that entries from {@link Iterator#next()} remain unaltered.
+     */
+    public void testIteratorNext() {
+        GridPartitionStateMap map = new GridPartitionStateMap();
+
+        initMap(map);
+
+        Iterator<Map.Entry<Integer, GridDhtPartitionState>> iter = map.entrySet().iterator();
+
+        for (int i = 0; i < map.size() + 1; i++)
+            assertTrue(iter.hasNext());
+
+        Map.Entry<Integer, GridDhtPartitionState> entry1 = iter.next();
+
+        for (int i = 0; i < map.size() + 1; i++)
+            assertTrue(iter.hasNext());
+
+        Map.Entry<Integer, GridDhtPartitionState> entry2 = iter.next();
+
+        iter.remove();
+
+        assertNotNull(entry1.getValue());
+        assertNotNull(entry2.getValue());
+
+        assertEquals(Integer.valueOf(0), entry1.getKey());
+        assertEquals(Integer.valueOf(1), entry2.getKey());
+
+        assertEquals(GridDhtPartitionState.MOVING, entry1.getValue());
+        assertEquals(GridDhtPartitionState.RENTING, entry2.getValue());
+    }
+
+    /**
+     * Tests {@link GridDhtPartitionState} compatibility with {@link TreeMap} on random operations.
+     */
+    public void testOnRandomOperations() {
+        ThreadLocalRandom rnd = ThreadLocalRandom.current();
+
+        Map<Integer, GridDhtPartitionState> treeMap = new TreeMap<>();
+        Map<Integer, GridDhtPartitionState> gridMap = new GridPartitionStateMap();
+
+        int statesNum = GridDhtPartitionState.values().length;
+
+        for (int i = 0; i < 10_000; i++) {
+            Integer part = rnd.nextInt(1024);
+
+            GridDhtPartitionState state = GridDhtPartitionState.fromOrdinal(rnd.nextInt(statesNum));
+
+            int rndOperation = rnd.nextInt(9);
+
+            if (rndOperation <= 5) {
+                treeMap.put(part, state);
+                gridMap.put(part, state);
+            }
+            else if (rndOperation == 6) {
+                treeMap.remove(part);
+                gridMap.remove(part);
+            }
+            else if (treeMap.size() > 0) {
+                int n = rnd.nextInt(0, treeMap.size());
+
+                Iterator<Map.Entry<Integer, GridDhtPartitionState>> iter1 = treeMap.entrySet().iterator();
+                Iterator<Map.Entry<Integer, GridDhtPartitionState>> iter2 = gridMap.entrySet().iterator();
+
+                Map.Entry<Integer, GridDhtPartitionState> entry1 = iter1.next();
+                Map.Entry<Integer, GridDhtPartitionState> entry2 = iter2.next();
+
+                for (int j = 1; j <= n; j++) {
+                    entry1 = iter1.next();
+                    entry2 = iter2.next();
+                }
+
+                if (rndOperation == 7) {
+                    entry1.setValue(state);
+                    entry2.setValue(state);
+                }
+                else {
+                    iter1.remove();
+                    iter2.remove();
+                }
+            }
+
+            assertEquals(treeMap.size(), gridMap.size());
+        }
+
+        assertEquals(treeMap, gridMap);
+    }
+
     /** */
     private GridPartitionStateMap initMap(GridPartitionStateMap map) {
         map.put(0, GridDhtPartitionState.MOVING);
@@ -159,4 +250,4 @@ private GridPartitionStateMap initMap(GridPartitionStateMap map) {
 
         return map;
     }
-}
\ No newline at end of file
+}