[Darwin] MTRDevice attribute storage should store values by cluster (#32765)

* [Darwin] MTRDevice attribute storage should store values by cluster

* Address review comments

* Only persist cluster data when there is something to persist
diff --git a/src/darwin/Framework/CHIP/MTRBaseDevice.h b/src/darwin/Framework/CHIP/MTRBaseDevice.h
index 5eb0047..fe1592a 100644
--- a/src/darwin/Framework/CHIP/MTRBaseDevice.h
+++ b/src/darwin/Framework/CHIP/MTRBaseDevice.h
@@ -26,6 +26,9 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
+typedef NSDictionary<NSString *, id> * MTRDeviceResponseValueDictionary;
+typedef NSDictionary<NSString *, id> * MTRDeviceDataValueDictionary;
+
 /**
  * Handler for read attribute response, write attribute response, invoke command response and reports.
  *
@@ -96,7 +99,7 @@
  *
  *                MTRDataKey : Data-value NSDictionary object.
  */
-typedef void (^MTRDeviceResponseHandler)(NSArray<NSDictionary<NSString *, id> *> * _Nullable values, NSError * _Nullable error);
+typedef void (^MTRDeviceResponseHandler)(NSArray<MTRDeviceResponseValueDictionary> * _Nullable values, NSError * _Nullable error);
 
 /**
  * Handler for -subscribeWithQueue: attribute and event reports
diff --git a/src/darwin/Framework/CHIP/MTRDevice.mm b/src/darwin/Framework/CHIP/MTRDevice.mm
index 36ce246..9640a4a 100644
--- a/src/darwin/Framework/CHIP/MTRDevice.mm
+++ b/src/darwin/Framework/CHIP/MTRDevice.mm
@@ -176,6 +176,7 @@
 @implementation MTRDeviceClusterData
 
 static NSString * const sDataVersionKey = @"dataVersion";
+static NSString * const sAttributesKey = @"attributes";
 
 + (BOOL)supportsSecureCoding
 {
@@ -184,7 +185,20 @@
 
 - (NSString *)description
 {
-    return [NSString stringWithFormat:@"<MTRDeviceClusterData: dataVersion %@>", _dataVersion];
+    return [NSString stringWithFormat:@"<MTRDeviceClusterData: dataVersion %@ attributes count %lu>", _dataVersion, static_cast<unsigned long>(_attributes.count)];
+}
+
+- (nullable instancetype)initWithDataVersion:(NSNumber * _Nullable)dataVersion attributes:(NSDictionary<NSNumber *, MTRDeviceDataValueDictionary> * _Nullable)attributes
+{
+    self = [super init];
+    if (self == nil) {
+        return nil;
+    }
+
+    _dataVersion = [dataVersion copy];
+    _attributes = [attributes copy];
+
+    return self;
 }
 
 - (nullable instancetype)initWithCoder:(NSCoder *)decoder
@@ -200,12 +214,25 @@
         return nil;
     }
 
+    static NSSet * const sAttributeValueClasses = [NSSet setWithObjects:[NSDictionary class], [NSArray class], [NSData class], [NSString class], [NSNumber class], nil];
+    _attributes = [decoder decodeObjectOfClasses:sAttributeValueClasses forKey:sAttributesKey];
+    if (_attributes != nil && ![_attributes isKindOfClass:[NSDictionary class]]) {
+        MTR_LOG_ERROR("MTRDeviceClusterData got %@ for attributes, not NSDictionary.", _attributes);
+        return nil;
+    }
+
     return self;
 }
 
 - (void)encodeWithCoder:(NSCoder *)coder
 {
     [coder encodeObject:self.dataVersion forKey:sDataVersionKey];
+    [coder encodeObject:self.attributes forKey:sAttributesKey];
+}
+
+- (id)copyWithZone:(NSZone *)zone
+{
+    return [[MTRDeviceClusterData alloc] initWithDataVersion:_dataVersion attributes:_attributes];
 }
 
 @end
@@ -285,8 +312,11 @@
     NSUInteger _unitTestAttributesReportedSinceLastCheck;
 #endif
     BOOL _delegateDeviceCachePrimedCalled;
+
+    // With MTRDeviceClusterData now able to hold attribute data, the plan is to move to using it
+    // as the read cache, should testing prove attribute storage by cluster is the better solution.
     NSMutableDictionary<MTRClusterPath *, MTRDeviceClusterData *> * _clusterData;
-    NSMutableDictionary<MTRClusterPath *, MTRDeviceClusterData *> * _clusterDataToPersist;
+    NSMutableSet<MTRClusterPath *> * _clustersToPersist;
 }
 
 - (instancetype)initWithNodeID:(NSNumber *)nodeID controller:(MTRDeviceController *)controller
@@ -836,6 +866,34 @@
     }
 }
 
+- (NSDictionary<NSNumber *, MTRDeviceDataValueDictionary> *)_attributesForCluster:(MTRClusterPath *)clusterPath
+{
+    os_unfair_lock_assert_owner(&self->_lock);
+    NSMutableDictionary * attributesToReturn = [NSMutableDictionary dictionary];
+    for (MTRAttributePath * attributePath in _readCache) {
+        if ([attributePath.endpoint isEqualToNumber:clusterPath.endpoint] && [attributePath.cluster isEqualToNumber:clusterPath.cluster]) {
+            attributesToReturn[attributePath.attribute] = _readCache[attributePath];
+        }
+    }
+    return attributesToReturn;
+}
+
+- (NSDictionary<MTRClusterPath *, MTRDeviceClusterData *> *)_clusterDataForPaths:(NSSet<MTRClusterPath *> *)clusterPaths
+{
+    os_unfair_lock_assert_owner(&self->_lock);
+    NSMutableDictionary * clusterDataToReturn = [NSMutableDictionary dictionary];
+    for (MTRClusterPath * clusterPath in clusterPaths) {
+        NSNumber * dataVersion = _clusterData[clusterPath].dataVersion;
+        NSDictionary<NSNumber *, MTRDeviceDataValueDictionary> * attributes = [self _attributesForCluster:clusterPath];
+        if (dataVersion || attributes) {
+            MTRDeviceClusterData * clusterData = [[MTRDeviceClusterData alloc] initWithDataVersion:dataVersion attributes:attributes];
+            clusterDataToReturn[clusterPath] = clusterData;
+        }
+    }
+
+    return clusterDataToReturn;
+}
+
 - (void)_handleReportEnd
 {
     std::lock_guard lock(_lock);
@@ -844,10 +902,11 @@
     _estimatedStartTimeFromGeneralDiagnosticsUpTime = nil;
 
     BOOL dataStoreExists = _deviceController.controllerDataStore != nil;
-    if (dataStoreExists && _clusterDataToPersist.count) {
-        MTR_LOG_DEFAULT("%@ Storing cluster information (data version) count: %lu", self, static_cast<unsigned long>(_clusterDataToPersist.count));
-        [_deviceController.controllerDataStore storeClusterData:_clusterDataToPersist forNodeID:_nodeID];
-        _clusterDataToPersist = nil;
+    if (dataStoreExists && _clustersToPersist.count) {
+        MTR_LOG_DEFAULT("%@ Storing cluster information (data version) count: %lu", self, static_cast<unsigned long>(_clustersToPersist.count));
+        NSDictionary<MTRClusterPath *, MTRDeviceClusterData *> * clusterData = [self _clusterDataForPaths:_clustersToPersist];
+        [_deviceController.controllerDataStore storeClusterData:clusterData forNodeID:_nodeID];
+        _clustersToPersist = nil;
     }
 
 // For unit testing only
@@ -1939,13 +1998,28 @@
 
 - (BOOL)_attributeDataValue:(NSDictionary *)one isEqualToDataValue:(NSDictionary *)theOther
 {
-    // Attribute data-value dictionaries are equal if type and value are equal
+    // Sanity check for nil cases
+    if (!one && !theOther) {
+        MTR_LOG_ERROR("%@ attribute data-value comparison does not expect comparing two nil dictionaries", self);
+        return YES;
+    }
+    if (!one || !theOther) {
+        MTR_LOG_ERROR("%@ attribute data-value comparison does not expect a dictionary to be nil: %@ %@", self, one, theOther);
+        return NO;
+    }
+
+    // Attribute data-value dictionaries are equal if type and value are equal, and specifically, this should return true if values are both nil
     return [one[MTRTypeKey] isEqual:theOther[MTRTypeKey]] && ((one[MTRValueKey] == theOther[MTRValueKey]) || [one[MTRValueKey] isEqual:theOther[MTRValueKey]]);
 }
 
 // Utility to return data value dictionary without data version
 - (NSDictionary *)_dataValueWithoutDataVersion:(NSDictionary *)attributeValue;
 {
+    // Sanity check for nil - return the same input to fail gracefully
+    if (!attributeValue || !attributeValue[MTRTypeKey]) {
+        return attributeValue;
+    }
+
     if (attributeValue[MTRValueKey]) {
         return @{ MTRTypeKey : attributeValue[MTRTypeKey], MTRValueKey : attributeValue[MTRValueKey] };
     } else {
@@ -1954,9 +2028,12 @@
 }
 
 // Update cluster data version and also note the change, so at onReportEnd it can be persisted
-- (void)_updateDataVersion:(NSNumber *)dataVersion forClusterPath:(MTRClusterPath *)clusterPath
+- (void)_noteDataVersion:(NSNumber *)dataVersion forClusterPath:(MTRClusterPath *)clusterPath
 {
+    os_unfair_lock_assert_owner(&self->_lock);
+
     BOOL dataVersionChanged = NO;
+    // Update data version used for subscription filtering
     MTRDeviceClusterData * clusterData = _clusterData[clusterPath];
     if (!clusterData) {
         clusterData = [[MTRDeviceClusterData alloc] init];
@@ -1969,17 +2046,25 @@
     if (dataVersionChanged) {
         clusterData.dataVersion = dataVersion;
 
-        // Set up for persisting if there is data store
+        // Mark cluster path as needing persistence if needed
         BOOL dataStoreExists = _deviceController.controllerDataStore != nil;
         if (dataStoreExists) {
-            if (!_clusterDataToPersist) {
-                _clusterDataToPersist = [NSMutableDictionary dictionary];
-            }
-            _clusterDataToPersist[clusterPath] = clusterData;
+            [self _noteChangeForClusterPath:clusterPath];
         }
     }
 }
 
+// Assuming data store exists, note that the cluster should be persisted at onReportEnd
+- (void)_noteChangeForClusterPath:(MTRClusterPath *)clusterPath
+{
+    os_unfair_lock_assert_owner(&self->_lock);
+
+    if (!_clustersToPersist) {
+        _clustersToPersist = [NSMutableSet set];
+    }
+    [_clustersToPersist addObject:clusterPath];
+}
+
 // assume lock is held
 - (NSArray *)_getAttributesToReportWithReportedValues:(NSArray<NSDictionary<NSString *, id> *> *)reportedAttributeValues
 {
@@ -1988,10 +2073,12 @@
     NSMutableArray * attributesToReport = [NSMutableArray array];
     NSMutableArray * attributePathsToReport = [NSMutableArray array];
     BOOL dataStoreExists = _deviceController.controllerDataStore != nil;
+#if !MTRDEVICE_ATTRIBUTE_CACHE_STORE_ATTRIBUTES_BY_CLUSTER
     NSMutableArray * attributesToPersist;
     if (dataStoreExists) {
         attributesToPersist = [NSMutableArray array];
     }
+#endif
     for (NSDictionary<NSString *, id> * attributeResponseValue in reportedAttributeValues) {
         MTRAttributePath * attributePath = attributeResponseValue[MTRAttributePathKey];
         NSDictionary * attributeDataValue = attributeResponseValue[MTRDataKey];
@@ -2023,9 +2110,9 @@
         } else {
             // First separate data version and restore data value to a form without data version
             NSNumber * dataVersion = attributeDataValue[MTRDataVersionKey];
+            MTRClusterPath * clusterPath = [MTRClusterPath clusterPathWithEndpointID:attributePath.endpoint clusterID:attributePath.cluster];
             if (dataVersion) {
-                MTRClusterPath * clusterPath = [MTRClusterPath clusterPathWithEndpointID:attributePath.endpoint clusterID:attributePath.cluster];
-                [self _updateDataVersion:dataVersion forClusterPath:clusterPath];
+                [self _noteDataVersion:dataVersion forClusterPath:clusterPath];
 
                 // Remove data version from what we cache in memory
                 attributeDataValue = [self _dataValueWithoutDataVersion:attributeDataValue];
@@ -2034,6 +2121,9 @@
             BOOL readCacheValueChanged = ![self _attributeDataValue:attributeDataValue isEqualToDataValue:_readCache[attributePath]];
             // Check if attribute needs to be persisted - compare only to read cache and disregard expected values
             if (dataStoreExists && readCacheValueChanged) {
+#if MTRDEVICE_ATTRIBUTE_CACHE_STORE_ATTRIBUTES_BY_CLUSTER
+                [self _noteChangeForClusterPath:clusterPath];
+#else
                 NSDictionary * attributeResponseValueToPersist;
                 if (dataVersion) {
                     // Remove data version from what we cache in memory and storage
@@ -2044,6 +2134,7 @@
                     attributeResponseValueToPersist = attributeResponseValue;
                 }
                 [attributesToPersist addObject:attributeResponseValueToPersist];
+#endif
             }
 
             // if expected values exists, purge and update read cache
@@ -2106,17 +2197,22 @@
 
     MTR_LOG_INFO("%@ report from reported values %@", self, attributePathsToReport);
 
+#if !MTRDEVICE_ATTRIBUTE_CACHE_STORE_ATTRIBUTES_BY_CLUSTER
     if (dataStoreExists && attributesToPersist.count) {
         [_deviceController.controllerDataStore storeAttributeValues:attributesToPersist forNodeID:_nodeID];
     }
+#endif
 
     return attributesToReport;
 }
 
-- (void)setAttributeValues:(NSArray<NSDictionary *> *)attributeValues reportChanges:(BOOL)reportChanges
+- (void)_setAttributeValues:(NSArray<NSDictionary *> *)attributeValues reportChanges:(BOOL)reportChanges
 {
-    MTR_LOG_INFO("%@ setAttributeValues count: %lu reportChanges: %d", self, static_cast<unsigned long>(attributeValues.count), reportChanges);
-    std::lock_guard lock(_lock);
+    os_unfair_lock_assert_owner(&self->_lock);
+
+    if (!attributeValues.count) {
+        return;
+    }
 
     if (reportChanges) {
         [self _reportAttributes:[self _getAttributesToReportWithReportedValues:attributeValues]];
@@ -2134,8 +2230,42 @@
     }
 }
 
+- (void)setAttributeValues:(NSArray<NSDictionary *> *)attributeValues reportChanges:(BOOL)reportChanges
+{
+    MTR_LOG_INFO("%@ setAttributeValues count: %lu reportChanges: %d", self, static_cast<unsigned long>(attributeValues.count), reportChanges);
+    std::lock_guard lock(_lock);
+    [self _setAttributeValues:attributeValues reportChanges:reportChanges];
+}
+
 - (void)setClusterData:(NSDictionary<MTRClusterPath *, MTRDeviceClusterData *> *)clusterData
 {
+    MTR_LOG_INFO("%@ setClusterData count: %lu", self, static_cast<unsigned long>(clusterData.count));
+    if (!clusterData.count) {
+        return;
+    }
+
+    std::lock_guard lock(_lock);
+
+#if MTRDEVICE_ATTRIBUTE_CACHE_STORE_ATTRIBUTES_BY_CLUSTER
+    // For each cluster, extract and create the attribute response-value for the read cache
+    // TODO: consider some optimization in how the read cache is structured so there's fewer conversions from this format to what's in the cache
+    for (MTRClusterPath * clusterPath in clusterData) {
+        MTRDeviceClusterData * data = clusterData[clusterPath];
+        // Build and set attributes one cluster at a time to avoid creating a ton of temporary objects at a time
+        @autoreleasepool {
+            NSMutableArray * attributeValues = [NSMutableArray array];
+            for (NSNumber * attributeID in data.attributes) {
+                MTRAttributePath * attributePath = [MTRAttributePath attributePathWithEndpointID:clusterPath.endpoint clusterID:clusterPath.cluster attributeID:attributeID];
+                NSDictionary * responseValue = @{ MTRAttributePathKey : attributePath, MTRDataKey : data.attributes[attributeID] };
+                [attributeValues addObject:responseValue];
+            }
+            if (attributeValues.count) {
+                [self _setAttributeValues:attributeValues reportChanges:NO];
+            }
+        }
+    }
+#endif
+
     [_clusterData addEntriesFromDictionary:clusterData];
 }
 
diff --git a/src/darwin/Framework/CHIP/MTRDeviceController.mm b/src/darwin/Framework/CHIP/MTRDeviceController.mm
index 1e52592..95c21e8 100644
--- a/src/darwin/Framework/CHIP/MTRDeviceController.mm
+++ b/src/darwin/Framework/CHIP/MTRDeviceController.mm
@@ -926,12 +926,15 @@
             _nodeIDToDeviceMap[nodeID] = deviceToReturn;
         }
 
+#if !MTRDEVICE_ATTRIBUTE_CACHE_STORE_ATTRIBUTES_BY_CLUSTER
         // Load persisted attributes if they exist.
         NSArray * attributesFromCache = [_controllerDataStore getStoredAttributesForNodeID:nodeID];
         MTR_LOG_INFO("Loaded %lu attributes from storage for %@", static_cast<unsigned long>(attributesFromCache.count), deviceToReturn);
         if (attributesFromCache.count) {
             [deviceToReturn setAttributeValues:attributesFromCache reportChanges:NO];
         }
+#endif
+        // Load persisted cluster data if they exist.
         NSDictionary * clusterData = [_controllerDataStore getStoredClusterDataForNodeID:nodeID];
         MTR_LOG_INFO("Loaded %lu cluster data from storage for %@", static_cast<unsigned long>(clusterData.count), deviceToReturn);
         if (clusterData.count) {
diff --git a/src/darwin/Framework/CHIP/MTRDeviceControllerDataStore.h b/src/darwin/Framework/CHIP/MTRDeviceControllerDataStore.h
index 90007db..c49bc89 100644
--- a/src/darwin/Framework/CHIP/MTRDeviceControllerDataStore.h
+++ b/src/darwin/Framework/CHIP/MTRDeviceControllerDataStore.h
@@ -68,8 +68,8 @@
  * Storage for MTRDevice attribute read cache. This is local-only storage as an optimization. New controller devices using MTRDevice API can prime their own local cache from devices directly.
  */
 - (nullable NSArray<NSDictionary *> *)getStoredAttributesForNodeID:(NSNumber *)nodeID;
-- (nullable NSDictionary<MTRClusterPath *, MTRDeviceClusterData *> *)getStoredClusterDataForNodeID:(NSNumber *)nodeID;
 - (void)storeAttributeValues:(NSArray<NSDictionary *> *)dataValues forNodeID:(NSNumber *)nodeID;
+- (nullable NSDictionary<MTRClusterPath *, MTRDeviceClusterData *> *)getStoredClusterDataForNodeID:(NSNumber *)nodeID;
 - (void)storeClusterData:(NSDictionary<MTRClusterPath *, MTRDeviceClusterData *> *)clusterData forNodeID:(NSNumber *)nodeID;
 - (void)clearStoredAttributesForNodeID:(NSNumber *)nodeID;
 - (void)clearAllStoredAttributes;
diff --git a/src/darwin/Framework/CHIP/MTRDeviceControllerDataStore.mm b/src/darwin/Framework/CHIP/MTRDeviceControllerDataStore.mm
index 38a089b..49a653e 100644
--- a/src/darwin/Framework/CHIP/MTRDeviceControllerDataStore.mm
+++ b/src/darwin/Framework/CHIP/MTRDeviceControllerDataStore.mm
@@ -371,16 +371,31 @@
 
 - (nullable NSArray<NSNumber *> *)_fetchEndpointIndexForNodeID:(NSNumber *)nodeID
 {
+    if (!nodeID) {
+        MTR_LOG_ERROR("%s: unexpected nil input", __func__);
+        return nil;
+    }
+
     return [self _fetchAttributeCacheValueForKey:[self _endpointIndexKeyForNodeID:nodeID] expectedClass:[NSArray class]];
 }
 
 - (BOOL)_storeEndpointIndex:(NSArray<NSNumber *> *)endpointIndex forNodeID:(NSNumber *)nodeID
 {
+    if (!nodeID) {
+        MTR_LOG_ERROR("%s: unexpected nil input", __func__);
+        return NO;
+    }
+
     return [self _storeAttributeCacheValue:endpointIndex forKey:[self _endpointIndexKeyForNodeID:nodeID]];
 }
 
 - (BOOL)_deleteEndpointIndexForNodeID:(NSNumber *)nodeID
 {
+    if (!nodeID) {
+        MTR_LOG_ERROR("%s: unexpected nil input", __func__);
+        return NO;
+    }
+
     return [self _removeAttributeCacheValueForKey:[self _endpointIndexKeyForNodeID:nodeID]];
 }
 
@@ -393,16 +408,31 @@
 
 - (nullable NSArray<NSNumber *> *)_fetchClusterIndexForNodeID:(NSNumber *)nodeID endpointID:(NSNumber *)endpointID
 {
+    if (!nodeID || !endpointID) {
+        MTR_LOG_ERROR("%s: unexpected nil input", __func__);
+        return nil;
+    }
+
     return [self _fetchAttributeCacheValueForKey:[self _clusterIndexKeyForNodeID:nodeID endpointID:endpointID] expectedClass:[NSArray class]];
 }
 
 - (BOOL)_storeClusterIndex:(NSArray<NSNumber *> *)clusterIndex forNodeID:(NSNumber *)nodeID endpointID:(NSNumber *)endpointID
 {
+    if (!nodeID || !endpointID) {
+        MTR_LOG_ERROR("%s: unexpected nil input", __func__);
+        return NO;
+    }
+
     return [self _storeAttributeCacheValue:clusterIndex forKey:[self _clusterIndexKeyForNodeID:nodeID endpointID:endpointID]];
 }
 
 - (BOOL)_deleteClusterIndexForNodeID:(NSNumber *)nodeID endpointID:(NSNumber *)endpointID
 {
+    if (!nodeID || !endpointID) {
+        MTR_LOG_ERROR("%s: unexpected nil input", __func__);
+        return NO;
+    }
+
     return [self _removeAttributeCacheValueForKey:[self _clusterIndexKeyForNodeID:nodeID endpointID:endpointID]];
 }
 
@@ -415,16 +445,31 @@
 
 - (nullable MTRDeviceClusterData *)_fetchClusterDataForNodeID:(NSNumber *)nodeID endpointID:(NSNumber *)endpointID clusterID:(NSNumber *)clusterID
 {
+    if (!nodeID || !endpointID || !clusterID) {
+        MTR_LOG_ERROR("%s: unexpected nil input", __func__);
+        return nil;
+    }
+
     return [self _fetchAttributeCacheValueForKey:[self _clusterDataKeyForNodeID:nodeID endpointID:endpointID clusterID:clusterID] expectedClass:[MTRDeviceClusterData class]];
 }
 
 - (BOOL)_storeClusterData:(MTRDeviceClusterData *)clusterData forNodeID:(NSNumber *)nodeID endpointID:(NSNumber *)endpointID clusterID:(NSNumber *)clusterID
 {
+    if (!nodeID || !endpointID || !clusterID || !clusterData) {
+        MTR_LOG_ERROR("%s: unexpected nil input", __func__);
+        return NO;
+    }
+
     return [self _storeAttributeCacheValue:clusterData forKey:[self _clusterDataKeyForNodeID:nodeID endpointID:endpointID clusterID:clusterID]];
 }
 
 - (BOOL)_deleteClusterDataForNodeID:(NSNumber *)nodeID endpointID:(NSNumber *)endpointID clusterID:(NSNumber *)clusterID
 {
+    if (!nodeID || !endpointID || !clusterID) {
+        MTR_LOG_ERROR("%s: unexpected nil input", __func__);
+        return NO;
+    }
+
     return [self _removeAttributeCacheValueForKey:[self _clusterDataKeyForNodeID:nodeID endpointID:endpointID clusterID:clusterID]];
 }
 
@@ -510,7 +555,7 @@
 #endif
 
         for (NSNumber * endpointID in endpointIndex) {
-            // Fetch endpoint index
+            // Fetch cluster index
             NSArray<NSNumber *> * clusterIndex = [self _fetchClusterIndexForNodeID:nodeID endpointID:endpointID];
 
 #if ATTRIBUTE_CACHE_VERBOSE_LOGGING
@@ -518,7 +563,7 @@
 #endif
 
             for (NSNumber * clusterID in clusterIndex) {
-                // Fetch endpoint index
+                // Fetch attribute index
                 NSArray<NSNumber *> * attributeIndex = [self _fetchAttributeIndexForNodeID:nodeID endpointID:endpointID clusterID:clusterID];
 
 #if ATTRIBUTE_CACHE_VERBOSE_LOGGING
@@ -576,12 +621,12 @@
         NSMutableArray<NSNumber *> * endpointIndexCopy = [endpointIndex mutableCopy];
 
         for (NSNumber * endpointID in endpointIndex) {
-            // Fetch endpoint index
+            // Fetch cluster index
             NSArray<NSNumber *> * clusterIndex = [self _fetchClusterIndexForNodeID:nodeID endpointID:endpointID];
             NSMutableArray<NSNumber *> * clusterIndexCopy = [clusterIndex mutableCopy];
 
             for (NSNumber * clusterID in clusterIndex) {
-                // Fetch endpoint index
+                // Fetch attribute index
                 NSArray<NSNumber *> * attributeIndex = [self _fetchAttributeIndexForNodeID:nodeID endpointID:endpointID clusterID:clusterID];
                 NSMutableArray<NSNumber *> * attributeIndexCopy = [attributeIndex mutableCopy];
 
@@ -857,6 +902,11 @@
 
 - (nullable NSDictionary<MTRClusterPath *, MTRDeviceClusterData *> *)getStoredClusterDataForNodeID:(NSNumber *)nodeID
 {
+    if (!nodeID) {
+        MTR_LOG_ERROR("%s: unexpected nil input", __func__);
+        return nil;
+    }
+
     __block NSMutableDictionary<MTRClusterPath *, MTRDeviceClusterData *> * clusterDataToReturn = nil;
     dispatch_sync(_storageDelegateQueue, ^{
         // Fetch node index
@@ -887,7 +937,7 @@
 #endif
 
         for (NSNumber * endpointID in endpointIndex) {
-            // Fetch endpoint index
+            // Fetch cluster index
             NSArray<NSNumber *> * clusterIndex = [self _fetchClusterIndexForNodeID:nodeID endpointID:endpointID];
 
 #if ATTRIBUTE_CACHE_VERBOSE_LOGGING
@@ -920,6 +970,16 @@
 
 - (void)storeClusterData:(NSDictionary<MTRClusterPath *, MTRDeviceClusterData *> *)clusterData forNodeID:(NSNumber *)nodeID
 {
+    if (!nodeID) {
+        MTR_LOG_ERROR("%s: unexpected nil input", __func__);
+        return;
+    }
+
+    if (!clusterData.count) {
+        MTR_LOG_ERROR("%s: nothing to store", __func__);
+        return;
+    }
+
     dispatch_async(_storageDelegateQueue, ^{
         NSUInteger storeFailures = 0;
 
diff --git a/src/darwin/Framework/CHIP/MTRDevice_Internal.h b/src/darwin/Framework/CHIP/MTRDevice_Internal.h
index cb6aa46..5d82847 100644
--- a/src/darwin/Framework/CHIP/MTRDevice_Internal.h
+++ b/src/darwin/Framework/CHIP/MTRDevice_Internal.h
@@ -28,13 +28,16 @@
 
 typedef void (^MTRDevicePerformAsyncBlock)(MTRBaseDevice * baseDevice);
 
+// Whether to store attributes by cluster instead of as individual entries for each attribute
+#define MTRDEVICE_ATTRIBUTE_CACHE_STORE_ATTRIBUTES_BY_CLUSTER 1
+
 /**
  * Information about a cluster, currently is just data version
  */
 MTR_TESTABLE
-@interface MTRDeviceClusterData : NSObject <NSSecureCoding>
+@interface MTRDeviceClusterData : NSObject <NSSecureCoding, NSCopying>
 @property (nonatomic) NSNumber * dataVersion;
-// TODO: add cluster attributes in this object, and remove direct attribute storage
+@property (nonatomic) NSDictionary<NSNumber *, MTRDeviceDataValueDictionary> * attributes; // attributeID => data-value dictionary
 @end
 
 @interface MTRDevice ()
diff --git a/src/darwin/Framework/CHIPTests/MTRDeviceTests.m b/src/darwin/Framework/CHIPTests/MTRDeviceTests.m
index e5539f4..1d71c7b 100644
--- a/src/darwin/Framework/CHIPTests/MTRDeviceTests.m
+++ b/src/darwin/Framework/CHIPTests/MTRDeviceTests.m
@@ -27,6 +27,7 @@
 #import "MTRCommandPayloadExtensions_Internal.h"
 #import "MTRDeviceControllerLocalTestStorage.h"
 #import "MTRDeviceTestDelegate.h"
+#import "MTRDevice_Internal.h"
 #import "MTRErrorTestUtils.h"
 #import "MTRTestDeclarations.h"
 #import "MTRTestKeys.h"
@@ -2870,8 +2871,13 @@
 
     NSUInteger attributesReportedWithFirstSubscription = [device unitTestAttributesReportedSinceLastCheck];
 
+#if MTRDEVICE_ATTRIBUTE_CACHE_STORE_ATTRIBUTES_BY_CLUSTER
+    NSDictionary * dataStoreClusterDataAfterFirstSubscription = [sController.controllerDataStore getStoredClusterDataForNodeID:@(kDeviceId)];
+    XCTAssertTrue(dataStoreClusterDataAfterFirstSubscription.count > 0);
+#else
     NSArray * dataStoreValuesAfterFirstSubscription = [sController.controllerDataStore getStoredAttributesForNodeID:@(kDeviceId)];
     XCTAssertTrue(dataStoreValuesAfterFirstSubscription.count > 0);
+#endif
 
     // Now remove device, resubscribe, and see that it succeeds
     [sController removeDevice:device];
diff --git a/src/darwin/Framework/CHIPTests/MTRPerControllerStorageTests.m b/src/darwin/Framework/CHIPTests/MTRPerControllerStorageTests.m
index 0fede08..fce1ff1 100644
--- a/src/darwin/Framework/CHIPTests/MTRPerControllerStorageTests.m
+++ b/src/darwin/Framework/CHIPTests/MTRPerControllerStorageTests.m
@@ -1249,7 +1249,21 @@
 
     [self waitForExpectations:@[ subscriptionExpectation ] timeout:60];
 
+    NSUInteger dataStoreValuesCount = 0;
+#if MTRDEVICE_ATTRIBUTE_CACHE_STORE_ATTRIBUTES_BY_CLUSTER
+    NSDictionary<MTRClusterPath *, MTRDeviceClusterData *> * dataStoreClusterData = [controller.controllerDataStore getStoredClusterDataForNodeID:deviceID];
+    for (MTRClusterPath * path in dataStoreClusterData) {
+        MTRDeviceClusterData * data = dataStoreClusterData[path];
+        for (NSNumber * attributeID in data.attributes) {
+            dataStoreValuesCount++;
+            NSDictionary * dataValue = data.attributes[attributeID];
+            NSDictionary * dataValueFromMTRDevice = [device readAttributeWithEndpointID:path.endpoint clusterID:path.cluster attributeID:attributeID params:nil];
+            XCTAssertTrue([device _attributeDataValue:dataValue isEqualToDataValue:dataValueFromMTRDevice]);
+        }
+    }
+#else
     NSArray * dataStoreValues = [controller.controllerDataStore getStoredAttributesForNodeID:deviceID];
+    dataStoreValuesCount = dataStoreValues.count;
 
     // Verify all values are stored into storage
     for (NSDictionary * responseValue in dataStoreValues) {
@@ -1261,6 +1275,7 @@
         NSDictionary * dataValueFromMTRDevice = [device readAttributeWithEndpointID:path.endpoint clusterID:path.cluster attributeID:path.attribute params:nil];
         XCTAssertTrue([device _attributeDataValue:dataValue isEqualToDataValue:dataValueFromMTRDevice]);
     }
+#endif
 
     // Now force the removal of the object from controller to test reloading read cache from storage
     [controller removeDevice:device];
@@ -1268,6 +1283,18 @@
     // Verify the new device is initialized with the same values
     __auto_type * newDevice = [MTRDevice deviceWithNodeID:deviceID controller:controller];
     NSUInteger storedAttributeDifferFromMTRDeviceCount = 0;
+#if MTRDEVICE_ATTRIBUTE_CACHE_STORE_ATTRIBUTES_BY_CLUSTER
+    for (MTRClusterPath * path in dataStoreClusterData) {
+        MTRDeviceClusterData * data = dataStoreClusterData[path];
+        for (NSNumber * attributeID in data.attributes) {
+            NSDictionary * dataValue = data.attributes[attributeID];
+            NSDictionary * dataValueFromMTRDevice = [newDevice readAttributeWithEndpointID:path.endpoint clusterID:path.cluster attributeID:attributeID params:nil];
+            if (![newDevice _attributeDataValue:dataValue isEqualToDataValue:dataValueFromMTRDevice]) {
+                storedAttributeDifferFromMTRDeviceCount++;
+            }
+        }
+    }
+#else
     for (NSDictionary * responseValue in dataStoreValues) {
         MTRAttributePath * path = responseValue[MTRAttributePathKey];
         XCTAssertNotNil(path);
@@ -1279,10 +1306,11 @@
             storedAttributeDifferFromMTRDeviceCount++;
         }
     }
+#endif
 
     // Only test that 90% of attributes are the same because there are some changing attributes each time (UTC time, for example)
     //   * With all-clusters-app as of 2024-02-10, about 1.476% of attributes change.
-    double storedAttributeDifferFromMTRDevicePercentage = storedAttributeDifferFromMTRDeviceCount * 100.0 / dataStoreValues.count;
+    double storedAttributeDifferFromMTRDevicePercentage = storedAttributeDifferFromMTRDeviceCount * 100.0 / dataStoreValuesCount;
     XCTAssertTrue(storedAttributeDifferFromMTRDevicePercentage < 10.0);
 
     // Now
@@ -1303,7 +1331,7 @@
     // 2) Some attributes do change on resubscribe
     //   * With all-clusts-app as of 2024-02-10, out of 1287 persisted attributes, still 450 attributes were reported with filter
     // And so conservatively, assert that data version filters save at least 300 entries.
-    NSUInteger storedAttributeCountDifferenceFromMTRDeviceReport = dataStoreValues.count - [device unitTestAttributesReportedSinceLastCheck];
+    NSUInteger storedAttributeCountDifferenceFromMTRDeviceReport = dataStoreValuesCount - [device unitTestAttributesReportedSinceLastCheck];
     XCTAssertTrue(storedAttributeCountDifferenceFromMTRDeviceReport > 300);
 
     // Reset our commissionee.