From 90acbda4a93b19c81725186f8e546ed7f4afd2cc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20G=C3=B6ransson?= <david.goransson90@gmail.com>
Date: Mon, 17 Oct 2022 17:01:47 +0200
Subject: [PATCH] Merge elements when getting new composition data

---
 .../mesh/transport/ProvisionedMeshNode.java   | 51 +++++++++++++++----
 1 file changed, 40 insertions(+), 11 deletions(-)

diff --git a/mesh/src/main/java/no/nordicsemi/android/mesh/transport/ProvisionedMeshNode.java b/mesh/src/main/java/no/nordicsemi/android/mesh/transport/ProvisionedMeshNode.java
index 3896f14f..030f6875 100644
--- a/mesh/src/main/java/no/nordicsemi/android/mesh/transport/ProvisionedMeshNode.java
+++ b/mesh/src/main/java/no/nordicsemi/android/mesh/transport/ProvisionedMeshNode.java
@@ -22,16 +22,11 @@
 
 package no.nordicsemi.android.mesh.transport;
 
+import static androidx.room.ForeignKey.CASCADE;
+
 import android.annotation.SuppressLint;
 import android.os.Parcel;
 
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-
 import androidx.annotation.NonNull;
 import androidx.annotation.RestrictTo;
 import androidx.annotation.VisibleForTesting;
@@ -39,6 +34,15 @@
 import androidx.room.ForeignKey;
 import androidx.room.Ignore;
 import androidx.room.Index;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
 import no.nordicsemi.android.mesh.ApplicationKey;
 import no.nordicsemi.android.mesh.Features;
 import no.nordicsemi.android.mesh.MeshNetwork;
@@ -54,8 +58,6 @@
 import no.nordicsemi.android.mesh.utils.SecureUtils;
 import no.nordicsemi.android.mesh.utils.SparseIntArrayParcelable;
 
-import static androidx.room.ForeignKey.CASCADE;
-
 @SuppressWarnings({"WeakerAccess"})
 @Entity(tableName = "nodes",
         foreignKeys = @ForeignKey(entity = MeshNetwork.class,
@@ -100,7 +102,7 @@ public ProvisionedMeshNode(final UnprovisionedMeshNode node) {
         ttl = node.getTtl();
         mTimeStampInMillis = node.getTimeStamp();
         // Here we add some dummy elements with empty models to occupy the addresses in use.
-        for(int i = 0; i < node.getProvisioningCapabilities().getNumberOfElements(); i++){
+        for (int i = 0; i < node.getProvisioningCapabilities().getNumberOfElements(); i++) {
             mElements.put(unicastAddress + i, new Element(unicastAddress + i, 0, new HashMap<>()));
         }
         security = node.isSecure() ? HIGH : LOW;
@@ -506,7 +508,34 @@ void setCompositionData(
                 lowPowerFeatureSupported ? Features.DISABLED : Features.UNSUPPORTED,
                 proxyFeatureSupported ? Features.DISABLED : Features.UNSUPPORTED,
                 relayFeatureSupported ? Features.DISABLED : Features.UNSUPPORTED);
-        mElements.putAll(configCompositionDataStatus.getElements());
+
+        LinkedHashMap<Integer, Element> modifiableNewElements = new LinkedHashMap<>(configCompositionDataStatus.getElements());
+
+        // For each old element
+        for (Map.Entry<Integer, Element> newElementEntry : configCompositionDataStatus.getElements().entrySet()) {
+            final Element newElement = newElementEntry.getValue();
+            // Fetch same element in new composition data
+            final Element oldElement = mElements.get(newElementEntry.getKey());
+
+            LinkedHashMap<Integer, MeshModel> updatedElementModels = new LinkedHashMap<>(newElement.getMeshModels());
+
+            // Old element existed previously, may have same mesh models
+            if (oldElement != null) {
+                // For each model in that element
+                for (Map.Entry<Integer, MeshModel> modelEntry : newElement.getMeshModels().entrySet()) {
+
+                    // Override new element models with old to ensure we keep mesh model info
+                    if (oldElement.getMeshModels().containsKey(modelEntry.getKey())) {
+                        updatedElementModels.put(modelEntry.getKey(), oldElement.meshModels.get(modelEntry.getKey()));
+                    }
+                }
+
+                Element newUpdatedElement = new Element(newElement.elementAddress, newElement.locationDescriptor, updatedElementModels, newElement.name);
+
+                modifiableNewElements.put(newElementEntry.getKey(), newUpdatedElement);
+            }
+        }
+        mElements.putAll(modifiableNewElements);
     }
 
     /**