From 80449d6eb05a75536208196d0ea95ab321dfc6f2 Mon Sep 17 00:00:00 2001
From: Andrei Koval <andrei.koval@neo4j.com>
Date: Mon, 29 Jan 2018 13:25:43 +0100
Subject: [PATCH] Removed TxState.LABEL_STATE

---
 .../kernel/impl/api/state/LabelState.java     | 95 -------------------
 .../neo4j/kernel/impl/api/state/TxState.java  | 44 ++++-----
 2 files changed, 21 insertions(+), 118 deletions(-)
 delete mode 100644 community/kernel/src/main/java/org/neo4j/kernel/impl/api/state/LabelState.java

diff --git a/community/kernel/src/main/java/org/neo4j/kernel/impl/api/state/LabelState.java b/community/kernel/src/main/java/org/neo4j/kernel/impl/api/state/LabelState.java
deleted file mode 100644
index 80e82b843aaaf..0000000000000
--- a/community/kernel/src/main/java/org/neo4j/kernel/impl/api/state/LabelState.java
+++ /dev/null
@@ -1,95 +0,0 @@
-/*
- * Copyright (c) 2002-2018 "Neo Technology,"
- * Network Engine for Objects in Lund AB [http://neotechnology.com]
- *
- * This file is part of Neo4j.
- *
- * Neo4j is free software: you can redistribute it and/or modify
- * it under the terms of the GNU General Public License as published by
- * the Free Software Foundation, either version 3 of the License, or
- * (at your option) any later version.
- *
- * This program is distributed in the hope that it will be useful,
- * but WITHOUT ANY WARRANTY; without even the implied warranty of
- * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- * GNU General Public License for more details.
- *
- * You should have received a copy of the GNU General Public License
- * along with this program.  If not, see <http://www.gnu.org/licenses/>.
- */
-package org.neo4j.kernel.impl.api.state;
-
-import org.neo4j.kernel.impl.util.diffsets.DiffSets;
-import org.neo4j.storageengine.api.txstate.ReadableDiffSets;
-
-/**
- * Represents the transactional changes that involve a particular label:
- * <ul>
- * <li>{@linkplain #nodeDiffSets() Nodes} where the label has been {@linkplain ReadableDiffSets#getAdded() added}
- * or {@linkplain ReadableDiffSets#getRemoved() removed}.</li>
- * </ul>
- */
-public abstract class LabelState
-{
-    public abstract ReadableDiffSets<Long> nodeDiffSets();
-
-    public static class Mutable extends LabelState
-    {
-        private DiffSets<Long> nodeDiffSets;
-        private final long labelId;
-
-        private Mutable( long labelId )
-        {
-            this.labelId = labelId;
-        }
-
-        public long getLabelId()
-        {
-            return labelId;
-        }
-
-        @Override
-        public ReadableDiffSets<Long> nodeDiffSets()
-        {
-            return ReadableDiffSets.Empty.ifNull( nodeDiffSets );
-        }
-
-        public DiffSets<Long> getOrCreateNodeDiffSets()
-        {
-            if ( nodeDiffSets == null )
-            {
-                nodeDiffSets = new DiffSets<>();
-            }
-            return nodeDiffSets;
-        }
-    }
-
-    abstract static class Defaults extends StateDefaults<LabelState, Mutable>
-    {
-        @Override
-        Mutable createValue( long key, TxState state )
-        {
-            return new Mutable( key );
-        }
-
-        @Override
-        LabelState defaultValue()
-        {
-            return DEFAULT;
-        }
-    }
-
-    private static final LabelState DEFAULT = new LabelState()
-    {
-        @Override
-        public ReadableDiffSets<Long> nodeDiffSets()
-        {
-            return ReadableDiffSets.Empty.instance();
-        }
-    };
-
-    private LabelState()
-    {
-        // limited subclasses
-    }
-}
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 b17fb04ec1ffe..b88550d530ad8 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
@@ -95,11 +95,10 @@
  */
 public class TxState implements TransactionState, RelationshipVisitor.Home
 {
-    private static final LabelState.Defaults LABEL_STATE = new TransactionLabelState();
     private static final NodeStateImpl.Defaults NODE_STATE = new TransactionNodeState();
     private static final RelationshipStateImpl.Defaults RELATIONSHIP_STATE = new TransactionRelationshipState();
 
-    private PrimitiveLongObjectMap<LabelState.Mutable> labelStatesMap;
+    private PrimitiveIntObjectMap<DiffSets<Long>> labelStatesMap;
     private PrimitiveLongObjectMap<NodeStateImpl> nodeStatesMap;
     private PrimitiveLongObjectMap<RelationshipStateImpl> relationshipStatesMap;
 
@@ -358,7 +357,21 @@ public Iterable<NodeState> modifiedNodes()
 
     private DiffSets<Long> getOrCreateLabelStateNodeDiffSets( int labelId )
     {
-        return LABEL_STATE.getOrCreate( this, labelId ).getOrCreateNodeDiffSets();
+        if ( labelStatesMap == null )
+        {
+            labelStatesMap = Primitive.intObjectMap();
+        }
+        return labelStatesMap.computeIfAbsent( labelId, unused -> new DiffSets<>() );
+    }
+
+    private ReadableDiffSets<Long> getLabelStateNodeDiffSets( int labelId )
+    {
+        if ( labelStatesMap == null )
+        {
+            return ReadableDiffSets.Empty.instance();
+        }
+        final DiffSets<Long> nodeDiffSets = labelStatesMap.get( labelId );
+        return ReadableDiffSets.Empty.ifNull( nodeDiffSets );
     }
 
     @Override
@@ -706,7 +719,7 @@ public Cursor<RelationshipItem> augmentRelationshipsGetAllCursor( Cursor<Relatio
     @Override
     public ReadableDiffSets<Long> nodesWithLabelChanged( int label )
     {
-        return LABEL_STATE.get( this, label ).nodeDiffSets();
+        return getLabelStateNodeDiffSets( label );
     }
 
     @Override
@@ -718,7 +731,7 @@ public ReadableDiffSets<Long> nodesWithAnyOfLabelsChanged( int... labels )
         Set<Long> removed = new HashSet<>();
         for ( int i = 0; i < labels.length; i++ )
         {
-            ReadableDiffSets<Long> nodeDiffSets = LABEL_STATE.get( this, labels[i] ).nodeDiffSets();
+            ReadableDiffSets<Long> nodeDiffSets = getLabelStateNodeDiffSets( labels[i] );
             if ( i == 0 )
             {
                 removed.addAll( nodeDiffSets.getRemoved() );
@@ -739,9 +752,9 @@ public ReadableDiffSets<Long> nodesWithAllLabelsChanged( int... labels )
         DiffSets<Long> changes = new DiffSets<>();
         for ( int label : labels )
         {
-            LabelState labelState = LABEL_STATE.get( this, label );
-            changes.addAll( labelState.nodeDiffSets().getAdded().iterator() );
-            changes.removeAll( labelState.nodeDiffSets().getRemoved().iterator() );
+            final ReadableDiffSets<Long> nodeDiffSets = getLabelStateNodeDiffSets( label );
+            changes.addAll( nodeDiffSets.getAdded().iterator() );
+            changes.removeAll( nodeDiffSets.getRemoved().iterator() );
         }
         return changes;
     }
@@ -1333,21 +1346,6 @@ public void visitRemoved( ConstraintDescriptor constraint ) throws ConstraintVal
         }
     }
 
-    private static class TransactionLabelState extends LabelState.Defaults
-    {
-        @Override
-        PrimitiveLongObjectMap<LabelState.Mutable> getMap( TxState state )
-        {
-            return state.labelStatesMap;
-        }
-
-        @Override
-        void setMap( TxState state, PrimitiveLongObjectMap<LabelState.Mutable> map )
-        {
-            state.labelStatesMap = map;
-        }
-    }
-
     private static class TransactionNodeState extends NodeStateImpl.Defaults
     {
         @Override