Speed up bluetooth matching (#122626)

- use a defaultdict to avoid lots of setdefault
- move the intersection outside of the genexpr
  to avoid entering the genexpr if there is no
  intersection
This commit is contained in:
J. Nick Koston 2024-07-26 02:22:56 -05:00 committed by GitHub
parent 78a98afb8d
commit e262f759af
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -2,6 +2,7 @@
from __future__ import annotations
from collections import defaultdict
from dataclasses import dataclass
from fnmatch import translate
from functools import lru_cache
@ -173,10 +174,10 @@ class BluetoothMatcherIndexBase[
def __init__(self) -> None:
"""Initialize the matcher index."""
self.local_name: dict[str, list[_T]] = {}
self.service_uuid: dict[str, list[_T]] = {}
self.service_data_uuid: dict[str, list[_T]] = {}
self.manufacturer_id: dict[int, list[_T]] = {}
self.local_name: defaultdict[str, list[_T]] = defaultdict(list)
self.service_uuid: defaultdict[str, list[_T]] = defaultdict(list)
self.service_data_uuid: defaultdict[str, list[_T]] = defaultdict(list)
self.manufacturer_id: defaultdict[int, list[_T]] = defaultdict(list)
self.service_uuid_set: set[str] = set()
self.service_data_uuid_set: set[str] = set()
self.manufacturer_id_set: set[int] = set()
@ -190,26 +191,22 @@ class BluetoothMatcherIndexBase[
"""
# Local name is the cheapest to match since its just a dict lookup
if LOCAL_NAME in matcher:
self.local_name.setdefault(
_local_name_to_index_key(matcher[LOCAL_NAME]), []
).append(matcher)
self.local_name[_local_name_to_index_key(matcher[LOCAL_NAME])].append(
matcher
)
return True
# Manufacturer data is 2nd cheapest since its all ints
if MANUFACTURER_ID in matcher:
self.manufacturer_id.setdefault(matcher[MANUFACTURER_ID], []).append(
matcher
)
self.manufacturer_id[matcher[MANUFACTURER_ID]].append(matcher)
return True
if SERVICE_UUID in matcher:
self.service_uuid.setdefault(matcher[SERVICE_UUID], []).append(matcher)
self.service_uuid[matcher[SERVICE_UUID]].append(matcher)
return True
if SERVICE_DATA_UUID in matcher:
self.service_data_uuid.setdefault(matcher[SERVICE_DATA_UUID], []).append(
matcher
)
self.service_data_uuid[matcher[SERVICE_DATA_UUID]].append(matcher)
return True
return False
@ -260,32 +257,38 @@ class BluetoothMatcherIndexBase[
if ble_device_matches(matcher, service_info)
)
if self.service_data_uuid_set and service_info.service_data:
if (
(service_data_uuid_set := self.service_data_uuid_set)
and (service_data := service_info.service_data)
and (matched_uuids := service_data_uuid_set.intersection(service_data))
):
matches.extend(
matcher
for service_data_uuid in self.service_data_uuid_set.intersection(
service_info.service_data
)
for service_data_uuid in matched_uuids
for matcher in self.service_data_uuid[service_data_uuid]
if ble_device_matches(matcher, service_info)
)
if self.manufacturer_id_set and service_info.manufacturer_data:
if (
(manufacturer_id_set := self.manufacturer_id_set)
and (manufacturer_data := service_info.manufacturer_data)
and (matched_ids := manufacturer_id_set.intersection(manufacturer_data))
):
matches.extend(
matcher
for manufacturer_id in self.manufacturer_id_set.intersection(
service_info.manufacturer_data
)
for manufacturer_id in matched_ids
for matcher in self.manufacturer_id[manufacturer_id]
if ble_device_matches(matcher, service_info)
)
if self.service_uuid_set and service_info.service_uuids:
if (
(service_uuid_set := self.service_uuid_set)
and (service_uuids := service_info.service_uuids)
and (matched_uuids := service_uuid_set.intersection(service_uuids))
):
matches.extend(
matcher
for service_uuid in self.service_uuid_set.intersection(
service_info.service_uuids
)
for service_uuid in matched_uuids
for matcher in self.service_uuid[service_uuid]
if ble_device_matches(matcher, service_info)
)
@ -310,7 +313,9 @@ class BluetoothCallbackMatcherIndex(
def __init__(self) -> None:
"""Initialize the matcher index."""
super().__init__()
self.address: dict[str, list[BluetoothCallbackMatcherWithCallback]] = {}
self.address: defaultdict[str, list[BluetoothCallbackMatcherWithCallback]] = (
defaultdict(list)
)
self.connectable: list[BluetoothCallbackMatcherWithCallback] = []
def add_callback_matcher(
@ -323,7 +328,7 @@ class BluetoothCallbackMatcherIndex(
We put them in the bucket that they are most likely to match.
"""
if ADDRESS in matcher:
self.address.setdefault(matcher[ADDRESS], []).append(matcher)
self.address[matcher[ADDRESS]].append(matcher)
return
if super().add(matcher):