Refactor restore state saving to avoid a dict lookup of ATTR_RESTORED (#106854)

homeassistant/helpers/restore_state.py

@@ -178,8 +178,8 @@ class RestoreStateData:
         now = dt_util.utcnow()
         all_states = self.hass.states.async_all()
         # Entities currently backed by an entity object
-        current_entity_ids = {
-            state.entity_id
+        current_states_by_entity_id = {
+            state.entity_id: state
             for state in all_states
             if not state.attributes.get(ATTR_RESTORED)
         }
@@ -187,13 +187,12 @@ class RestoreStateData:
         # Start with the currently registered states
         stored_states = [
             StoredState(
-                state, self.entities[state.entity_id].extra_restore_state_data, now
+                current_states_by_entity_id[entity_id],
+                entity.extra_restore_state_data,
+                now,
             )
-            for state in all_states
-            if state.entity_id in self.entities
-            and
-            # Ignore all states that are entity registry placeholders
-            not state.attributes.get(ATTR_RESTORED)
+            for entity_id, entity in self.entities.items()
+            if entity_id in current_states_by_entity_id
         ]
         expiration_time = now - STATE_EXPIRATION
 
@@ -201,7 +200,7 @@ class RestoreStateData:
             # Don't save old states that have entities in the current run
             # They are either registered and already part of stored_states,
             # or no longer care about restoring.
-            if entity_id in current_entity_ids:
+            if entity_id in current_states_by_entity_id:
                 continue
 
             # Don't save old states that have expired