Optimize MapUtils slightly

smithy-lang · Dec 28, 2021 · 0a1a825 · 0a1a825
1 parent 331ab01
commit 0a1a825
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 8 deletions.
diff --git a/smithy-utils/src/main/java/software/amazon/smithy/utils/MapUtils.java b/smithy-utils/src/main/java/software/amazon/smithy/utils/MapUtils.java
@@ -76,9 +76,7 @@ public static <K, V> Map<K, V> of() {
      * @throws NullPointerException if the key or the value is {@code null}
      */
     public static <K, V> Map<K, V> of(K k1, V v1) {
-        Map<K, V> result = new HashMap<>(1);
-        result.put(k1, v1);
-        return Collections.unmodifiableMap(result);
+        return Collections.singletonMap(k1, v1);
     }
 
     /**
@@ -399,11 +397,17 @@ public static <K, V> Map.Entry<K, V> entry(K key, V value) {
     @SafeVarargs
     @SuppressWarnings("varargs")
     public static <K, V> Map<K, V> ofEntries(Map.Entry<? extends K, ? extends V>... entries) {
-        Map<K, V> result = new HashMap<>(entries.length);
-        for (Map.Entry<? extends K, ? extends V> entry : entries) {
-            result.put(entry.getKey(), entry.getValue());
+        if (entries.length == 0) {
+            return MapUtils.of();
+        } else if (entries.length == 1) {
+            return MapUtils.of(entries[0].getKey(), entries[0].getValue());
+        } else {
+            Map<K, V> result = new HashMap<>(entries.length);
+            for (Map.Entry<? extends K, ? extends V> entry : entries) {
+                result.put(entry.getKey(), entry.getValue());
+            }
+            return Collections.unmodifiableMap(result);
         }
-        return Collections.unmodifiableMap(result);
     }
 
     /**

diff --git a/smithy-utils/src/test/java/software/amazon/smithy/utils/MapUtilsTest.java b/smithy-utils/src/test/java/software/amazon/smithy/utils/MapUtilsTest.java
@@ -18,8 +18,10 @@
 import static java.util.function.Function.identity;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.anEmptyMap;
+import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasEntry;
+import static org.hamcrest.Matchers.hasSize;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
 import java.util.Collections;
@@ -182,12 +184,29 @@ public void suppliesEntry() {
     }
 
     @Test
-    public void buildsFromEntries() {
+    public void buildsFromTwoEntries() {
         Map<String, String> map = MapUtils.ofEntries(
                 MapUtils.entry("1", "A"),
                 MapUtils.entry("2", "B"));
+
         assertThat(map, hasEntry("1", "A"));
         assertThat(map, hasEntry("2", "B"));
+        assertThat(map.entrySet(), hasSize(2));
+    }
+
+    @Test
+    public void buildsFromEntriesSingleEntry() {
+        Map<String, String> map = MapUtils.ofEntries(MapUtils.entry("1", "A"));
+
+        assertThat(map, hasEntry("1", "A"));
+        assertThat(map.entrySet(), hasSize(1));
+    }
+
+    @Test
+    public void buildsFromEntriesEmpty() {
+        Map<String, String> map = MapUtils.ofEntries();
+
+        assertThat(map.entrySet(), empty());
     }
 
     @Test