From 9689cb448db3ab5ce635d899fec3feacb4618875 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Fri, 9 Feb 2024 08:42:30 -0600
Subject: [PATCH] Avoid linear search of entity registry in
 async_extract_referenced_entity_ids (#109667)

* Index area_ids in the entity registry

I missed that these are used in _resolve_area in search.

Eventually we can make async_extract_referenced_entity_ids
a bit faster with this as well

* Avoid linear search of entity registry in async_extract_referenced_entity_ids

needs https://github.com/home-assistant/core/pull/109660
---
 homeassistant/helpers/service.py | 43 ++++++++++++++++++++------------
 1 file changed, 27 insertions(+), 16 deletions(-)

diff --git a/homeassistant/helpers/service.py b/homeassistant/helpers/service.py
index d397764c1be..b170026f375 100644
--- a/homeassistant/helpers/service.py
+++ b/homeassistant/helpers/service.py
@@ -496,26 +496,37 @@ def async_extract_referenced_entity_ids(
     if not selector.area_ids and not selected.referenced_devices:
         return selected
 
-    for ent_entry in ent_reg.entities.values():
+    entities = ent_reg.entities
+    # Add indirectly referenced by area
+    selected.indirectly_referenced.update(
+        entry.entity_id
+        for area_id in selector.area_ids
+        # The entity's area matches a targeted area
+        for entry in entities.get_entries_for_area_id(area_id)
+        # Do not add entities which are hidden or which are config
+        # or diagnostic entities.
+        if entry.entity_category is None and entry.hidden_by is None
+    )
+    # Add indirectly referenced by device
+    selected.indirectly_referenced.update(
+        entry.entity_id
+        for device_id in selected.referenced_devices
+        for entry in entities.get_entries_for_device_id(device_id)
         # Do not add entities which are hidden or which are config
         # or diagnostic entities.
-        if ent_entry.entity_category is not None or ent_entry.hidden_by is not None:
-            continue
-
         if (
-            # The entity's area matches a targeted area
-            ent_entry.area_id in selector.area_ids
-            # The entity's device matches a device referenced by an area and the entity
-            # has no explicitly set area
-            or (
-                not ent_entry.area_id
-                and ent_entry.device_id in selected.referenced_devices
+            entry.entity_category is None
+            and entry.hidden_by is None
+            and (
+                # The entity's device matches a device referenced
+                # by an area and the entity
+                # has no explicitly set area
+                not entry.area_id
+                # The entity's device matches a targeted device
+                or device_id in selector.device_ids
             )
-            # The entity's device matches a targeted device
-            or ent_entry.device_id in selector.device_ids
-        ):
-            selected.indirectly_referenced.add(ent_entry.entity_id)
-
+        )
+    )
     return selected