Fix keymapper reverse mapping (#9535)

Fixes #9470

5 files changed, 104 insertions, 51 deletions

diff --git a/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java b/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
index 2f46b6526a..0971c36755 100644
--- a/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
+++ b/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
@@ -30,6 +30,7 @@ import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
+import com.vaadin.data.ValueProvider;
 import com.vaadin.data.provider.DataChangeEvent.DataRefreshEvent;
 import com.vaadin.server.AbstractExtension;
 import com.vaadin.server.KeyMapper;
@@ -167,7 +168,7 @@ public class DataCommunicator<T> extends AbstractExtension {
         @Override
         public void generateData(T data, JsonObject jsonObject) {
             // Make sure KeyMapper is up to date
-            getKeyMapper().refresh(data, dataProvider::getId);
+            getKeyMapper().refresh(data);
 
             // Write the key string for given data object
             jsonObject.put(DataCommunicatorConstants.KEY,
@@ -193,7 +194,7 @@ public class DataCommunicator<T> extends AbstractExtension {
     private final ActiveDataHandler handler = new ActiveDataHandler();
 
     /** Empty default data provider. */
-    protected DataProvider<T, ?> dataProvider = new CallbackDataProvider<>(
+    private DataProvider<T, ?> dataProvider = new CallbackDataProvider<>(
             q -> Stream.empty(), q -> 0);
     private final DataKeyMapper<T> keyMapper;
 
@@ -212,7 +213,7 @@ public class DataCommunicator<T> extends AbstractExtension {
         addDataGenerator(handler);
         rpc = getRpcProxy(DataCommunicatorClientRpc.class);
         registerRpc(createRpc());
-        keyMapper = createKeyMapper();
+        keyMapper = createKeyMapper(dataProvider::getId);
     }
 
     @Override
@@ -357,13 +358,13 @@ public class DataCommunicator<T> extends AbstractExtension {
 
     /**
      * Fetches a list of items from the DataProvider.
-     * 
+     *
      * @param offset
      *            the starting index of the range
      * @param limit
      *            the max number of results
      * @return the list of items in given range
-     * 
+     *
      * @since 8.1
      */
     @SuppressWarnings({ "rawtypes", "unchecked" })
@@ -585,10 +586,16 @@ public class DataCommunicator<T> extends AbstractExtension {
      * <p>
      * This method is called from the constructor.
      *
+     * @param identifierGetter has to return a unique key for every bean, and the returned key has to
+     *                         follow general {@code hashCode()} and {@code equals()} contract,
+     *                         see {@link Object#hashCode()} for details.
      * @return key mapper
+     *
+     * @since 8.1
+     *
      */
-    protected DataKeyMapper<T> createKeyMapper() {
-        return new KeyMapper<>();
+    protected DataKeyMapper<T> createKeyMapper(ValueProvider<T,Object> identifierGetter) {
+        return new KeyMapper<T>(identifierGetter);
     }
 
     /**
@@ -633,9 +640,7 @@ public class DataCommunicator<T> extends AbstractExtension {
             DataProvider<T, F> dataProvider, F initialFilter) {
         Objects.requireNonNull(dataProvider, "data provider cannot be null");
         filter = initialFilter;
-        detachDataProviderListener();
-        dropAllData();
-        this.dataProvider = dataProvider;
+        setDataProvider(dataProvider);
 
         /*
          * This introduces behavior which influence on the client-server
@@ -672,13 +677,13 @@ public class DataCommunicator<T> extends AbstractExtension {
      * Sets the filter for this DataCommunicator. This method is used by user
      * through the consumer method from {@link #setDataProvider} and should not
      * be called elsewhere.
-     * 
+     *
      * @param filter
      *            the filter
-     * 
+     *
      * @param <F>
      *            the filter type
-     * 
+     *
      * @since 8.1
      */
     protected <F> void setFilter(F filter) {
@@ -725,7 +730,7 @@ public class DataCommunicator<T> extends AbstractExtension {
     /**
      * Getter method for finding the size of DataProvider. Can be overridden by
      * a subclass that uses a specific type of DataProvider and/or query.
-     * 
+     *
      * @return the size of data provider with current filter
      */
     @SuppressWarnings({ "unchecked", "rawtypes" })
@@ -749,7 +754,7 @@ public class DataCommunicator<T> extends AbstractExtension {
                     getUI().access(() -> {
                         if (event instanceof DataRefreshEvent) {
                             T item = ((DataRefreshEvent<T>) event).getItem();
-                            keyMapper.refresh(item, dataProvider::getId);
+                            getKeyMapper().refresh(item);
                             generators.forEach(g -> g.refreshData(item));
                             refresh(item);
                         } else {
@@ -773,4 +778,17 @@ public class DataCommunicator<T> extends AbstractExtension {
             dataProviderUpdateRegistration = null;
         }
     }
+
+    /**
+     * Sets a new {@code DataProvider} and refreshes all the internal structures
+     *
+     * @param dataProvider
+     * @since 8.1
+     */
+    protected void setDataProvider(DataProvider<T, ?> dataProvider) {
+        detachDataProviderListener();
+        dropAllData();
+        this.dataProvider = dataProvider;
+        keyMapper.setIdentifierGetter(dataProvider::getId);
+    }
 }
diff --git a/server/src/main/java/com/vaadin/data/provider/DataKeyMapper.java b/server/src/main/java/com/vaadin/data/provider/DataKeyMapper.java
index f5be454f86..6db4a491fe 100644
--- a/server/src/main/java/com/vaadin/data/provider/DataKeyMapper.java
+++ b/server/src/main/java/com/vaadin/data/provider/DataKeyMapper.java
@@ -79,8 +79,19 @@ public interface DataKeyMapper<T> extends Serializable {
      *
      * @param dataObject
      *            the data object to update
-     * @param identifierGetter
-     *            the function to get an identifier from a data object
+     *
+     * @since 8.1
+     *
+     */
+    void refresh(T dataObject);
+
+    /**
+     * Takes identifier getter into use and updates existing mappings
+     *
+     * @param identifierGetter has to return a unique key for every bean, and the returned key has to
+     *                         follow general {@code hashCode()} and {@code equals()} contract,
+     *                         see {@link Object#hashCode()} for details.
+     * @since 8.1
      */
-    void refresh(T dataObject, ValueProvider<T, Object> identifierGetter);
+    void setIdentifierGetter(ValueProvider<T, Object> identifierGetter);
 }
diff --git a/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java b/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
index 3e9eb01921..f53287b1df 100644
--- a/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
+++ b/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
@@ -52,7 +52,7 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
      */
     public HierarchicalDataCommunicator() {
         super();
-        setDataProvider(new TreeDataProvider<T>(new TreeData<>()), null);
+        setDataProvider(new TreeDataProvider<>(new TreeData<>()), null);
     }
 
     @Override
@@ -220,7 +220,7 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
 
     /**
      * Returns whether given item has children.
-     * 
+     *
      * @param item
      *            the item to test
      * @return {@code true} if item has children; {@code false} if not
@@ -231,7 +231,7 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
 
     /**
      * Returns whether given item is expanded.
-     * 
+     *
      * @param item
      *            the item to test
      * @return {@code true} if item is expanded; {@code false} if not
diff --git a/server/src/main/java/com/vaadin/server/KeyMapper.java b/server/src/main/java/com/vaadin/server/KeyMapper.java
index 58058c6f69..f7b878dd27 100644
--- a/server/src/main/java/com/vaadin/server/KeyMapper.java
+++ b/server/src/main/java/com/vaadin/server/KeyMapper.java
@@ -18,6 +18,7 @@ package com.vaadin.server;
 
 import java.io.Serializable;
 import java.util.HashMap;
+import java.util.Map;
 
 import com.vaadin.data.ValueProvider;
 import com.vaadin.data.provider.DataKeyMapper;
@@ -33,15 +34,35 @@ public class KeyMapper<V> implements DataKeyMapper<V>, Serializable {
 
     private int lastKey = 0;
 
-    private final HashMap<V, String> objectKeyMap = new HashMap<>();
+    private final HashMap<Object, String> objectIdKeyMap = new HashMap<>();
 
     private final HashMap<String, V> keyObjectMap = new HashMap<>();
 
+    private ValueProvider<V, Object> identifierGetter;
+
+    /**
+     * Constructs a new mapper
+     *
+     * @param identifierGetter has to return a unique key for every bean, and the returned key has to
+     *                         follow general {@code hashCode()} and {@code equals()} contract,
+     *                         see {@link Object#hashCode()} for details.
+     * @since 8.1
+     */
+    public KeyMapper(ValueProvider<V, Object> identifierGetter) {
+        this.identifierGetter = identifierGetter;
+    }
+
+    /**
+     * Constructs a new mapper with trivial {@code identifierGetter}
+     */
+    public KeyMapper() {
+        this(v -> v);
+    }
+
     /**
      * Gets key for an object.
      *
-     * @param o
-     *            the object.
+     * @param o the object.
      */
     @Override
     public String key(V o) {
@@ -51,14 +72,15 @@ public class KeyMapper<V> implements DataKeyMapper<V>, Serializable {
         }
 
         // If the object is already mapped, use existing key
-        String key = objectKeyMap.get(o);
+        Object id = identifierGetter.apply(o);
+        String key = objectIdKeyMap.get(id);
         if (key != null) {
             return key;
         }
 
         // If the object is not yet mapped, map it
         key = String.valueOf(++lastKey);
-        objectKeyMap.put(o, key);
+        objectIdKeyMap.put(id, key);
         keyObjectMap.put(key, o);
 
         return key;
@@ -66,14 +88,13 @@ public class KeyMapper<V> implements DataKeyMapper<V>, Serializable {
 
     @Override
     public boolean has(V o) {
-        return objectKeyMap.containsKey(o);
+        return objectIdKeyMap.containsKey(identifierGetter.apply(o));
     }
 
     /**
      * Retrieves object with the key.
      *
-     * @param key
-     *            the name with the desired value.
+     * @param key the name with the desired value.
      * @return the object with the key.
      */
     @Override
@@ -84,15 +105,12 @@ public class KeyMapper<V> implements DataKeyMapper<V>, Serializable {
     /**
      * Removes object from the mapper.
      *
-     * @param removeobj
-     *            the object to be removed.
+     * @param removeobj the object to be removed.
      */
     @Override
     public void remove(V removeobj) {
-        final String key = objectKeyMap.get(removeobj);
-
+        final String key = objectIdKeyMap.remove(identifierGetter.apply(removeobj));
         if (key != null) {
-            objectKeyMap.remove(removeobj);
             keyObjectMap.remove(key);
         }
     }
@@ -102,34 +120,39 @@ public class KeyMapper<V> implements DataKeyMapper<V>, Serializable {
      */
     @Override
     public void removeAll() {
-        objectKeyMap.clear();
+        objectIdKeyMap.clear();
         keyObjectMap.clear();
     }
 
     /**
      * Checks if the given key is mapped to an object.
      *
-     * @since 7.7
-     *
-     * @param key
-     *            the key to check
+     * @param key the key to check
      * @return <code>true</code> if the key is currently mapped,
-     *         <code>false</code> otherwise
+     * <code>false</code> otherwise
+     * @since 7.7
      */
     public boolean containsKey(String key) {
         return keyObjectMap.containsKey(key);
     }
 
     @Override
-    public void refresh(V dataObject,
-            ValueProvider<V, Object> identifierGetter) {
+    public void refresh(V dataObject) {
         Object id = identifierGetter.apply(dataObject);
-        objectKeyMap.entrySet().stream()
-                .filter(e -> identifierGetter.apply(e.getKey()).equals(id))
-                .findAny().ifPresent(e -> {
-                    String key = objectKeyMap.remove(e.getKey());
-                    objectKeyMap.put(dataObject, key);
-                    keyObjectMap.put(key, dataObject);
-                });
+        String key = objectIdKeyMap.get(id);
+        if (key != null) {
+            keyObjectMap.put(key, dataObject);
+        }
+    }
+
+    @Override
+    public void setIdentifierGetter(ValueProvider<V, Object> identifierGetter) {
+        if (this.identifierGetter != identifierGetter) {
+            this.identifierGetter = identifierGetter;
+            objectIdKeyMap.clear();
+            for (Map.Entry<String, V> entry : keyObjectMap.entrySet()) {
+                objectIdKeyMap.put(identifierGetter.apply(entry.getValue()),entry.getKey());
+            }
+        }
     }
 }
diff --git a/server/src/main/java/com/vaadin/ui/ComboBox.java b/server/src/main/java/com/vaadin/ui/ComboBox.java
index 62e7ef7751..73d9ad84cf 100644
--- a/server/src/main/java/com/vaadin/ui/ComboBox.java
+++ b/server/src/main/java/com/vaadin/ui/ComboBox.java
@@ -25,6 +25,7 @@ import java.util.Objects;
 import java.util.Set;
 import java.util.stream.Stream;
 
+import com.vaadin.data.ValueProvider;
 import org.jsoup.nodes.Element;
 
 import com.vaadin.data.HasFilterableDataProvider;
@@ -188,8 +189,8 @@ public class ComboBox<T> extends AbstractSingleSelect<T>
     public ComboBox() {
         super(new DataCommunicator<T>() {
             @Override
-            protected DataKeyMapper<T> createKeyMapper() {
-                return new KeyMapper<T>() {
+            protected DataKeyMapper<T> createKeyMapper(ValueProvider<T,Object> identifierGetter) {
+                return new KeyMapper<T>(identifierGetter) {
                     @Override
                     public void remove(T removeobj) {
                         // never remove keys from ComboBox to support selection