From 62afb4ba663d9cbd7c4d210270afeb909a38ad2d Mon Sep 17 00:00:00 2001
From: Googler <jhorvitz@google.com>
Date: Wed, 28 Jun 2023 10:05:39 -0700
Subject: [PATCH] Avoid creating an identical `ImmutableMap` when the input to
 `Dict#copyOf` is already an `ImmutableMap`.

Also use `forEach` for iteration to avoid creating entry set views.

PiperOrigin-RevId: 544088895
Change-Id: Iac12fe7754b5c95f5939e4f2124935abc26d955d
---
 .../java/net/starlark/java/eval/Dict.java     | 36 ++++++++++---------
 1 file changed, 20 insertions(+), 16 deletions(-)

diff --git a/src/main/java/net/starlark/java/eval/Dict.java b/src/main/java/net/starlark/java/eval/Dict.java
index 649ec3520943fb..6ff9a7eb66f474 100644
--- a/src/main/java/net/starlark/java/eval/Dict.java
+++ b/src/main/java/net/starlark/java/eval/Dict.java
@@ -441,31 +441,35 @@ public static <K, V> Dict<K, V> copyOf(@Nullable Mutability mu, Map<? extends K,
       mu = Mutability.IMMUTABLE;
     }
 
-    if (mu == Mutability.IMMUTABLE && m instanceof Dict && ((Dict) m).isImmutable()) {
-      @SuppressWarnings("unchecked")
-      Dict<K, V> dict = (Dict<K, V>) m; // safe
-      return dict;
-    }
-
     if (mu == Mutability.IMMUTABLE) {
       if (m.isEmpty()) {
         return empty();
       }
+
+      if (m instanceof ImmutableMap) {
+        m.forEach(
+            (k, v) -> {
+              Starlark.checkValid(k);
+              Starlark.checkValid(v);
+            });
+        @SuppressWarnings("unchecked")
+        var immutableMap = (ImmutableMap<K, V>) m;
+        return new Dict<>(immutableMap);
+      }
+
+      if (m instanceof Dict && ((Dict<?, ?>) m).isImmutable()) {
+        @SuppressWarnings("unchecked")
+        var dict = (Dict<K, V>) m;
+        return dict;
+      }
+
       ImmutableMap.Builder<K, V> immutableMapBuilder =
           ImmutableMap.builderWithExpectedSize(m.size());
-      for (Map.Entry<? extends K, ? extends V> e : m.entrySet()) {
-        immutableMapBuilder.put(
-            Starlark.checkValid(e.getKey()), //
-            Starlark.checkValid(e.getValue()));
-      }
+      m.forEach((k, v) -> immutableMapBuilder.put(Starlark.checkValid(k), Starlark.checkValid(v)));
       return new Dict<>(immutableMapBuilder.buildOrThrow());
     } else {
       LinkedHashMap<K, V> linkedHashMap = Maps.newLinkedHashMapWithExpectedSize(m.size());
-      for (Map.Entry<? extends K, ? extends V> e : m.entrySet()) {
-        linkedHashMap.put(
-            Starlark.checkValid(e.getKey()), //
-            Starlark.checkValid(e.getValue()));
-      }
+      m.forEach((k, v) -> linkedHashMap.put(Starlark.checkValid(k), Starlark.checkValid(v)));
       return new Dict<>(mu, linkedHashMap);
     }
   }