From 35aef8bb3677b2df3ba37f25474475daeb6d555e Mon Sep 17 00:00:00 2001
From: Igor Motov <igor@motovs.org>
Date: Tue, 5 May 2020 08:02:26 -0400
Subject: [PATCH] Simplify the ValuesSourceRegistry structure (#56154)

Follow up to #55747.
---
 .../support/ValuesSourceRegistry.java           | 17 +++++++----------
 1 file changed, 7 insertions(+), 10 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSourceRegistry.java b/server/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSourceRegistry.java
index 6be2f97f9f749..7941583af6686 100644
--- a/server/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSourceRegistry.java
+++ b/server/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSourceRegistry.java
@@ -98,19 +98,21 @@ public ValuesSourceRegistry build() {
 
     /** Maps Aggregation names to (ValuesSourceType, Supplier) pairs, keyed by ValuesSourceType */
     private final AggregationUsageService usageService;
-    private Map<String, List<Map.Entry<ValuesSourceType, AggregatorSupplier>>> aggregatorRegistry;
+    private Map<String, Map<ValuesSourceType, AggregatorSupplier>> aggregatorRegistry;
     public ValuesSourceRegistry(Map<String, List<Map.Entry<ValuesSourceType, AggregatorSupplier>>> aggregatorRegistry,
                                 AggregationUsageService usageService) {
         /*
          Make an immutatble copy of our input map. Since this is write once, read many, we'll spend a bit of extra time to shape this
          into a Map.of(), which is more read optimized than just using a hash map.
          */
-        Map.Entry[] copiedEntries = new Map.Entry[aggregatorRegistry.size()];
+        @SuppressWarnings("unchecked")
+        Map.Entry<String, Map<ValuesSourceType, AggregatorSupplier>>[] copiedEntries = new Map.Entry[aggregatorRegistry.size()];
         int i = 0;
         for (Map.Entry<String, List<Map.Entry<ValuesSourceType, AggregatorSupplier>>> entry : aggregatorRegistry.entrySet()) {
             String aggName = entry.getKey();
             List<Map.Entry<ValuesSourceType, AggregatorSupplier>> values = entry.getValue();
-            Map.Entry newEntry = Map.entry(aggName, List.of(values.toArray()));
+            @SuppressWarnings("unchecked") Map.Entry<String, Map<ValuesSourceType, AggregatorSupplier>> newEntry =
+                Map.entry(aggName, Map.ofEntries(values.toArray(new Map.Entry[0])));
             copiedEntries[i++] = newEntry;
         }
         this.aggregatorRegistry = Map.ofEntries(copiedEntries);
@@ -118,13 +120,8 @@ public ValuesSourceRegistry(Map<String, List<Map.Entry<ValuesSourceType, Aggrega
     }
 
     private AggregatorSupplier findMatchingSuppier(ValuesSourceType valuesSourceType,
-                                                   List<Map.Entry<ValuesSourceType, AggregatorSupplier>> supportedTypes) {
-        for (Map.Entry<ValuesSourceType, AggregatorSupplier> candidate : supportedTypes) {
-            if (candidate.getKey().equals(valuesSourceType)) {
-                return candidate.getValue();
-            }
-        }
-        return null;
+                                                   Map<ValuesSourceType, AggregatorSupplier> supportedTypes) {
+        return supportedTypes.get(valuesSourceType);
     }
 
     public AggregatorSupplier getAggregator(ValuesSourceType valuesSourceType, String aggregationName) {