diff --git a/community/kernel/src/main/java/org/neo4j/kernel/impl/api/state/TxState.java b/community/kernel/src/main/java/org/neo4j/kernel/impl/api/state/TxState.java
index 6adcdff4c604..958b6a23e72b 100644
--- a/community/kernel/src/main/java/org/neo4j/kernel/impl/api/state/TxState.java
+++ b/community/kernel/src/main/java/org/neo4j/kernel/impl/api/state/TxState.java
@@ -28,6 +28,8 @@
 import java.util.function.Consumer;
 
 import org.neo4j.collection.primitive.Primitive;
+import org.neo4j.collection.primitive.PrimitiveIntObjectMap;
+import org.neo4j.collection.primitive.PrimitiveIntObjectVisitor;
 import org.neo4j.collection.primitive.PrimitiveIntSet;
 import org.neo4j.collection.primitive.PrimitiveLongIterator;
 import org.neo4j.collection.primitive.PrimitiveLongSet;
@@ -133,9 +135,9 @@ void setMap( TxState state, Map<Long, RelationshipStateImpl> map )
         }
     };
 
-    private Map<Integer/*Token ID*/, String> createdLabelTokens;
-    private Map<Integer/*Token ID*/, String> createdPropertyKeyTokens;
-    private Map<Integer/*Token ID*/, String> createdRelationshipTypeTokens;
+    private PrimitiveIntObjectMap<String> createdLabelTokens;
+    private PrimitiveIntObjectMap<String> createdPropertyKeyTokens;
+    private PrimitiveIntObjectMap<String> createdRelationshipTypeTokens;
 
     private GraphState graphState;
     private DiffSets<IndexDescriptor> indexChanges;
@@ -267,26 +269,17 @@ public void accept( final TxStateVisitor visitor )
 
         if ( createdLabelTokens != null )
         {
-            for ( Map.Entry<Integer, String> entry : createdLabelTokens.entrySet() )
-            {
-                visitor.visitCreatedLabelToken( entry.getValue(), entry.getKey() );
-            }
+            createdLabelTokens.visitEntries( new LabelTokenStateVisitor( visitor ) );
         }
 
         if ( createdPropertyKeyTokens != null )
         {
-            for ( Map.Entry<Integer, String> entry : createdPropertyKeyTokens.entrySet() )
-            {
-                visitor.visitCreatedPropertyKeyToken( entry.getValue(), entry.getKey() );
-            }
+            createdPropertyKeyTokens.visitEntries( new PropertyKeyTokenStateVisitor( visitor ) );
         }
 
         if ( createdRelationshipTypeTokens != null )
         {
-            for ( Map.Entry<Integer, String> entry : createdRelationshipTypeTokens.entrySet() )
-            {
-                visitor.visitCreatedRelationshipTypeToken( entry.getValue(), entry.getKey() );
-            }
+            createdRelationshipTypeTokens.visitEntries( new RelationshipTypeTokenStateVisitor( visitor ) );
         }
     }
 
@@ -673,7 +666,7 @@ public void labelDoCreateForName( String labelName, int id )
     {
         if ( createdLabelTokens == null )
         {
-            createdLabelTokens = new HashMap<>();
+            createdLabelTokens = Primitive.intObjectMap();
         }
         createdLabelTokens.put( id, labelName );
         changed();
@@ -684,7 +677,7 @@ public void propertyKeyDoCreateForName( String propertyKeyName, int id )
     {
         if ( createdPropertyKeyTokens == null )
         {
-            createdPropertyKeyTokens = new HashMap<>();
+            createdPropertyKeyTokens = Primitive.intObjectMap();
         }
         createdPropertyKeyTokens.put( id, propertyKeyName );
         changed();
@@ -695,7 +688,7 @@ public void relationshipTypeDoCreateForName( String labelName, int id )
     {
         if ( createdRelationshipTypeTokens == null )
         {
-            createdRelationshipTypeTokens = new HashMap<>();
+            createdRelationshipTypeTokens = Primitive.intObjectMap();
         }
         createdRelationshipTypeTokens.put( id, labelName );
         changed();
@@ -1352,4 +1345,55 @@ private void recordRelationshipDeleted( long id )
         }
         relationshipsDeletedInTx.add( id );
     }
+
+    private static class LabelTokenStateVisitor implements PrimitiveIntObjectVisitor<String,RuntimeException>
+    {
+        private final TxStateVisitor visitor;
+
+        LabelTokenStateVisitor( TxStateVisitor visitor )
+        {
+            this.visitor = visitor;
+        }
+
+        @Override
+        public boolean visited( int key, String value )
+        {
+            visitor.visitCreatedLabelToken( value, key );
+            return false;
+        }
+    }
+
+    private static class PropertyKeyTokenStateVisitor implements PrimitiveIntObjectVisitor<String,RuntimeException>
+    {
+        private final TxStateVisitor visitor;
+
+        PropertyKeyTokenStateVisitor( TxStateVisitor visitor )
+        {
+            this.visitor = visitor;
+        }
+
+        @Override
+        public boolean visited( int key, String value ) throws RuntimeException
+        {
+            visitor.visitCreatedPropertyKeyToken( value, key );
+            return false;
+        }
+    }
+
+    private static class RelationshipTypeTokenStateVisitor implements PrimitiveIntObjectVisitor<String,RuntimeException>
+    {
+        private final TxStateVisitor visitor;
+
+        RelationshipTypeTokenStateVisitor( TxStateVisitor visitor )
+        {
+            this.visitor = visitor;
+        }
+
+        @Override
+        public boolean visited( int key, String value ) throws RuntimeException
+        {
+            visitor.visitCreatedRelationshipTypeToken( value, key );
+            return false;
+        }
+    }
 }
diff --git a/community/kernel/src/test/java/org/neo4j/kernel/impl/newapi/IndexTxStateUpdaterTest.java b/community/kernel/src/test/java/org/neo4j/kernel/impl/newapi/IndexTxStateUpdaterTest.java
index baa4dc63bceb..fe224ff60ff9 100644
--- a/community/kernel/src/test/java/org/neo4j/kernel/impl/newapi/IndexTxStateUpdaterTest.java
+++ b/community/kernel/src/test/java/org/neo4j/kernel/impl/newapi/IndexTxStateUpdaterTest.java
@@ -85,11 +85,19 @@ public void setup()
 
         StoreReadLayer storeReadLayer = mock( StoreReadLayer.class );
         when( storeReadLayer.indexesGetAll() ).thenAnswer( x -> indexes.iterator() );
-        when( storeReadLayer.indexesGetForLabel(anyInt() ) )
-                .thenAnswer( x -> filter( hasLabel( x.getArgument( 0 ) ), indexes.iterator() ) );
+        when( storeReadLayer.indexesGetForLabel( anyInt() ) )
+                .thenAnswer( x ->
+                {
+                    Integer argument = x.getArgument( 0 );
+                    return filter( hasLabel( argument ), indexes.iterator() );
+                } );
 
         when( storeReadLayer.indexesGetRelatedToProperty( anyInt() ) )
-                .thenAnswer( x -> filter( hasProperty( x.getArgument( 0 ) ), indexes.iterator() ) );
+                .thenAnswer( x ->
+                {
+                    Integer argument = x.getArgument( 0 );
+                    return filter( hasProperty( argument ), indexes.iterator() );
+                } );
 
         HashMap<Integer,Value> map = new HashMap<>();
         map.put( propId1, Values.of( "hi1" ) );