From 1bcc8203bb6dd8bd1df76da3b3439cf30a671092 Mon Sep 17 00:00:00 2001
From: Nikolas Havrikov <havrikov@google.com>
Date: Thu, 20 Jan 2022 17:08:34 +0100
Subject: [PATCH] Use mAddressMap in cdm getAssociationsByAddress

This CL leverages the mAddressMap which was specifically introduces
for the kind of address-based lookup that is needed in the
getAssociationsByAddress method.

Test: atest CtsCompanionDeviceManagerCoreTestCases
Test: atest CtsCompanionDeviceManagerUiAutomationTestCases
Change-Id: Ifb00e184bd491e2385214b7c89d64da33dc36ee3
---
 .../android/server/companion/AssociationStoreImpl.java    | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/services/companion/java/com/android/server/companion/AssociationStoreImpl.java b/services/companion/java/com/android/server/companion/AssociationStoreImpl.java
index dbcdd0f877a1d..38f8633eaa0a0 100644
--- a/services/companion/java/com/android/server/companion/AssociationStoreImpl.java
+++ b/services/companion/java/com/android/server/companion/AssociationStoreImpl.java
@@ -213,11 +213,9 @@ class AssociationStoreImpl implements AssociationStore {
             final Set<Integer> ids = mAddressMap.get(address);
             if (ids == null) return Collections.emptyList();
 
-            final List<AssociationInfo> associations = new ArrayList<>();
-            for (AssociationInfo association : mIdMap.values()) {
-                if (address.equals(association.getDeviceMacAddress())) {
-                    associations.add(association);
-                }
+            final List<AssociationInfo> associations = new ArrayList<>(ids.size());
+            for (Integer id : ids) {
+                associations.add(mIdMap.get(id));
             }
 
             return Collections.unmodifiableList(associations);