Use a simpler representation for NodeMap with only one value

flow-server/src/main/java/com/vaadin/flow/internal/nodefeature/NodeMap.java

@@ -25,6 +25,7 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.function.Consumer;
+import java.util.stream.Stream;
 
 import com.vaadin.flow.internal.StateNode;
 import com.vaadin.flow.internal.change.EmptyChange;
@@ -43,7 +44,107 @@ public abstract class NodeMap extends NodeFeature {
     private static final Serializable REMOVED_MARKER = new UniqueSerializable() {
     };
 
-    private Map<String, Serializable> values;
+    private interface Values extends Serializable {
+        int size();
+
+        Serializable get(String key);
+
+        Set<String> keySet();
+
+        boolean containsKey(String key);
+
+        default boolean isEmpty() {
+            return size() == 0;
+        }
+
+        Stream<Serializable> streamValues();
+
+        // Named set instead of put to avoid incompatibility with HashMap where
+        // put returns the previous value
+        void set(String key, Serializable value);
+    }
+
+    private static class SingleValue implements Values {
+
+        private final String key;
+
+        private Serializable value;
+
+        public SingleValue(String key, Serializable value) {
+            assert key != null;
+            this.key = key;
+            this.value = value;
+        }
+
+        @Override
+        public int size() {
+            return 1;
+        }
+
+        @Override
+        public Serializable get(String key) {
+            if (containsKey(key)) {
+                return value;
+            } else {
+                return null;
+            }
+        }
+
+        @Override
+        public Set<String> keySet() {
+            return Collections.singleton(key);
+        }
+
+        @Override
+        public boolean containsKey(String key) {
+            return this.key.equals(key);
+        }
+
+        @Override
+        public Stream<Serializable> streamValues() {
+            return Stream.of(value);
+        }
+
+        @Override
+        public void set(String key, Serializable value) {
+            assert key.equals(this.key);
+            this.value = value;
+        }
+    }
+
+    private static class HashMapValues extends HashMap<String, Serializable>
+            implements Values {
+
+        public HashMapValues(Values previousValues) {
+            super(previousValues == null ? 0 : previousValues.size());
+            if (previousValues != null) {
+                previousValues.keySet().forEach(
+                        key -> super.put(key, previousValues.get(key)));
+            }
+        }
+
+        @Override
+        public Serializable get(String key) {
+            return super.get(key);
+        }
+
+        @Override
+        public void set(String key, Serializable value) {
+            super.put(key, value);
+        }
+
+        @Override
+        public boolean containsKey(String key) {
+            return super.containsKey(key);
+        }
+
+        @Override
+        public Stream<Serializable> streamValues() {
+            return super.values().stream();
+        }
+    }
+
+    private Values values;
 
     private boolean isPopulated;
 
@@ -71,12 +172,6 @@ protected void put(String key, Serializable value) {
         put(key, value, true);
     }
 
-    private void ensureValues() {
-        if (values == null) {
-            values = new HashMap<>();
-        }
-    }
-
     /**
      * Stores a value with the given key, replacing any value previously stored
      * with the same key.
@@ -101,8 +196,16 @@ protected Serializable put(String key, Serializable value,
         } else {
             setUnChanged(key);
         }
-        ensureValues();
-        values.put(key, value);
+
+        // Optimize memory use when there's only one key
+        if (values == null) {
+            values = new SingleValue(key, value);
+        } else {
+            if (values instanceof SingleValue && !values.containsKey(key)) {
+                values = new HashMapValues(values);
+            }
+            values.set(key, value);
+        }
 
         detatchPotentialChild(oldValue);
 
@@ -236,14 +339,27 @@ protected boolean contains(String key) {
      */
     protected Serializable remove(String key) {
         setChanged(key);
+        Serializable oldValue;
+
         if (values == null) {
             return null;
+        } else if (values instanceof SingleValue) {
+            oldValue = values.get(key);
+            if (values.containsKey(key)) {
+                values = null;
+            }
+        } else {
+            assert values instanceof HashMapValues;
+            HashMapValues hashMapValues = (HashMapValues) values;
+            oldValue = hashMapValues.remove(key);
+
+            if (hashMapValues.isEmpty()) {
+                values = null;
+            }
         }
-        Serializable oldValue = values.remove(key);
+
         detatchPotentialChild(oldValue);
-        if (values.isEmpty()) {
-            values = null;
-        }
+
         return oldValue;
     }
 
@@ -343,7 +459,7 @@ public void forEachChild(Consumer<StateNode> action) {
         }
         assert !values.isEmpty();
 
-        values.values().stream().filter(v -> v instanceof StateNode)
+        values.streamValues().filter(v -> v instanceof StateNode)
                 .forEach(v -> action.accept((StateNode) v));
     }
 
@@ -385,4 +501,9 @@ protected boolean mayUpdateFromClient(String key, Serializable value) {
         return false;
     }
 
+    // Exposed for testing purposes
+    boolean usesSingleMap() {
+        return values instanceof SingleValue;
+    }
+
 }

flow-server/src/test/java/com/vaadin/flow/internal/nodefeature/NodeMapTest.java

@@ -349,4 +349,27 @@ public void put_sameValue_hasNoEffect() {
         // The attach listener is not called one more time
         nodeMap.put("foo", child);
     }
+
+    @Test
+    public void put_replaceSingleValue_stillUseSingleValue() {
+        nodeMap.put("foo", "bar");
+
+        Assert.assertTrue(nodeMap.usesSingleMap());
+
+        nodeMap.put("foo", "baz");
+
+        Assert.assertTrue(nodeMap.usesSingleMap());
+    }
+
+    @Test
+    public void streamSingleNullValue() {
+        nodeMap.put("foo", null);
+
+        Assert.assertTrue(nodeMap.usesSingleMap());
+
+        nodeMap.forEachChild(child -> {
+            Assert.fail(
+                    "Should not happen, but forEachChild shouldn't explode either");
+        });
+    }
 }