Addressed review comments

Author: nivi-apple

src/darwin/Framework/CHIP/MTRDevice.mm

@@ -675,7 +675,7 @@ - (void)_setDSTOffsets:(NSArray<MTRTimeSynchronizationClusterDSTOffsetStruct *>
                                  completion:setDSTOffsetResponseHandler];
 }
 
-- (NSMutableArray<NSNumber *> *)arrayOfNumbersFromAttributeValue:(NSDictionary *)dataDictionary
+- (NSMutableArray<NSNumber *> *)arrayOfNumbersFromAttributeValue:(MTRDeviceDataValueDictionary)dataDictionary
 {
     if (![MTRArrayValueType isEqual:dataDictionary[MTRTypeKey]]) {
         return nil;
@@ -3040,35 +3040,43 @@ - (void)_removeClusters:(NSSet<MTRClusterPath *> *)clusterPathsToRemove
 
     for (MTRClusterPath * path in clusterPathsToRemove) {
         [_persistedClusterData removeObjectForKey:path];
-        [_persistedClusters removeObject:path];
         [_clusterDataToPersist removeObjectForKey:path];
         [self.deviceController.controllerDataStore clearStoredClusterDataForNodeID:self.nodeID endpointID:path.endpoint clusterID:path.cluster];
     }
 }
 
-- (void)_removeAttributes:(NSSet<NSNumber *> *)attributes fromCluster:(MTRClusterPath *)clusterPath
+- (void)_removeClustersFromCache:(NSSet<MTRClusterPath *> *)clusterPathsToRemove
 {
-    if (toBeRemovedAttributes == nil || clusterPathToRemoveAttributesFrom == nil) {
-        return;
+    os_unfair_lock_assert_owner(&self->_lock);
+
+    [_persistedClusters minusSet:clusterPathsToRemove];
+
+    for (MTRClusterPath * path in clusterPathsToRemove) {
+        [_persistedClusterData removeObjectForKey:path];
+        [_clusterDataToPersist removeObjectForKey:path];
     }
+}
+
+- (void)_removeAttributes:(NSSet<NSNumber *> *)attributes fromCluster:(MTRClusterPath *)clusterPath
+{
     os_unfair_lock_assert_owner(&self->_lock);
 
-    for (NSNumber * attribute in toBeRemovedAttributes) {
-        [self _removeCachedAttribute:attribute fromCluster:clusterPathToRemoveAttributesFrom];
+    for (NSNumber * attribute in attributes) {
+        [self _removeCachedAttribute:attribute fromCluster:clusterPath];
     }
     // Just clear out the NSCache entry for this cluster, so we'll load it from storage as needed.
-    [_persistedClusterData removeObjectForKey:clusterPathToRemoveAttributesFrom];
-    [self.deviceController.controllerDataStore clearStoredClusterDataForNodeID:self.nodeID endpointID:clusterPathToRemoveAttributesFrom.endpoint clusterID:clusterPathToRemoveAttributesFrom.cluster];
+    [_persistedClusterData removeObjectForKey:clusterPath];
+    [self.deviceController.controllerDataStore removeAttributes:attributes fromCluster:clusterPath forNodeID:self.nodeID];
 }
 
-- (void)_pruneEndpointsIn:(NSDictionary *)previousPartsListValue
-              missingFrom:(NSDictionary *)newPartsListValue
+- (void)_pruneEndpointsIn:(MTRDeviceDataValueDictionary)previousPartsListValue
+              missingFrom:(MTRDeviceDataValueDictionary)newPartsListValue
 {
     // If the parts list changed and one or more endpoints were removed, remove all the
     // clusters for all those endpoints from our data structures.
     // Also remove those endpoints from the data store.
-    NSMutableSet * toBeRemovedEndpoints = [NSMutableSet setWithArray:[self arrayOfNumbersFromAttributeValue:[self _dataValueWithoutDataVersion:previousPartsListValue]]];
-    NSSet * endpointsOnDevice = [NSSet setWithArray:[self arrayOfNumbersFromAttributeValue:newPartsListValue]];
+    NSMutableSet<NSNumber*> * toBeRemovedEndpoints = [NSMutableSet setWithArray:[self arrayOfNumbersFromAttributeValue:previousPartsListValue]];
+    NSSet<NSNumber *> * endpointsOnDevice = [NSSet setWithArray:[self arrayOfNumbersFromAttributeValue:newPartsListValue]];
     [toBeRemovedEndpoints minusSet:endpointsOnDevice];
 
     for (NSNumber * endpoint in toBeRemovedEndpoints) {
@@ -3078,75 +3086,71 @@ - (void)_pruneEndpointsIn:(NSDictionary *)previousPartsListValue
                 [clusterPathsToRemove addObject:path];
             }
         }
-        [self _removeClusters:[clusterPathsToRemove copy]];
-        [self.deviceController.controllerDataStore removeEndpointFromEndpointIndex:endpoint forNodeID:self.nodeID];
+        [self _removeClustersFromCache:clusterPathsToRemove];
+        [self.deviceController.controllerDataStore clearStoredClusterDataForNodeID:self.nodeID endpointID:endpoint];
     }
 }
 
-- (void)_pruneOrphanedClusters:(MTRAttributePath *)attributePath
-       previousServerListValue:(NSDictionary *)previousServerListValue
-            newServerListValue:(NSDictionary *)newServerListValue
+- (void)_pruneClustersIn:(MTRDeviceDataValueDictionary)previousServerListValue
+             missingFrom:(MTRDeviceDataValueDictionary)newServerListValue
+              forEndpoint:(NSNumber *)endpointID
 {
     // If the server list changed and clusters were removed, remove those clusters from our data structures.
     // Also remove it from the data store.
-    NSMutableSet<NSNumber *> * toBeRemovedClusters = [NSMutableSet setWithArray:[self arrayOfNumbersFromAttributeValue:[self _dataValueWithoutDataVersion:previousServerListValue]]];
+    NSMutableSet<NSNumber *> * toBeRemovedClusters = [NSMutableSet setWithArray:[self arrayOfNumbersFromAttributeValue:previousServerListValue]];
     NSSet<NSNumber *> * clustersStillOnEndpoint = [NSSet setWithArray:[self arrayOfNumbersFromAttributeValue:newServerListValue]];
     [toBeRemovedClusters minusSet:clustersStillOnEndpoint];
 
     NSMutableSet<MTRClusterPath *> * clusterPathsToRemove = [[NSMutableSet alloc] init];
     for (NSNumber * cluster in toBeRemovedClusters) {
         for (MTRClusterPath * path in _persistedClusters) {
-            if ([path.endpoint isEqualToNumber:attributePath.endpoint] && [path.cluster isEqualToNumber:cluster]) {
+            if ([path.endpoint isEqualToNumber:endpointID] && [path.cluster isEqualToNumber:cluster]) {
                 [clusterPathsToRemove addObject:path];
             }
         }
     }
-    [self _removeClusters:[clusterPathsToRemove copy]];
+    [self _removeClusters:clusterPathsToRemove];
 }
 
-- (void)_pruneOrphanedAttributes:(MTRAttributePath *)attributePath
-           newAttributeListValue:(NSDictionary *)newAttributeListValue
+- (void)_pruneAttributesIn:(MTRDeviceDataValueDictionary)previousAttributeListValue
+               missingFrom:(MTRDeviceDataValueDictionary)newAttributeListValue
+                forCluster:(MTRClusterPath *)clusterPath
 {
     // If the attribute list changed and attributes were removed, remove the attributes from our
     // data structures.
-    NSMutableSet<NSNumber *> * toBeRemovedAttributes = [NSMutableSet setWithArray:[self arrayOfNumbersFromAttributeValue:[self _cachedAttributeValueForPath:attributePath]]];
+    NSMutableSet<NSNumber *> * toBeRemovedAttributes = [NSMutableSet setWithArray:[self arrayOfNumbersFromAttributeValue:previousAttributeListValue]];
     NSSet<NSNumber *> * attributesStillInCluster = [NSSet setWithArray:[self arrayOfNumbersFromAttributeValue:newAttributeListValue]];
 
     [toBeRemovedAttributes minusSet:attributesStillInCluster];
-    MTRClusterPath * clusterPathToRemoveAttributesFrom;
-    for (MTRClusterPath * path in _persistedClusters) {
-        if ([path.endpoint isEqualToNumber:attributePath.endpoint] && [path.cluster isEqualToNumber:attributePath.cluster]) {
-            clusterPathToRemoveAttributesFrom = path;
-            break;
-        }
-    }
-    [self _removeAttributes:[toBeRemovedAttributes copy] fromCluster:clusterPathToRemoveAttributesFrom];
-    [self.deviceController.controllerDataStore removeAttributes:[toBeRemovedAttributes copy] fromCluster:clusterPathToRemoveAttributesFrom forNodeID:self.nodeID];
+    [self _removeAttributes:toBeRemovedAttributes fromCluster:clusterPath];
 }
 
-- (void)_pruneOrphanedEndpointsAndClusters:(MTRAttributePath *)attributePath
-                             previousValue:(NSDictionary *)previousValue
-                        attributeDataValue:(NSDictionary *)attributeDataValue
+- (void)_pruneStoredDataForPath:(MTRAttributePath *)attributePath
+                    missingFrom:(MTRDeviceDataValueDictionary)newAttributeDataValue
 {
     os_unfair_lock_assert_owner(&self->_lock);
 
-    NSNumber * rootEndpoint = @0;
-
-    if (_persistedClusters == nil || _persistedClusterData == nil || !previousValue.count) {
+    if (![self _dataStoreExists] && !_clusterDataToPersist.count) {
+        MTR_LOG_DEBUG("%@ No data store to prune from", self);
         return;
     }
+    
     // Check if parts list changed or server list changed for the descriptor cluster or the attribute list changed for a cluster.
     // If yes, we might need to prune any deleted endpoints, clusters or attributes from the storage and persisted cluster data.
     if (attributePath.cluster.unsignedLongValue == MTRClusterIDTypeDescriptorID) {
-        if (attributePath.attribute.unsignedLongValue == MTRAttributeIDTypeClusterDescriptorAttributePartsListID && [attributePath.endpoint isEqualToNumber:rootEndpoint]) {
-            [self _pruneOrphanedEndpoints:previousValue newPartsListValue:attributeDataValue];
-        } else if (attributePath.attribute.unsignedLongValue == MTRAttributeIDTypeClusterDescriptorAttributeServerListID) {
-            [self _pruneOrphanedClusters:attributePath previousServerListValue:previousValue newServerListValue:attributeDataValue];
+        if (attributePath.attribute.unsignedLongValue == MTRAttributeIDTypeClusterDescriptorAttributePartsListID && [attributePath.endpoint isEqualToNumber:@(kRootEndpointId)]) {
+            [self _pruneEndpointsIn:[self _cachedAttributeValueForPath:attributePath] missingFrom:newAttributeDataValue];
+            return;
+        }
+        
+        if (attributePath.attribute.unsignedLongValue == MTRAttributeIDTypeClusterDescriptorAttributeServerListID) {
+            [self _pruneClustersIn:[self _cachedAttributeValueForPath:attributePath] missingFrom:newAttributeDataValue forEndpoint:attributePath.endpoint];
+            return;
         }
     }
 
     if (attributePath.attribute.unsignedLongValue == MTRAttributeIDTypeGlobalAttributeAttributeListID) {
-        [self _pruneOrphanedAttributes:attributePath newAttributeListValue:attributeDataValue];
+        [self _pruneAttributesIn:[self _cachedAttributeValueForPath:attributePath] missingFrom:newAttributeDataValue forCluster:[MTRClusterPath clusterPathWithEndpointID:attributePath.endpoint clusterID:attributePath.cluster]];
     }
 }
 
@@ -3203,7 +3207,7 @@ - (NSArray *)_getAttributesToReportWithReportedValues:(NSArray<NSDictionary<NSSt
                     [self _noteDataVersion:dataVersion forClusterPath:clusterPath];
                 }
 
-                [self _pruneOrphanedEndpointsAndClusters:attributePath previousValue:previousValue attributeDataValue:attributeDataValue];
+                [self _pruneStoredDataForPath:attributePath missingFrom:attributeDataValue];
 
                 if (!_deviceConfigurationChanged) {
                     _deviceConfigurationChanged = [self _attributeAffectsDeviceConfiguration:attributePath];

src/darwin/Framework/CHIP/MTRDeviceControllerDataStore.h

@@ -76,8 +76,8 @@ typedef void (^MTRDeviceControllerDataStoreClusterDataHandler)(NSDictionary<NSNu
 - (nullable MTRDeviceClusterData *)getStoredClusterDataForNodeID:(NSNumber *)nodeID endpointID:(NSNumber *)endpointID clusterID:(NSNumber *)clusterID;
 - (void)storeClusterData:(NSDictionary<MTRClusterPath *, MTRDeviceClusterData *> *)clusterData forNodeID:(NSNumber *)nodeID;
 - (void)clearStoredClusterDataForNodeID:(NSNumber *)nodeID;
+- (void)clearStoredClusterDataForNodeID:(NSNumber *)nodeID endpointID:(NSNumber *)endpointID;
 - (void)clearStoredClusterDataForNodeID:(NSNumber *)nodeID endpointID:(NSNumber *)endpointID clusterID:(NSNumber *)clusterID;
-- (void)removeEndpointFromEndpointIndex:(NSNumber *)endpointID forNodeID:(NSNumber *)nodeID;
 - (void)removeAttributes:(NSSet<NSNumber *> *)attributes fromCluster:(MTRClusterPath *)path forNodeID:(NSNumber *)nodeID;
 - (void)clearAllStoredClusterData;
 

src/darwin/Framework/CHIP/MTRDeviceControllerDataStore.mm

@@ -446,24 +446,21 @@ - (BOOL)_storeEndpointIndex:(NSArray<NSNumber *> *)endpointIndex forNodeID:(NSNu
     return [self _storeAttributeCacheValue:endpointIndex forKey:[self _endpointIndexKeyForNodeID:nodeID]];
 }
 
-- (void)removeEndpointFromEndpointIndex:(NSNumber *)endpointID forNodeID:(NSNumber *)nodeID
+- (BOOL)_removeEndpointFromEndpointIndex:(NSNumber *)endpointID forNodeID:(NSNumber *)nodeID
 {
+    dispatch_assert_queue(_storageDelegateQueue);
     if (!endpointID || !nodeID) {
         MTR_LOG_ERROR("%s: unexpected nil input", __func__);
-        return;
+        return NO;
     }
-    dispatch_async(_storageDelegateQueue, ^{
-        NSMutableArray * endpointIndex = [[self _fetchEndpointIndexForNodeID:nodeID] mutableCopy];
-        if (endpointIndex == nil) {
-            return;
-        }
 
-        [endpointIndex removeObject:endpointID];
-        BOOL success = [self _storeEndpointIndex:endpointIndex forNodeID:nodeID];
-        if (!success) {
-            MTR_LOG_ERROR("removeEndpointFromEndpointIndex: _storeEndpointIndex for node 0x%016llX", nodeID.unsignedLongLongValue);
-        }
-    });
+    NSMutableArray<NSNumber *> * endpointIndex = [[self _fetchEndpointIndexForNodeID:nodeID] mutableCopy];
+    if (endpointIndex == nil) {
+        return NO;
+    }
+
+    [endpointIndex removeObject:endpointID];
+    return [self _storeEndpointIndex:endpointIndex forNodeID:nodeID];
 }
 
 - (BOOL)_deleteEndpointIndexForNodeID:(NSNumber *)nodeID
@@ -517,6 +514,7 @@ - (BOOL)_deleteClusterIndexForNodeID:(NSNumber *)nodeID endpointID:(NSNumber *)e
         MTR_LOG_ERROR("%s: unexpected nil input", __func__);
         return NO;
     }
+
     return [self _removeAttributeCacheValueForKey:[self _clusterIndexKeyForNodeID:nodeID endpointID:endpointID]];
 }
 
@@ -718,25 +716,54 @@ - (void)clearStoredClusterDataForNodeID:(NSNumber *)nodeID
     });
 }
 
-- (void)clearStoredClusterDataForNodeID:(NSNumber *)nodeID endpointID:(NSNumber *)endpointID clusterID:(NSNumber *)clusterID
+- (void)clearStoredClusterDataForNodeID:(NSNumber *)nodeID endpointID:(NSNumber *)endpointID
 {
     dispatch_async(_storageDelegateQueue, ^{
-        BOOL success = [self _deleteClusterDataForNodeID:nodeID endpointID:endpointID clusterID:clusterID];
+        BOOL success = [self _removeEndpointFromEndpointIndex:endpointID forNodeID:nodeID];
         if (!success) {
-            MTR_LOG_ERROR("clearStoredClusterDataForNodeID: _deleteClusterDataForNodeID failed for node 0x%016llX endpoint %u cluster 0x%08lX", nodeID.unsignedLongLongValue, endpointID.unsignedShortValue, clusterID.unsignedLongValue);
-            return;
+            MTR_LOG_ERROR("removeEndpointFromEndpointIndex for endpointID %u failed for node 0x%016llX", endpointID.unsignedShortValue, nodeID.unsignedLongLongValue);
+        }
+        
+        NSArray<NSNumber *> * clusterIndex = [self _fetchClusterIndexForNodeID:nodeID endpointID:endpointID];
+        
+        for (NSNumber * cluster in clusterIndex)
+        {
+            success = [self _deleteClusterDataForNodeID:nodeID endpointID:endpointID clusterID:cluster];
+            if (!success) {
+                MTR_LOG_ERROR("Delete failed for clusterData for node 0x%016llX endpoint %u cluster 0x%08lX", nodeID.unsignedLongLongValue, endpointID.unsignedShortValue, cluster.unsignedLongValue);
+            }
+        }
+        
+        success = [self _deleteClusterIndexForNodeID:nodeID endpointID:endpointID];
+        if (!success) {
+            MTR_LOG_ERROR("Delete failed for clusterIndex for node 0x%016llX endpoint %u", nodeID.unsignedLongLongValue, endpointID.unsignedShortValue);
         }
 
+        MTR_LOG("clearStoredClusterDataForNodeID: Deleted endpoint %u for node 0x%016llX successfully", endpointID.unsignedShortValue, nodeID.unsignedLongLongValue);
+    });
+}
+
+- (void)clearStoredClusterDataForNodeID:(NSNumber *)nodeID endpointID:(NSNumber *)endpointID clusterID:(NSNumber *)clusterID
+{
+    dispatch_async(_storageDelegateQueue, ^{
         NSArray<NSNumber *> * clusterIndex = [self _fetchClusterIndexForNodeID:nodeID endpointID:endpointID];
         NSMutableArray<NSNumber *> * clusterIndexCopy = [clusterIndex mutableCopy];
         [clusterIndexCopy removeObject:clusterID];
 
+        BOOL success;
         if (clusterIndexCopy.count != clusterIndex.count) {
             success = [self _storeClusterIndex:clusterIndexCopy forNodeID:nodeID endpointID:endpointID];
             if (!success) {
                 MTR_LOG_ERROR("clearStoredClusterDataForNodeID: _storeClusterIndex failed for node 0x%016llX endpoint %u", nodeID.unsignedLongLongValue, endpointID.unsignedShortValue);
             }
         }
+
+        success = [self _deleteClusterDataForNodeID:nodeID endpointID:endpointID clusterID:clusterID];
+        if (!success) {
+            MTR_LOG_ERROR("clearStoredClusterDataForNodeID: _deleteClusterDataForNodeID failed for node 0x%016llX endpoint %u cluster 0x%08lX", nodeID.unsignedLongLongValue, endpointID.unsignedShortValue, clusterID.unsignedLongValue);
+            return;
+        }
+
         MTR_LOG("clearStoredClusterDataForNodeID: Deleted endpoint %u cluster 0x%08lX for node 0x%016llX successfully", endpointID.unsignedShortValue, clusterID.unsignedLongValue, nodeID.unsignedLongLongValue);
     });
 }
@@ -750,11 +777,14 @@ - (void)removeAttributes:(NSSet<NSNumber *> *)attributes fromCluster:(MTRCluster
     for (NSNumber * attribute in attributes) {
         [clusterData removeValueForAttribute:attribute];
     }
-    BOOL success = [self _storeClusterData:clusterData forNodeID:nodeID endpointID:path.endpoint clusterID:path.cluster];
-    if (!success) {
-        MTR_LOG_ERROR("removeAttributes: _storeClusterData failed for node 0x%016llX endpoint %u", nodeID.unsignedLongLongValue, path.endpoint.unsignedShortValue);
-    }
-    MTR_LOG("removeAttributes: Deleted attributes %@ from endpoint %u cluster 0x%08lX for node 0x%016llX successfully", attributes, path.endpoint.unsignedShortValue, path.cluster.unsignedLongValue, nodeID.unsignedLongLongValue);
+    
+    dispatch_async(_storageDelegateQueue, ^{
+        BOOL success = [self _storeClusterData:clusterData forNodeID:nodeID endpointID:path.endpoint clusterID:path.cluster];
+        if (!success) {
+            MTR_LOG_ERROR("removeAttributes: _storeClusterData failed for node 0x%016llX endpoint %u", nodeID.unsignedLongLongValue, path.endpoint.unsignedShortValue);
+        }
+        MTR_LOG("removeAttributes: Deleted attributes %@ from endpoint %u cluster 0x%08lX for node 0x%016llX successfully", attributes, path.endpoint.unsignedShortValue, path.cluster.unsignedLongValue, nodeID.unsignedLongLongValue);
+    });
 }
 
 - (void)clearAllStoredClusterData